[incubator-superset] branch master updated: Hook for auditing queries (#6484)

[be...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

beto pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-superset.git


The following commit(s) were added to refs/heads/master by this push:
     new 5168c69  Hook for auditing queries (#6484)
5168c69 is described below

commit 5168c69a25bcbcb3b06b950286b63fa373113675
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Fri Dec 7 12:19:19 2018 -0800

    Hook for auditing queries (#6484)
    
    * Hook for auditing queries
    
    * Address comments
    
    * Use __name__
---
 superset/config.py      | 6 ++++++
 superset/models/core.py | 8 ++++++++
 superset/sql_lab.py     | 9 +++++++++
 3 files changed, 23 insertions(+)

diff --git a/superset/config.py b/superset/config.py
index 73e5a2c..1613e75 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -281,6 +281,12 @@ ROLLOVER = 'midnight'
 INTERVAL = 1
 BACKUP_COUNT = 30
 
+# Custom logger for auditing queries. This can be used to send ran queries to a
+# structured immutable store for auditing purposes. The function is called for
+# every query ran, in both SQL Lab and charts/dashboards.
+# def QUERY_LOGGER(database, query, schema=None, user=None, client=None):
+#     pass
+
 # Set this API key to enable Mapbox visualizations
 MAPBOX_API_KEY = os.environ.get('MAPBOX_API_KEY', '')
 
diff --git a/superset/models/core.py b/superset/models/core.py
index 1323eed..6f520aa 100644
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -43,6 +43,7 @@ from urllib import parse  # noqa
 config = app.config
 custom_password_store = config.get('SQLALCHEMY_CUSTOM_PASSWORD_STORE')
 stats_logger = config.get('STATS_LOGGER')
+log_query = config.get('QUERY_LOGGER')
 metadata = Model.metadata  # pylint: disable=no-member
 
 PASSWORD_MASK = 'X' * 10
@@ -801,6 +802,7 @@ class Database(Model, AuditMixinNullable, ImportMixin):
     def get_df(self, sql, schema):
         sqls = [str(s).strip().strip(';') for s in sqlparse.parse(sql)]
         engine = self.get_sqla_engine(schema=schema)
+        username = utils.get_username()
 
         def needs_conversion(df_series):
             if df_series.empty:
@@ -809,12 +811,18 @@ class Database(Model, AuditMixinNullable, ImportMixin):
                 return True
             return False
 
+        def _log_query(sql):
+            if log_query:
+                log_query(engine.url, sql, schema, username, __name__)
+
         with closing(engine.raw_connection()) as conn:
             with closing(conn.cursor()) as cursor:
                 for sql in sqls[:-1]:
+                    _log_query(sql)
                     self.db_engine_spec.execute(cursor, sql)
                     cursor.fetchall()
 
+                _log_query(sqls[-1])
                 self.db_engine_spec.execute(cursor, sqls[-1])
 
                 if cursor.description is not None:
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index 1d4171b..0969147 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -26,6 +26,7 @@ config = app.config
 celery_app = get_celery_app(config)
 stats_logger = app.config.get('STATS_LOGGER')
 SQLLAB_TIMEOUT = config.get('SQLLAB_ASYNC_TIME_LIMIT_SEC', 600)
+log_query = config.get('QUERY_LOGGER')
 
 
 class SqlLabException(Exception):
@@ -180,6 +181,14 @@ def execute_sql(
         logging.info('Running query: \n{}'.format(executed_sql))
         logging.info(query.executed_sql)
         query_start_time = now_as_float()
+        if log_query:
+            log_query(
+                query.database.sqlalchemy_uri,
+                query.executed_sql,
+                query.schema,
+                user_name,
+                __name__,
+            )
         db_engine_spec.execute(cursor, query.executed_sql, async_=True)
         logging.info('Handling cursor')
         db_engine_spec.handle_cursor(cursor, query, session)