fix(ingest): sort by last modified not working in the UI (#11343)

Author: sid-acryl

metadata-ingestion/src/datahub/ingestion/api/auto_work_units/__init__.py

@@ -0,0 +1,133 @@
+import dataclasses
+import json
+from typing import Dict, Iterable, Optional
+
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.metadata.schema_classes import (
+    ChangeTypeClass,
+    DatasetPropertiesClass,
+    GenericAspectClass,
+    MetadataChangeProposalClass,
+    OperationClass,
+    TimeStampClass,
+)
+from datahub.specific.dataset import DatasetPatchBuilder
+from datahub.utilities.urns.urn import guess_entity_type
+
+
+@dataclasses.dataclass
+class TimestampPair:
+    last_modified_dataset_props: Optional[
+        TimeStampClass
+    ]  # last_modified of datasetProperties aspect
+    last_updated_timestamp_dataset_props: Optional[
+        int
+    ]  # lastUpdatedTimestamp of the operation aspect
+
+
+def try_aspect_from_metadata_change_proposal_class(
+    wu: MetadataWorkUnit,
+) -> Optional[DatasetPropertiesClass]:
+    if (
+        isinstance(wu.metadata, MetadataChangeProposalClass)
+        and wu.metadata.aspectName == "datasetProperties"
+        and wu.metadata.changeType == ChangeTypeClass.PATCH
+        and isinstance(wu.metadata.aspect, GenericAspectClass)
+    ):
+        patch_dataset_properties = json.loads(wu.metadata.aspect.value)
+        for operation in patch_dataset_properties:
+            if operation.get("path") == "/lastModified":
+                # Deserializing `lastModified` as the `auto_patch_last_modified` function relies on this property
+                # to decide if a patch aspect for the datasetProperties aspect should be generated
+                return DatasetPropertiesClass(
+                    lastModified=TimeStampClass(time=operation["value"]["time"])
+                )
+
+    return None
+
+
+def auto_patch_last_modified(
+    stream: Iterable[MetadataWorkUnit],
+) -> Iterable[MetadataWorkUnit]:
+    """
+    Generate a patch request for datasetProperties aspect in-case
+        1. `lastModified` of datasetProperties is not set
+        2. And there are operation aspects
+    in this case set the `lastModified` of datasetProperties to max value of operation aspects `lastUpdatedTimestamp`.
+
+    We need this functionality to support sort by `last modified` on UI.
+    """
+    candidate_dataset_for_patch: Dict[str, TimestampPair] = {}
+
+    for wu in stream:
+        if (
+            guess_entity_type(wu.get_urn()) != "dataset"
+        ):  # we are only processing datasets
+            yield wu
+            continue
+
+        dataset_properties_aspect = wu.get_aspect_of_type(
+            DatasetPropertiesClass
+        ) or try_aspect_from_metadata_change_proposal_class(wu)
+        dataset_operation_aspect = wu.get_aspect_of_type(OperationClass)
+
+        timestamp_pair = candidate_dataset_for_patch.get(wu.get_urn())
+
+        if timestamp_pair:
+            # Update the timestamp_pair
+            if dataset_properties_aspect and dataset_properties_aspect.lastModified:
+                timestamp_pair.last_modified_dataset_props = (
+                    dataset_properties_aspect.lastModified
+                )
+
+            if (
+                dataset_operation_aspect
+                and dataset_operation_aspect.lastUpdatedTimestamp
+            ):
+                timestamp_pair.last_updated_timestamp_dataset_props = max(
+                    timestamp_pair.last_updated_timestamp_dataset_props or 0,
+                    dataset_operation_aspect.lastUpdatedTimestamp,
+                )
+
+        else:
+            # Create new TimestampPair
+            last_modified_dataset_props: Optional[TimeStampClass] = None
+            last_updated_timestamp_dataset_props: Optional[int] = None
+
+            if dataset_properties_aspect:
+                last_modified_dataset_props = dataset_properties_aspect.lastModified
+
+            if dataset_operation_aspect:
+                last_updated_timestamp_dataset_props = (
+                    dataset_operation_aspect.lastUpdatedTimestamp
+                )
+
+            candidate_dataset_for_patch[wu.get_urn()] = TimestampPair(
+                last_modified_dataset_props=last_modified_dataset_props,
+                last_updated_timestamp_dataset_props=last_updated_timestamp_dataset_props,
+            )
+
+        yield wu
+
+    # Emit a patch datasetProperties aspect for dataset where last_modified is None
+    for entity_urn, timestamp_pair in candidate_dataset_for_patch.items():
+        # Emit patch if last_modified is not set and last_updated_timestamp is set
+        if (
+            timestamp_pair.last_modified_dataset_props is None
+            and timestamp_pair.last_updated_timestamp_dataset_props
+        ):
+            dataset_patch_builder = DatasetPatchBuilder(urn=entity_urn)
+
+            dataset_patch_builder.set_last_modified(
+                timestamp=TimeStampClass(
+                    time=timestamp_pair.last_updated_timestamp_dataset_props
+                )
+            )
+
+            yield from [
+                MetadataWorkUnit(
+                    id=MetadataWorkUnit.generate_workunit_id(mcp),
+                    mcp_raw=mcp,
+                )
+                for mcp in dataset_patch_builder.build()
+            ]

metadata-ingestion/src/datahub/ingestion/api/auto_work_units/auto_dataset_properties_aspect.py

@@ -28,6 +28,9 @@
 from datahub.configuration.common import ConfigModel
 from datahub.configuration.source_common import PlatformInstanceConfigMixin
 from datahub.emitter.mcp_builder import mcps_from_mce
+from datahub.ingestion.api.auto_work_units.auto_dataset_properties_aspect import (
+    auto_patch_last_modified,
+)
 from datahub.ingestion.api.closeable import Closeable
 from datahub.ingestion.api.common import PipelineContext, RecordEnvelope, WorkUnit
 from datahub.ingestion.api.report import Report
@@ -443,6 +446,7 @@ def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
             ),
             browse_path_processor,
             partial(auto_workunit_reporter, self.get_report()),
+            auto_patch_last_modified,
         ]
 
     @staticmethod

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

@@ -498,5 +498,27 @@
         "runId": "bigquery-2022_02_03-07_00_00",
         "lastRunId": "no-run-id-provided"
     }
+},
+{
+    "entityType": "dataset",
+    "entityUrn": "urn:li:dataset:(urn:li:dataPlatform:bigquery,project-id-1.bigquery-dataset-1.view-1,PROD)",
+    "changeType": "PATCH",
+    "aspectName": "datasetProperties",
+    "aspect": {
+        "json": [
+            {
+                "op": "add",
+                "path": "/lastModified",
+                "value": {
+                    "time": 1643871600000
+                }
+            }
+        ]
+    },
+    "systemMetadata": {
+        "lastObserved": 1643871600000,
+        "runId": "bigquery-2022_02_03-07_00_00",
+        "lastRunId": "no-run-id-provided"
+    }
 }
 ]