fix(ingest/redshift): patch instead of replace redshift custom properties

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

@@ -25,12 +25,14 @@
     BrowsePathsV2Class,
     ChangeTypeClass,
     ContainerClass,
+    DatasetPropertiesClass,
     DatasetUsageStatisticsClass,
     MetadataChangeEventClass,
     MetadataChangeProposalClass,
     StatusClass,
     TimeWindowSizeClass,
 )
+from datahub.specific.dataset import DatasetPatchBuilder
 from datahub.telemetry import telemetry
 from datahub.utilities.urns.dataset_urn import DatasetUrn
 from datahub.utilities.urns.tag_urn import TagUrn
@@ -58,6 +60,22 @@ def auto_workunit(
             yield item.as_workunit()
 
 
+def create_dataset_props_patch_builder(
+    dataset_urn: str,
+    dataset_properties: DatasetPropertiesClass,
+) -> DatasetPatchBuilder:
+    """Creates a patch builder with a table's or view's attributes and dataset properties"""
+    patch_builder = DatasetPatchBuilder(dataset_urn)
+    patch_builder.set_display_name(dataset_properties.name)
+    patch_builder.set_description(dataset_properties.description)
+    patch_builder.set_created(dataset_properties.created)
+    patch_builder.set_last_modified(dataset_properties.lastModified)
+    patch_builder.set_qualified_name(dataset_properties.qualifiedName)
+    patch_builder.add_custom_properties(dataset_properties.customProperties)
+
+    return patch_builder
+
+
 def auto_status_aspect(
     stream: Iterable[MetadataWorkUnit],
 ) -> Iterable[MetadataWorkUnit]:
@@ -371,9 +389,9 @@ def auto_empty_dataset_usage_statistics(
                     userCounts=[],
                     fieldCounts=[],
                 ),
-                changeType=ChangeTypeClass.CREATE
-                if all_buckets
-                else ChangeTypeClass.UPSERT,
+                changeType=(
+                    ChangeTypeClass.CREATE if all_buckets else ChangeTypeClass.UPSERT
+                ),
             ).as_workunit()
 
 

metadata-ingestion/src/datahub/ingestion/source/redshift/config.py

@@ -142,6 +142,11 @@ class RedshiftConfig(
         description="When enabled, emits lineage as incremental to existing lineage already in DataHub. When disabled, re-states lineage on each run.  This config works with rest-sink only.",
     )
 
+    patch_custom_properties: bool = Field(
+        default=True,
+        description="Whether to patch custom properties on existing datasets rather than replace.",
+    )
+
     resolve_temp_table_in_lineage: bool = Field(
         default=True,
         description="Whether to resolve temp table appear in lineage to upstream permanent tables.",

metadata-ingestion/src/datahub/ingestion/source/redshift/redshift.py

@@ -33,6 +33,7 @@
     TestableSource,
     TestConnectionReport,
 )
