bito-code-review[bot] commented on code in PR #44694: URL: https://github.com/apache/superset/pull/44694#discussion_r4114430997
########## tests/unit_tests/mcp_service/catalog/tool/test_get_catalog.py: ########## @@ -0,0 +1,896 @@ +# 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 permission-filtered ``get_catalog`` MCP tool. + +These run the real DAO ``base_filter`` implementations (DatabaseFilter, +DatasourceFilter, ChartFilter, DashboardAccessFilter) against an in-memory +database. Roles are simulated at the security-manager grant lookups those +filters consult (``can_access``, ``user_view_menu_names``, ...), so the grant +set can be changed between calls to model a revoked permission. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import ExitStack +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any +from unittest.mock import patch + +import pytest +from fastmcp import Client +from fastmcp.exceptions import ToolError +from flask import g +from pydantic import ValidationError +from sqlalchemy.orm.session import Session + +from superset.mcp_service.app import mcp +from superset.mcp_service.auth import MCPPermissionDeniedError +from superset.mcp_service.catalog.schemas import ( + CATALOG_MAX_RESPONSE_BYTES, + CatalogItem, + GetCatalogRequest, +) +from superset.mcp_service.catalog.tool.get_catalog import ( + build_catalog_page, + decode_cursor, + encode_cursor, +) +from superset.utils import json + +# Values that must never appear in catalog output, planted in fields that the +# regular list/get serializers do expose (SQL, extra, params, metrics, ...). +MARKER_SQL = "SELECT secret_sql_marker FROM hidden" +MARKER_EXTRA = '{"secret_extra_marker": true}' +MARKER_TEMPLATE = '{"secret_template_marker": 1}' +MARKER_METRIC = "SUM(secret_metric_marker)" +MARKER_PARAMS = '{"secret_params_marker": "x"}' +MARKER_URI = "postgresql://secret_user:[email protected]/secret_db" +MARKER_JSON = '{"secret_position_marker": {}}' +FORBIDDEN_MARKERS = ( + "secret_sql_marker", + "secret_extra_marker", + "secret_template_marker", + "secret_metric_marker", + "secret_params_marker", + "secret_password", + "secret_position_marker", +) +FORBIDDEN_KEYS = { + "sql", + "extra", + "template_params", + "params", + "metrics", + "expression", + "sqlalchemy_uri", + "query_context", + "form_data", + "json_metadata", + "position_json", + "owners", + "schema", + "database_name", + "datasource_name", +} +ALLOWED_ITEM_KEYS = set(CatalogItem.model_fields) + +TABLE_A_PERM = "[examples].[table_a](id:1)" +TABLE_B_PERM = "[examples].[table_b](id:2)" +DATABASE_PERM = "[examples].(id:1)" + + +@dataclass +class Role: + """A simulated principal: class permissions plus data-access grants.""" + + name: str + permissions: set[tuple[str, str]] + all_datasource_access: bool = False + all_database_access: bool = False + is_admin: bool = False + grants: dict[str, set[str]] = field(default_factory=dict) + + def can_access(self, permission_name: str, view_name: str) -> bool: + return self.is_admin or (permission_name, view_name) in self.permissions + + def view_menu_names(self, permission_name: str) -> set[str]: + return set(self.grants.get(permission_name, set())) + + +READ_ALL = { + ("can_read", "Chart"), + ("can_read", "Dashboard"), + ("can_read", "Dataset"), + ("can_read", "Database"), +} + + +def admin_role() -> Role: + return Role( + name="Admin", + permissions=set(), + all_datasource_access=True, + all_database_access=True, + is_admin=True, + ) + + +def gamma_with_table_a_grant() -> Role: + """Gamma-like role with a datasource-level grant on table_a only.""" + return Role( + name="Gamma+table_a", + permissions=READ_ALL | {("can_get_drill_info", "Dataset")}, + grants={"datasource_access": {TABLE_A_PERM}}, + ) + + +def viewer_without_data_model_access() -> Role: + """Can read charts/dashboards on table_a, but has no data-model access.""" + return Role( + name="Viewer", + permissions=set(READ_ALL), + grants={"datasource_access": {TABLE_A_PERM}}, + ) + + [email protected](autouse=True) +def enable_mcp_rbac(app: Any, disable_mcp_rbac: None) -> None: + """Run with RBAC on: the package conftest turns it off for mock users.""" + app.config["MCP_RBAC_ENABLED"] = True + + [email protected] +def act_as() -> Iterator[Any]: + """Return a function that makes ``role`` the current principal.""" + from superset import security_manager + + stack = ExitStack() + current: dict[str, Role] = {} + + def _role() -> Role: + return current["role"] + + stack.enter_context( + patch.object( + security_manager, + "can_access", + side_effect=lambda perm, view: _role().can_access(perm, view), + ) + ) + stack.enter_context( + patch.object( + security_manager, + "can_access_all_datasources", + side_effect=lambda: _role().all_datasource_access, + ) + ) + stack.enter_context( + patch.object( + security_manager, + "can_access_all_databases", + side_effect=lambda: _role().all_database_access, + ) + ) + stack.enter_context( + patch.object(security_manager, "is_admin", side_effect=lambda: _role().is_admin) + ) + stack.enter_context( + patch.object( + security_manager, + "user_view_menu_names", + side_effect=lambda perm: _role().view_menu_names(perm), + ) + ) + stack.enter_context( + patch.object(security_manager, "get_accessible_databases", return_value=[]) + ) + stack.enter_context( + patch.object(security_manager, "is_guest_user", return_value=False) + ) + stack.enter_context( + patch.object( + security_manager, "get_current_guest_user_if_guest", return_value=None + ) + ) + + def _act_as(role: Role, user_id: int | None = None) -> Role: + current["role"] = role + g.user = SimpleNamespace( + id=user_id, username=role.name, is_authenticated=True, roles=[] + ) + return role + + with stack: + yield _act_as + + [email protected] +def catalog_fixtures(session: Session) -> SimpleNamespace: + """One database, two datasets, two charts and two dashboards. + + Every data-model field the catalog must never expose is populated with a + recognizable marker. + """ + # pylint: disable=import-outside-toplevel + from superset.connectors.sqla.models import SqlaTable, SqlMetric + from superset.models.core import Database + from superset.models.dashboard import Dashboard + from superset.models.slice import Slice + + engine = session.get_bind() + Dashboard.metadata.create_all(engine) # pylint: disable=no-member + + database = Database( + id=1, + database_name="examples", + sqlalchemy_uri=MARKER_URI, + extra=MARKER_EXTRA, + ) + session.add(database) + session.flush() + + table_a = SqlaTable( + id=1, + table_name="table_a", + database_id=database.id, + schema="secret_schema", + sql=MARKER_SQL, + extra=MARKER_EXTRA, + template_params=MARKER_TEMPLATE, + description="Orders fact table", + perm=TABLE_A_PERM, + metrics=[SqlMetric(metric_name="revenue", expression=MARKER_METRIC)], + ) + table_b = SqlaTable( + id=2, + table_name="table_b", + database_id=database.id, + sql=MARKER_SQL, + description="Payroll", + perm=TABLE_B_PERM, + ) + session.add_all([table_a, table_b]) + session.flush() + + chart_a = Slice( + id=1, + slice_name="Revenue by month", + description="Monthly revenue", + datasource_id=table_a.id, + datasource_type="table", + datasource_name="table_a", + viz_type="table", + params=MARKER_PARAMS, + query_context=MARKER_PARAMS, + perm=TABLE_A_PERM, + ) + chart_b = Slice( + id=2, + slice_name="Salaries", + datasource_id=table_b.id, + datasource_type="table", + datasource_name="table_b", + viz_type="table", + params=MARKER_PARAMS, + perm=TABLE_B_PERM, + ) + session.add_all([chart_a, chart_b]) + session.flush() + # set_related_perm recomputes perms on flush; pin them for the filters. + chart_a.perm = TABLE_A_PERM + chart_b.perm = TABLE_B_PERM + + dashboard_a = Dashboard( + id=1, + dashboard_title="Sales", + slug="sales", + description="Sales overview", + published=True, + slices=[chart_a], + json_metadata=MARKER_JSON, + position_json=MARKER_JSON, + ) + dashboard_b = Dashboard( + id=2, + dashboard_title="HR", + published=True, + slices=[chart_b], + json_metadata=MARKER_JSON, + ) + session.add_all([dashboard_a, dashboard_b]) + session.flush() + + return SimpleNamespace( + session=session, + database=database, + table_a=table_a, + table_b=table_b, + chart_a=chart_a, + chart_b=chart_b, + dashboard_a=dashboard_a, + dashboard_b=dashboard_b, + ) + + +def _page(asset_type: str, **kwargs: Any) -> Any: + return build_catalog_page(GetCatalogRequest(asset_type=asset_type, **kwargs)) + + +def _names(asset_type: str, **kwargs: Any) -> list[str]: + return [item.name for item in _page(asset_type, **kwargs).items] + + +def _assert_allowlisted(payload: dict[str, Any]) -> None: + text = json.dumps(payload) + for marker in FORBIDDEN_MARKERS: + assert marker not in text, marker + assert "secret_schema" not in text + for item in payload["items"]: + assert set(item) <= ALLOWED_ITEM_KEYS + assert not FORBIDDEN_KEYS & set(item) Review Comment: <!-- Bito Reply --> The addition of the docstring to the `_assert_allowlisted` helper function addresses the requirement to document new test helpers. This change is appropriate and follows the project's documentation standards. ########## tests/unit_tests/mcp_service/catalog/tool/test_get_catalog.py: ########## @@ -0,0 +1,896 @@ +# 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 permission-filtered ``get_catalog`` MCP tool. + +These run the real DAO ``base_filter`` implementations (DatabaseFilter, +DatasourceFilter, ChartFilter, DashboardAccessFilter) against an in-memory +database. Roles are simulated at the security-manager grant lookups those +filters consult (``can_access``, ``user_view_menu_names``, ...), so the grant +set can be changed between calls to model a revoked permission. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import ExitStack +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any +from unittest.mock import patch + +import pytest +from fastmcp import Client +from fastmcp.exceptions import ToolError +from flask import g +from pydantic import ValidationError +from sqlalchemy.orm.session import Session + +from superset.mcp_service.app import mcp +from superset.mcp_service.auth import MCPPermissionDeniedError +from superset.mcp_service.catalog.schemas import ( + CATALOG_MAX_RESPONSE_BYTES, + CatalogItem, + GetCatalogRequest, +) +from superset.mcp_service.catalog.tool.get_catalog import ( + build_catalog_page, + decode_cursor, + encode_cursor, +) +from superset.utils import json + +# Values that must never appear in catalog output, planted in fields that the +# regular list/get serializers do expose (SQL, extra, params, metrics, ...). +MARKER_SQL = "SELECT secret_sql_marker FROM hidden" +MARKER_EXTRA = '{"secret_extra_marker": true}' +MARKER_TEMPLATE = '{"secret_template_marker": 1}' +MARKER_METRIC = "SUM(secret_metric_marker)" +MARKER_PARAMS = '{"secret_params_marker": "x"}' +MARKER_URI = "postgresql://secret_user:[email protected]/secret_db" +MARKER_JSON = '{"secret_position_marker": {}}' +FORBIDDEN_MARKERS = ( + "secret_sql_marker", + "secret_extra_marker", + "secret_template_marker", + "secret_metric_marker", + "secret_params_marker", + "secret_password", + "secret_position_marker", +) +FORBIDDEN_KEYS = { + "sql", + "extra", + "template_params", + "params", + "metrics", + "expression", + "sqlalchemy_uri", + "query_context", + "form_data", + "json_metadata", + "position_json", + "owners", + "schema", + "database_name", + "datasource_name", +} +ALLOWED_ITEM_KEYS = set(CatalogItem.model_fields) + +TABLE_A_PERM = "[examples].[table_a](id:1)" +TABLE_B_PERM = "[examples].[table_b](id:2)" +DATABASE_PERM = "[examples].(id:1)" + + +@dataclass +class Role: + """A simulated principal: class permissions plus data-access grants.""" + + name: str + permissions: set[tuple[str, str]] + all_datasource_access: bool = False + all_database_access: bool = False + is_admin: bool = False + grants: dict[str, set[str]] = field(default_factory=dict) + + def can_access(self, permission_name: str, view_name: str) -> bool: + return self.is_admin or (permission_name, view_name) in self.permissions + + def view_menu_names(self, permission_name: str) -> set[str]: + return set(self.grants.get(permission_name, set())) + + +READ_ALL = { + ("can_read", "Chart"), + ("can_read", "Dashboard"), + ("can_read", "Dataset"), + ("can_read", "Database"), +} + + +def admin_role() -> Role: + return Role( + name="Admin", + permissions=set(), + all_datasource_access=True, + all_database_access=True, + is_admin=True, + ) + + +def gamma_with_table_a_grant() -> Role: + """Gamma-like role with a datasource-level grant on table_a only.""" + return Role( + name="Gamma+table_a", + permissions=READ_ALL | {("can_get_drill_info", "Dataset")}, + grants={"datasource_access": {TABLE_A_PERM}}, + ) + + +def viewer_without_data_model_access() -> Role: + """Can read charts/dashboards on table_a, but has no data-model access.""" + return Role( + name="Viewer", + permissions=set(READ_ALL), + grants={"datasource_access": {TABLE_A_PERM}}, + ) + + [email protected](autouse=True) +def enable_mcp_rbac(app: Any, disable_mcp_rbac: None) -> None: + """Run with RBAC on: the package conftest turns it off for mock users.""" + app.config["MCP_RBAC_ENABLED"] = True + + [email protected] +def act_as() -> Iterator[Any]: + """Return a function that makes ``role`` the current principal.""" + from superset import security_manager + + stack = ExitStack() + current: dict[str, Role] = {} + + def _role() -> Role: + return current["role"] + + stack.enter_context( + patch.object( + security_manager, + "can_access", + side_effect=lambda perm, view: _role().can_access(perm, view), + ) + ) + stack.enter_context( + patch.object( + security_manager, + "can_access_all_datasources", + side_effect=lambda: _role().all_datasource_access, + ) + ) + stack.enter_context( + patch.object( + security_manager, + "can_access_all_databases", + side_effect=lambda: _role().all_database_access, + ) + ) + stack.enter_context( + patch.object(security_manager, "is_admin", side_effect=lambda: _role().is_admin) + ) + stack.enter_context( + patch.object( + security_manager, + "user_view_menu_names", + side_effect=lambda perm: _role().view_menu_names(perm), + ) + ) + stack.enter_context( + patch.object(security_manager, "get_accessible_databases", return_value=[]) + ) + stack.enter_context( + patch.object(security_manager, "is_guest_user", return_value=False) + ) + stack.enter_context( + patch.object( + security_manager, "get_current_guest_user_if_guest", return_value=None + ) + ) + + def _act_as(role: Role, user_id: int | None = None) -> Role: + current["role"] = role + g.user = SimpleNamespace( + id=user_id, username=role.name, is_authenticated=True, roles=[] + ) + return role + + with stack: + yield _act_as + + [email protected] +def catalog_fixtures(session: Session) -> SimpleNamespace: + """One database, two datasets, two charts and two dashboards. + + Every data-model field the catalog must never expose is populated with a + recognizable marker. + """ + # pylint: disable=import-outside-toplevel + from superset.connectors.sqla.models import SqlaTable, SqlMetric + from superset.models.core import Database + from superset.models.dashboard import Dashboard + from superset.models.slice import Slice + + engine = session.get_bind() + Dashboard.metadata.create_all(engine) # pylint: disable=no-member + + database = Database( + id=1, + database_name="examples", + sqlalchemy_uri=MARKER_URI, + extra=MARKER_EXTRA, + ) + session.add(database) + session.flush() + + table_a = SqlaTable( + id=1, + table_name="table_a", + database_id=database.id, + schema="secret_schema", + sql=MARKER_SQL, + extra=MARKER_EXTRA, + template_params=MARKER_TEMPLATE, + description="Orders fact table", + perm=TABLE_A_PERM, + metrics=[SqlMetric(metric_name="revenue", expression=MARKER_METRIC)], + ) + table_b = SqlaTable( + id=2, + table_name="table_b", + database_id=database.id, + sql=MARKER_SQL, + description="Payroll", + perm=TABLE_B_PERM, + ) + session.add_all([table_a, table_b]) + session.flush() + + chart_a = Slice( + id=1, + slice_name="Revenue by month", + description="Monthly revenue", + datasource_id=table_a.id, + datasource_type="table", + datasource_name="table_a", + viz_type="table", + params=MARKER_PARAMS, + query_context=MARKER_PARAMS, + perm=TABLE_A_PERM, + ) + chart_b = Slice( + id=2, + slice_name="Salaries", + datasource_id=table_b.id, + datasource_type="table", + datasource_name="table_b", + viz_type="table", + params=MARKER_PARAMS, + perm=TABLE_B_PERM, + ) + session.add_all([chart_a, chart_b]) + session.flush() + # set_related_perm recomputes perms on flush; pin them for the filters. + chart_a.perm = TABLE_A_PERM + chart_b.perm = TABLE_B_PERM + + dashboard_a = Dashboard( + id=1, + dashboard_title="Sales", + slug="sales", + description="Sales overview", + published=True, + slices=[chart_a], + json_metadata=MARKER_JSON, + position_json=MARKER_JSON, + ) + dashboard_b = Dashboard( + id=2, + dashboard_title="HR", + published=True, + slices=[chart_b], + json_metadata=MARKER_JSON, + ) + session.add_all([dashboard_a, dashboard_b]) + session.flush() + + return SimpleNamespace( + session=session, + database=database, + table_a=table_a, + table_b=table_b, + chart_a=chart_a, + chart_b=chart_b, + dashboard_a=dashboard_a, + dashboard_b=dashboard_b, + ) + + +def _page(asset_type: str, **kwargs: Any) -> Any: + return build_catalog_page(GetCatalogRequest(asset_type=asset_type, **kwargs)) + + +def _names(asset_type: str, **kwargs: Any) -> list[str]: + return [item.name for item in _page(asset_type, **kwargs).items] + + +def _assert_allowlisted(payload: dict[str, Any]) -> None: + text = json.dumps(payload) + for marker in FORBIDDEN_MARKERS: + assert marker not in text, marker + assert "secret_schema" not in text + for item in payload["items"]: + assert set(item) <= ALLOWED_ITEM_KEYS + assert not FORBIDDEN_KEYS & set(item) + + +# --------------------------------------------------------------------------- +# Role-based visibility (same DAO filters as the list tools) +# --------------------------------------------------------------------------- + + [email protected]( + ("asset_type", "expected"), + [ + ("databases", ["examples"]), + ("datasets", ["table_a", "table_b"]), + ("charts", ["Revenue by month", "Salaries"]), + ("dashboards", ["Sales", "HR"]), + ], +) +def test_admin_sees_every_asset( + catalog_fixtures: SimpleNamespace, act_as: Any, asset_type: str, expected: list[str] +) -> None: + act_as(admin_role()) + page = _page(asset_type) + assert [item.name for item in page.items] == expected + assert page.restricted is False + assert page.next_cursor is None Review Comment: <!-- Bito Reply --> The addition of the docstring to the test function addresses the reviewer's suggestion. The updated code now includes a docstring describing the test scenario and expected behavior, consistent with the project's documentation standards. -- 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]
