From 34a081b9267d0473ba9726d057dfcdd90011fa45 Mon Sep 17 00:00:00 2001 From: Maxime Beauchemin Date: Fri, 9 Mar 2018 11:27:36 -0800 Subject: [PATCH] [sql lab] comment injection hook (#4585) --- 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 3f19d5797d..87da811fbb 100644 --- a/superset/config.py +++ b/superset/config.py @@ -384,9 +384,10 @@ ENABLE_JAVASCRIPT_CONTROLS = False # 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 @@ ENABLE_JAVASCRIPT_CONTROLS = False # 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 41d8742b65..9caa7adce4 100644 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -677,7 +677,7 @@ class Database(Model, AuditMixinNullable, ImportMixin): 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 4dae72720d..f98231ed77 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -17,7 +17,7 @@ import sqlalchemy 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 execute_sql( 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()