fix: OAuth2 in async DBs (#29461)

This commit is contained in:
Beto Dealmeida 2024-07-02 21:12:07 -04:00 committed by GitHub
parent fb1f2c4f18
commit d5c0506faa
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 20 additions and 19 deletions

View File

@ -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,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

View File

@ -273,9 +273,9 @@ def test_get_sql_results_oauth2(mocker: MockerFixture, app) -> None:
"error_type": SupersetErrorType.OAUTH2_REDIRECT,
"level": ErrorLevel.WARNING,
"extra": {
"url": "https://abcd1234.snowflakecomputing.com/oauth/authorize?scope=refresh_token+session%3Arole%3ASYSADMIN&access_type=offline&include_granted_scopes=false&response_type=code&state=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9%252EeyJleHAiOjE2MTcyMzU1MDAsImRhdGFiYXNlX2lkIjoxLCJ1c2VyX2lkIjo0MiwiZGVmYXVsdF9yZWRpcmVjdF91cmkiOiJodHRwOi8vZXhhbXBsZS5jb20vYXBpL3YxL2RhdGFiYXNlL29hdXRoMi8iLCJ0YWJfaWQiOiJmYjExZjUyOC02ZWJhLTRhOGEtODM3ZS02YjBkMzllZTkxODcifQ%252Ec_m_35xwwSrLgCXwV4aKO1928flOEFQIqqg9ctiXjcM&redirect_uri=http%3A%2F%2Fexample.com%2Fapi%2Fv1%2Fdatabase%2Foauth2%2F&client_id=my_client_id&prompt=consent",
"url": "https://abcd1234.snowflakecomputing.com/oauth/authorize?scope=refresh_token+session%3Arole%3ASYSADMIN&access_type=offline&include_granted_scopes=false&response_type=code&state=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9%252EeyJleHAiOjE2MTcyMzU1MDAsImRhdGFiYXNlX2lkIjoxLCJ1c2VyX2lkIjo0MiwiZGVmYXVsdF9yZWRpcmVjdF91cmkiOiJodHRwOi8vbG9jYWxob3N0L2FwaS92MS9kYXRhYmFzZS9vYXV0aDIvIiwidGFiX2lkIjoiZmIxMWY1MjgtNmViYS00YThhLTgzN2UtNmIwZDM5ZWU5MTg3In0%252E7nLkei6-V8sVk_Pgm8cFhk0tnKRKayRE1Vc7RxuM9mw&redirect_uri=http%3A%2F%2Flocalhost%2Fapi%2Fv1%2Fdatabase%2Foauth2%2F&client_id=my_client_id&prompt=consent",
"tab_id": "fb11f528-6eba-4a8a-837e-6b0d39ee9187",
"redirect_uri": "http://example.com/api/v1/database/oauth2/",
"redirect_uri": "http://localhost/api/v1/database/oauth2/",
},
}
],