EnxDev commented on code in PR #43990:
URL: https://github.com/apache/superset/pull/43990#discussion_r4069225310
##########
tests/unit_tests/dao/dashboard_test.py:
##########
@@ -168,3 +172,221 @@ def
test_set_dash_metadata_updates_refresh_frequency_when_present(
assert md["refresh_frequency"] == 0, (
"refresh_frequency should be updated when present in data"
)
+
+
+def _make_dashboard(slug: str, n_charts: int) -> Dashboard:
+ """A dashboard with n_charts member charts, each with one editor and
viewer."""
+ editor = Subject(label=f"editor-{slug}", type=SubjectType.ROLE)
+ viewer = Subject(label=f"viewer-{slug}", type=SubjectType.ROLE)
+ dashboard = Dashboard(dashboard_title=slug, slug=slug)
+ for i in range(n_charts):
+ dashboard.slices.append(
+ Slice(
+ slice_name=f"{slug}-chart-{i}",
+ datasource_type="table",
+ datasource_id=1,
+ viz_type="table",
+ editors=[editor],
+ viewers=[viewer],
+ )
+ )
+ db.session.add(dashboard)
+ db.session.flush()
+ return dashboard
+
+
+def _count_statements(session: Session, fn) -> int:
+ """Number of SQL statements fn issues."""
+ seen = []
+
+ def before(conn, cursor, statement, params, ctx, many): # noqa: PLR0913
+ seen.append(statement)
+
+ bind = session.get_bind()
+ event.listen(bind, "before_cursor_execute", before)
+ try:
+ fn()
+ finally:
+ event.remove(bind, "before_cursor_execute", before)
+ return len(seen)
+
+
+def test_prefetch_chart_access_loads_editors_and_viewers(
+ session: Session,
+) -> None:
+ """The per-chart access check reads editors and viewers on every slice.
+
+ Without the prefetch those are two lazy loads per chart, so the dashboard
+ GET issues a pair of queries for each member chart it narrows.
+ """
+ Dashboard.metadata.create_all(session.get_bind())
+
+ editor = Subject(label="editor", type=SubjectType.ROLE)
+ viewer = Subject(label="viewer", type=SubjectType.ROLE)
+ dashboard = Dashboard(dashboard_title="prefetch", slug="prefetch")
+ for i in range(3):
+ dashboard.slices.append(
+ Slice(
+ slice_name=f"chart-{i}",
+ datasource_type="table",
+ datasource_id=1,
+ viz_type="table",
+ editors=[editor],
+ viewers=[viewer],
+ )
+ )
+ session.add(dashboard)
+ session.flush()
+
+ # Drop everything from the identity map so the relationships start
unloaded.
+ session.expire_all()
+ dashboard = session.query(Dashboard).filter_by(slug="prefetch").one()
+ assert all("editors" in inspect(slc).unloaded for slc in dashboard.slices)
+
+ # The access check short-circuits for an admin without reading either
+ # relationship, so the prefetch is a non-admin path.
+ with patch.object(security_manager, "is_admin", return_value=False):
+ DashboardDAO.prefetch_chart_access(dashboard)
+
+ for slc in dashboard.slices:
+ unloaded = inspect(slc).unloaded
+ assert "editors" not in unloaded
+ assert "viewers" not in unloaded
+ assert [s.label for s in slc.editors] == ["editor"]
+ assert [s.label for s in slc.viewers] == ["viewer"]
+
+
+def test_prefetch_chart_access_does_no_per_chart_work(
+ session: Session,
+) -> None:
+ """Adding charts must not add statements.
+
+ Asserting only that the relationships end up loaded would also pass for an
+ implementation that walks the slices and touches each one, which is the 2N
+ behaviour this is meant to remove. So count the statements at two sizes and
+ require the same number.
+
+ Not literally constant forever: selectinload batches its own IN lists at
500,
+ so each relationship adds one statement per 500 charts (1200 charts is 7
+ statements, against 2400 before). Both sizes here sit inside the first
batch,
+ which is what makes the equality check meaningful.
+ """
+ Dashboard.metadata.create_all(session.get_bind())
+
+ small = _make_dashboard("count-small", 3)
+ large = _make_dashboard("count-large", 25)
+ session.expire_all()
+
+ with patch.object(security_manager, "is_admin", return_value=False):
+ small = session.query(Dashboard).filter_by(slug="count-small").one()
+ small_n = _count_statements(
+ session, lambda: DashboardDAO.prefetch_chart_access(small)
+ )
+ large = session.query(Dashboard).filter_by(slug="count-large").one()
+ large_n = _count_statements(
+ session, lambda: DashboardDAO.prefetch_chart_access(large)
+ )
+
+ assert small_n == large_n, (
+ f"prefetch scaled with chart count: {small_n} statements for 3 charts,
"
+ f"{large_n} for 25"
+ )
+ assert small_n == 3, f"expected slices + editors + viewers, got {small_n}"
+
+
+def test_prefetch_chart_access_skips_the_query_for_admins(
+ session: Session,
+) -> None:
+ """An admin's access check never reads editors or viewers, so don't
prefetch.
+
+ is_editor and is_viewer both return True on is_admin() before touching
+ either relationship. Prefetching for an admin would turn a zero-query
+ access check into three extra statements.
+ """
+ Dashboard.metadata.create_all(session.get_bind())
+
+ dashboard = _make_dashboard("admin-skip", 5)
+ session.expire_all()
+ dashboard = session.query(Dashboard).filter_by(slug="admin-skip").one()
+
+ with patch.object(security_manager, "is_admin", return_value=True):
+ n = _count_statements(
+ session, lambda: DashboardDAO.prefetch_chart_access(dashboard)
+ )
+
+ assert n == 0, f"prefetch issued {n} statements for an admin"
+
+
+def test_get_charts_for_dashboard_returns_the_prefetched_charts(
+ session: Session,
+) -> None:
+ """The charts endpoint must read editors/viewers without per-chart queries.
+
+ get_charts_for_dashboard looks the dashboard up with its slices unloaded,
+ so the prefetched charts and the returned collection have to be the same
+ instances. If the method reloaded dashboard.slices after prefetching, the
+ weak identity map could drop the prefetched charts and reading their
+ editors/viewers would fall back to a query per chart -- the 2N cost this is
+ meant to remove.
+ """
+ Dashboard.metadata.create_all(session.get_bind())
+
+ _make_dashboard("charts-endpoint", 3)
+ # Detach everything so the lookup returns a dashboard with slices unloaded,
+ # matching the real request path.
+ session.expunge_all()
+ dashboard =
session.query(Dashboard).filter_by(slug="charts-endpoint").one()
+
+ with (
+ patch.object(security_manager, "is_admin", return_value=False),
+ patch.object(DashboardDAO, "get_by_id_or_slug",
return_value=dashboard),
+ ):
+ charts = DashboardDAO.get_charts_for_dashboard("charts-endpoint")
+
+ def read_relationships() -> None:
+ for slc in charts:
+ assert [s.label for s in slc.editors] ==
["editor-charts-endpoint"]
+ assert [s.label for s in slc.viewers] ==
["viewer-charts-endpoint"]
+
+ n = _count_statements(session, read_relationships)
+
+ assert len(charts) == 3
+ assert n == 0, f"reading the returned charts issued {n} statements"
+
+
+def test_prefetch_before_the_main_get_access_loop_reads_no_extra_sql(
+ session: Session,
+) -> None:
+ """The dashboard GET narrows ``charts`` with the same per-slice access
check.
+
+ ``DashboardApi.get`` prefetches, then keeps only the slices the caller can
+ access by reading each one's editors and viewers
(superset/dashboards/api.py).
+ The /charts test above protects ``get_charts_for_dashboard``; this protects
+ the main GET call site, which relies on the prefetch running before the
loop.
+ If that call were removed or moved after serialization the loop would fall
+ back to two lazy loads per chart -- the 2N cost the prefetch removes.
+ """
+ Dashboard.metadata.create_all(session.get_bind())
+
+ _make_dashboard("main-get", 3)
+ # Detach everything so the dashboard comes back with its slices unloaded,
+ # matching the request path.
+ session.expunge_all()
+ dashboard = session.query(Dashboard).filter_by(slug="main-get").one()
+
+ with patch.object(security_manager, "is_admin", return_value=False):
+ # ``schema.dump(dash)`` reads ``Dashboard.charts`` and materializes the
+ # slices before the prefetch runs; hold the same instances so the loop
+ # below reads their relationships rather than reloading the collection.
+ slices = list(dashboard.slices)
+ DashboardDAO.prefetch_chart_access(dashboard)
Review Comment:
Thanks for picking this up. Could we invoke the actual
`DashboardRestApi.get` handler here instead of reproducing its sequence in the
test? I temporarily removed `DashboardDAO.prefetch_chart_access(dash)` from
`superset/dashboards/api.py` and all eight tests still passed, including this
one. Calling the helper directly here means the test can't detect removal or
reordering at the API call site.
A focused unit test could use `inspect.unwrap(DashboardRestApi.get)` inside
a request context, keep the real schema serialization, and have the
`can_access_chart` stub read editors/viewers while counting those SQL
statements. That would exercise the production ordering and fail when the
handler's prefetch is removed. This remains a non-blocking test suggestion.
--
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]
