From 9d70c348d3403fa4682383c868fcb1d7571fded9 Mon Sep 17 00:00:00 2001
From: timifasubaa <30888507+timifasubaa@users.noreply.github.com>
Date: Thu, 10 Jan 2019 17:30:32 -0800
Subject: [PATCH] pass source to db api mutator (#6497)

---
 superset/models/core.py | 16 ++++++++++++----
 superset/sql_lab.py     |  2 ++
 superset/utils/core.py  |  6 ++++++
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/superset/models/core.py b/superset/models/core.py
index 97dd54a400..e9b1918d1c 100644
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -757,7 +757,7 @@ class Database(Model, AuditMixinNullable, ImportMixin):
 
     @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)
@@ -790,7 +790,7 @@ class Database(Model, AuditMixinNullable, ImportMixin):
         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):
@@ -801,7 +801,14 @@ class Database(Model, AuditMixinNullable, ImportMixin):
 
     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):
@@ -860,7 +867,8 @@ class Database(Model, AuditMixinNullable, ImportMixin):
             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,
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index 63de7830fa..28ad921908 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -20,6 +20,7 @@ from superset.tasks.celery_app import app as celery_app
 from superset.utils.core import (
     json_iso_dttm_ser,
     QueryStatus,
+    sources,
     zlib_compress,
 )
 from superset.utils.dates import now_as_float
@@ -226,6 +227,7 @@ def execute_sql_statements(
         schema=query.schema,
         nullpool=True,
         user_name=user_name,
+        source=sources.get('sql_lab', None),
     )
     # Sharing a single connection and cursor across the
     # execution of all statements (if many)
diff --git a/superset/utils/core.py b/superset/utils/core.py
index b1b81e56d2..5bf3054651 100644
--- a/superset/utils/core.py
+++ b/superset/utils/core.py
@@ -54,6 +54,12 @@ ADHOC_METRIC_EXPRESSION_TYPES = {
 
 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"""