From b95a1e9b0ac4b8e93013f008a91059c84aff3309 Mon Sep 17 00:00:00 2001
From: Maxime Beauchemin <maximebeauchemin@gmail.com>
Date: Fri, 9 Mar 2018 11:27:36 -0800
Subject: [PATCH] [sql lab] comment injection hook (#4585)

(cherry picked from commit 34a081b9267d0473ba9726d057dfcdd90011fa45)
---
 superset/config.py      | 14 ++++++++++++--
 superset/models/core.py |  2 +-
 superset/sql_lab.py     |  7 ++++++-
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/superset/config.py b/superset/config.py
index c9b8607e3e6ff..5e5a822f7a947 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -382,9 +382,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
@@ -393,6 +394,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
diff --git a/superset/models/core.py b/superset/models/core.py
index 41d8742b65b7c..9caa7adce4bfb 100644
--- a/superset/models/core.py
+++ b/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):
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index 4dae72720dc97..f98231ed775c9 100644
--- a/superset/sql_lab.py
+++ b/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()