Apply SQL_QUERY_MUTATOR to explore & dashboard (#5493)

* Apply SQL_QUERY_MUTATOR kn explore & dashboard

* Add unit test

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

setup.py

@@ -67,7 +67,6 @@ setup(
         'flask-caching',
         'flask-compress',
         'flask-migrate',
-        'flask-testing',
         'flask-wtf',
         'flower',  # deprecated
         'future>=0.16.0, <0.17',

superset/connectors/sqla/models.py

@@ -23,7 +23,7 @@ from sqlalchemy.sql import column, literal_column, table, text
 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 QueryResult
 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 @@ class SqlaTable(Model, BaseDatasource):
         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 @@ class SqlaTable(Model, BaseDatasource):
         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

@@ -28,7 +28,7 @@ import bleach
 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

tests/model_tests.py

@@ -8,7 +8,7 @@ import textwrap
 
 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 @@ class SqlaTableModelTestCase(SupersetTestCase):
         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