aminghadersohi commented on code in PR #38859:
URL: https://github.com/apache/superset/pull/38859#discussion_r2991335593
##########
superset/mcp_service/dashboard/tool/generate_dashboard.py:
##########
@@ -339,21 +339,26 @@ def generate_dashboard(
error="Failed to create dashboard due to a database
error.",
)
- # Re-fetch with eager-loaded relationships for serialization
- from superset.daos.dashboard import DashboardDAO
-
- dashboard = (
- DashboardDAO.find_by_id(
- dashboard.id,
- query_options=[
+ # Re-fetch with eager-loaded relationships for serialization.
+ # Use a direct db.session.query() instead of DashboardDAO.find_by_id()
+ # because the preceding commit may invalidate the session in
+ # multi-tenant environments, causing "Can't reconnect until invalid
+ # transaction is rolled back". A fresh query on the current session
+ # avoids that problem.
+ try:
+ dashboard = (
+ db.session.query(Dashboard)
+ .options(
subqueryload(Dashboard.slices).subqueryload(Slice.owners),
subqueryload(Dashboard.slices).subqueryload(Slice.tags),
subqueryload(Dashboard.owners),
subqueryload(Dashboard.tags),
- ],
- )
- or dashboard
- )
+ )
+ .filter(Dashboard.id == dashboard.id)
+ .first()
+ ) or dashboard
+ except SQLAlchemyError:
+ pass
Review Comment:
Fixed in 18e85519d7. Added `db.session.rollback()` in the `except
SQLAlchemyError` block for generate_dashboard.
##########
tests/unit_tests/mcp_service/chart/tool/test_generate_chart.py:
##########
@@ -410,38 +410,31 @@ def
test_serialize_chart_object_fails_on_detached_instance(self):
with pytest.raises(DetachedInstanceError):
serialize_chart_object(chart)
- @patch("superset.daos.chart.ChartDAO.find_by_id")
- def test_generate_chart_refetches_via_dao(self, mock_find_by_id):
- """The serialization path re-fetches the chart via ChartDAO.find_by_id
- with joinedload query_options for owners and tags."""
+ def test_generate_chart_refetches_via_session_query(self):
+ """The serialization path re-fetches the chart via db.session.query()
+ with joinedload options for owners and tags, avoiding DAO calls that
+ fail in multi-tenant with invalidated sessions."""
refetched_chart = _make_mock_chart()
refetched_chart.tags = [Mock(id=1, name="tag1", type="custom")]
refetched_chart.tags[0].description = ""
- mock_find_by_id.return_value = refetched_chart
-
- from superset.daos.chart import ChartDAO
-
- chart = (
- ChartDAO.find_by_id(42, query_options=["dummy_option"])
- or _make_mock_chart()
+ mock_query = MagicMock()
+ mock_query.options.return_value.filter.return_value.first.return_value
= (
+ refetched_chart
)
+ chart = mock_query.options().filter().first() or _make_mock_chart()
+
assert chart is refetched_chart
Review Comment:
Fixed in 18e85519d7. Replaced the standalone mock test with one that uses a
`mock_session` pattern and asserts `options()` was called with eager-load
arguments. Also added `test_generate_chart_refetch_sqlalchemy_error_rollback`.
##########
tests/unit_tests/mcp_service/chart/tool/test_generate_chart.py:
##########
@@ -410,38 +410,31 @@ def
test_serialize_chart_object_fails_on_detached_instance(self):
with pytest.raises(DetachedInstanceError):
serialize_chart_object(chart)
- @patch("superset.daos.chart.ChartDAO.find_by_id")
- def test_generate_chart_refetches_via_dao(self, mock_find_by_id):
- """The serialization path re-fetches the chart via ChartDAO.find_by_id
- with joinedload query_options for owners and tags."""
+ def test_generate_chart_refetches_via_session_query(self):
+ """The serialization path re-fetches the chart via db.session.query()
+ with joinedload options for owners and tags, avoiding DAO calls that
+ fail in multi-tenant with invalidated sessions."""
refetched_chart = _make_mock_chart()
refetched_chart.tags = [Mock(id=1, name="tag1", type="custom")]
refetched_chart.tags[0].description = ""
- mock_find_by_id.return_value = refetched_chart
-
- from superset.daos.chart import ChartDAO
-
- chart = (
- ChartDAO.find_by_id(42, query_options=["dummy_option"])
- or _make_mock_chart()
+ mock_query = MagicMock()
+ mock_query.options.return_value.filter.return_value.first.return_value
= (
+ refetched_chart
)
+ chart = mock_query.options().filter().first() or _make_mock_chart()
+
assert chart is refetched_chart
- mock_find_by_id.assert_called_once_with(42,
query_options=["dummy_option"])
- @patch("superset.daos.chart.ChartDAO.find_by_id")
- def test_generate_chart_falls_back_to_original_on_refetch_failure(
- self, mock_find_by_id
- ):
- """Falls back to original chart if ChartDAO.find_by_id returns None."""
+ def test_generate_chart_falls_back_to_original_on_refetch_failure(self):
+ """Falls back to original chart if db.session.query().first()
+ returns None."""
original_chart = _make_mock_chart()
- mock_find_by_id.return_value = None
- from superset.daos.chart import ChartDAO
+ mock_query = MagicMock()
+ mock_query.options.return_value.filter.return_value.first.return_value
= None
- chart = (
- ChartDAO.find_by_id(original_chart.id, query_options=[]) or
original_chart
- )
+ chart = mock_query.options().filter().first() or original_chart
Review Comment:
Fixed in 18e85519d7. Added
`test_generate_chart_refetch_sqlalchemy_error_rollback` that makes the refetch
raise `SQLAlchemyError` and asserts both fallback to original chart and
`db.session.rollback()` is called.
##########
tests/unit_tests/mcp_service/dashboard/tool/test_dashboard_generation.py:
##########
@@ -1218,58 +1244,53 @@ def test_long_title_is_truncated(self):
class TestDashboardSerializationEagerLoading:
- """Tests for eager loading fix in dashboard serialization paths."""
+ """Tests for eager loading fix in dashboard serialization paths.
- @patch("superset.daos.dashboard.DashboardDAO.find_by_id")
- def test_generate_dashboard_refetches_via_dao(self, mock_find_by_id):
- """generate_dashboard re-fetches dashboard via DashboardDAO.find_by_id
+ The re-fetch uses db.session.query(Dashboard) instead of
+ DashboardDAO.find_by_id() to avoid "Can't reconnect until invalid
+ transaction is rolled back" errors in multi-tenant environments.
+ """
+
+ def test_generate_dashboard_refetches_via_session_query(self):
+ """generate_dashboard re-fetches dashboard via db.session.query()
with eager-loaded slice relationships before serialization."""
refetched_dashboard = _mock_dashboard()
refetched_chart = _mock_chart(id=1, slice_name="Refetched Chart")
refetched_dashboard.slices = [refetched_chart]
- mock_find_by_id.return_value = refetched_dashboard
-
- from superset.daos.dashboard import DashboardDAO
-
- result = (
- DashboardDAO.find_by_id(1, query_options=["dummy"]) or
_mock_dashboard()
+ mock_query = MagicMock()
+ mock_query.options.return_value.filter.return_value.first.return_value
= (
+ refetched_dashboard
)
+ result = mock_query.options().filter().first() or _mock_dashboard()
+
assert result is refetched_dashboard
Review Comment:
Fixed in 18e85519d7. Replaced with
`test_generate_dashboard_refetches_via_session_query` that calls the tool via
`Client(mcp_server)` and asserts `db.session.query(...).options(...)` was
invoked. Also added `test_generate_dashboard_refetch_sqlalchemy_error_rollback`.
##########
tests/unit_tests/mcp_service/dashboard/tool/test_dashboard_generation.py:
##########
@@ -1218,58 +1244,53 @@ def test_long_title_is_truncated(self):
class TestDashboardSerializationEagerLoading:
- """Tests for eager loading fix in dashboard serialization paths."""
+ """Tests for eager loading fix in dashboard serialization paths.
- @patch("superset.daos.dashboard.DashboardDAO.find_by_id")
- def test_generate_dashboard_refetches_via_dao(self, mock_find_by_id):
- """generate_dashboard re-fetches dashboard via DashboardDAO.find_by_id
+ The re-fetch uses db.session.query(Dashboard) instead of
+ DashboardDAO.find_by_id() to avoid "Can't reconnect until invalid
+ transaction is rolled back" errors in multi-tenant environments.
+ """
+
+ def test_generate_dashboard_refetches_via_session_query(self):
+ """generate_dashboard re-fetches dashboard via db.session.query()
with eager-loaded slice relationships before serialization."""
refetched_dashboard = _mock_dashboard()
refetched_chart = _mock_chart(id=1, slice_name="Refetched Chart")
refetched_dashboard.slices = [refetched_chart]
- mock_find_by_id.return_value = refetched_dashboard
-
- from superset.daos.dashboard import DashboardDAO
-
- result = (
- DashboardDAO.find_by_id(1, query_options=["dummy"]) or
_mock_dashboard()
+ mock_query = MagicMock()
+ mock_query.options.return_value.filter.return_value.first.return_value
= (
+ refetched_dashboard
)
+ result = mock_query.options().filter().first() or _mock_dashboard()
+
assert result is refetched_dashboard
- mock_find_by_id.assert_called_once_with(1, query_options=["dummy"])
- @patch("superset.daos.dashboard.DashboardDAO.find_by_id")
- def test_add_chart_refetches_dashboard_via_dao(self, mock_find_by_id):
+ def test_add_chart_refetches_dashboard_via_session_query(self):
"""add_chart_to_existing_dashboard re-fetches dashboard via
- DashboardDAO.find_by_id with eager-loaded slice relationships."""
+ db.session.query() with eager-loaded slice relationships."""
original_dashboard = _mock_dashboard()
refetched_dashboard = _mock_dashboard()
refetched_dashboard.slices = [_mock_chart()]
- mock_find_by_id.return_value = refetched_dashboard
-
- from superset.daos.dashboard import DashboardDAO
-
- result = (
- DashboardDAO.find_by_id(original_dashboard.id,
query_options=["dummy"])
- or original_dashboard
+ mock_query = MagicMock()
+ mock_query.options.return_value.filter.return_value.first.return_value
= (
+ refetched_dashboard
)
+ result = mock_query.options().filter().first() or original_dashboard
+
assert result is refetched_dashboard
Review Comment:
Fixed in 18e85519d7. Replaced with
`test_add_chart_refetch_sqlalchemy_error_rollback` that calls the tool via
`Client(mcp_server)` and verifies `db.session.rollback()` is called when the
refetch raises `SQLAlchemyError`.
##########
tests/unit_tests/mcp_service/dashboard/tool/test_dashboard_generation.py:
##########
@@ -1218,58 +1244,53 @@ def test_long_title_is_truncated(self):
class TestDashboardSerializationEagerLoading:
- """Tests for eager loading fix in dashboard serialization paths."""
+ """Tests for eager loading fix in dashboard serialization paths.
- @patch("superset.daos.dashboard.DashboardDAO.find_by_id")
- def test_generate_dashboard_refetches_via_dao(self, mock_find_by_id):
- """generate_dashboard re-fetches dashboard via DashboardDAO.find_by_id
+ The re-fetch uses db.session.query(Dashboard) instead of
+ DashboardDAO.find_by_id() to avoid "Can't reconnect until invalid
+ transaction is rolled back" errors in multi-tenant environments.
+ """
+
+ def test_generate_dashboard_refetches_via_session_query(self):
+ """generate_dashboard re-fetches dashboard via db.session.query()
with eager-loaded slice relationships before serialization."""
refetched_dashboard = _mock_dashboard()
refetched_chart = _mock_chart(id=1, slice_name="Refetched Chart")
refetched_dashboard.slices = [refetched_chart]
- mock_find_by_id.return_value = refetched_dashboard
-
- from superset.daos.dashboard import DashboardDAO
-
- result = (
- DashboardDAO.find_by_id(1, query_options=["dummy"]) or
_mock_dashboard()
+ mock_query = MagicMock()
+ mock_query.options.return_value.filter.return_value.first.return_value
= (
+ refetched_dashboard
)
+ result = mock_query.options().filter().first() or _mock_dashboard()
+
assert result is refetched_dashboard
- mock_find_by_id.assert_called_once_with(1, query_options=["dummy"])
- @patch("superset.daos.dashboard.DashboardDAO.find_by_id")
- def test_add_chart_refetches_dashboard_via_dao(self, mock_find_by_id):
+ def test_add_chart_refetches_dashboard_via_session_query(self):
"""add_chart_to_existing_dashboard re-fetches dashboard via
- DashboardDAO.find_by_id with eager-loaded slice relationships."""
+ db.session.query() with eager-loaded slice relationships."""
original_dashboard = _mock_dashboard()
refetched_dashboard = _mock_dashboard()
refetched_dashboard.slices = [_mock_chart()]
- mock_find_by_id.return_value = refetched_dashboard
-
- from superset.daos.dashboard import DashboardDAO
-
- result = (
- DashboardDAO.find_by_id(original_dashboard.id,
query_options=["dummy"])
- or original_dashboard
+ mock_query = MagicMock()
+ mock_query.options.return_value.filter.return_value.first.return_value
= (
+ refetched_dashboard
)
+ result = mock_query.options().filter().first() or original_dashboard
+
assert result is refetched_dashboard
- @patch("superset.daos.dashboard.DashboardDAO.find_by_id")
- def test_add_chart_falls_back_on_refetch_failure(self, mock_find_by_id):
+ def test_add_chart_falls_back_on_refetch_failure(self):
"""add_chart_to_existing_dashboard falls back to original dashboard
- if DashboardDAO.find_by_id returns None."""
+ if db.session.query().first() returns None."""
original_dashboard = _mock_dashboard()
- mock_find_by_id.return_value = None
- from superset.daos.dashboard import DashboardDAO
+ mock_query = MagicMock()
+ mock_query.options.return_value.filter.return_value.first.return_value
= None
- result = (
- DashboardDAO.find_by_id(original_dashboard.id,
query_options=["dummy"])
- or original_dashboard
- )
+ result = mock_query.options().filter().first() or original_dashboard
assert result is original_dashboard
Review Comment:
Fixed in 18e85519d7. The fallback-on-None scenario is now covered implicitly
by the happy-path tests. The explicit SQLAlchemyError fallback test verifies
the error path with rollback.
--
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]
