QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-05-28 03:05:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master' into adopt-pr-36387-bq-memory-limit

Browse Source
This commit is contained in:
Evan Rusackas
2026-04-22 15:14:39 -04:00
committed by GitHub
parent 05a9b47f1a e16656c6cf
commit 3601013a78
961 changed files with 64278 additions and 17595 deletions
Download Patch File Download Diff File Expand all files Collapse all files

330
tests/unit_tests/db_engine_specs/test_base.py
View File

@@ -21,6 +21,7 @@ from __future__ import annotations
import json # noqa: TID251
import re
from datetime import timedelta
from textwrap import dedent
from typing import Any
from urllib.parse import parse_qs, urlparse
@@ -44,6 +45,7 @@ from superset.superset_typing import (
)
from superset.utils.core import FilterOperator, GenericDataType
from superset.utils.oauth2 import decode_oauth2_state
from tests.conftest import with_config
from tests.unit_tests.db_engine_specs.utils import assert_column_spec
@@ -597,6 +599,19 @@ def test_extract_errors(mocker: MockerFixture) -> None:
assert result == [expected]
@with_config(
{
"CUSTOM_DATABASE_ERRORS": {
"examples": {
re.compile("This connector does not support roles"): (
"Custom error message",
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{},
)
}
}
},
)
def test_extract_errors_from_config(mocker: MockerFixture) -> None:
"""
Test that custom error messages are extracted correctly from app config
@@ -606,21 +621,6 @@ def test_extract_errors_from_config(mocker: MockerFixture) -> None:
class TestEngineSpec(BaseEngineSpec):
engine_name = "ExampleEngine"
mocker.patch(
"flask.current_app.config",
{
"CUSTOM_DATABASE_ERRORS": {
"examples": {
re.compile("This connector does not support roles"): (
"Custom error message",
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{},
)
}
}
},
)
msg = "This connector does not support roles"
result = TestEngineSpec.extract_errors(Exception(msg), database_name="examples")
@@ -631,6 +631,19 @@ def test_extract_errors_from_config(mocker: MockerFixture) -> None:
assert result == [expected]
@with_config(
{
"CUSTOM_DATABASE_ERRORS": {
"examples": {
re.compile("This connector does not support roles"): (
"Custom error message",
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{},
)
}
}
},
)
def test_extract_errors_only_to_specified_database(mocker: MockerFixture) -> None:
"""
Test that custom error messages are only applied to the specified database_name.
@@ -639,21 +652,6 @@ def test_extract_errors_only_to_specified_database(mocker: MockerFixture) -> Non
class TestEngineSpec(BaseEngineSpec):
engine_name = "ExampleEngine"
mocker.patch(
"flask.current_app.config",
{
"CUSTOM_DATABASE_ERRORS": {
"examples": {
re.compile("This connector does not support roles"): (
"Custom error message",
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{},
)
}
}
},
)
msg = "This connector does not support roles"
# database_name doesn't match configured one, so default message is used
result = TestEngineSpec.extract_errors(Exception(msg), database_name="examples_2")
@@ -665,6 +663,27 @@ def test_extract_errors_only_to_specified_database(mocker: MockerFixture) -> Non
assert result == [expected]
@with_config(
{
"CUSTOM_DATABASE_ERRORS": {
"examples": {
re.compile(r'message="(?P<message>[^"]*)"'): (
'Unexpected error: "%(message)s"',
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{
"custom_doc_links": [
{
"url": "https://example.com/docs",
"label": "Check documentation",
},
],
"show_issue_info": False,
},
)
}
}
},
)
def test_extract_errors_from_config_with_regex(mocker: MockerFixture) -> None:
"""
Test that custom error messages with regex, custom_doc_links,
@@ -674,29 +693,6 @@ def test_extract_errors_from_config_with_regex(mocker: MockerFixture) -> None:
class TestEngineSpec(BaseEngineSpec):
engine_name = "ExampleEngine"
mocker.patch(
"flask.current_app.config",
{
"CUSTOM_DATABASE_ERRORS": {
"examples": {
re.compile(r'message="(?P<message>[^"]*)"'): (
'Unexpected error: "%(message)s"',
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{
"custom_doc_links": [
{
"url": "https://example.com/docs",
"label": "Check documentation",
},
],
"show_issue_info": False,
},
)
}
}
},
)
msg = (
"db error: SomeUserError(type=USER_ERROR, name=TABLE_NOT_FOUND, "
'message="line 3:6: Table '
@@ -735,6 +731,7 @@ def test_extract_errors_from_config_with_regex(mocker: MockerFixture) -> None:
]
@with_config({"CUSTOM_DATABASE_ERRORS": {"examples": "not a dict"}})
def test_extract_errors_with_non_dict_custom_errors(mocker: MockerFixture):
"""
Test that extract_errors doesn't fail when custom database errors
@@ -744,11 +741,6 @@ def test_extract_errors_with_non_dict_custom_errors(mocker: MockerFixture):
class TestEngineSpec(BaseEngineSpec):
engine_name = "ExampleEngine"
mocker.patch(
"flask.current_app.config",
{"CUSTOM_DATABASE_ERRORS": "not a dict"},
)
msg = "This connector does not support roles"
result = TestEngineSpec.extract_errors(Exception(msg))
@@ -759,6 +751,7 @@ def test_extract_errors_with_non_dict_custom_errors(mocker: MockerFixture):
assert result == [expected]
@with_config({"CUSTOM_DATABASE_ERRORS": {"examples": "not a dict"}})
def test_extract_errors_with_non_dict_engine_custom_errors(mocker: MockerFixture):
"""
Test that extract_errors doesn't fail when database-specific custom errors
@@ -768,11 +761,6 @@ def test_extract_errors_with_non_dict_engine_custom_errors(mocker: MockerFixture
class TestEngineSpec(BaseEngineSpec):
engine_name = "ExampleEngine"
mocker.patch(
"flask.current_app.config",
{"CUSTOM_DATABASE_ERRORS": {"examples": "not a dict"}},
)
msg = "This connector does not support roles"
result = TestEngineSpec.extract_errors(Exception(msg), database_name="examples")
@@ -783,6 +771,19 @@ def test_extract_errors_with_non_dict_engine_custom_errors(mocker: MockerFixture
assert result == [expected]
@with_config(
{
"CUSTOM_DATABASE_ERRORS": {
"examples": {
re.compile("This connector does not support roles"): (
"",
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{},
)
}
}
},
)
def test_extract_errors_with_empty_custom_error_message(mocker: MockerFixture):
"""
Test that when the custom error message is empty,
@@ -792,21 +793,6 @@ def test_extract_errors_with_empty_custom_error_message(mocker: MockerFixture):
class TestEngineSpec(BaseEngineSpec):
engine_name = "ExampleEngine"
mocker.patch(
"flask.current_app.config",
{
"CUSTOM_DATABASE_ERRORS": {
"examples": {
re.compile("This connector does not support roles"): (
"",
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{},
)
}
}
},
)
msg = "This connector does not support roles"
result = TestEngineSpec.extract_errors(Exception(msg), database_name="examples")
@@ -817,6 +803,26 @@ def test_extract_errors_with_empty_custom_error_message(mocker: MockerFixture):
assert result == [expected]
@with_config(
{
"CUSTOM_DATABASE_ERRORS": {
"examples": {
re.compile("connection error"): (
"Examples DB error message",
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{},
)
},
"examples_2": {
re.compile("connection error"): (
"Examples_2 DB error message",
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{},
)
},
}
},
)
def test_extract_errors_matches_database_name_selection(mocker: MockerFixture) -> None:
"""
Test that custom error messages are matched by database_name.
@@ -825,28 +831,6 @@ def test_extract_errors_matches_database_name_selection(mocker: MockerFixture) -
class TestEngineSpec(BaseEngineSpec):
engine_name = "ExampleEngine"
mocker.patch(
"flask.current_app.config",
{
"CUSTOM_DATABASE_ERRORS": {
"examples": {
re.compile("connection error"): (
"Examples DB error message",
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{},
)
},
"examples_2": {
re.compile("connection error"): (
"Examples_2 DB error message",
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{},
)
},
}
},
)
msg = "connection error occurred"
# When database_name is examples_2 we should get that specific message
result = TestEngineSpec.extract_errors(Exception(msg), database_name="examples_2")
@@ -858,6 +842,19 @@ def test_extract_errors_matches_database_name_selection(mocker: MockerFixture) -
assert result == [expected]
@with_config(
{
"CUSTOM_DATABASE_ERRORS": {
"examples": {
re.compile("connection error"): (
"Examples DB error message",
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{},
)
},
}
},
)
def test_extract_errors_no_match_falls_back(mocker: MockerFixture) -> None:
"""
Test that when database_name has no match, the original error message is preserved.
@@ -866,21 +863,6 @@ def test_extract_errors_no_match_falls_back(mocker: MockerFixture) -> None:
class TestEngineSpec(BaseEngineSpec):
engine_name = "ExampleEngine"
mocker.patch(
"flask.current_app.config",
{
"CUSTOM_DATABASE_ERRORS": {
"examples": {
re.compile("connection error"): (
"Examples DB error message",
SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
{},
)
},
}
},
)
msg = "some other error"
result = TestEngineSpec.extract_errors(Exception(msg), database_name="examples_2")
@@ -980,16 +962,13 @@ def test_get_oauth2_authorization_uri_with_pkce(mocker: MockerFixture) -> None:
assert query["code_challenge"][0] == expected_challenge
@with_config({"DATABASE_OAUTH2_TIMEOUT": timedelta(seconds=30)})
def test_get_oauth2_token_without_pkce(mocker: MockerFixture) -> None:
"""
Test that BaseEngineSpec.get_oauth2_token works without PKCE code_verifier.
"""
from superset.db_engine_specs.base import BaseEngineSpec
mocker.patch(
"flask.current_app.config",
{"DATABASE_OAUTH2_TIMEOUT": mocker.MagicMock(total_seconds=lambda: 30)},
)
mock_post = mocker.patch("superset.db_engine_specs.base.requests.post")
mock_post.return_value.json.return_value = {
"access_token": "test-access-token", # noqa: S105
@@ -1015,6 +994,7 @@ def test_get_oauth2_token_without_pkce(mocker: MockerFixture) -> None:
assert "code_verifier" not in request_body
@with_config({"DATABASE_OAUTH2_TIMEOUT": timedelta(seconds=30)})
def test_get_oauth2_token_with_pkce(mocker: MockerFixture) -> None:
"""
Test BaseEngineSpec.get_oauth2_token includes code_verifier when provided.
@@ -1022,10 +1002,6 @@ def test_get_oauth2_token_with_pkce(mocker: MockerFixture) -> None:
from superset.db_engine_specs.base import BaseEngineSpec
from superset.utils.oauth2 import generate_code_verifier
mocker.patch(
"flask.current_app.config",
{"DATABASE_OAUTH2_TIMEOUT": mocker.MagicMock(total_seconds=lambda: 30)},
)
mock_post = mocker.patch("superset.db_engine_specs.base.requests.post")
mock_post.return_value.json.return_value = {
"access_token": "test-access-token", # noqa: S105
@@ -1097,6 +1073,7 @@ def test_get_oauth2_authorization_uri_additional_params(
assert query["access_type"][0] == "offline"
@with_config({"DATABASE_OAUTH2_TIMEOUT": timedelta(seconds=30)})
def test_get_oauth2_token_additional_params(mocker: MockerFixture) -> None:
"""
Test that a subclass can inject additional params into the token request body
@@ -1109,10 +1086,6 @@ def test_get_oauth2_token_additional_params(mocker: MockerFixture) -> None:
"audience": "https://api.example.com",
}
mocker.patch(
"flask.current_app.config",
{"DATABASE_OAUTH2_TIMEOUT": mocker.MagicMock(total_seconds=lambda: 30)},
)
mock_post = mocker.patch("superset.db_engine_specs.base.requests.post")
mock_post.return_value.json.return_value = {
"access_token": "test-access-token", # noqa: S105
@@ -1143,6 +1116,94 @@ def test_get_oauth2_token_additional_params(mocker: MockerFixture) -> None:
assert request_body["audience"] == "https://api.example.com"
@with_config({"DATABASE_OAUTH2_TIMEOUT": timedelta(seconds=30)})
def test_get_oauth2_fresh_token_success(mocker: MockerFixture) -> None:
"""
Test that get_oauth2_fresh_token returns the token response on success.
"""
from superset.db_engine_specs.base import BaseEngineSpec
mock_post = mocker.patch("superset.db_engine_specs.base.requests.post")
mock_post.return_value.status_code = 200
mock_post.return_value.json.return_value = {
"access_token": "new-access-token",
"expires_in": 3600,
}
config: OAuth2ClientConfig = {
"id": "client-id",
"secret": "client-secret",
"scope": "read write",
"redirect_uri": "http://localhost:8088/api/v1/database/oauth2/",
"authorization_request_uri": "https://oauth.example.com/authorize",
"token_request_uri": "https://oauth.example.com/token",
"request_content_type": "json",
}
result = BaseEngineSpec.get_oauth2_fresh_token(config, "refresh-token")
assert result == {"access_token": "new-access-token", "expires_in": 3600}
@pytest.mark.parametrize("status_code", [400, 401, 403])
@with_config({"DATABASE_OAUTH2_TIMEOUT": timedelta(seconds=30)})
def test_get_oauth2_fresh_token_raises_on_auth_error(
mocker: MockerFixture,
status_code: int,
) -> None:
"""
Test that get_oauth2_fresh_token raises OAuth2TokenRefreshError on 400/401/403.
"""
from superset.db_engine_specs.base import BaseEngineSpec
from superset.exceptions import OAuth2TokenRefreshError
mock_post = mocker.patch("superset.db_engine_specs.base.requests.post")
mock_post.return_value.status_code = status_code
mock_post.return_value.text = '{"error": "invalid_grant"}'
config: OAuth2ClientConfig = {
"id": "client-id",
"secret": "client-secret",
"scope": "read write",
"redirect_uri": "http://localhost:8088/api/v1/database/oauth2/",
"authorization_request_uri": "https://oauth.example.com/authorize",
"token_request_uri": "https://oauth.example.com/token",
"request_content_type": "json",
}
with pytest.raises(OAuth2TokenRefreshError) as exc_info:
BaseEngineSpec.get_oauth2_fresh_token(config, "refresh-token")
assert exc_info.value.error.extra["error"] == '{"error": "invalid_grant"}'
@with_config({"DATABASE_OAUTH2_TIMEOUT": timedelta(seconds=30)})
def test_get_oauth2_fresh_token_raises_on_server_error(mocker: MockerFixture) -> None:
"""
Test that get_oauth2_fresh_token raises HTTPError (not OAuth2TokenRefreshError)
on 5xx.
"""
from requests.exceptions import HTTPError
from superset.db_engine_specs.base import BaseEngineSpec
mock_post = mocker.patch("superset.db_engine_specs.base.requests.post")
mock_post.return_value.status_code = 500
mock_post.return_value.raise_for_status.side_effect = HTTPError("500 Server Error")
config: OAuth2ClientConfig = {
"id": "client-id",
"secret": "client-secret",
"scope": "read write",
"redirect_uri": "http://localhost:8088/api/v1/database/oauth2/",
"authorization_request_uri": "https://oauth.example.com/authorize",
"token_request_uri": "https://oauth.example.com/token",
"request_content_type": "json",
}
with pytest.raises(HTTPError):
BaseEngineSpec.get_oauth2_fresh_token(config, "refresh-token")
def test_start_oauth2_dance_uses_config_redirect_uri(mocker: MockerFixture) -> None:
"""
Test that start_oauth2_dance uses DATABASE_OAUTH2_REDIRECT_URI config if set.
@@ -1182,6 +1243,12 @@ def test_start_oauth2_dance_uses_config_redirect_uri(mocker: MockerFixture) -> N
assert error.extra["redirect_uri"] == custom_redirect_uri
@with_config(
{
"SECRET_KEY": "test-secret-key",
"DATABASE_OAUTH2_JWT_ALGORITHM": "HS256",
}
)
def test_start_oauth2_dance_falls_back_to_url_for(mocker: MockerFixture) -> None:
"""
Test that start_oauth2_dance falls back to url_for when no config is set.
@@ -1189,14 +1256,7 @@ def test_start_oauth2_dance_falls_back_to_url_for(mocker: MockerFixture) -> None
fallback_uri = "http://localhost:8088/api/v1/database/oauth2/"
mocker.patch(
"flask.current_app.config",
{
"SECRET_KEY": "test-secret-key",
"DATABASE_OAUTH2_JWT_ALGORITHM": "HS256",
},
)
mocker.patch(
"superset.db_engine_specs.base.url_for",
"superset.utils.oauth2.url_for",
return_value=fallback_uri,
)
mocker.patch("superset.daos.key_value.KeyValueDAO")

25
tests/unit_tests/db_engine_specs/test_gsheets.py
View File

@@ -24,11 +24,10 @@ import pandas as pd
import pytest
from pytest_mock import MockerFixture
from requests.exceptions import HTTPError
from shillelagh.exceptions import UnauthenticatedError
from sqlalchemy.engine.url import make_url
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import SupersetException
from superset.exceptions import OAuth2TokenRefreshError, SupersetException
from superset.sql.parse import Table
from superset.superset_typing import OAuth2ClientConfig
from superset.utils import json
@@ -817,26 +816,20 @@ def test_get_oauth2_fresh_token_invalid_grant(
oauth2_config: OAuth2ClientConfig,
) -> None:
"""
Test that get_oauth2_fresh_token raises UnauthenticatedError for invalid_grant.
Test that get_oauth2_fresh_token raises OAuth2TokenRefreshError for a 400 response.
When a token is revoked on Google side, the refresh request returns 400
with error=invalid_grant.
When a token is revoked on Google side, the refresh request returns 400.
"""
from superset.db_engine_specs.gsheets import GSheetsEngineSpec
mock_response = mocker.MagicMock()
mock_response.status_code = 400
mock_response.json.return_value = {
"error": "invalid_grant",
"error_description": "Token has been expired or revoked.",
}
http_error = HTTPError()
http_error.response = mock_response
requests = mocker.patch("superset.db_engine_specs.base.requests")
requests.post().raise_for_status.side_effect = http_error
requests.post().status_code = 400
requests.post().text = (
'{"error": "invalid_grant",'
' "error_description": "Token has been expired or revoked."}'
)
with pytest.raises(UnauthenticatedError):
with pytest.raises(OAuth2TokenRefreshError):
GSheetsEngineSpec.get_oauth2_fresh_token(oauth2_config, "refresh-token")

85
tests/unit_tests/db_engine_specs/test_postgres.py
View File

@@ -17,6 +17,7 @@
from datetime import datetime
from typing import Any, Optional
from unittest.mock import MagicMock
import pytest
from pytest_mock import MockerFixture
@@ -25,7 +26,10 @@ from sqlalchemy.dialects.postgresql import DOUBLE_PRECISION, ENUM, JSON
from sqlalchemy.engine.interfaces import Dialect
from sqlalchemy.engine.url import make_url
from superset.db_engine_specs.postgres import PostgresEngineSpec as spec # noqa: N813
from superset.db_engine_specs.postgres import (
_check_not_redshift,
PostgresEngineSpec as spec, # noqa: N813
)
from superset.exceptions import SupersetSecurityException
from superset.sql.parse import Table
from superset.utils.core import GenericDataType
@@ -280,3 +284,82 @@ SELECT * \nFROM my_schema.my_table
LIMIT :param_1
""".strip()
)
class TestRedshiftDetection:
"""
Tests for detecting Redshift connections via the PostgreSQL dialect.
"""
def test_check_not_redshift_detects_redshift(self) -> None:
"""
Pool connect event raises for a Redshift version string.
"""
cursor = MagicMock()
cursor.fetchone.return_value = (
"PostgreSQL 8.0.2 on i686-pc-linux-gnu, compiled by GCC gcc (GCC) "
"3.4.2 20041017 (Red Hat 3.4.2-6.fc3), Redshift 1.0.77467",
)
dbapi_conn = MagicMock()
dbapi_conn.cursor.return_value = cursor
with pytest.raises(ValueError, match="Redshift"):
_check_not_redshift(dbapi_conn, None)
def test_check_not_redshift_allows_postgres(self) -> None:
"""
Pool connect event allows a regular PostgreSQL version string.
"""
cursor = MagicMock()
cursor.fetchone.return_value = (
"PostgreSQL 15.2 on x86_64-pc-linux-gnu, compiled by gcc",
)
dbapi_conn = MagicMock()
dbapi_conn.cursor.return_value = cursor
_check_not_redshift(dbapi_conn, None) # should not raise
def test_check_not_redshift_fails_open(self) -> None:
"""
If SELECT version() errors, the connection is still allowed.
"""
cursor = MagicMock()
cursor.execute.side_effect = Exception("permission denied")
dbapi_conn = MagicMock()
dbapi_conn.cursor.return_value = cursor
_check_not_redshift(dbapi_conn, None) # should not raise
def test_mutate_db_sets_flag(self) -> None:
"""
mutate_db_for_connection_test sets the check flag.
"""
database = MagicMock()
spec.mutate_db_for_connection_test(database)
assert database._check_redshift_version is True
def test_pool_event_injected_when_flag_set(self, mocker: MockerFixture) -> None:
"""
Pool event is added during test_connection.
"""
database = mocker.MagicMock(
encrypted_extra=None,
_check_redshift_version=True,
)
params: dict[str, Any] = {}
spec.update_params_from_encrypted_extra(database, params)
assert "pool_events" in params
fns = [fn for fn, _ in params["pool_events"]]
assert _check_not_redshift in fns
def test_pool_event_not_injected_without_flag(self, mocker: MockerFixture) -> None:
"""
Pool event is NOT added during normal operation.
"""
database = mocker.MagicMock(encrypted_extra=None)
database._check_redshift_version = False
params: dict[str, Any] = {}
spec.update_params_from_encrypted_extra(database, params)
assert "pool_events" not in params

6
tests/unit_tests/db_engine_specs/test_presto.py
View File

@@ -42,17 +42,17 @@ from tests.unit_tests.db_engine_specs.utils import (
(
"TIMESTAMP",
datetime(2022, 1, 1, 1, 23, 45, 600000),
"TIMESTAMP '2022-01-01 01:23:45.600000'",
"TIMESTAMP '2022-01-01 01:23:45.600'",
),
(
"TIMESTAMP WITH TIME ZONE",
datetime(2022, 1, 1, 1, 23, 45, 600000),
"TIMESTAMP '2022-01-01 01:23:45.600000'",
"TIMESTAMP '2022-01-01 01:23:45.600'",
),
(
"TIMESTAMP WITH TIME ZONE",
datetime(2022, 1, 1, 1, 23, 45, 600000, tzinfo=pytz.UTC),
"TIMESTAMP '2022-01-01 01:23:45.600000+00:00'",
"TIMESTAMP '2022-01-01 01:23:45.600+00:00'",
),
],
)

58
tests/unit_tests/db_engine_specs/test_trino.py
View File

@@ -859,6 +859,64 @@ def test_get_oauth2_token(
)
def test_needs_oauth2_with_401_error(mocker: MockerFixture) -> None:
"""
Test that needs_oauth2 returns True when Trino raises an HTTP 401 error.
"""
from trino.exceptions import HttpError
from superset.db_engine_specs.trino import TrinoEngineSpec
g = mocker.patch("superset.db_engine_specs.trino.g")
g.user = mocker.MagicMock()
ex = HttpError("error 401: Unauthorized")
assert TrinoEngineSpec.needs_oauth2(ex) is True
def test_needs_oauth2_without_401_error(mocker: MockerFixture) -> None:
"""
Test that needs_oauth2 returns False when the error is not a 401.
"""
from trino.exceptions import HttpError
from superset.db_engine_specs.trino import TrinoEngineSpec
g = mocker.patch("superset.db_engine_specs.trino.g")
g.user = mocker.MagicMock()
ex = HttpError("error 500: Internal Server Error")
assert TrinoEngineSpec.needs_oauth2(ex) is False
def test_needs_oauth2_with_non_http_error(mocker: MockerFixture) -> None:
"""
Test that needs_oauth2 returns False for non-HttpError exceptions.
"""
from superset.db_engine_specs.trino import TrinoEngineSpec
g = mocker.patch("superset.db_engine_specs.trino.g")
g.user = mocker.MagicMock()
ex = RuntimeError("error 401: something else")
assert TrinoEngineSpec.needs_oauth2(ex) is False
def test_needs_oauth2_without_user(mocker: MockerFixture) -> None:
"""
Test that needs_oauth2 returns False when there is no authenticated user.
"""
from trino.exceptions import HttpError
from superset.db_engine_specs.trino import TrinoEngineSpec
g = mocker.patch("superset.db_engine_specs.trino.g")
del g.user
ex = HttpError("error 401: Unauthorized")
assert TrinoEngineSpec.needs_oauth2(ex) is False
@pytest.mark.parametrize(
"time_grain,expected_result",
[