QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-07-27 17:12:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
(cherry picked from commit 68067d7f44)
This commit is contained in:
Amin Ghadersohi
2026-04-09 15:47:00 -04:00
committed by Joe Li
parent ece8b8ff5b
commit 3cda699864
11 changed files with 288 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
superset/db_engine_specs/base.py
View File

@@ -43,7 +43,7 @@ import requests
from apispec import APISpec
from apispec.ext.marshmallow import MarshmallowPlugin
from deprecation import deprecated
from flask import current_app as app, g, url_for
from flask import current_app as app, g
from flask_appbuilder.security.sqla.models import User
from flask_babel import gettext as __, lazy_gettext as _
from marshmallow import fields, Schema
@@ -88,6 +88,7 @@ from superset.utils.oauth2 import (
encode_oauth2_state,
generate_code_challenge,
generate_code_verifier,
get_oauth2_redirect_uri,
)
if TYPE_CHECKING:
@@ -520,10 +521,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
from superset.daos.key_value import KeyValueDAO
tab_id = str(uuid4())
default_redirect_uri = app.config.get(
"DATABASE_OAUTH2_REDIRECT_URI",
url_for("DatabaseRestApi.oauth2", _external=True),
)
default_redirect_uri = get_oauth2_redirect_uri()
# Generate PKCE code verifier (RFC 7636)
code_verifier = generate_code_verifier()
@@ -586,10 +584,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
return None
db_engine_spec_config = oauth2_config[cls.engine_name]
redirect_uri = app.config.get(
"DATABASE_OAUTH2_REDIRECT_URI",
url_for("DatabaseRestApi.oauth2", _external=True),
)
redirect_uri = get_oauth2_redirect_uri()
config: OAuth2ClientConfig = {
"id": db_engine_spec_config["id"],

36
superset/mcp_service/chart/tool/generate_chart.py
View File

@@ -29,6 +29,7 @@ from mcp.types import ToolAnnotations
from sqlalchemy.exc import SQLAlchemyError
from superset.commands.exceptions import CommandException
from superset.exceptions import OAuth2Error, OAuth2RedirectError
from superset.extensions import event_logger
from superset.mcp_service.app import mcp
from superset.mcp_service.auth import has_dataset_access, mcp_auth_hook
@@ -46,6 +47,10 @@ from superset.mcp_service.chart.schemas import (
GenerateChartResponse,
PerformanceMetadata,
)
from superset.mcp_service.utils.oauth2_utils import (
build_oauth2_redirect_message,
OAUTH2_CONFIG_ERROR_MESSAGE,
)
from superset.mcp_service.utils.url_utils import get_superset_base_url
from superset.utils import json
@@ -805,6 +810,37 @@ async def generate_chart( # noqa: C901
)
return GenerateChartResponse.model_validate(result)
except OAuth2RedirectError as ex:
await ctx.error(
"Chart generation requires OAuth authentication: dataset_id=%s"
% request.dataset_id
)
return GenerateChartResponse.model_validate(
{
"chart": None,
"success": False,
"error": {
"error_type": "OAUTH2_REDIRECT",
"message": build_oauth2_redirect_message(ex),
"details": "OAuth2 authentication required",
},
}
)
except OAuth2Error:
await ctx.error(
"OAuth2 configuration error: dataset_id=%s" % request.dataset_id
)
return GenerateChartResponse.model_validate(
{
"chart": None,
"success": False,
"error": {
"error_type": "OAUTH2_REDIRECT_ERROR",
"message": OAUTH2_CONFIG_ERROR_MESSAGE,
"details": "OAuth2 configuration or provider error",
},
}
)
except (CommandException, SQLAlchemyError, KeyError, ValueError) as e:
from superset import db

23
superset/mcp_service/chart/tool/get_chart_data.py
View File

@@ -35,7 +35,7 @@ if TYPE_CHECKING:
from superset.commands.exceptions import CommandException
from superset.commands.explore.form_data.parameters import CommandParameters
from superset.exceptions import SupersetException
from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
from superset.mcp_service.auth import mcp_auth_hook
from superset.mcp_service.chart.chart_utils import validate_chart_dataset
from superset.mcp_service.chart.schemas import (
@@ -46,6 +46,10 @@ from superset.mcp_service.chart.schemas import (
PerformanceMetadata,
)
from superset.mcp_service.utils.cache_utils import get_cache_status_from_result
from superset.mcp_service.utils.oauth2_utils import (
build_oauth2_redirect_message,
OAUTH2_CONFIG_ERROR_MESSAGE,
)
from superset.utils.core import merge_extra_filters
logger = logging.getLogger(__name__)
@@ -728,6 +732,23 @@ async def get_chart_data( # noqa: C901
error_type="DataError",
)
except OAuth2RedirectError as ex:
await ctx.error(
"Chart data requires OAuth authentication: identifier=%s"
% request.identifier
)
return ChartError(
error=build_oauth2_redirect_message(ex),
error_type="OAUTH2_REDIRECT",
)
except OAuth2Error:
await ctx.error(
"OAuth2 configuration error: identifier=%s" % request.identifier
)
return ChartError(
error=OAUTH2_CONFIG_ERROR_MESSAGE,
error_type="OAUTH2_REDIRECT_ERROR",
)
except Exception as e:
await ctx.error(
"Chart data retrieval failed: identifier=%s, error=%s, error_type=%s"

23
superset/mcp_service/chart/tool/get_chart_preview.py
View File

@@ -26,7 +26,7 @@ from fastmcp import Context
from mcp.types import ToolAnnotations
from superset.commands.exceptions import CommandException
from superset.exceptions import SupersetException
from superset.exceptions import OAuth2Error, OAuth2RedirectError, SupersetException
from superset.extensions import event_logger
from superset.mcp_service.app import mcp
from superset.mcp_service.auth import mcp_auth_hook
@@ -43,6 +43,10 @@ from superset.mcp_service.chart.schemas import (
URLPreview,
VegaLitePreview,
)
from superset.mcp_service.utils.oauth2_utils import (
build_oauth2_redirect_message,
OAUTH2_CONFIG_ERROR_MESSAGE,
)
from superset.mcp_service.utils.url_utils import get_superset_base_url
logger = logging.getLogger(__name__)
@@ -2127,6 +2131,23 @@ async def get_chart_preview(
)
return result
except OAuth2RedirectError as ex:
await ctx.error(
"Chart preview requires OAuth authentication: identifier=%s"
% request.identifier
)
return ChartError(
error=build_oauth2_redirect_message(ex),
error_type="OAUTH2_REDIRECT",
)
except OAuth2Error:
await ctx.error(
"OAuth2 configuration error: identifier=%s" % request.identifier
)
return ChartError(
error=OAUTH2_CONFIG_ERROR_MESSAGE,
error_type="OAUTH2_REDIRECT_ERROR",
)
except Exception as e:
await ctx.error(
"Chart preview generation failed: identifier=%s, error=%s, error_type=%s"

36
superset/mcp_service/chart/tool/update_chart.py
View File

@@ -27,6 +27,7 @@ from mcp.types import ToolAnnotations
from sqlalchemy.exc import SQLAlchemyError
from superset.commands.exceptions import CommandException
from superset.exceptions import OAuth2Error, OAuth2RedirectError
from superset.extensions import event_logger
from superset.mcp_service.app import mcp
from superset.mcp_service.auth import mcp_auth_hook
@@ -42,6 +43,10 @@ from superset.mcp_service.chart.schemas import (
PerformanceMetadata,
UpdateChartRequest,
)
from superset.mcp_service.utils.oauth2_utils import (
build_oauth2_redirect_message,
OAUTH2_CONFIG_ERROR_MESSAGE,
)
from superset.mcp_service.utils.url_utils import get_superset_base_url
from superset.utils import json
@@ -57,7 +62,7 @@ logger = logging.getLogger(__name__)
),
)
@mcp_auth_hook(class_permission_name="Chart", method_permission_name="write")
async def update_chart(
async def update_chart( # noqa: C901
request: UpdateChartRequest, ctx: Context
) -> GenerateChartResponse:
"""Update existing chart with new configuration.
@@ -271,6 +276,35 @@ async def update_chart(
}
return GenerateChartResponse.model_validate(result)
except OAuth2RedirectError as ex:
await ctx.error(
"Chart update requires OAuth authentication: identifier=%s"
% request.identifier
)
return GenerateChartResponse.model_validate(
{
"chart": None,
"success": False,
"error": {
"error_type": "OAUTH2_REDIRECT",
"message": build_oauth2_redirect_message(ex),
"details": "OAuth2 authentication required",
},
}
)
except OAuth2Error:
await ctx.error("OAuth2 configuration error: chart_id=%s" % request.identifier)
return GenerateChartResponse.model_validate(
{
"chart": None,
"success": False,
"error": {
"error_type": "OAUTH2_REDIRECT_ERROR",
"message": OAUTH2_CONFIG_ERROR_MESSAGE,
"details": "OAuth2 configuration or provider error",
},
}
)
except (
CommandException,
SQLAlchemyError,

24
superset/mcp_service/chart/tool/update_chart_preview.py
View File

@@ -26,6 +26,7 @@ from typing import Any, Dict
from fastmcp import Context
from mcp.types import ToolAnnotations
from superset.exceptions import OAuth2Error, OAuth2RedirectError
from superset.extensions import event_logger
from superset.mcp_service.app import mcp
from superset.mcp_service.auth import mcp_auth_hook
@@ -41,6 +42,10 @@ from superset.mcp_service.chart.schemas import (
PerformanceMetadata,
UpdateChartPreviewRequest,
)
from superset.mcp_service.utils.oauth2_utils import (
build_oauth2_redirect_message,
OAUTH2_CONFIG_ERROR_MESSAGE,
)
logger = logging.getLogger(__name__)
@@ -143,6 +148,25 @@ def update_chart_preview(
}
return result
except OAuth2RedirectError as ex:
logger.warning(
"Chart preview update requires OAuth authentication: form_data_key=%s",
request.form_data_key,
)
return {
"chart": None,
"error": build_oauth2_redirect_message(ex),
"success": False,
}
except OAuth2Error:
logger.warning(
"OAuth2 configuration error: form_data_key=%s", request.form_data_key
)
return {
"chart": None,
"error": OAUTH2_CONFIG_ERROR_MESSAGE,
"success": False,
}
except Exception as e:
execution_time = int((time.time() - start_time) * 1000)
return {

47
superset/mcp_service/utils/oauth2_utils.py Normal file
View File

@@ -0,0 +1,47 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""
Utilities for handling OAuth2 errors in MCP tools.
"""
from superset.exceptions import OAuth2RedirectError
def build_oauth2_redirect_message(ex: OAuth2RedirectError) -> str:
"""
Build a user-facing message for OAuth2RedirectError.
Extracts the authorization URL from the exception and includes it
so the MCP client can present it to the user for authentication.
"""
# extra is always set by OAuth2RedirectError.__init__
assert ex.error.extra is not None # noqa: S101
oauth_url = ex.error.extra["url"]
return (
"This database uses OAuth for authentication. "
"Please open the following URL in your browser to "
"authorize access, then retry this request:\n\n"
f"{oauth_url}"
)
OAUTH2_CONFIG_ERROR_MESSAGE = (
"OAuth authentication failed due to a configuration "
"or provider error. "
"Please contact your Superset administrator."
)

6
superset/sql/execution/executor.py
View File

@@ -69,6 +69,8 @@ from flask import current_app as app, g, has_app_context
from superset import db
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import (
OAuth2Error,
OAuth2RedirectError,
SupersetSecurityException,
SupersetTimeoutException,
)
@@ -318,6 +320,10 @@ class SQLExecutor:
)
except SupersetSecurityException as ex:
return self._create_error_result(QueryStatus.FAILED, str(ex), start_time)
except (OAuth2RedirectError, OAuth2Error):
# Let OAuth2 exceptions propagate so callers (MCP, API) can
# handle them with context-appropriate responses.
raise
except Exception as ex:
error_msg = self.database.db_engine_spec.extract_error_message(ex)
return self._create_error_result(QueryStatus.FAILED, error_msg, start_time)

32
superset/utils/oauth2.py
View File

@@ -29,10 +29,14 @@ import backoff
import jwt
from flask import current_app as app, url_for
from marshmallow import EXCLUDE, fields, post_load, Schema, validate
from werkzeug.routing import BuildError
from superset import db
from superset.distributed_lock import KeyValueDistributedLock
from superset.exceptions import CreateKeyValueDistributedLockFailedException
from superset.exceptions import (
CreateKeyValueDistributedLockFailedException,
OAuth2Error,
)
from superset.superset_typing import OAuth2ClientConfig, OAuth2State
if TYPE_CHECKING:
@@ -243,16 +247,34 @@ def decode_oauth2_state(encoded_state: str) -> OAuth2State:
return state
def get_oauth2_redirect_uri() -> str:
"""
Return the OAuth2 redirect URI.
Tries the explicit config first, then falls back to url_for().
If url_for() fails (e.g. in headless/MCP contexts where the
DatabaseRestApi blueprint may not be registered), raises
OAuth2Error so callers don't silently proceed with an invalid URI.
"""
if configured := app.config.get("DATABASE_OAUTH2_REDIRECT_URI"):
return configured
try:
return url_for("DatabaseRestApi.oauth2", _external=True)
except (BuildError, RuntimeError):
raise OAuth2Error(
"Unable to determine the OAuth2 redirect URI. "
"Set DATABASE_OAUTH2_REDIRECT_URI in the configuration."
) from None
class OAuth2ClientConfigSchema(Schema):
id = fields.String(required=True)
secret = fields.String(required=True)
scope = fields.String(required=True)
redirect_uri = fields.String(
required=False,
load_default=lambda: app.config.get(
"DATABASE_OAUTH2_REDIRECT_URI",
url_for("DatabaseRestApi.oauth2", _external=True),
),
load_default=get_oauth2_redirect_uri,
)
authorization_request_uri = fields.String(required=True)
token_request_uri = fields.String(required=True)

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

@@ -1200,7 +1200,7 @@ def test_start_oauth2_dance_falls_back_to_url_for(mocker: MockerFixture) -> None
},
)
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")

