feat: pass proxy user to sqlalchemy executor (#1303)

pinterest · Aug 3, 2023 · 18e15b7 · 18e15b7
1 parent 4cec5c1
commit 18e15b7
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/querybook/server/lib/query_executor/clients/sqlalchemy.py b/querybook/server/lib/query_executor/clients/sqlalchemy.py
@@ -1,13 +1,25 @@
 import sqlalchemy
+from sqlalchemy.engine import make_url
 
 from lib.query_executor.base_client import ClientBaseClass, CursorBaseClass
 from lib.query_executor.connection_string.sqlalchemy import create_sqlalchemy_engine
 
 
 class SqlAlchemyClient(ClientBaseClass):
     def __init__(
-        self, connection_string=None, connect_args=[], proxy_user=None, *args, **kwargs
+        self,
+        connection_string=None,
+        connect_args=[],
+        proxy_user=None,
+        impersonate=False,
+        *args,
+        **kwargs
     ):
+        if impersonate and proxy_user:
+            parsed_connection_url = make_url(connection_string)
+            new_connection_url = parsed_connection_url.set(username=proxy_user)
+            connection_string = str(new_connection_url)
+
         self._engine = create_sqlalchemy_engine(
             {
                 "connection_string": connection_string,

diff --git a/querybook/server/lib/query_executor/executor_template/templates.py b/querybook/server/lib/query_executor/executor_template/templates.py
@@ -97,6 +97,7 @@
 </p>""",
         ),
     ),
+    ("impersonate", FormField(field_type=FormFieldType.Boolean)),
     (
         "connect_args",
         ExpandableFormField(