Apply SQL_QUERY_MUTATOR to explore & dashboard (#5493)

* Apply SQL_QUERY_MUTATOR kn explore & dashboard

* Add unit test

(cherry picked from commit 94cb20c)

superset/connectors/sqla/models.py

@@ -24,7 +24,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
@@ -33,6 +33,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
@@ -411,10 +413,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):

superset/utils.py

@@ -27,7 +27,8 @@
 import bleach
 import celery
 from dateutil.parser import parse
-from flask import flash, Markup, render_template
+from dateutil.relativedelta import relativedelta
+from flask import flash, g, Markup, render_template
 from flask_babel import gettext as __
 from flask_caching import Cache
 import markdown as md
@@ -885,3 +886,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

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