gabotorresruiz commented on code in PR #45535: URL: https://github.com/apache/superset/pull/45535#discussion_r4340459537
########## tests/unit_tests/mcp_service/dashboard/tool/test_chart_access_query_count.py: ########## @@ -0,0 +1,150 @@ +# 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. + +"""Chart permission relationships are loaded in batches before access checks.""" + +import inspect +from collections.abc import Iterator +from unittest.mock import AsyncMock, MagicMock + +import pytest +from pytest_mock import MockerFixture +from sqlalchemy import create_engine, event +from sqlalchemy.engine import Connection, Engine +from sqlalchemy.orm import Session +from sqlalchemy.pool import StaticPool + +from superset import security_manager +from superset.mcp_service.chart.chart_utils import DatasetValidationResult +from superset.mcp_service.dashboard.schemas import ( + GenerateDashboardRequest, + GenerateDashboardResponse, +) +from superset.mcp_service.dashboard.tool.generate_dashboard import generate_dashboard +from superset.models.slice import Slice +from superset.subjects.models import chart_editors, chart_viewers, Subject + + [email protected] +def session_engine() -> Iterator[Engine]: + """Let the tool's worker thread use the isolated in-memory test database.""" + engine: Engine = create_engine( + "sqlite://", poolclass=StaticPool, connect_args={"check_same_thread": False} + ) + try: + yield engine + finally: + engine.dispose() + + [email protected] [email protected]("chart_count", [1, 5]) [email protected]("access", ["editor", "viewer", "denied"]) +async def test_generate_dashboard_batches_chart_access_relationships( + session: Session, + mocker: MockerFixture, + chart_count: int, + access: str, +) -> None: + """Real permission checks retain their decisions without per-chart SELECTs.""" + Slice.metadata.create_all(session.get_bind()) + session.execute( + Subject.__table__.insert(), + [ + {"id": 1, "label": "reader", "type": 1}, + {"id": 2, "label": "other", "type": 1}, + ], + ) + chart_ids: list[int] = list(range(1, chart_count + 1)) + session.execute( + Slice.__table__.insert(), + [{"id": chart_id, "slice_name": f"Chart {chart_id}"} for chart_id in chart_ids], + ) + session.execute( + chart_editors.insert(), + [ + {"chart_id": chart_id, "subject_id": 1 if access == "editor" else 2} + for chart_id in chart_ids + ], + ) + session.execute( + chart_viewers.insert(), + [ + { + "chart_id": chart_id, + "subject_id": 2 + if access == "denied" and chart_id == chart_count + else 1, + } + for chart_id in chart_ids + ], + ) + mocker.patch.object(security_manager, "is_admin", return_value=False) + mocker.patch.object(security_manager, "is_guest_user", return_value=False) + mocker.patch("superset.security.manager.get_user_id", return_value=1) + mocker.patch("superset.subjects.utils.get_user_subject_ids", return_value=[1]) + mocker.patch( + "superset.security.manager.get_extra_editor_subject_ids", return_value=[] + ) + mocker.patch( + "superset.mcp_service.auth.check_chart_data_access", + return_value=DatasetValidationResult( + is_valid=True, dataset_id=1, dataset_name="allowed", warnings=[], error=None + ), + ) + layout: MagicMock = mocker.patch( + "superset.mcp_service.dashboard.tool.generate_dashboard._create_dashboard_layout", + side_effect=RuntimeError("access checks completed"), Review Comment: Not a blocker, but this stub raises before anything reads a column off the charts, and that read is what triggers the refresh. If I make the stub do what the real builder does, read `slice_name` and `uuid` per chart and then raise, this same test reports 12 relationship SELECTs for 5 charts instead of `<= 2`. Letting the tool run through the real layout step would make the bound mean what it says. ########## tests/unit_tests/mcp_service/dashboard/tool/test_add_chart_to_existing_dashboard.py: ########## @@ -136,6 +144,66 @@ def _mock_dashboard( # --------------------------------------------------------------------------- [email protected]("datasource_type", ["table", "semantic_view"]) Review Comment: Just a small NIT: I think the `semantic_view` half of this parametrization is inert. `resolved_datasource` is set to a `Mock()` on line 171 and `can_access_datasource` is patched, so `Slice.resolved_datasource`, the only place `datasource_type` is read, never runs. I swapped both values for nonsense strings and all 87 tests across this file and `test_dashboard_generation.py` still passed. Real semantic-view coverage needs a chart whose `resolved_datasource` actually resolves, which `test_chart_access_query_count.py` is already wired for. Same shape at `test_dashboard_generation.py:210`. ########## superset/mcp_service/dashboard/tool/generate_dashboard.py: ########## @@ -244,6 +245,7 @@ def generate_dashboard( # noqa: C901 with event_logger.log_context(action="mcp.generate_dashboard.chart_validation"): chart_objects = ( db.session.query(Slice) + .options(selectinload(Slice.editors), selectinload(Slice.viewers)) Review Comment: This block worries me a bit: measured, the `selectinload` is a net regression rather than a saving. The `log_context` wrapping this query commits on exit (`DBEventLogger` is the default `EVENT_LOGGER`), which expires every chart, and SQLAlchemy re-applies the saved loader options when `_create_dashboard_layout` reads `chart.slice_name`, so `chart_editors` and `chart_viewers` end up selected per chart anyway. One `generate_dashboard` call with 20 charts on this branch: 392 statements as an Admin and 680 as an explicit chart viewer, versus 350 and 679 with just this line removed (`master` is 350 and 599). Would you consider dropping it, or moving the chart query and the access loop out of the committing `log_context` so the batched load actually survives? -- 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]
