[sql lab] comment injection hook (apache#4585)

superset/config.py

@@ -384,9 +384,10 @@ class CeleryConfig(object):
 # arbitrary logic. For instance you can wire different users to
 # use different connection parameters, or pass their email address as the
 # username. The function receives the connection uri object, connection
-# params, and user object, and returns the mutated uri and params objects.
+# params, the username, and returns the mutated uri and params objects.
 # Example:
-#   def DB_CONNECTION_MUTATOR(uri, params, user):
+#   def DB_CONNECTION_MUTATOR(uri, params, username, security_manager):
+#       user = security_manager.find_user(username=username)
 #       if user and user.email:
 #           uri.username = user.email
 #       return uri, params
@@ -395,6 +396,15 @@ class CeleryConfig(object):
 # as such `create_engine(url, **params)`
 DB_CONNECTION_MUTATOR = None
 
+# A function that intercepts the SQL to be executed and can alter it.
+# The use case is can be around adding some sort of comment header
+# with information such as the username and worker node information
+#
+#    def SQL_QUERY_MUTATOR(sql, username, security_manager):
+#        dttm = datetime.now().isoformat()
+#        return "-- [SQL LAB] {username} {dttm}\n sql"(**locals())
+SQL_QUERY_MUTATOR = None
+
 try:
     if CONFIG_PATH_ENV_VAR in os.environ:
         # Explicitly import config module that is not in pythonpath; useful

superset/models/core.py

@@ -677,7 +677,7 @@ def get_sqla_engine(self, schema=None, nullpool=True, user_name=None):
 
         DB_CONNECTION_MUTATOR = config.get('DB_CONNECTION_MUTATOR')
         if DB_CONNECTION_MUTATOR:
-            url, params = DB_CONNECTION_MUTATOR(url, params, g.user)
+            url, params = DB_CONNECTION_MUTATOR(url, params, user_name, sm)
         return create_engine(url, **params)
 
     def get_reserved_words(self):

superset/sql_lab.py

@@ -17,7 +17,7 @@
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy.pool import NullPool
 
-from superset import app, dataframe, db, results_backend, utils
+from superset import app, dataframe, db, results_backend, sm, utils
 from superset.db_engine_specs import LimitMethod
 from superset.jinja_context import get_template_processor
 from superset.models.sql_lab import Query
@@ -194,6 +194,11 @@ def handle_error(msg):
         msg = 'Template rendering failed: ' + utils.error_msg_from_exception(e)
         return handle_error(msg)
 
+    # Hook to allow environment-specific mutation (usually comments) to the SQL
+    SQL_QUERY_MUTATOR = config.get('SQL_QUERY_MUTATOR')
+    if SQL_QUERY_MUTATOR:
+        executed_sql = SQL_QUERY_MUTATOR(executed_sql, user_name, sm, database)
+
     query.executed_sql = executed_sql
     query.status = QueryStatus.RUNNING
     query.start_running_time = utils.now_as_float()