fix chardEdges destination type validation in DashboardPatchBuilder + add missing case for dashboards

Author: sgomezvillamor

metadata-ingestion/src/datahub/ingestion/api/incremental_lineage_helper.py

@@ -102,6 +102,10 @@ def convert_dashboard_info_to_patch(
     if aspect.datasets:
         patch_builder.add_datasets(aspect.datasets)
 
+    if aspect.dashboards:
+        for dashboard in aspect.dashboards:
+            patch_builder.add_dashboard(dashboard)
+
     if aspect.access:
         patch_builder.set_access(aspect.access)
 

metadata-ingestion/src/datahub/specific/dashboard.py

@@ -161,7 +161,7 @@ def add_chart_edge(self, chart: Union[Edge, Urn, str]) -> "DashboardPatchBuilder
                 lastModified=self._mint_auditstamp(),
             )
 
-        self._ensure_urn_type("dataset", [chart_edge], "add_chart_edge")
+        self._ensure_urn_type("chart", [chart_edge], "add_chart_edge")
         self._add_patch(
             DashboardInfo.ASPECT_NAME,
             "add",
@@ -271,6 +271,48 @@ def add_datasets(
 
         return self
 
+    def add_dashboard(
+        self, dashboard: Union[Edge, Urn, str]
+    ) -> "DashboardPatchBuilder":
+        """
+        Adds an dashboard to the DashboardPatchBuilder.
+
+        Args:
+            dashboard: The dashboard, which can be an Edge object, Urn object, or a string.
+
+        Returns:
+            The DashboardPatchBuilder instance.
+
+        Raises:
+            ValueError: If the dashboard is not a Dashboard urn.
+
+        Notes:
+            If `dashboard` is an Edge object, it is used directly. If `dashboard` is a Urn object or string,
+            it is converted to an Edge object and added with default audit stamps.
+        """
+        if isinstance(dashboard, Edge):
+            dashboard_urn: str = dashboard.destinationUrn
+            dashboard_edge: Edge = dashboard
+        elif isinstance(dashboard, (Urn, str)):
+            dashboard_urn = str(dashboard)
+            if not dashboard_urn.startswith("urn:li:dashboard:"):
+                raise ValueError(f"Input {dashboard} is not a Dashboard urn")
+
+            dashboard_edge = Edge(
+                destinationUrn=dashboard_urn,
+                created=self._mint_auditstamp(),
+                lastModified=self._mint_auditstamp(),
+            )
+
+        self._ensure_urn_type("dashboard", [dashboard_edge], "add_dashboard")
+        self._add_patch(
+            DashboardInfo.ASPECT_NAME,
+            "add",
+            path=("dashboards", dashboard_urn),
+            value=dashboard_edge,
+        )
+        return self
+
     def set_dashboard_url(
         self, dashboard_url: Optional[str]
     ) -> "DashboardPatchBuilder":

metadata-ingestion/tests/unit/api/source_helpers/test_incremental_lineage_helper.py

@@ -1,11 +1,19 @@
-from typing import List
+import re
+import unittest
+from typing import List, Optional
+
+import pytest
 
 import datahub.emitter.mce_builder as builder
 import datahub.metadata.schema_classes as models
 from datahub.emitter.mce_builder import make_dataset_urn, make_schema_field_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
-from datahub.ingestion.api.incremental_lineage_helper import auto_incremental_lineage
+from datahub.ingestion.api.incremental_lineage_helper import (
+    auto_incremental_lineage,
+    convert_dashboard_info_to_patch,
+)
 from datahub.ingestion.api.source_helpers import auto_workunit
+from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.sink.file import write_metadata_file
 from tests.test_helpers import mce_helpers
 
@@ -183,3 +191,129 @@ def test_incremental_lineage_pass_through(tmp_path, pytestconfig):
     mce_helpers.check_golden_file(
         pytestconfig=pytestconfig, output_path=test_file, golden_path=golden_file
     )
+
+
+class TestConvertDashboardInfoToPatch(unittest.TestCase):
+    def setUp(self):
+        self.urn = builder.make_dashboard_urn(
+            platform="platform", name="test", platform_instance="instance"
+        )
+        self.system_metadata = models.SystemMetadataClass()
+        self.aspect = models.DashboardInfoClass(
+            title="Test Dashboard",
+            description="This is a test dashboard",
+            lastModified=models.ChangeAuditStampsClass(),
+        )
+
+    def _validate_dashboard_info_patch(
+        self, result: Optional[MetadataWorkUnit], expected_aspect_value: bytes
+    ) -> None:
+        assert (
+            result
+            and result.metadata
+            and isinstance(result.metadata, models.MetadataChangeProposalClass)
+        )
+        mcp: models.MetadataChangeProposalClass = result.metadata
+        assert (
+            mcp.entityUrn == self.urn
+            and mcp.entityType == "dashboard"
+            and mcp.aspectName == "dashboardInfo"
+            and mcp.changeType == "PATCH"
+        )
+        assert mcp.aspect and isinstance(mcp.aspect, models.GenericAspectClass)
+        aspect: models.GenericAspectClass = mcp.aspect
+        assert aspect.value == expected_aspect_value
+
+    def test_convert_dashboard_info_to_patch_with_no_additional_values(self):
+        result = convert_dashboard_info_to_patch(
+            self.urn, self.aspect, self.system_metadata
+        )
+        self._validate_dashboard_info_patch(
+            result=result,
+            expected_aspect_value=b'[{"op": "add", "path": "/title", "value": "Test Dashboard"}, {"op": "add", "path": "/description", "value": "This is a test dashboard"}, {"op": "add", "path": "/lastModified", "value": {"created": {"time": 0, "actor": "urn:li:corpuser:unknown"}, "lastModified": {"time": 0, "actor": "urn:li:corpuser:unknown"}}}]',
+        )
+
+    def test_convert_dashboard_info_to_patch_with_custom_properties(self):
+        self.aspect.customProperties = {
+            "key1": "value1",
+            "key2": "value2",
+        }
+        result = convert_dashboard_info_to_patch(
+            self.urn, self.aspect, self.system_metadata
+        )
+        self._validate_dashboard_info_patch(
+            result=result,
+            expected_aspect_value=b'[{"op": "add", "path": "/customProperties/key1", "value": "value1"}, {"op": "add", "path": "/customProperties/key2", "value": "value2"}, {"op": "add", "path": "/title", "value": "Test Dashboard"}, {"op": "add", "path": "/description", "value": "This is a test dashboard"}, {"op": "add", "path": "/lastModified", "value": {"created": {"time": 0, "actor": "urn:li:corpuser:unknown"}, "lastModified": {"time": 0, "actor": "urn:li:corpuser:unknown"}}}]',
+        )
+
+    def test_convert_dashboard_info_to_patch_with_chart_edges(self):
+        self.aspect.chartEdges = [
+            models.EdgeClass(
+                destinationUrn=builder.make_chart_urn(
+                    platform="platform",
+                    name="target-test",
+                    platform_instance="instance",
+                ),
+            )
+        ]
+        result = convert_dashboard_info_to_patch(
+            self.urn, self.aspect, self.system_metadata
+        )
+        self._validate_dashboard_info_patch(
+            result,
+            b'[{"op": "add", "path": "/chartEdges/urn:li:chart:(platform,instance.target-test)", "value": {"destinationUrn": "urn:li:chart:(platform,instance.target-test)"}}, {"op": "add", "path": "/title", "value": "Test Dashboard"}, {"op": "add", "path": "/description", "value": "This is a test dashboard"}, {"op": "add", "path": "/lastModified", "value": {"created": {"time": 0, "actor": "urn:li:corpuser:unknown"}, "lastModified": {"time": 0, "actor": "urn:li:corpuser:unknown"}}}]',
+        )
+
+    def test_convert_dashboard_info_to_patch_with_chart_edges_no_chart(self):
+        self.aspect.chartEdges = [
+            models.EdgeClass(
+                destinationUrn=builder.make_dashboard_urn(
+                    platform="platform",
+                    name="target-test",
+                    platform_instance="instance",
+                ),
+            )
+        ]
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "add_chart_edge: urn:li:dashboard:(platform,instance.target-test) is not of type chart"
+            ),
+        ):
+            convert_dashboard_info_to_patch(self.urn, self.aspect, self.system_metadata)
+
+    def test_convert_dashboard_info_to_patch_with_dashboards(self):
+        self.aspect.dashboards = [
+            models.EdgeClass(
+                destinationUrn=builder.make_dashboard_urn(
+                    platform="platform",
+                    name="target-test",
+                    platform_instance="instance",
+                ),
+            )
+        ]
+        result = convert_dashboard_info_to_patch(
+            self.urn, self.aspect, self.system_metadata
+        )
+        self._validate_dashboard_info_patch(
+            result,
+            b'[{"op": "add", "path": "/title", "value": "Test Dashboard"}, {"op": "add", "path": "/description", "value": "This is a test dashboard"}, {"op": "add", "path": "/dashboards/urn:li:dashboard:(platform,instance.target-test)", "value": {"destinationUrn": "urn:li:dashboard:(platform,instance.target-test)"}}, {"op": "add", "path": "/lastModified", "value": {"created": {"time": 0, "actor": "urn:li:corpuser:unknown"}, "lastModified": {"time": 0, "actor": "urn:li:corpuser:unknown"}}}]',
+        )
+
+    def test_convert_dashboard_info_to_patch_with_dashboards_no_dashboard(self):
+        self.aspect.dashboards = [
+            models.EdgeClass(
+                destinationUrn=builder.make_chart_urn(
+                    platform="platform",
+                    name="target-test",
+                    platform_instance="instance",
+                ),
+            )
+        ]
+        with pytest.raises(
+            ValueError,
+            match=re.escape(
+                "add_dashboard: urn:li:chart:(platform,instance.target-test) is not of type dashboard"
+            ),
+        ):
+            convert_dashboard_info_to_patch(self.urn, self.aspect, self.system_metadata)