Feature: "Impersonate user" setting on Datasource (apache#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

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')

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):

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))

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 = (

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)