Further null checks

Author: treff7es

metadata-ingestion/src/datahub/ingestion/source/gc/dataprocess_cleanup.py

@@ -262,7 +262,6 @@ def keep_last_n_dpi(
                     self.report.report_failure(
                         f"Exception while deleting DPI: {e}", exc=e
                     )
-
             if deleted_count_last_n % self.config.batch_size == 0:
                 logger.info(f"Deleted {deleted_count_last_n} DPIs from {job.urn}")
                 if self.config.delay:
@@ -296,7 +295,7 @@ def delete_dpi_from_datajobs(self, job: DataJobEntity) -> None:
         dpis = self.fetch_dpis(job.urn, self.config.batch_size)
         dpis.sort(
             key=lambda x: x["created"]["time"]
-            if "created" in x and "time" in x["created"]
+            if x.get("created") and x["created"].get("time")
             else 0,
             reverse=True,
         )

metadata-ingestion/tests/unit/test_gc.py

@@ -84,6 +84,54 @@ def test_delete_dpi_from_datajobs_without_dpi_created_time(self, mock_fetch_dpis
         self.cleanup.delete_dpi_from_datajobs(job)
         self.assertEqual(10, self.report.num_aspects_removed)
 
+    @patch(
+        "datahub.ingestion.source.gc.dataprocess_cleanup.DataProcessCleanup.fetch_dpis"
+    )
+    def test_delete_dpi_from_datajobs_without_dpi_null_created_time(
+        self, mock_fetch_dpis
+    ):
+        job = DataJobEntity(
+            urn="urn:li:dataJob:1",
+            flow_urn="urn:li:dataFlow:1",
+            lastIngested=int(datetime.now(timezone.utc).timestamp()),
+            jobId="job1",
+            dataPlatformInstance="urn:li:dataPlatformInstance:1",
+            total_runs=10,
+        )
+        mock_fetch_dpis.return_value = [
+            {"urn": f"urn:li:dataprocessInstance:{i}"} for i in range(10)
+        ] + [
+            {
+                "urn": "urn:li:dataprocessInstance:11",
+                "created": {"time": None},
+            }
+        ]
+        self.cleanup.delete_dpi_from_datajobs(job)
+        self.assertEqual(11, self.report.num_aspects_removed)
+
+    @patch(
+        "datahub.ingestion.source.gc.dataprocess_cleanup.DataProcessCleanup.fetch_dpis"
+    )
+    def test_delete_dpi_from_datajobs_without_dpi_without_time(self, mock_fetch_dpis):
+        job = DataJobEntity(
+            urn="urn:li:dataJob:1",
+            flow_urn="urn:li:dataFlow:1",
+            lastIngested=int(datetime.now(timezone.utc).timestamp()),
+            jobId="job1",
+            dataPlatformInstance="urn:li:dataPlatformInstance:1",
+            total_runs=10,
+        )
+        mock_fetch_dpis.return_value = [
+            {"urn": f"urn:li:dataprocessInstance:{i}"} for i in range(10)
+        ] + [
+            {
+                "urn": "urn:li:dataprocessInstance:11",
+                "created": None,
+            }
+        ]
+        self.cleanup.delete_dpi_from_datajobs(job)
+        self.assertEqual(11, self.report.num_aspects_removed)
+
     def test_fetch_dpis(self):
         assert self.cleanup.ctx.graph
         self.cleanup.ctx.graph = MagicMock()