aminghadersohi commented on code in PR #41606: URL: https://github.com/apache/superset/pull/41606#discussion_r3604744498
########## tests/unit_tests/mcp_service/dashboard/tool/test_manage_dashboard_certification.py: ########## @@ -0,0 +1,363 @@ +# 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. + +"""Tests for the manage_dashboard_certification MCP tool.""" + +from typing import Any +from unittest.mock import Mock, patch + +import pytest +from fastmcp import Client + +from superset.utils import json + +DAO_GET: str = "superset.daos.dashboard.DashboardDAO.get_by_id_or_slug" + + +def _mock_dashboard( + id: int = 42, + title: str = "Test Dashboard", + slug: str | None = "test-slug", + certified_by: str | None = None, + certification_details: str | None = None, +) -> Mock: + dashboard: Mock = Mock() + dashboard.id = id + dashboard.dashboard_title = title + dashboard.slug = slug + dashboard.certified_by = certified_by + dashboard.certification_details = certification_details + return dashboard + + +class TestManageDashboardCertification: + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_set_certification( + self, mock_session: Mock, mock_get: Mock, mcp_server: object + ) -> None: + dash: Mock = _mock_dashboard() + mock_get.return_value = dash + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + { + "request": { + "identifier": 42, + "certified_by": "Data Platform Team", + "certification_details": "Verified against source-of-truth.", + } + }, + ) + + assert dash.certified_by == "Data Platform Team" + assert dash.certification_details == "Verified against source-of-truth." + assert mock_session.commit.call_count >= 1 + payload: dict[str, Any] = json.loads(result.content[0].text) + # Response text is wrapped for LLM context (mirrors the read-path + # sanitization on DashboardInfo.certified_by), so check substring. + assert "Data Platform Team" in payload["certified_by"] + assert set(payload["changed_fields"]) == { + "certified_by", + "certification_details", + } + + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_set_certified_by_only( + self, mock_session: Mock, mock_get: Mock, mcp_server: object + ) -> None: Review Comment: Fixed in cde36518a5: `dash: Mock` annotated. ########## tests/unit_tests/mcp_service/dashboard/tool/test_manage_dashboard_certification.py: ########## @@ -0,0 +1,363 @@ +# 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. + +"""Tests for the manage_dashboard_certification MCP tool.""" + +from typing import Any +from unittest.mock import Mock, patch + +import pytest +from fastmcp import Client + +from superset.utils import json + +DAO_GET: str = "superset.daos.dashboard.DashboardDAO.get_by_id_or_slug" + + +def _mock_dashboard( + id: int = 42, + title: str = "Test Dashboard", + slug: str | None = "test-slug", + certified_by: str | None = None, + certification_details: str | None = None, +) -> Mock: + dashboard: Mock = Mock() + dashboard.id = id + dashboard.dashboard_title = title + dashboard.slug = slug + dashboard.certified_by = certified_by + dashboard.certification_details = certification_details + return dashboard + + +class TestManageDashboardCertification: + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_set_certification( + self, mock_session: Mock, mock_get: Mock, mcp_server: object + ) -> None: + dash: Mock = _mock_dashboard() + mock_get.return_value = dash + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + { + "request": { + "identifier": 42, + "certified_by": "Data Platform Team", + "certification_details": "Verified against source-of-truth.", + } + }, + ) + + assert dash.certified_by == "Data Platform Team" + assert dash.certification_details == "Verified against source-of-truth." + assert mock_session.commit.call_count >= 1 + payload: dict[str, Any] = json.loads(result.content[0].text) + # Response text is wrapped for LLM context (mirrors the read-path + # sanitization on DashboardInfo.certified_by), so check substring. + assert "Data Platform Team" in payload["certified_by"] + assert set(payload["changed_fields"]) == { + "certified_by", + "certification_details", + } + + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_set_certified_by_only( + self, mock_session: Mock, mock_get: Mock, mcp_server: object + ) -> None: + dash = _mock_dashboard(certification_details="Existing details") + mock_get.return_value = dash + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + {"request": {"identifier": 42, "certified_by": "Analytics Guild"}}, + ) + + assert dash.certified_by == "Analytics Guild" + assert dash.certification_details == "Existing details" + payload = json.loads(result.content[0].text) Review Comment: Fixed in cde36518a5: `payload: dict[str, Any]` annotated. ########## tests/unit_tests/mcp_service/dashboard/tool/test_manage_dashboard_certification.py: ########## @@ -0,0 +1,363 @@ +# 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. + +"""Tests for the manage_dashboard_certification MCP tool.""" + +from typing import Any +from unittest.mock import Mock, patch + +import pytest +from fastmcp import Client + +from superset.utils import json + +DAO_GET: str = "superset.daos.dashboard.DashboardDAO.get_by_id_or_slug" + + +def _mock_dashboard( + id: int = 42, + title: str = "Test Dashboard", + slug: str | None = "test-slug", + certified_by: str | None = None, + certification_details: str | None = None, +) -> Mock: + dashboard: Mock = Mock() + dashboard.id = id + dashboard.dashboard_title = title + dashboard.slug = slug + dashboard.certified_by = certified_by + dashboard.certification_details = certification_details + return dashboard + + +class TestManageDashboardCertification: + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_set_certification( + self, mock_session: Mock, mock_get: Mock, mcp_server: object + ) -> None: + dash: Mock = _mock_dashboard() + mock_get.return_value = dash + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + { + "request": { + "identifier": 42, + "certified_by": "Data Platform Team", + "certification_details": "Verified against source-of-truth.", + } + }, + ) + + assert dash.certified_by == "Data Platform Team" + assert dash.certification_details == "Verified against source-of-truth." + assert mock_session.commit.call_count >= 1 + payload: dict[str, Any] = json.loads(result.content[0].text) + # Response text is wrapped for LLM context (mirrors the read-path + # sanitization on DashboardInfo.certified_by), so check substring. + assert "Data Platform Team" in payload["certified_by"] + assert set(payload["changed_fields"]) == { + "certified_by", + "certification_details", + } + + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_set_certified_by_only( + self, mock_session: Mock, mock_get: Mock, mcp_server: object + ) -> None: + dash = _mock_dashboard(certification_details="Existing details") + mock_get.return_value = dash + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + {"request": {"identifier": 42, "certified_by": "Analytics Guild"}}, + ) + + assert dash.certified_by == "Analytics Guild" + assert dash.certification_details == "Existing details" + payload = json.loads(result.content[0].text) + assert payload["changed_fields"] == ["certified_by"] + + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_clear_with_empty_string( + self, mock_session: Mock, mock_get: Mock, mcp_server: object + ) -> None: + dash = _mock_dashboard( + certified_by="Old Team", certification_details="Old details" + ) Review Comment: Fixed in cde36518a5: `dash: Mock` annotated. ########## tests/unit_tests/mcp_service/dashboard/tool/test_manage_dashboard_certification.py: ########## @@ -0,0 +1,363 @@ +# 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. + +"""Tests for the manage_dashboard_certification MCP tool.""" + +from typing import Any +from unittest.mock import Mock, patch + +import pytest +from fastmcp import Client + +from superset.utils import json + +DAO_GET: str = "superset.daos.dashboard.DashboardDAO.get_by_id_or_slug" + + +def _mock_dashboard( + id: int = 42, + title: str = "Test Dashboard", + slug: str | None = "test-slug", + certified_by: str | None = None, + certification_details: str | None = None, +) -> Mock: + dashboard: Mock = Mock() + dashboard.id = id + dashboard.dashboard_title = title + dashboard.slug = slug + dashboard.certified_by = certified_by + dashboard.certification_details = certification_details + return dashboard + + +class TestManageDashboardCertification: + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_set_certification( + self, mock_session: Mock, mock_get: Mock, mcp_server: object + ) -> None: + dash: Mock = _mock_dashboard() + mock_get.return_value = dash + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + { + "request": { + "identifier": 42, + "certified_by": "Data Platform Team", + "certification_details": "Verified against source-of-truth.", + } + }, + ) + + assert dash.certified_by == "Data Platform Team" + assert dash.certification_details == "Verified against source-of-truth." + assert mock_session.commit.call_count >= 1 + payload: dict[str, Any] = json.loads(result.content[0].text) + # Response text is wrapped for LLM context (mirrors the read-path + # sanitization on DashboardInfo.certified_by), so check substring. + assert "Data Platform Team" in payload["certified_by"] + assert set(payload["changed_fields"]) == { + "certified_by", + "certification_details", + } + + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_set_certified_by_only( + self, mock_session: Mock, mock_get: Mock, mcp_server: object + ) -> None: + dash = _mock_dashboard(certification_details="Existing details") + mock_get.return_value = dash + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + {"request": {"identifier": 42, "certified_by": "Analytics Guild"}}, + ) + + assert dash.certified_by == "Analytics Guild" + assert dash.certification_details == "Existing details" + payload = json.loads(result.content[0].text) + assert payload["changed_fields"] == ["certified_by"] + + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_clear_with_empty_string( + self, mock_session: Mock, mock_get: Mock, mcp_server: object + ) -> None: + dash = _mock_dashboard( + certified_by="Old Team", certification_details="Old details" + ) + mock_get.return_value = dash + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + { + "request": { + "identifier": 42, + "certified_by": "", + "certification_details": "", + } + }, + ) + + assert dash.certified_by is None + assert dash.certification_details is None + payload = json.loads(result.content[0].text) + assert payload["certified_by"] is None + assert payload["certification_details"] is None + + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_refresh_failure_after_commit_returns_captured_values( + self, mock_session: Mock, mock_get: Mock, mcp_server: object + ) -> None: + """A failed post-commit refresh() must not fail the call: the + response returns the values captured before the commit and appends + an advisory warning.""" + from sqlalchemy.exc import SQLAlchemyError + + dash = _mock_dashboard() + mock_get.return_value = dash + mock_session.refresh.side_effect = SQLAlchemyError("stale connection") + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + { + "request": { + "identifier": 42, + "certified_by": "Data Platform Team", + } + }, + ) + + assert mock_session.commit.call_count >= 1 + payload = json.loads(result.content[0].text) + assert payload.get("error") is None + assert "Data Platform Team" in payload["certified_by"] + assert payload["changed_fields"] == ["certified_by"] + assert any("refresh failed" in w.lower() for w in payload["warnings"]) + + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_rollback_failure_still_returns_structured_error( + self, mock_session: Mock, mock_get: Mock, mcp_server: object + ) -> None: + """Even when the rollback itself fails on the already-broken + session, the caller must still get the structured database-error + response, not an unhandled exception.""" + from sqlalchemy.exc import SQLAlchemyError + + dash = _mock_dashboard() + mock_get.return_value = dash + mock_session.commit.side_effect = SQLAlchemyError("connection lost") + mock_session.rollback.side_effect = SQLAlchemyError("still broken") + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + {"request": {"identifier": 42, "certified_by": "Data Platform Team"}}, + ) + + payload = json.loads(result.content[0].text) + assert "database error" in (payload.get("error") or "").lower() + + @patch(DAO_GET) + @pytest.mark.asyncio + async def test_no_fields_is_noop(self, mock_get: Mock, mcp_server: object) -> None: + dash = _mock_dashboard(certified_by="Existing") + mock_get.return_value = dash + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + {"request": {"identifier": 42}}, + ) + + assert dash.certified_by == "Existing" + payload = json.loads(result.content[0].text) + assert payload["changed_fields"] == [] + assert any("No fields provided" in w for w in payload["warnings"]) + + @patch(DAO_GET) + @pytest.mark.asyncio + async def test_dashboard_not_found( + self, mock_get: Mock, mcp_server: object + ) -> None: + from superset.commands.dashboard.exceptions import DashboardNotFoundError + + mock_get.side_effect = DashboardNotFoundError() + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + {"request": {"identifier": 999999, "certified_by": "Team"}}, + ) + + payload = json.loads(result.content[0].text) + assert "not found" in (payload.get("error") or "").lower() + + @pytest.mark.asyncio + async def test_rejects_boolean_identifier(self, mcp_server: object) -> None: + """bool subclasses int; identifier=true must not coerce to dashboard ID 1.""" + from fastmcp.exceptions import ToolError + + async with Client(mcp_server) as client: + with pytest.raises(ToolError): + await client.call_tool( + "manage_dashboard_certification", + {"request": {"identifier": True, "certified_by": "Team"}}, + ) + + @patch(DAO_GET) + @pytest.mark.asyncio + async def test_lookup_database_error_is_not_masked_as_not_found( + self, mock_get: Mock, mcp_server: object + ) -> None: + """A real DB/infra failure during lookup must surface as a distinct + database error, not be collapsed into "not found".""" + from sqlalchemy.exc import SQLAlchemyError + + mock_get.side_effect = SQLAlchemyError("connection to server lost") + + with patch( + "superset.mcp_service.dashboard.tool.governance_utils.logger" + ) as mock_logger: + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + {"request": {"identifier": 999999, "certified_by": "Team"}}, + ) + + payload = json.loads(result.content[0].text) + assert "not found" not in (payload.get("error") or "").lower() + assert "database error" in (payload.get("error") or "").lower() + assert "connection to server lost" not in (payload.get("error") or "") + mock_logger.exception.assert_called_once() + + @patch(DAO_GET) + @pytest.mark.asyncio + async def test_view_only_caller_gets_permission_denied( + self, mock_get: Mock, mcp_server: object + ) -> None: + """get_by_id_or_slug itself re-checks view access and raises + DashboardAccessDeniedError for dashboards the caller cannot see; + that must surface as permission_denied, not an unhandled error.""" + from superset.commands.dashboard.exceptions import ( + DashboardAccessDeniedError, + ) + + mock_get.side_effect = DashboardAccessDeniedError() + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + {"request": {"identifier": 42, "certified_by": "Team"}}, + ) + + payload = json.loads(result.content[0].text) + assert payload.get("permission_denied") is True + + @patch(DAO_GET) + @pytest.mark.asyncio + async def test_non_owner_gets_permission_denied( + self, mock_get: Mock, mcp_server: object + ) -> None: + from superset.exceptions import SupersetSecurityException + + dash = _mock_dashboard() + mock_get.return_value = dash + + with patch( + "superset.security_manager.raise_for_editorship", + side_effect=SupersetSecurityException(Mock(message="forbidden")), + ): + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_certification", + {"request": {"identifier": 42, "certified_by": "Team"}}, + ) + + payload = json.loads(result.content[0].text) + assert payload.get("permission_denied") is True + assert dash.certified_by is None + + @pytest.mark.asyncio + async def test_xss_certified_by_is_sanitized(self, mcp_server: object) -> None: + from superset.mcp_service.dashboard.schemas import ( + ManageDashboardCertificationRequest, + ) + + request: ManageDashboardCertificationRequest = ( + ManageDashboardCertificationRequest( + identifier=1, certified_by="<script>alert(1)</script>Team" + ) + ) + assert "<script>" not in (request.certified_by or "") + + @pytest.mark.asyncio + async def test_xss_certification_details_is_sanitized( + self, mcp_server: object + ) -> None: + from superset.mcp_service.dashboard.schemas import ( + ManageDashboardCertificationRequest, + ) + + request = ManageDashboardCertificationRequest( + identifier=1, + certification_details="<script>alert(1)</script>Verified details", + ) Review Comment: Fixed in cde36518a5: `request: ManageDashboardCertificationRequest` annotated. ########## tests/unit_tests/mcp_service/dashboard/tool/test_manage_dashboard_roles.py: ########## @@ -0,0 +1,538 @@ +# 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. + +"""Tests for the manage_dashboard_roles MCP tool. + +"Roles" are ROLE-type Subjects in the dashboard's ``viewers`` list (the +Subject-based model apache/superset#38831 introduced, replacing the legacy +``roles``/``DASHBOARD_RBAC`` relationship), gated by ``ENABLE_VIEWERS``. +""" + +from typing import Callable +from unittest.mock import Mock, patch + +import pytest +from fastmcp import Client + +from superset.subjects.types import SubjectType +from superset.utils import json + +DAO_GET: str = "superset.daos.dashboard.DashboardDAO.get_by_id_or_slug" +GET_OR_CREATE_ROLE_SUBJECT: str = "superset.subjects.utils.get_or_create_role_subject" +IS_FEATURE_ENABLED: str = "superset.is_feature_enabled" + + +def _get_or_create_side_effect( + mapping: dict[int, Mock], +) -> Callable[[int], Mock | None]: + """Typed stand-in for ``get_or_create_role_subject``'s ``side_effect``: + resolves a role ID to its mocked Subject, or ``None`` if unknown.""" + + def _lookup(role_id: int) -> Mock | None: + return mapping.get(role_id) + + return _lookup + + +def _mock_subject(id: int, role_id: int, label: str = "role") -> Mock: + subject = Mock() + subject.id = id + subject.type = SubjectType.ROLE + subject.user_id = None + subject.role_id = role_id + subject.label = label + subject.active = True + return subject + + +def _mock_dashboard( + id: int = 42, + title: str = "Test Dashboard", + slug: str | None = "test-slug", + viewers: list[Mock] | None = None, +) -> Mock: + dashboard = Mock() + dashboard.id = id + dashboard.dashboard_title = title + dashboard.slug = slug + dashboard.viewers = viewers if viewers is not None else [] + return dashboard + + +class TestManageDashboardRoles: + @patch(IS_FEATURE_ENABLED, return_value=True) + @patch(GET_OR_CREATE_ROLE_SUBJECT) + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_add_role( + self, + mock_session: Mock, + mock_get: Mock, + mock_get_or_create: Mock, + mock_flag: Mock, + mcp_server: object, + ) -> None: + new_subject = _mock_subject(200, 5, "Analyst") + dash = _mock_dashboard(viewers=[]) + mock_get.return_value = dash + mock_get_or_create.side_effect = _get_or_create_side_effect({5: new_subject}) + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_roles", + {"request": {"identifier": 42, "add_role_ids": [5]}}, + ) + + mock_get_or_create.assert_called_once_with(5) + assert dash.viewers == [new_subject] + assert mock_session.commit.call_count >= 1 + payload = json.loads(result.content[0].text) + assert [r["id"] for r in payload["roles"]] == [200] + assert payload["added_role_ids"] == [5] + assert payload["viewers_enabled"] is True + assert payload["warnings"] == [] + + @patch(IS_FEATURE_ENABLED, return_value=True) + @patch(GET_OR_CREATE_ROLE_SUBJECT) + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_commit_failure_rolls_back_and_returns_error( + self, + mock_session: Mock, + mock_get: Mock, + mock_get_or_create: Mock, + mock_flag: Mock, + mcp_server: object, + ) -> None: + """A DB fault during commit must be caught, trigger a rollback, and + surface as a structured error response rather than an unhandled + exception.""" + from sqlalchemy.exc import SQLAlchemyError + + new_subject = _mock_subject(200, 5, "Analyst") + dash = _mock_dashboard(viewers=[]) + mock_get.return_value = dash + mock_get_or_create.side_effect = _get_or_create_side_effect({5: new_subject}) + mock_session.commit.side_effect = SQLAlchemyError("connection lost") + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_roles", + {"request": {"identifier": 42, "add_role_ids": [5]}}, + ) + + mock_session.rollback.assert_called_once() + payload = json.loads(result.content[0].text) + assert "database error" in (payload.get("error") or "").lower() + + @patch(IS_FEATURE_ENABLED, return_value=True) + @patch(GET_OR_CREATE_ROLE_SUBJECT) + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_rollback_failure_still_returns_structured_error( + self, + mock_session: Mock, + mock_get: Mock, + mock_get_or_create: Mock, + mock_flag: Mock, + mcp_server: object, + ) -> None: + """Even when the rollback itself fails on the already-broken + session, the caller must still get the structured database-error + response, not an unhandled exception.""" + from sqlalchemy.exc import SQLAlchemyError + + new_subject = _mock_subject(200, 5, "Analyst") + dash = _mock_dashboard(viewers=[]) + mock_get.return_value = dash + mock_get_or_create.side_effect = _get_or_create_side_effect({5: new_subject}) + mock_session.commit.side_effect = SQLAlchemyError("connection lost") + mock_session.rollback.side_effect = SQLAlchemyError("still broken") + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_roles", + {"request": {"identifier": 42, "add_role_ids": [5]}}, + ) + + payload = json.loads(result.content[0].text) + assert "database error" in (payload.get("error") or "").lower() + + @patch(IS_FEATURE_ENABLED, return_value=True) + @patch(GET_OR_CREATE_ROLE_SUBJECT) + @patch(DAO_GET) + @patch("superset.extensions.db.session") + @pytest.mark.asyncio + async def test_refresh_failure_after_commit_returns_captured_values( + self, + mock_session: Mock, + mock_get: Mock, + mock_get_or_create: Mock, + mock_flag: Mock, + mcp_server: object, + ) -> None: + """A failed post-commit refresh() must not fail the call: the + response is built from values captured before the commit.""" + from sqlalchemy.exc import SQLAlchemyError + + new_subject = _mock_subject(200, 5, "Analyst") + dash = _mock_dashboard(viewers=[]) + mock_get.return_value = dash + mock_get_or_create.side_effect = _get_or_create_side_effect({5: new_subject}) + mock_session.refresh.side_effect = SQLAlchemyError("stale connection") + + async with Client(mcp_server) as client: + result = await client.call_tool( + "manage_dashboard_roles", + {"request": {"identifier": 42, "add_role_ids": [5]}}, + ) + + assert mock_session.commit.call_count >= 1 + payload = json.loads(result.content[0].text) + assert payload.get("error") is None + assert [r["id"] for r in payload["roles"]] == [200] + assert payload["added_role_ids"] == [5] + + @patch(IS_FEATURE_ENABLED, return_value=True) + @patch(DAO_GET) + @pytest.mark.asyncio + async def test_viewers_lazy_load_failure_returns_error( + self, mock_get: Mock, mock_flag: Mock, mcp_server: object + ) -> None: + """A DB fault while lazy-loading dashboard.viewers must surface as + the structured load error, not an unhandled exception.""" + from sqlalchemy.exc import SQLAlchemyError + + class _ViewersRaise: + id = 42 + dashboard_title = "Test Dashboard" + slug = "test-slug" Review Comment: Fixed in cde36518a5: `_ViewersRaise`'s `id`/`dashboard_title`/`slug` class attributes are now annotated. ########## tests/unit_tests/subjects/test_utils.py: ########## @@ -87,6 +88,59 @@ def test_get_current_user_subject_ids_guest_user(app_context): mock_get_user_id.assert_not_called() +# -------------------------------------------------------------------------- +# get_or_create_role_subject +# -------------------------------------------------------------------------- + + +@patch("superset.subjects.utils.get_role_subject") +def test_get_or_create_role_subject_returns_existing(mock_get_role_subject): Review Comment: Fixed in cde36518a5: added `mock_get_role_subject: MagicMock` and `-> None`. ########## tests/unit_tests/subjects/test_utils.py: ########## @@ -87,6 +88,59 @@ def test_get_current_user_subject_ids_guest_user(app_context): mock_get_user_id.assert_not_called() +# -------------------------------------------------------------------------- +# get_or_create_role_subject +# -------------------------------------------------------------------------- + + +@patch("superset.subjects.utils.get_role_subject") +def test_get_or_create_role_subject_returns_existing(mock_get_role_subject): + """An already-synced role resolves to its Subject with no sync.""" + existing = _make_subject(10, SubjectType.ROLE) + mock_get_role_subject.return_value = existing + + assert get_or_create_role_subject(5) is existing + mock_get_role_subject.assert_called_once_with(5) + + +def test_get_or_create_role_subject_syncs_unsynced_role(): Review Comment: Fixed in cde36518a5: added `-> None`. ########## tests/unit_tests/subjects/test_utils.py: ########## @@ -87,6 +88,59 @@ def test_get_current_user_subject_ids_guest_user(app_context): mock_get_user_id.assert_not_called() +# -------------------------------------------------------------------------- +# get_or_create_role_subject +# -------------------------------------------------------------------------- + + +@patch("superset.subjects.utils.get_role_subject") +def test_get_or_create_role_subject_returns_existing(mock_get_role_subject): + """An already-synced role resolves to its Subject with no sync.""" + existing = _make_subject(10, SubjectType.ROLE) + mock_get_role_subject.return_value = existing + + assert get_or_create_role_subject(5) is existing + mock_get_role_subject.assert_called_once_with(5) + + +def test_get_or_create_role_subject_syncs_unsynced_role(): + """A role that exists but has no Subject row yet is synced on demand + rather than treated as nonexistent.""" + created = _make_subject(11, SubjectType.ROLE) + role = MagicMock() + role.id = 5 + + with ( + patch( + "superset.subjects.utils.get_role_subject", + side_effect=[None, created], + ) as mock_get_role_subject, + patch("superset.subjects.utils.db") as mock_db, + patch("superset.subjects.sync.sync_role_subject") as mock_sync, + ): + mock_db.session.get.return_value = role + + assert get_or_create_role_subject(5) is created + + mock_sync.assert_called_once_with(role) + mock_db.session.flush.assert_called_once() + assert mock_get_role_subject.call_count == 2 + + +def test_get_or_create_role_subject_missing_role_returns_none(): Review Comment: Fixed in cde36518a5: added `-> None`. ########## superset/mcp_service/dashboard/tool/governance_utils.py: ########## @@ -0,0 +1,107 @@ +# 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. + +""" +Shared helpers for the dashboard governance tools +(``manage_dashboard_owners`` / ``manage_dashboard_roles`` / +``manage_dashboard_certification``). + +``update_dashboard`` keeps its own variant of the lookup/authorization +helper because its not-found contract differs — it returns a +``DashboardError`` carrying an ``error_type`` rather than the tool's own +response schema; unifying that shape is left to a follow-up. +""" + +import logging +from typing import Any, TypeVar + +from sqlalchemy.exc import SQLAlchemyError + +from superset.commands.dashboard.exceptions import ( + DashboardAccessDeniedError, + DashboardNotFoundError, +) +from superset.exceptions import SupersetSecurityException +from superset.mcp_service.dashboard.schemas import DashboardMutationErrorFields +from superset.mcp_service.utils.url_utils import get_superset_base_url + +logger: logging.Logger = logging.getLogger(__name__) + +ResponseT = TypeVar("ResponseT", bound=DashboardMutationErrorFields) + + +def find_and_authorize_dashboard( + identifier: int | str, + response_cls: type[ResponseT], +) -> tuple[Any, ResponseT | None]: + """Return (dashboard, None) on success or (None, error_response) on failure. + + ``response_cls`` is the calling tool's response schema; every failure + mode is reported through it so the caller has a single pre-condition + branch. Mirrors the helper in ``update_dashboard``: avoids ImportError + before Flask app initialisation by co-locating the imports it needs + with the call site rather than importing them at module load time. + """ + from superset import security_manager + from superset.daos.dashboard import DashboardDAO + + try: + dashboard = DashboardDAO.get_by_id_or_slug(identifier) Review Comment: Fixed in cde36518a5: annotated as `Dashboard | None` (via a `TYPE_CHECKING`-guarded import, matching the pattern already used in `delete_dashboard.py`/`list_dashboards.py` in this package). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
