From 1a083d3ebb4fb8ea22390ede662986ab7243f16d Mon Sep 17 00:00:00 2001 From: Maxime Beauchemin Date: Fri, 9 Mar 2018 17:14:22 +0000 Subject: [PATCH] [sql lab] comment injection hook --- 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()