This is an automated email from the ASF dual-hosted git repository. beto pushed a commit to branch async-oauth2-refresh in repository https://gitbox.apache.org/repos/asf/superset.git
commit d9897f6f7865df8f7c9bb960e645572c5573a62f Author: Beto Dealmeida <[email protected]> AuthorDate: Tue Jul 2 19:21:25 2024 -0400 fix: OAuth2 in async DBs --- superset/sql_lab.py | 40 ++++++++++++++++++++-------------------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 9712ab47ab..9a5c01f54f 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -14,7 +14,6 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# pylint: disable=consider-using-transaction import dataclasses import logging import uuid @@ -26,6 +25,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 +128,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 +174,23 @@ 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) + 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 @@ -314,9 +314,9 @@ def execute_sql_statement( # pylint: disable=too-many-statements level=ErrorLevel.ERROR, ) ) from ex - except OAuth2RedirectError: + except OAuth2RedirectError as ex: # user needs to authenticate with OAuth2 in order to run query - raise + raise ex except Exception as ex: # query is stopped in another thread/worker # stopping raises expected exceptions which we should skip
