pass source to db api mutator (#6497)

(cherry picked from commit 9d70c34)

superset/models/core.py

@@ -756,7 +756,7 @@ def get_effective_user(self, url, user_name=None):
 
     @utils.memoized(
         watch=('impersonate_user', 'sqlalchemy_uri_decrypted', 'extra'))
-    def get_sqla_engine(self, schema=None, nullpool=True, user_name=None):
+    def get_sqla_engine(self, schema=None, nullpool=True, user_name=None, source=None):
         extra = self.get_extra()
         url = make_url(self.sqlalchemy_uri_decrypted)
         url = self.db_engine_spec.adjust_database_uri(url, schema)
@@ -789,7 +789,7 @@ def get_sqla_engine(self, schema=None, nullpool=True, user_name=None):
         DB_CONNECTION_MUTATOR = config.get('DB_CONNECTION_MUTATOR')
         if DB_CONNECTION_MUTATOR:
             url, params = DB_CONNECTION_MUTATOR(
-                url, params, effective_username, security_manager)
+                url, params, effective_username, security_manager, source)
         return create_engine(url, **params)
 
     def get_reserved_words(self):
@@ -800,7 +800,15 @@ def get_quoter(self):
 
     def get_df(self, sql, schema):
         sqls = [str(s).strip().strip(';') for s in sqlparse.parse(sql)]
-        engine = self.get_sqla_engine(schema=schema)
+        source_key = None
+        if request and request.referrer:
+            if '/superset/dashboard/' in request.referrer:
+                source_key = 'dashboard'
+            elif '/superset/explore/' in request.referrer:
+                source_key = 'chart'
+        engine = self.get_sqla_engine(
+            schema=schema, source=utils.sources.get(source_key, None))
+        username = utils.get_username()
 
         def needs_conversion(df_series):
             if df_series.empty:
@@ -852,7 +860,8 @@ def select_star(
             self, table_name, schema=None, limit=100, show_cols=False,
             indent=True, latest_partition=False, cols=None):
         """Generates a ``select *`` statement in the proper dialect"""
-        eng = self.get_sqla_engine(schema=schema)
+        eng = self.get_sqla_engine(
+            schema=schema, source=utils.sources.get('sql_lab', None))
         return self.db_engine_spec.select_star(
             self, table_name, schema=schema, engine=eng,
             limit=limit, show_cols=show_cols,

superset/sql_lab.py

@@ -19,6 +19,7 @@
     json_iso_dttm_ser,
     now_as_float,
     QueryStatus,
+    sources,
     zlib_compress,
 )
 
@@ -174,6 +175,7 @@ def handle_error(msg):
             schema=query.schema,
             nullpool=True,
             user_name=user_name,
+            source=sources.get('sql_lab', None),
         )
         conn = engine.raw_connection()
         cursor = conn.cursor()

superset/utils/core.py

@@ -54,6 +54,12 @@
 
 JS_MAX_INTEGER = 9007199254740991   # Largest int Java Script can handle 2^53-1
 
+sources = {
+    'chart': 0,
+    'dashboard': 1,
+    'sql_lab': 2,
+}
+
 
 def flasher(msg, severity=None):
     """Flask's flash if available, logging call if not"""