From f5855dbcf6ee1d4aeef671a033fac1f1870bed88 Mon Sep 17 00:00:00 2001
From: Dmitry Goryunov <d.f.goryunov@gmail.com>
Date: Mon, 18 Sep 2017 19:52:29 +0300
Subject: [PATCH] Feature: "Impersonate user" setting on Datasource (#3404)

* Add "Impersonate user" setting to Datasource

* Add tests

* Use g.user.username for all the sync cases

* use uri.username instead of uri.user

* Small refactoring
---
 dump.rdb                                      | Bin 0 -> 157 bytes
 ...9c47e2c1547_add_impersonate_user_to_dbs.py |  22 ++++++++++++++++++
 superset/models/core.py                       |   5 +++-
 superset/sql_lab.py                           |  10 ++++----
 superset/views/core.py                        |   8 +++++--
 tests/model_tests.py                          |  13 +++++++++++
 6 files changed, 50 insertions(+), 8 deletions(-)
 create mode 100644 dump.rdb
 create mode 100644 superset/migrations/versions/a9c47e2c1547_add_impersonate_user_to_dbs.py

diff --git a/dump.rdb b/dump.rdb
new file mode 100644
index 0000000000000000000000000000000000000000..5e6dde24da9b5e5e19bad9186563ac55820af381
GIT binary patch
literal 157
zcmWG?b@2=~Ffg$A#aWb^l3A=<mRiJWqGzCI_=^i7kd#?ce8AxsYjR0uZt5Wqt&Net
zI7*9CQ*?7va}Nb<;$!&5lbD~TTTqmmn46T7dVt{<dr@jZj&5d(hOvpIk(qIdv4Lrt
zVVZH8X{u?mWm2-SxrMo9Qc_BiMXF^=vZ+OK+AnUfj{LN=;?$A@4FAPNF6F*me8~s^
Dgbz93

literal 0
HcmV?d00001

diff --git a/superset/migrations/versions/a9c47e2c1547_add_impersonate_user_to_dbs.py b/superset/migrations/versions/a9c47e2c1547_add_impersonate_user_to_dbs.py
new file mode 100644
index 0000000000000..e0cf1e286cbfd
--- /dev/null
+++ b/superset/migrations/versions/a9c47e2c1547_add_impersonate_user_to_dbs.py
@@ -0,0 +1,22 @@
+"""add impersonate_user to dbs
+
+Revision ID: a9c47e2c1547
+Revises: ca69c70ec99b
+Create Date: 2017-08-31 17:35:58.230723
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = 'a9c47e2c1547'
+down_revision = 'ca69c70ec99b'
+
+from alembic import op
+import sqlalchemy as sa
+
+
+def upgrade():
+    op.add_column('dbs', sa.Column('impersonate_user', sa.Boolean(), nullable=True))
+
+
+def downgrade():
+    op.drop_column('dbs', 'impersonate_user')
diff --git a/superset/models/core.py b/superset/models/core.py
index 637ed091d950a..b9daa72fdcdf2 100644
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -562,6 +562,7 @@ class Database(Model, AuditMixinNullable):
     """))
     perm = Column(String(1000))
     custom_password_store = config.get('SQLALCHEMY_CUSTOM_PASSWORD_STORE')
+    impersonate_user = Column(Boolean, default=False)
 
     def __repr__(self):
         return self.verbose_name if self.verbose_name else self.database_name
@@ -588,13 +589,15 @@ def set_sqlalchemy_uri(self, uri):
         conn.password = password_mask if conn.password else None
         self.sqlalchemy_uri = str(conn)  # hides the password
 
-    def get_sqla_engine(self, schema=None, nullpool=False):
+    def get_sqla_engine(self, schema=None, nullpool=False, user_name=None):
         extra = self.get_extra()
         uri = make_url(self.sqlalchemy_uri_decrypted)
         params = extra.get('engine_params', {})
         if nullpool:
             params['poolclass'] = NullPool
         uri = self.db_engine_spec.adjust_database_uri(uri, schema)
+        if self.impersonate_user:
+            uri.username = user_name if user_name else g.user.username
         return create_engine(uri, **params)
 
     def get_reserved_words(self):
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index 0bfca71c81749..0f2f1fefa5e23 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -86,11 +86,11 @@ def get_session(nullpool):
 
 @celery_app.task(bind=True, soft_time_limit=SQLLAB_TIMEOUT)
 def get_sql_results(
-        ctask, query_id, return_results=True, store_results=False):
+        ctask, query_id, return_results=True, store_results=False, user_name=None):
     """Executes the sql query returns the results."""
     try:
         return execute_sql(
-            ctask, query_id, return_results, store_results)
+            ctask, query_id, return_results, store_results, user_name)
     except Exception as e:
         logging.exception(e)
         stats_logger.incr('error_sqllab_unhandled')
@@ -103,7 +103,7 @@ def get_sql_results(
         raise
 
 
-def execute_sql(ctask, query_id, return_results=True, store_results=False):
+def execute_sql(ctask, query_id, return_results=True, store_results=False, user_name=None):
     """Executes the sql query returns the results."""
     session = get_session(not ctask.request.called_directly)
 
@@ -170,10 +170,10 @@ def handle_error(msg):
     logging.info("Set query to 'running'")
 
     engine = database.get_sqla_engine(
-            schema=query.schema, nullpool=not ctask.request.called_directly)
+            schema=query.schema, nullpool=not ctask.request.called_directly, user_name=user_name)
     try:
         engine = database.get_sqla_engine(
-            schema=query.schema, nullpool=not ctask.request.called_directly)
+            schema=query.schema, nullpool=not ctask.request.called_directly, user_name=user_name)
         conn = engine.raw_connection()
         cursor = conn.cursor()
         logging.info("Running query: \n{}".format(executed_sql))
