mirror of https://github.com/apache/superset.git
feat(db_engine): Add custom_user_agent when connecting to MotherDuck (#27665)
This commit is contained in:
parent
8ae4662f17
commit
fcf90dffa8
|
@ -25,7 +25,8 @@ from flask_babel import gettext as __
|
||||||
from sqlalchemy import types
|
from sqlalchemy import types
|
||||||
from sqlalchemy.engine.reflection import Inspector
|
from sqlalchemy.engine.reflection import Inspector
|
||||||
|
|
||||||
from superset.constants import TimeGrain
|
from superset.config import VERSION_STRING
|
||||||
|
from superset.constants import TimeGrain, USER_AGENT
|
||||||
from superset.db_engine_specs.base import BaseEngineSpec
|
from superset.db_engine_specs.base import BaseEngineSpec
|
||||||
from superset.errors import SupersetErrorType
|
from superset.errors import SupersetErrorType
|
||||||
|
|
||||||
|
@ -41,6 +42,8 @@ class DuckDBEngineSpec(BaseEngineSpec):
|
||||||
engine = "duckdb"
|
engine = "duckdb"
|
||||||
engine_name = "DuckDB"
|
engine_name = "DuckDB"
|
||||||
|
|
||||||
|
sqlalchemy_uri_placeholder = "duckdb:////path/to/duck.db"
|
||||||
|
|
||||||
_time_grain_expressions = {
|
_time_grain_expressions = {
|
||||||
None: "{col}",
|
None: "{col}",
|
||||||
TimeGrain.SECOND: "DATE_TRUNC('second', {col})",
|
TimeGrain.SECOND: "DATE_TRUNC('second', {col})",
|
||||||
|
@ -81,9 +84,28 @@ class DuckDBEngineSpec(BaseEngineSpec):
|
||||||
) -> set[str]:
|
) -> set[str]:
|
||||||
return set(inspector.get_table_names(schema))
|
return set(inspector.get_table_names(schema))
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def get_extra_params(database: Database) -> dict[str, Any]:
|
||||||
|
"""
|
||||||
|
Add a user agent to be used in the requests.
|
||||||
|
"""
|
||||||
|
extra: dict[str, Any] = BaseEngineSpec.get_extra_params(database)
|
||||||
|
engine_params: dict[str, Any] = extra.setdefault("engine_params", {})
|
||||||
|
connect_args: dict[str, Any] = engine_params.setdefault("connect_args", {})
|
||||||
|
config: dict[str, Any] = connect_args.setdefault("config", {})
|
||||||
|
custom_user_agent = config.pop("custom_user_agent", "")
|
||||||
|
delim = " " if custom_user_agent else ""
|
||||||
|
user_agent = USER_AGENT.replace(" ", "-").lower()
|
||||||
|
user_agent = f"{user_agent}/{VERSION_STRING}{delim}{custom_user_agent}"
|
||||||
|
config.setdefault("custom_user_agent", user_agent)
|
||||||
|
|
||||||
|
return extra
|
||||||
|
|
||||||
|
|
||||||
class MotherDuckEngineSpec(DuckDBEngineSpec):
|
class MotherDuckEngineSpec(DuckDBEngineSpec):
|
||||||
engine = "duckdb"
|
engine = "duckdb"
|
||||||
engine_name = "MotherDuck"
|
engine_name = "MotherDuck"
|
||||||
|
|
||||||
sqlalchemy_uri_placeholder = "duckdb:///md:{SERVICE_TOKEN}@{database_name}"
|
sqlalchemy_uri_placeholder = (
|
||||||
|
"duckdb:///md:{database_name}?motherduck_token={SERVICE_TOKEN}"
|
||||||
|
)
|
||||||
|
|
|
@ -15,11 +15,14 @@
|
||||||
# specific language governing permissions and limitations
|
# specific language governing permissions and limitations
|
||||||
# under the License.
|
# under the License.
|
||||||
|
|
||||||
|
import json
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
from pytest_mock import MockerFixture
|
||||||
|
|
||||||
|
from superset.config import VERSION_STRING
|
||||||
from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm
|
from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm
|
||||||
from tests.unit_tests.fixtures.common import dttm
|
from tests.unit_tests.fixtures.common import dttm
|
||||||
|
|
||||||
|
@ -38,3 +41,34 @@ def test_convert_dttm(
|
||||||
from superset.db_engine_specs.duckdb import DuckDBEngineSpec as spec
|
from superset.db_engine_specs.duckdb import DuckDBEngineSpec as spec
|
||||||
|
|
||||||
assert_convert_dttm(spec, target_type, expected_result, dttm)
|
assert_convert_dttm(spec, target_type, expected_result, dttm)
|
||||||
|
|
||||||
|
|
||||||
|
def test_get_extra_params(mocker: MockerFixture) -> None:
|
||||||
|
"""
|
||||||
|
Test the ``get_extra_params`` method.
|
||||||
|
"""
|
||||||
|
from superset.db_engine_specs.duckdb import DuckDBEngineSpec
|
||||||
|
|
||||||
|
database = mocker.MagicMock()
|
||||||
|
|
||||||
|
database.extra = {}
|
||||||
|
assert DuckDBEngineSpec.get_extra_params(database) == {
|
||||||
|
"engine_params": {
|
||||||
|
"connect_args": {
|
||||||
|
"config": {"custom_user_agent": f"apache-superset/{VERSION_STRING}"}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
database.extra = json.dumps(
|
||||||
|
{"engine_params": {"connect_args": {"config": {"custom_user_agent": "my-app"}}}}
|
||||||
|
)
|
||||||
|
assert DuckDBEngineSpec.get_extra_params(database) == {
|
||||||
|
"engine_params": {
|
||||||
|
"connect_args": {
|
||||||
|
"config": {
|
||||||
|
"custom_user_agent": f"apache-superset/{VERSION_STRING} my-app"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
Loading…
Reference in New Issue