64
tests/unit_tests/utils/oauth2_tests.py
View File

@@ -33,6 +33,7 @@ from superset.utils.oauth2 import (
generate_code_challenge,
generate_code_verifier,
get_oauth2_access_token,
get_oauth2_redirect_uri,
refresh_oauth2_token,
)
@@ -335,3 +336,66 @@ def test_encode_decode_oauth2_state(
assert "code_verifier" not in decoded
assert decoded["database_id"] == 1
assert decoded["user_id"] == 2
def test_get_oauth2_redirect_uri_from_config(mocker: MockerFixture) -> None:
"""
Test that get_oauth2_redirect_uri returns the configured value when set.
"""
custom_uri = "https://proxy.example.com/oauth2/"
mocker.patch(
"flask.current_app.config",
{"DATABASE_OAUTH2_REDIRECT_URI": custom_uri},
)
assert get_oauth2_redirect_uri() == custom_uri
def test_get_oauth2_redirect_uri_falls_back_to_url_for(mocker: MockerFixture) -> None:
"""
Test that get_oauth2_redirect_uri falls back to url_for when config is not set.
"""
fallback_uri = "http://localhost:8088/api/v1/database/oauth2/"
mocker.patch("flask.current_app.config", {})
mocker.patch(
"superset.utils.oauth2.url_for",
return_value=fallback_uri,
)
assert get_oauth2_redirect_uri() == fallback_uri
def test_get_oauth2_redirect_uri_raises_on_build_error(
mocker: MockerFixture,
) -> None:
"""
Test that get_oauth2_redirect_uri raises OAuth2Error when url_for raises
BuildError (e.g. in headless/MCP contexts).
"""
from werkzeug.routing import BuildError
from superset.exceptions import OAuth2Error
mocker.patch("flask.current_app.config", {})
mocker.patch(
"superset.utils.oauth2.url_for",
side_effect=BuildError("DatabaseRestApi.oauth2", {}, ("GET",)),
)
with pytest.raises(OAuth2Error):
get_oauth2_redirect_uri()
def test_get_oauth2_redirect_uri_raises_on_runtime_error(
mocker: MockerFixture,
) -> None:
"""
Test that get_oauth2_redirect_uri raises OAuth2Error when url_for raises
RuntimeError (e.g. no request context and no SERVER_NAME).
"""
from superset.exceptions import OAuth2Error
mocker.patch("flask.current_app.config", {})
mocker.patch(
"superset.utils.oauth2.url_for",
side_effect=RuntimeError("Unable to build URL outside of request context"),
)
with pytest.raises(OAuth2Error):
get_oauth2_redirect_uri()