[incubator-superset] branch master updated: pass source to db api mutator (#6497)

[ti...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

timi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-superset.git


The following commit(s) were added to refs/heads/master by this push:
     new 9d70c34  pass source to db api mutator (#6497)
9d70c34 is described below

commit 9d70c348d3403fa4682383c868fcb1d7571fded9
Author: timifasubaa <30...@users.noreply.github.com>
AuthorDate: Thu Jan 10 17:30:32 2019 -0800

    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 97dd54a..e9b1918 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 63de783..28ad921 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 b1b81e5..5bf3054 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"""