diff --git a/superset/db_engine_specs/drill.py b/superset/db_engine_specs/drill.py index 3c06017fa2d..de8c8397f66 100644 --- a/superset/db_engine_specs/drill.py +++ b/superset/db_engine_specs/drill.py @@ -21,6 +21,7 @@ from urllib import parse from sqlalchemy.engine.url import URL from superset.db_engine_specs.base import BaseEngineSpec +from superset.db_engine_specs.exceptions import SupersetDBAPIProgrammingError from superset.utils import core as utils @@ -84,7 +85,9 @@ class DrillEngineSpec(BaseEngineSpec): if impersonate_user and username is not None: if url.drivername == "drill+odbc": url.query["DelegationUID"] = username - elif url.drivername == "drill+jdbc": + elif url.drivername in ["drill+sadrill", "drill+jdbc"]: url.query["impersonation_target"] = username else: - url.username = username + raise SupersetDBAPIProgrammingError( + f"impersonation is not supported for {url.drivername}" + ) diff --git a/tests/unit_tests/db_engine_specs/test_drill.py b/tests/unit_tests/db_engine_specs/test_drill.py index 08f7972d80b..ad7254870f6 100644 --- a/tests/unit_tests/db_engine_specs/test_drill.py +++ b/tests/unit_tests/db_engine_specs/test_drill.py @@ -17,6 +17,7 @@ # pylint: disable=unused-argument, import-outside-toplevel, protected-access from flask.ctx import AppContext +from pytest import raises def test_odbc_impersonation(app_context: AppContext) -> None: @@ -55,7 +56,7 @@ def test_sadrill_impersonation(app_context: AppContext) -> None: """ Test ``modify_url_for_impersonation`` method when driver == sadrill. - The method changes the username of URL Object. + The method adds the parameter ``impersonation_target`` to the query string. """ from sqlalchemy.engine.url import URL @@ -64,4 +65,23 @@ def test_sadrill_impersonation(app_context: AppContext) -> None: url = URL("drill+sadrill") username = "DoAsUser" DrillEngineSpec.modify_url_for_impersonation(url, True, username) - assert url.username == username + assert url.query["impersonation_target"] == username + + +def test_invalid_impersonation(app_context: AppContext) -> None: + """ + Test ``modify_url_for_impersonation`` method when driver == foobar. + + The method raises an exception because impersonation is not supported + for drill+foobar. + """ + from sqlalchemy.engine.url import URL + + from superset.db_engine_specs.drill import DrillEngineSpec + from superset.db_engine_specs.exceptions import SupersetDBAPIProgrammingError + + url = URL("drill+foobar") + username = "DoAsUser" + + with raises(SupersetDBAPIProgrammingError): + DrillEngineSpec.modify_url_for_impersonation(url, True, username)