Skip to content

perf(dashboard): batch dataset lookups during dashboard export #43017

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ColtenOuO wants to merge 1 commit into
base: master
Choose a base branch
from
+185 −21
Open

perf(dashboard): batch dataset lookups during dashboard export #43017

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 35 additions & 12 deletions superset/commands/dashboard/export.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -309,17 +309,37 @@ def _file_content(model: Dashboard) -> str:
logger.info("Unable to decode `%s` field: %s", key, value)
payload[new_name] = {}

referenced_dataset_ids = {
target["datasetId"]
for native_filter in payload.get("metadata", {}).get(
"native_filter_configuration", []
)
for target in native_filter.get("targets", [])
if target.get("datasetId") is not None
} | {
target["datasetId"]
for customization in (
payload.get("metadata", {}).get("chart_customization_config") or []
)
for target in customization.get("targets") or []
if target.get("datasetId") is not None
}
datasets_by_id = {
dataset.id: dataset
for dataset in DatasetDAO.find_by_ids(list(referenced_dataset_ids))
}

# Extract all native filter datasets and replace native
# filter dataset references with uuid
for native_filter in payload.get("metadata", {}).get(
"native_filter_configuration", []
):
for target in native_filter.get("targets", []):
dataset_id = target.pop("datasetId", None)
if dataset_id is not None:
dataset = DatasetDAO.find_by_id(dataset_id)
if dataset:
target["datasetUuid"] = str(dataset.uuid)
if dataset_id is not None and (
dataset := datasets_by_id.get(dataset_id)
):
target["datasetUuid"] = str(dataset.uuid)

