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

@@ -1093,3 +1093,74 @@ class TestSanitizeRowValues:
assert rows[0]["name"] == "test"
assert rows[0]["price"] == 9.99
assert rows[0]["blob"] == "000102ff"
class TestExecuteSqlOAuth2:
"""Tests for OAuth2 error handling in execute_sql."""
@patch("superset.security_manager")
@patch("superset.db")
@pytest.mark.asyncio
async def test_execute_sql_oauth2_redirect_error(
self, mock_db, mock_security_manager, mcp_server
):
"""Test that OAuth2RedirectError is caught and returns a clear message."""
from superset.exceptions import OAuth2RedirectError
mock_database = _mock_database()
mock_database.execute.side_effect = OAuth2RedirectError(
url="https://oauth.example.com/authorize",
tab_id="test-tab-id",
redirect_uri="https://superset.example.com/callback",
)
mock_db.session.query.return_value.filter_by.return_value.first.return_value = (
mock_database
)
mock_security_manager.can_access_database.return_value = True
request = {
"database_id": 1,
"sql": "SELECT 1",
"limit": 100,
}
async with Client(mcp_server) as client:
result = await client.call_tool("execute_sql", {"request": request})
data = result.structured_content
assert data["success"] is False
assert "OAuth" in data["error"]
assert "https://oauth.example.com/authorize" in data["error"]
assert data["error_type"] == "OAUTH2_REDIRECT"
@patch("superset.security_manager")
@patch("superset.db")
@pytest.mark.asyncio
async def test_execute_sql_oauth2_error(
self, mock_db, mock_security_manager, mcp_server
):
"""Test that OAuth2Error is caught and returns a clear message."""
from superset.exceptions import OAuth2Error
mock_database = _mock_database()
mock_database.execute.side_effect = OAuth2Error(
"Unable to determine the OAuth2 redirect URI."
)
mock_db.session.query.return_value.filter_by.return_value.first.return_value = (
mock_database
)
mock_security_manager.can_access_database.return_value = True
request = {
"database_id": 1,
"sql": "SELECT 1",
"limit": 100,
}
async with Client(mcp_server) as client:
result = await client.call_tool("execute_sql", {"request": request})
data = result.structured_content
assert data["success"] is False
assert "configuration" in data["error"]
assert data["error_type"] == "OAUTH2_REDIRECT_ERROR"