gabotorresruiz commented on code in PR #45535:
URL: https://github.com/apache/superset/pull/45535#discussion_r4361945057
##########
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:
Thanks for turning this around so fast. Confirmed on `2f93648c` with the
same harness: a 20 chart `generate_dashboard` is 300 statements as an Admin,
which is exactly what `master` costs at this merge base, and the explicit chart
viewer path now pays 2 relationship reads per chart instead of those plus a
refresh each. Both columns moved down from the 350 and 599 I quoted, because of
the tag work that came in with the master merge rather than anything here.
Doing the batched load as a separate follow-up is fine by me, and the number to
beat there is those 40 reads for 20 charts.
##########
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:
This went further than I asked for, thanks. I put the `selectinload` back on
top of this head and 9 of the 12 cases fail; with it off, the actual counts
land exactly on the bound in every case (0 for admins, 1 per chart for editors,
2 per chart for viewers and for the denied caller), so the `<=` has no slack
left in it.
##########
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:
Dropping it from both files rather than keeping a case that could not fail
is the right call, thanks.
--
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]
