[superset] 14/24: fix(drill): specify an SA URL parm of `impersonation_target` for drill+sadrill (#19252)

[vi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

villebro pushed a commit to tag 1.5.0rc1
in repository https://gitbox.apache.org/repos/asf/superset.git

commit e4cbbdc65399de33c44002fc2ae7996942123762
Author: James Turton <91...@users.noreply.github.com>
AuthorDate: Thu Mar 31 18:42:27 2022 +0200

    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.
    
    (cherry picked from commit 85e330e94bc490d5a35fa0536f44e7189e299d20)
---
 superset/db_engine_specs/drill.py              |  7 +++++--
 tests/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 3c06017fa2..de8c8397f6 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 08f7972d80..ad7254870f 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)