From 94cb20cf9673cfe7b8806ee933fe9f1107eb9a65 Mon Sep 17 00:00:00 2001
From: Maxime Beauchemin <maximebeauchemin@gmail.com>
Date: Thu, 26 Jul 2018 15:20:23 -0700
Subject: [PATCH] Apply SQL_QUERY_MUTATOR to explore & dashboard (#5493)

* Apply SQL_QUERY_MUTATOR kn explore & dashboard

* Add unit test
---
 requirements.txt                   |  1 -
 setup.py                           |  1 -
 superset/connectors/sqla/models.py | 16 +++++++++++++++-
 superset/utils.py                  | 10 +++++++++-
 tests/model_tests.py               | 26 +++++++++++++++++++++++++-
 5 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index b81e1789484b4..b076fc0c5aec9 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -9,7 +9,6 @@ flask-appbuilder==1.10.0
 flask-caching==1.4.0
 flask-compress==1.4.0
 flask-migrate==2.1.1
-flask-testing==0.7.1
 flask-wtf==0.14.2
 flower==0.9.2
 future==0.16.0
diff --git a/setup.py b/setup.py
index 38c39a8921428..84c0c48405ef3 100644
--- a/setup.py
+++ b/setup.py
@@ -67,7 +67,6 @@ def get_git_sha():
         'flask-caching',
         'flask-compress',
         'flask-migrate',
-        'flask-testing',
         'flask-wtf',
         'flower',  # deprecated
         'future>=0.16.0, <0.17',
diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
index c86d4eaddad5b..4a4e024b88fd7 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -23,7 +23,7 @@
 from sqlalchemy.sql.expression import TextAsFrom
 import sqlparse
 
-from superset import db, import_util, security_manager, utils
+from superset import app, db, import_util, security_manager, utils
 from superset.connectors.base.models import BaseColumn, BaseDatasource, BaseMetric
 from superset.jinja_context import get_template_processor
 from superset.models.annotations import Annotation
@@ -32,6 +32,8 @@
 from superset.models.helpers import set_perm
 from superset.utils import DTTM_ALIAS, QueryStatus
 
+config = app.config
+
 
 class AnnotationDatasource(BaseDatasource):
     """ Dummy object so we can query annotations using 'Viz' objects just like
@@ -417,10 +419,21 @@ def values_for_column(self, column_name, limit=10000):
         sql = '{}'.format(
             qry.compile(engine, compile_kwargs={'literal_binds': True}),
         )
+        sql = self.mutate_query_from_config(sql)
 
         df = pd.read_sql_query(sql=sql, con=engine)
         return [row[0] for row in df.to_records(index=False)]
 
+    def mutate_query_from_config(self, sql):
+        """Apply config's SQL_QUERY_MUTATOR
+
+        Typically adds comments to the query with context"""
+        SQL_QUERY_MUTATOR = config.get('SQL_QUERY_MUTATOR')
+        if SQL_QUERY_MUTATOR:
+            username = utils.get_username()
+            sql = SQL_QUERY_MUTATOR(sql, username, security_manager, self.database)
+        return sql
+
     def get_template_processor(self, **kwargs):
         return get_template_processor(
             table=self, database=self.database, **kwargs)
@@ -432,6 +445,7 @@ def get_query_str(self, query_obj):
         sql = sqlparse.format(sql, reindent=True)
         if query_obj['is_prequery']:
             query_obj['prequeries'].append(sql)
+        sql = self.mutate_query_from_config(sql)
         return sql
 
     def get_sqla_table(self):
diff --git a/superset/utils.py b/superset/utils.py
index 8fd8337ccaf59..04885a6bb0efd 100644
--- a/superset/utils.py
+++ b/superset/utils.py
@@ -28,7 +28,7 @@
 import celery
 from dateutil.parser import parse
 from dateutil.relativedelta import relativedelta
-from flask import flash, Markup, render_template
+from flask import flash, g, Markup, render_template
 from flask_babel import gettext as __
 from flask_caching import Cache
 import markdown as md
@@ -959,3 +959,11 @@ def split_adhoc_filters_into_base_filters(fd):
         fd['having'] = ' AND '.join(['({})'.format(sql) for sql in sql_having_filters])
         fd['having_filters'] = simple_having_filters
         fd['filters'] = simple_where_filters
+
+
+def get_username():
+    """Get username if within the flask context, otherwise return noffin'"""
+    try:
+        return g.user.username
+    except Exception:
+        pass
diff --git a/tests/model_tests.py b/tests/model_tests.py
index 411fdd5ed38ae..74dc822645d00 100644
--- a/tests/model_tests.py
+++ b/tests/model_tests.py
@@ -8,7 +8,7 @@
 
 from sqlalchemy.engine.url import make_url
 
-from superset import db
+from superset import app, db
 from superset.models.core import Database
 from .base_tests import SupersetTestCase
 
@@ -186,3 +186,27 @@ def test_get_timestamp_expression_backward(self):
         compiled = '{}'.format(sqla_literal.compile())
         if tbl.database.backend == 'mysql':
             self.assertEquals(compiled, 'ds')
+
+    def test_sql_mutator(self):
+        tbl = self.get_table_by_name('birth_names')
+        query_obj = dict(
+            groupby=[],
+            metrics=[],
+            filter=[],
+            is_timeseries=False,
+            columns=['name'],
+            granularity=None,
+            from_dttm=None, to_dttm=None,
+            is_prequery=False,
+            extras={},
+        )
+        sql = tbl.get_query_str(query_obj)
+        self.assertNotIn('--COMMENT', sql)
+
+        def mutator(*args):
+            return '--COMMENT\n' + args[0]
+        app.config['SQL_QUERY_MUTATOR'] = mutator
+        sql = tbl.get_query_str(query_obj)
+        self.assertIn('--COMMENT', sql)
+
+        app.config['SQL_QUERY_MUTATOR'] = None