diff --git a/superset/views/core.py b/superset/views/core.py
index faa61ad4374a4..65d61bab16b4f 100755
--- a/superset/views/core.py
+++ b/superset/views/core.py
@@ -182,7 +182,7 @@ class DatabaseView(SupersetModelView, DeleteMixin):  # noqa
     add_columns = [
         'database_name', 'sqlalchemy_uri', 'cache_timeout', 'extra',
         'expose_in_sqllab', 'allow_run_sync', 'allow_run_async',
-        'allow_ctas', 'allow_dml', 'force_ctas_schema']
+        'allow_ctas', 'allow_dml', 'force_ctas_schema', 'impersonate_user']
     search_exclude_columns = (
         'password', 'tables', 'created_by', 'changed_by', 'queries',
         'saved_queries', )
@@ -235,6 +235,9 @@ class DatabaseView(SupersetModelView, DeleteMixin):  # noqa
             "gets unpacked into the [sqlalchemy.MetaData]"
             "(http://docs.sqlalchemy.org/en/rel_1_0/core/metadata.html"
             "#sqlalchemy.schema.MetaData) call. ", True),
+        'impersonate_user': _(
+            "All the queries in Sql Lab are going to be executed "
+            "on behalf of currently authorized user."),
     }
     label_columns = {
         'expose_in_sqllab': _("Expose in SQL Lab"),
@@ -249,6 +252,7 @@ class DatabaseView(SupersetModelView, DeleteMixin):  # noqa
         'extra': _("Extra"),
         'allow_run_sync': _("Allow Run Sync"),
         'allow_run_async': _("Allow Run Async"),
+        'impersonate_user': _("Impersonate queries to the database"),
     }
 
     def pre_add(self, db):
@@ -2057,7 +2061,7 @@ def sql_json(self):
             try:
                 sql_lab.get_sql_results.delay(
                     query_id=query_id, return_results=False,
-                    store_results=not query.select_as_cta)
+                    store_results=not query.select_as_cta, user_name=g.user.username)
             except Exception as e:
                 logging.exception(e)
                 msg = (
diff --git a/tests/model_tests.py b/tests/model_tests.py
index df2902b107b62..94a5358807135 100644
--- a/tests/model_tests.py
+++ b/tests/model_tests.py
@@ -54,3 +54,16 @@ def test_database_schema_mysql(self):
 
         db = make_url(model.get_sqla_engine(schema='staging').url).database
         self.assertEquals('staging', db)
+
+    def test_database_impersonate_user(self):
+        uri = 'mysql://root@localhost'
+        example_user = 'giuseppe'
+        model = Database(sqlalchemy_uri=uri)
+
+        model.impersonate_user = True
+        user_name = make_url(model.get_sqla_engine(user_name=example_user).url).username
+        self.assertEquals(example_user, user_name)
+
+        model.impersonate_user = False
+        user_name = make_url(model.get_sqla_engine(user_name=example_user).url).username
+        self.assertNotEquals(example_user, user_name)