diff --git a/superset/sql_lab.py b/superset/sql_lab.py index fb5f6aff79..eb9303455e 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -151,6 +151,7 @@ def get_sql_results( # pylint: disable=too-many-arguments user_name=None, start_time=None, expand_data=False, + log_params=None, ): """Executes the sql query returns the results.""" with session_scope(not ctask.request.called_directly) as session: @@ -165,6 +166,7 @@ def get_sql_results( # pylint: disable=too-many-arguments session=session, start_time=start_time, expand_data=expand_data, + log_params=log_params, ) except Exception as e: # pylint: disable=broad-except logger.exception(f"Query {query_id}: {e}") @@ -173,7 +175,8 @@ def get_sql_results( # pylint: disable=too-many-arguments return handle_query_error(str(e), query, session) -def execute_sql_statement(sql_statement, query, user_name, session, cursor): +# pylint: disable=too-many-arguments +def execute_sql_statement(sql_statement, query, user_name, session, cursor, log_params): """Executes a single SQL statement""" database = query.database db_engine_spec = database.db_engine_spec @@ -218,6 +221,7 @@ def execute_sql_statement(sql_statement, query, user_name, session, cursor): user_name, __name__, security_manager, + log_params, ) query.executed_sql = sql session.commit() @@ -303,6 +307,7 @@ def execute_sql_statements( session=None, start_time=None, expand_data=False, + log_params=None, ): # pylint: disable=too-many-arguments, too-many-locals, too-many-statements """Executes the sql query returns the results.""" if store_results and start_time: @@ -352,7 +357,7 @@ def execute_sql_statements( session.commit() try: cdf = execute_sql_statement( - statement, query, user_name, session, cursor + statement, query, user_name, session, cursor, log_params ) except Exception as e: # pylint: disable=broad-except msg = str(e) diff --git a/superset/views/core.py b/superset/views/core.py index ec247169b7..81314ff8a6 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -2439,7 +2439,12 @@ class Superset(BaseSupersetView): return json_error_response(f"{msg}") def _sql_json_async( - self, session: Session, rendered_query: str, query: Query, expand_data: bool + self, + session: Session, + rendered_query: str, + query: Query, + expand_data: bool, + log_params: Dict[str, Any], ) -> str: """ Send SQL JSON query to celery workers @@ -2460,6 +2465,7 @@ class Superset(BaseSupersetView): user_name=g.user.username if g.user else None, start_time=now_as_float(), expand_data=expand_data, + log_params=log_params, ) except Exception as e: logging.exception(f"Query {query.id}: {e}") @@ -2484,7 +2490,12 @@ class Superset(BaseSupersetView): return resp def _sql_json_sync( - self, session: Session, rendered_query: str, query: Query, expand_data: bool + self, + session: Session, + rendered_query: str, + query: Query, + expand_data: bool, + log_params: Dict[str, Any], ) -> str: """ Execute SQL query (sql json) @@ -2509,6 +2520,7 @@ class Superset(BaseSupersetView): store_results=store_results, user_name=g.user.username if g.user else None, expand_data=expand_data, + log_params=log_params, ) payload = json.dumps( @@ -2528,9 +2540,12 @@ class Superset(BaseSupersetView): @expose("/sql_json/", methods=["POST"]) @event_logger.log_this def sql_json(self): - return self.sql_json_exec(request.json) + log_params = { + "USER_AGENT": cast(Optional[str], request.headers.get("USER_AGENT")) + } + return self.sql_json_exec(request.json, log_params) - def sql_json_exec(self, query_params: dict): + def sql_json_exec(self, query_params: dict, log_params: dict): """Runs arbitrary sql and returns data as json""" # Collect Values database_id: int = cast(int, query_params.get("database_id")) @@ -2636,9 +2651,13 @@ class Superset(BaseSupersetView): # Async request. if async_flag: - return self._sql_json_async(session, rendered_query, query, expand_data) + return self._sql_json_async( + session, rendered_query, query, expand_data, log_params + ) # Sync request. - return self._sql_json_sync(session, rendered_query, query, expand_data) + return self._sql_json_sync( + session, rendered_query, query, expand_data, log_params + ) @has_access @expose("/csv/")