Compare commits

...

1 Commits

Author SHA1 Message Date
Beto Dealmeida
4a03e80fcb fix: pass request context to async queries 2024-07-01 17:11:11 -04:00

View File

@@ -26,6 +26,7 @@ from typing import Any, cast, Optional, Union
import backoff import backoff
import msgpack import msgpack
from celery.exceptions import SoftTimeLimitExceeded from celery.exceptions import SoftTimeLimitExceeded
from flask import current_app
from flask_babel import gettext as __ from flask_babel import gettext as __
from superset import ( from superset import (
@@ -128,7 +129,6 @@ def handle_query_error(
def get_query_backoff_handler(details: dict[Any, Any]) -> None: def get_query_backoff_handler(details: dict[Any, Any]) -> None:
print(details)
query_id = details["kwargs"]["query_id"] query_id = details["kwargs"]["query_id"]
logger.error( logger.error(
"Query with id `%s` could not be retrieved", str(query_id), exc_info=True "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, log_params: Optional[dict[str, Any]] = None,
) -> Optional[dict[str, Any]]: ) -> Optional[dict[str, Any]]:
"""Executes the sql query returns the results.""" """Executes the sql query returns the results."""
with override_user(security_manager.find_user(username)): # Make sure to pass a request context, so that DB engine specs that support OAuth2
try: # can compute the redirect URI using `url_for`.
return execute_sql_statements( with current_app.test_request_context():
query_id, with override_user(security_manager.find_user(username)):
rendered_query, try:
return_results, return execute_sql_statements(
store_results, query_id,
start_time=start_time, rendered_query,
expand_data=expand_data, return_results,
log_params=log_params, store_results,
) start_time=start_time,
except Exception as ex: # pylint: disable=broad-except expand_data=expand_data,
logger.debug("Query %d: %s", query_id, ex) log_params=log_params,
stats_logger.incr("error_sqllab_unhandled") )
query = get_query(query_id) except Exception as ex: # pylint: disable=broad-except
return handle_query_error(ex, query) 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 def execute_sql_statement( # pylint: disable=too-many-statements