# Replace display control dataset references with uuid.
# datasetId is intentionally preserved alongside datasetUuid so that
Expand All @@ -331,8 +351,7 @@ def _file_content(model: Dashboard) -> str:
for target in customization.get("targets") or []:
dataset_id = target.get("datasetId")
if dataset_id is not None:
dataset = DatasetDAO.find_by_id(dataset_id)
if dataset:
if dataset := datasets_by_id.get(dataset_id):
target["datasetUuid"] = str(dataset.uuid)
else:
logger.warning(
Expand Down Expand Up @@ -422,15 +441,14 @@ def _export(

if export_related:
# Extract all native filter datasets and export referenced datasets
referenced_dataset_ids: set[int] = set()
for native_filter in payload.get("metadata", {}).get(
"native_filter_configuration", []
):
for target in native_filter.get("targets", []):
dataset_id = target.pop("datasetId", None)
if dataset_id is not None:
dataset = DatasetDAO.find_by_id(dataset_id)
if dataset:
yield from ExportDatasetsCommand([dataset_id]).run()
referenced_dataset_ids.add(dataset_id)

# Export datasets referenced by display controls
for customization in (
Expand All @@ -439,6 +457,11 @@ def _export(
for target in customization.get("targets") or []:
dataset_id = target.get("datasetId")
if dataset_id is not None:
dataset = DatasetDAO.find_by_id(dataset_id)
if dataset:
yield from ExportDatasetsCommand([dataset_id]).run()
referenced_dataset_ids.add(dataset_id)

found_dataset_ids = [
dataset.id
for dataset in DatasetDAO.find_by_ids(list(referenced_dataset_ids))
]
if found_dataset_ids:
yield from ExportDatasetsCommand(found_dataset_ids).run()
159 changes: 150 additions & 9 deletions tests/unit_tests/commands/dashboard/export_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,12 +82,13 @@ def test_file_content_replaces_dataset_id_with_uuid_in_display_controls():
)

mock_dataset = MagicMock()
mock_dataset.id = 99
mock_dataset.uuid = dataset_uuid

with (
patch(
"superset.commands.dashboard.export.DatasetDAO.find_by_id",
return_value=mock_dataset,
"superset.commands.dashboard.export.DatasetDAO.find_by_ids",
return_value=[mock_dataset],
),
patch(
"superset.commands.dashboard.export.feature_flag_manager.is_feature_enabled",
Expand All @@ -108,6 +109,144 @@ def test_file_content_replaces_dataset_id_with_uuid_in_display_controls():
assert customizations[1]["targets"] == []


def test_file_content_batches_dataset_lookup_across_targets():
"""
Regression test: dataset lookups must go through a single batched
DatasetDAO.find_by_ids call, not one DatasetDAO.find_by_id call per
target. Multiple filters/customizations referencing the same dataset
must not trigger redundant DB round-trips.
"""
from superset.commands.dashboard.export import ExportDashboardsCommand

dataset_uuid_1 = str(uuid.uuid4())
dataset_uuid_2 = str(uuid.uuid4())

mock_dashboard = _make_mock_dashboard(
{
"native_filter_configuration": [
{
"id": "FILTER-1",
"targets": [{"datasetId": 1}, {"datasetId": 2}],
},
{
"id": "FILTER-2",
"targets": [{"datasetId": 1}],
},
],
"chart_customization_config": [
{
"id": "CUSTOMIZATION-1",
"type": "CHART_CUSTOMIZATION",
"targets": [{"datasetId": 1}],
},
],
}
)

mock_dataset_1 = MagicMock()
mock_dataset_1.id = 1
mock_dataset_1.uuid = dataset_uuid_1
mock_dataset_2 = MagicMock()
mock_dataset_2.id = 2
mock_dataset_2.uuid = dataset_uuid_2

with (
patch(
"superset.commands.dashboard.export.DatasetDAO.find_by_ids",
return_value=[mock_dataset_1, mock_dataset_2],
) as mock_find_by_ids,
patch(
"superset.commands.dashboard.export.DatasetDAO.find_by_id"
) as mock_find_by_id,
patch(
"superset.commands.dashboard.export.feature_flag_manager.is_feature_enabled",
return_value=False,
),
):
content = ExportDashboardsCommand._file_content(mock_dashboard)

mock_find_by_id.assert_not_called()
mock_find_by_ids.assert_called_once()
(called_ids,), _ = mock_find_by_ids.call_args
assert set(called_ids) == {1, 2}

result = yaml.safe_load(content)
native_filters = result["metadata"]["native_filter_configuration"]
assert native_filters[0]["targets"][0]["datasetUuid"] == dataset_uuid_1
assert native_filters[0]["targets"][1]["datasetUuid"] == dataset_uuid_2
assert native_filters[1]["targets"][0]["datasetUuid"] == dataset_uuid_1

customization_target = result["metadata"]["chart_customization_config"][0][
"targets"
][0]
assert customization_target["datasetUuid"] == dataset_uuid_1


def test_export_batches_dataset_export_across_targets():
"""
Regression test: _export must batch dataset exports into a single
ExportDatasetsCommand call, not one call per target. Multiple
filters/customizations referencing the same dataset must only trigger
a single find_by_ids lookup and a single export command.
"""
from superset.commands.dashboard.export import ExportDashboardsCommand

mock_dashboard = _make_mock_dashboard(
{
"native_filter_configuration": [
{
"id": "FILTER-1",
"targets": [{"datasetId": 1}, {"datasetId": 2}],
},
],
"chart_customization_config": [
{
"id": "CUSTOMIZATION-1",
"type": "CHART_CUSTOMIZATION",
"targets": [{"datasetId": 1}],
},
],
}
)

mock_dataset_1 = MagicMock()
mock_dataset_1.id = 1
mock_dataset_2 = MagicMock()
mock_dataset_2.id = 2
mock_datasets_cmd = MagicMock()
mock_datasets_cmd.run.return_value = iter([])

with (
patch(
"superset.commands.dashboard.export.DatasetDAO.find_by_ids",
return_value=[mock_dataset_1, mock_dataset_2],
) as mock_find_by_ids,
patch(
"superset.commands.dashboard.export.DatasetDAO.find_by_id"
) as mock_find_by_id,
patch(
"superset.commands.dashboard.export.ExportDatasetsCommand",
return_value=mock_datasets_cmd,
) as mock_datasets_cls,
patch(
"superset.commands.dashboard.export.ExportChartsCommand"
) as mock_charts_cls,
patch(
"superset.commands.dashboard.export.feature_flag_manager.is_feature_enabled",
return_value=False,
),
):
mock_charts_cls.return_value.run.return_value = iter([])
list(ExportDashboardsCommand._export(mock_dashboard))

mock_find_by_id.assert_not_called()
mock_find_by_ids.assert_called_once()
mock_datasets_cls.assert_called_once()
mock_datasets_cmd.run.assert_called_once()
(called_ids,), _ = mock_datasets_cls.call_args
assert set(called_ids) == {1, 2}


def test_export_yields_dataset_files_for_display_controls():
"""
_export must yield dataset files for datasets referenced by display controls.
Expand All @@ -134,14 +273,15 @@ def test_export_yields_dataset_files_for_display_controls():
)

mock_dataset = MagicMock()
mock_dataset.id = dataset_id
sentinel_file = ("datasets/my_dataset.yaml", lambda: "dataset_content")
mock_datasets_cmd = MagicMock()
mock_datasets_cmd.run.return_value = iter([sentinel_file])

with (
patch(
"superset.commands.dashboard.export.DatasetDAO.find_by_id",
return_value=mock_dataset,
"superset.commands.dashboard.export.DatasetDAO.find_by_ids",
return_value=[mock_dataset],
),
patch(
"superset.commands.dashboard.export.ExportDatasetsCommand",
Expand Down Expand Up @@ -686,9 +826,10 @@ def test_stabilize_chart_ids_remaps_expanded_slices() -> None:

def test_file_content_missing_dataset_preserves_dataset_id() -> None:
"""
When DatasetDAO.find_by_id returns None for a display control target,
datasetId is preserved (dual-write: it was never popped) and no
datasetUuid is added — the target is not silently emptied.
When DatasetDAO.find_by_ids does not return a match for a display
control target's dataset, datasetId is preserved (dual-write: it was
never popped) and no datasetUuid is added — the target is not silently
emptied.
"""
from superset.commands.dashboard.export import ExportDashboardsCommand

Expand All @@ -706,8 +847,8 @@ def test_file_content_missing_dataset_preserves_dataset_id() -> None:

with (
patch(
"superset.commands.dashboard.export.DatasetDAO.find_by_id",
return_value=None,
"superset.commands.dashboard.export.DatasetDAO.find_by_ids",
return_value=[],
),
patch(
"superset.commands.dashboard.export.feature_flag_manager.is_feature_enabled",
Expand Down
Loading