fix(mcp): handle OAuth-authenticated databases in execute_sql (#39166)

This commit is contained in:
Amin Ghadersohi
2026-04-09 15:47:00 -04:00
committed by GitHub
parent 5815665cc6
commit 68067d7f44
14 changed files with 452 additions and 18 deletions

View File

@@ -644,6 +644,78 @@ def test_execute_error(
assert "Database error" in result.error_message
def test_execute_oauth2_redirect_error_propagates(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""Test that OAuth2RedirectError propagates instead of being swallowed."""
from superset.exceptions import OAuth2RedirectError
mock_conn = MagicMock()
mock_cursor = MagicMock()
mock_conn.cursor.return_value = mock_cursor
mock_conn.__enter__ = MagicMock(return_value=mock_conn)
mock_conn.__exit__ = MagicMock(return_value=False)
mocker.patch.object(database, "get_raw_connection", return_value=mock_conn)
mocker.patch.object(
database, "mutate_sql_based_on_config", side_effect=lambda sql, **kw: sql
)
mocker.patch.object(
database.db_engine_spec,
"execute",
side_effect=OAuth2RedirectError(
url="https://oauth.example.com/authorize",
tab_id="test-tab",
redirect_uri="https://superset.example.com/callback",
),
)
mocker.patch.dict(
current_app.config,
{
"SQL_QUERY_MUTATOR": None,
"SQLLAB_TIMEOUT": 30,
"SQL_MAX_ROW": None,
"QUERY_LOGGER": None,
},
)
with pytest.raises(OAuth2RedirectError):
database.execute("SELECT 1")
def test_execute_oauth2_error_propagates(
mocker: MockerFixture, database: Database, app_context: None
) -> None:
"""Test that OAuth2Error propagates instead of being swallowed."""
from superset.exceptions import OAuth2Error
mock_conn = MagicMock()
mock_cursor = MagicMock()
mock_conn.cursor.return_value = mock_cursor
mock_conn.__enter__ = MagicMock(return_value=mock_conn)
mock_conn.__exit__ = MagicMock(return_value=False)
mocker.patch.object(database, "get_raw_connection", return_value=mock_conn)
mocker.patch.object(
database, "mutate_sql_based_on_config", side_effect=lambda sql, **kw: sql
)
mocker.patch.object(
database.db_engine_spec,
"execute",
side_effect=OAuth2Error("No configuration found for OAuth2"),
)
mocker.patch.dict(
current_app.config,
{
"SQL_QUERY_MUTATOR": None,
"SQLLAB_TIMEOUT": 30,
"SQL_MAX_ROW": None,
"QUERY_LOGGER": None,
},
)
with pytest.raises(OAuth2Error):
database.execute("SELECT 1")
# =============================================================================
# Async Execution Tests
# =============================================================================