+from datahub.ingestion.api.source_helpers import create_dataset_props_patch_builder
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.subtypes import (
     DatasetContainerSubTypes,
@@ -756,24 +757,36 @@ def gen_dataset_workunits(
 
         dataset_properties = DatasetProperties(
             name=table.name,
-            created=TimeStamp(time=int(table.created.timestamp() * 1000))
-            if table.created
-            else None,
-            lastModified=TimeStamp(time=int(table.last_altered.timestamp() * 1000))
-            if table.last_altered
-            else TimeStamp(time=int(table.created.timestamp() * 1000))
-            if table.created
-            else None,
+            created=(
+                TimeStamp(time=int(table.created.timestamp() * 1000))
+                if table.created
+                else None
+            ),
+            lastModified=(
+                TimeStamp(time=int(table.last_altered.timestamp() * 1000))
+                if table.last_altered
+                else (
+                    TimeStamp(time=int(table.created.timestamp() * 1000))
+                    if table.created
+                    else None
+                )
+            ),
             description=table.comment,
             qualifiedName=str(datahub_dataset_name),
+            customProperties=custom_properties,
         )
-
-        if custom_properties:
-            dataset_properties.customProperties = custom_properties
-
-        yield MetadataChangeProposalWrapper(
-            entityUrn=dataset_urn, aspect=dataset_properties
-        ).as_workunit()
+        if self.config.patch_custom_properties:
+            patch_builder = create_dataset_props_patch_builder(
+                dataset_urn, dataset_properties
+            )
+            for patch_mcp in patch_builder.build():
+                yield MetadataWorkUnit(
+                    id=f"{dataset_urn}-{patch_mcp.aspectName}", mcp_raw=patch_mcp
+                )
+        else:
+            yield MetadataChangeProposalWrapper(
+                entityUrn=dataset_urn, aspect=dataset_properties
+            ).as_workunit()
 
         # TODO: Check if needed
         # if tags_to_add:

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

@@ -1,6 +1,7 @@
 from typing import Dict, Generic, List, Optional, Tuple, TypeVar, Union
 
 from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
+from datahub.metadata.com.linkedin.pegasus2avro.common import TimeStamp
 from datahub.metadata.schema_classes import (
     DatasetPropertiesClass as DatasetProperties,
     EditableDatasetPropertiesClass as EditableDatasetProperties,
@@ -258,16 +259,19 @@ def for_field(
         )
 
     def set_description(
-        self, description: str, editable: bool = False
+        self, description: Optional[str] = None, editable: bool = False
     ) -> "DatasetPatchBuilder":
-        self._add_patch(
-            DatasetProperties.ASPECT_NAME
-            if not editable
-            else EditableDatasetProperties.ASPECT_NAME,
-            "add",
-            path="/description",
-            value=description,
-        )
+        if description is not None:
+            self._add_patch(
+                (
+                    DatasetProperties.ASPECT_NAME
+                    if not editable
+                    else EditableDatasetProperties.ASPECT_NAME
+                ),
+                "add",
+                path="/description",
+                value=description,
+            )
         return self
 
     def set_custom_properties(
@@ -285,11 +289,21 @@ def add_custom_property(self, key: str, value: str) -> "DatasetPatchBuilder":
         self.custom_properties_patch_helper.add_property(key, value)
         return self
 
+    def add_custom_properties(
+        self, custom_properties: Optional[Dict[str, str]] = None
+    ) -> "DatasetPatchBuilder":
+        if custom_properties is not None:
+            for key, value in custom_properties.items():
+                self.custom_properties_patch_helper.add_property(key, value)
+        return self
+
     def remove_custom_property(self, key: str) -> "DatasetPatchBuilder":
         self.custom_properties_patch_helper.remove_property(key)
         return self
 
-    def set_display_name(self, display_name: str) -> "DatasetPatchBuilder":
+    def set_display_name(
+        self, display_name: Optional[str] = None
+    ) -> "DatasetPatchBuilder":
         if display_name is not None:
             self._add_patch(
                 DatasetProperties.ASPECT_NAME,
@@ -299,6 +313,42 @@ def set_display_name(self, display_name: str) -> "DatasetPatchBuilder":
             )
         return self
 
+    def set_qualified_name(
+        self, qualified_name: Optional[str] = None
+    ) -> "DatasetPatchBuilder":
+        if qualified_name is not None:
+            self._add_patch(
+                DatasetProperties.ASPECT_NAME,
+                "add",
+                path="/qualifiedName",
+                value=qualified_name,
+            )
+        return self
+
+    def set_created(
+        self, timestamp: Optional[TimeStamp] = None
+    ) -> "DatasetPatchBuilder":
+        if timestamp is not None:
+            self._add_patch(
+                DatasetProperties.ASPECT_NAME,
+                "add",
+                path="/created",
+                value=timestamp,
+            )
+        return self
+
+    def set_last_modified(
+        self, timestamp: Optional[TimeStamp] = None
+    ) -> "DatasetPatchBuilder":
+        if timestamp is not None:
+            self._add_patch(
+                DatasetProperties.ASPECT_NAME,
+                "add",
+                path="/lastModified",
+                value=timestamp,
+            )
+        return self
+
     def set_structured_property(
         self, property_name: str, value: Union[str, float, List[Union[str, float]]]
     ) -> "DatasetPatchBuilder":

metadata-ingestion/tests/unit/test_redshift_source.py

@@ -0,0 +1,63 @@
+from typing import Iterable
+
+from datahub.emitter.mcp import (
+    MetadataChangeProposalClass,
+    MetadataChangeProposalWrapper,
+)
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.source.redshift.config import RedshiftConfig
+from datahub.ingestion.source.redshift.redshift import RedshiftSource
+from datahub.ingestion.source.redshift.redshift_schema import RedshiftTable
+from datahub.metadata.schema_classes import MetadataChangeEventClass
+
+
+def redshift_source_setup(custom_props_flag: bool) -> Iterable[MetadataWorkUnit]:
+    config = RedshiftConfig(
+        host_port="localhost:5439",
+        database="test",
+        patch_custom_properties=custom_props_flag,
+    )
+    source: RedshiftSource = RedshiftSource(config, ctx=PipelineContext(run_id="test"))
+    gen = source.gen_dataset_workunits(
+        table=RedshiftTable(
+            name="category",
+            columns=[],
+            created=None,
+            comment="",
+        ),
+        database="dev",
+        schema="public",
+        sub_type="test_sub_type",
+        custom_properties={"my_key": "my_value"},
+    )
+    return gen
+
+
+def test_gen_dataset_workunits_patch_custom_properties_patch():
+    gen = redshift_source_setup(True)
+    custom_props_exist = False
+    for item in gen:
+        mcp = item.metadata
+        assert not isinstance(mcp, MetadataChangeEventClass)
+        if mcp.aspectName == "datasetProperties":
+            assert isinstance(mcp, MetadataChangeProposalClass)
+            assert mcp.changeType == "PATCH"
+            custom_props_exist = True
+        else:
+            assert isinstance(mcp, MetadataChangeProposalWrapper)
+
+    assert custom_props_exist
+
+
+def test_gen_dataset_workunits_patch_custom_properties_upsert():
+    gen = redshift_source_setup(False)
+    custom_props_exist = False
+    for item in gen:
+        assert isinstance(item.metadata, MetadataChangeProposalWrapper)
+        mcp = item.metadata
+        if mcp.aspectName == "datasetProperties":
+            assert mcp.changeType == "UPSERT"
+            custom_props_exist = True
+
+    assert custom_props_exist