fix(drill): specify an SA URL parm of impersonation_target for dril…

…l+sadrill (#19252)

* Update drill+sadrill to specify an SA URL parm of "impersonation_target".

Sqlalchemy-drill is being updated to support impersonation with the
drill+sadrill driver, where previously it did not.  The way that callers
should specify impersonation matches that for the drill+jdbc driver in that
a SA URL parameter of impersonation_target should be set to the username
of the user to be impersonated, while the stadard SA username and password
should be those of the proxy user.

* Remove lint.

* Address review comments.

* Use idiomatic pytest to test for a raised exception.

* Fix import statement order in drill.py.

superset/db_engine_specs/drill.py

@@ -21,6 +21,7 @@
 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 @@ def modify_url_for_impersonation(
         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}"
+                )

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)