From 18e15b7bc4eec88fa5541d30525dc4766be83e25 Mon Sep 17 00:00:00 2001
From: "J.C. Zhong" <jianchengzhong@pinterest.com>
Date: Thu, 3 Aug 2023 10:23:13 -0700
Subject: [PATCH] feat: pass proxy user to sqlalchemy executor (#1303)

---
 .../lib/query_executor/clients/sqlalchemy.py       | 14 +++++++++++++-
 .../query_executor/executor_template/templates.py  |  1 +
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/querybook/server/lib/query_executor/clients/sqlalchemy.py b/querybook/server/lib/query_executor/clients/sqlalchemy.py
index 34614e749..0f170cedf 100644
--- a/querybook/server/lib/query_executor/clients/sqlalchemy.py
+++ b/querybook/server/lib/query_executor/clients/sqlalchemy.py
@@ -1,4 +1,5 @@
 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
@@ -6,8 +7,19 @@
 
 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
index d57c6ab7f..452d7f6ff 100644
--- 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(