From 85e330e94bc490d5a35fa0536f44e7189e299d20 Mon Sep 17 00:00:00 2001
From: James Turton <9107319+jnturton@users.noreply.github.com>
Date: Thu, 31 Mar 2022 18:42:27 +0200
Subject: [PATCH] fix(drill): specify an SA URL parm of `impersonation_target`
 for drill+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             |  7 ++++--
 .../unit_tests/db_engine_specs/test_drill.py  | 24 +++++++++++++++++--
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/superset/db_engine_specs/drill.py b/superset/db_engine_specs/drill.py
index 3c06017fa2daa..de8c8397f66b3 100644
--- a/superset/db_engine_specs/drill.py
+++ b/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}"
+                )
diff --git a/tests/unit_tests/db_engine_specs/test_drill.py b/tests/unit_tests/db_engine_specs/test_drill.py
index 08f7972d80b78..ad7254870f6ef 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)