This is an automated email from the ASF dual-hosted git repository. beto pushed a commit to branch request-context-oauth in repository https://gitbox.apache.org/repos/asf/superset.git
commit 4a03e80fcbc73bc57dadaf972ad773870642459a Author: Beto Dealmeida <[email protected]> AuthorDate: Mon Jul 1 17:10:22 2024 -0400 fix: pass request context to async queries --- superset/sql_lab.py | 37 ++++++++++++++++++++----------------- 1 file changed, 20 insertions(+), 17 deletions(-) diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 9712ab47ab..bd1b1ec3a4 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -26,6 +26,7 @@ from typing import Any, cast, Optional, Union import backoff import msgpack from celery.exceptions import SoftTimeLimitExceeded +from flask import current_app from flask_babel import gettext as __ from superset import ( @@ -128,7 +129,6 @@ def handle_query_error( def get_query_backoff_handler(details: dict[Any, Any]) -> None: - print(details) query_id = details["kwargs"]["query_id"] logger.error( "Query with id `%s` could not be retrieved", str(query_id), exc_info=True @@ -175,22 +175,25 @@ def get_sql_results( # pylint: disable=too-many-arguments log_params: Optional[dict[str, Any]] = None, ) -> Optional[dict[str, Any]]: """Executes the sql query returns the results.""" - with override_user(security_manager.find_user(username)): - try: - return execute_sql_statements( - query_id, - rendered_query, - return_results, - store_results, - start_time=start_time, - expand_data=expand_data, - log_params=log_params, - ) - except Exception as ex: # pylint: disable=broad-except - logger.debug("Query %d: %s", query_id, ex) - stats_logger.incr("error_sqllab_unhandled") - query = get_query(query_id) - return handle_query_error(ex, query) + # Make sure to pass a request context, so that DB engine specs that support OAuth2 + # can compute the redirect URI using `url_for`. + with current_app.test_request_context(): + with override_user(security_manager.find_user(username)): + try: + return execute_sql_statements( + query_id, + rendered_query, + return_results, + store_results, + start_time=start_time, + expand_data=expand_data, + log_params=log_params, + ) + except Exception as ex: # pylint: disable=broad-except + logger.debug("Query %d: %s", query_id, ex) + stats_logger.incr("error_sqllab_unhandled") + query = get_query(query_id) + return handle_query_error(ex, query) def execute_sql_statement( # pylint: disable=too-many-statements
