Copilot commented on code in PR #42705:
URL: https://github.com/apache/superset/pull/42705#discussion_r3705991206
##########
superset/db_engine_specs/duckdb.py:
##########
@@ -369,6 +369,38 @@ def get_extra_params(
return extra
+ @classmethod
+ def impersonate_user(
+ cls,
+ database: Database,
+ username: str | None,
+ user_token: str | None,
+ url: URL,
+ engine_kwargs: dict[str, Any],
+ ) -> tuple[URL, dict[str, Any]]:
+ # Local DuckDB has no user concept, and the base implementation's
+ # username-in-URL is passed by duckdb-engine as a connect() kwarg
+ # that duckdb rejects -- so never call it. Note that md: databases
+ # also resolve to this spec (spec lookup is by backend name only,
+ # and MotherDuck shares the duckdb:// scheme).
+ if username is None or not (url.database or "").startswith("md:"):
+ return url, engine_kwargs
+
+ # Identify the logged-in user to MotherDuck for per-user session
+ # attribution. The parameter is embedded in the md: path rather than
+ # the SQLAlchemy URL query: the setting must reach MotherDuck at
+ # connection initialization, and duckdb-engine applies unknown URL
+ # query parameters via SET after connecting, which is rejected.
+ # MotherDuck reads path parameter values literally (no URL decoding),
+ # so only the characters structural to the path are escaped.
+ session_name = (
+ username.replace("%", "%25").replace("&", "%26").replace("=",
"%3D")
+ )
+ separator = "&" if "?" in url.database else "?"
+ url =
url.set(database=f"{url.database}{separator}session_name={session_name}")
Review Comment:
This unconditionally appends `session_name=...` to `url.database`. If the
configured `md:` path already includes a `session_name` parameter, this will
create duplicates (e.g. `...&session_name=old&session_name=new`), which can
lead to incorrect or ambiguous session attribution. Consider parsing the
existing `md:` parameter string and replacing/removing any existing
`session_name` entry before adding the impersonated one (keeping parsing
literal—no URL decoding—since MotherDuck reads values literally).
##########
tests/unit_tests/db_engine_specs/test_duckdb.py:
##########
@@ -163,3 +163,98 @@ def test_column_type_recognition() -> None:
col_spec = DuckDBEngineSpec.get_column_spec("TINYINT")
# TINYINT matches the pattern "^int" so it should be recognized
assert col_spec is None, "TINYINT doesn't match any patterns"
+
+
+def test_motherduck_impersonation(mocker: MockerFixture) -> None:
+ """
+ Test ``impersonate_user`` embeds the username in the md: path.
+
+ The hook lives on DuckDBEngineSpec because engine spec resolution is by
+ backend name, so md: databases resolve to the DuckDB spec.
+ """
+ from sqlalchemy.engine.url import URL
+
+ from superset.db_engine_specs.duckdb import DuckDBEngineSpec
+
+ database = mocker.MagicMock()
+
+ url = URL.create("duckdb", database="md:my_db", query={"motherduck_token":
"abc"})
+ url, engine_kwargs = DuckDBEngineSpec.impersonate_user(
+ database=database,
+ username="alice",
+ user_token=None,
+ url=url,
+ engine_kwargs={},
+ )
+ assert url.database == "md:my_db?session_name=alice"
+ assert url.username is None
+ assert url.query["motherduck_token"] == "abc"
+ assert engine_kwargs == {}
+
+
+def test_duckdb_local_impersonation_is_a_noop(mocker: MockerFixture) -> None:
+ """
+ Test ``impersonate_user`` leaves local DuckDB URLs alone.
+
+ The base implementation puts the username in the URL, which duckdb-engine
+ forwards as a ``connect()`` kwarg that duckdb rejects.
+ """
+ from sqlalchemy.engine.url import URL
+
+ from superset.db_engine_specs.duckdb import DuckDBEngineSpec
+
+ database = mocker.MagicMock()
+
+ url = URL.create("duckdb", database="/path/to/duck.db")
+ url, _ = DuckDBEngineSpec.impersonate_user(
+ database=database,
+ username="alice",
+ user_token=None,
+ url=url,
+ engine_kwargs={},
+ )
+ assert url == URL.create("duckdb", database="/path/to/duck.db")
+
+
+def test_motherduck_impersonation_escapes_structural_characters(
+ mocker: MockerFixture,
+) -> None:
+ """
+ Test ``impersonate_user`` escapes characters that would inject parameters.
+ """
+ from sqlalchemy.engine.url import URL
+
+ from superset.db_engine_specs.duckdb import MotherDuckEngineSpec
+
+ database = mocker.MagicMock()
+
+ url = URL.create("duckdb", database="md:my_db?attach_mode=single")
+ url, _ = MotherDuckEngineSpec.impersonate_user( # inherited from
DuckDBEngineSpec
+ database=database,
+ username="a&host=evil",
+ user_token=None,
+ url=url,
+ engine_kwargs={},
+ )
+ assert url.database ==
"md:my_db?attach_mode=single&session_name=a%26host%3Devil"
+
+
+def test_motherduck_impersonation_without_username(mocker: MockerFixture) ->
None:
+ """
+ Test ``impersonate_user`` leaves the URL alone when there is no username.
+ """
+ from sqlalchemy.engine.url import URL
+
+ from superset.db_engine_specs.duckdb import MotherDuckEngineSpec
+
+ database = mocker.MagicMock()
+
+ url = URL.create("duckdb", database="md:my_db")
+ url, _ = MotherDuckEngineSpec.impersonate_user(
+ database=database,
+ username=None,
+ user_token=None,
+ url=url,
+ engine_kwargs={},
+ )
+ assert "motherduck_session_name" not in url.query
Review Comment:
This assertion doesn’t validate the intended behavior: the implementation
writes `session_name` into `url.database` (not `url.query`), and `url.query` is
empty here anyway—so this test can pass even if the URL is modified
incorrectly. Update the test to assert the URL is unchanged (e.g.,
`url.database == \"md:my_db\"` or `url == URL.create(\"duckdb\",
database=\"md:my_db\")`) and, if checking a key, use `session_name` rather than
`motherduck_session_name`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
