bito-code-review[bot] commented on code in PR #44694: URL: https://github.com/apache/superset/pull/44694#discussion_r4114143113
########## 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: <div> <div id="suggestion"> <div id="issue"><b>Missing docstring on test</b></div> <div id="fix"> `test_admin_sees_every_asset` lacks a docstring, while sibling tests in the same section (`test_datasource_grant_limits_visibility`, `test_non_admin_cannot_see_another_users_private_content`) document their scenario. BITO adaptive rule 12148/12490 requires a docstring on every new test function describing scenario and expected behavior. </div> </div> <small><i>Code Review Run #760a42</i></small> </div> --- Should Bito avoid suggestions like this for future reviews? (<a href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>) - [ ] Yes, avoid them ########## 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: <div> <div id="suggestion"> <div id="issue"><b>Missing docstring on helper</b></div> <div id="fix"> `_assert_allowlisted` (helper used by `test_output_is_allowlisted`, `test_response_size_is_bounded`, and `test_get_catalog_over_mcp`) has no docstring. BITO adaptive rule 12147/12490 requires docstrings on all new functions, including test helpers. Add a one-line docstring describing the allowlist contract it enforces. </div> </div> <small><i>Code Review Run #760a42</i></small> </div> --- Should Bito avoid suggestions like this for future reviews? (<a href=https://alpha.bito.ai/home/ai-agents/review-rules>Manage Rules</a>) - [ ] Yes, avoid them -- 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]
