complain more loudly for extractor failures

datahub-project · Jul 12, 2024 · 853bbc1 · 853bbc1
1 parent 2837ef5
commit 853bbc1
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 32 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/run/pipeline.py b/metadata-ingestion/src/datahub/ingestion/run/pipeline.py
@@ -436,6 +436,8 @@ def run(self) -> None:
                     )
                 )
 
+            stack.enter_context(self.sink)
+
             self.final_status = PipelineStatus.UNKNOWN
             self._notify_reporters_on_ingestion_start()
             callback = None
@@ -460,7 +462,17 @@ def run(self) -> None:
                     if not self.dry_run:
                         self.sink.handle_work_unit_start(wu)
                     try:
-                        record_envelopes = self.extractor.get_records(wu)
+                        # Most of this code is meant to be fully stream-based instead of generating all records into memory.
+                        # However, the extractor in particular will never generate a particularly large list. We want the
+                        # exception reporting to be associated with the source, and not the transformer. As such, we
+                        # need to materialize the generator returned by get_records().
+                        record_envelopes = list(self.extractor.get_records(wu))
+                    except Exception as e:
+                        self.source.get_report().failure(
+                            "Source produced bad metadata", context=wu.id, exc=e
+                        )
+                        continue
+                    try:
                         for record_envelope in self.transform(record_envelopes):
                             if not self.dry_run:
                                 try:
@@ -482,9 +494,9 @@ def run(self) -> None:
                         )
                         # TODO: Transformer errors should cause the pipeline to fail.
 
-                    self.extractor.close()
                     if not self.dry_run:
                         self.sink.handle_work_unit_end(wu)
+                self.extractor.close()
                 self.source.close()
                 # no more data is coming, we need to let the transformers produce any additional records if they are holding on to state
                 for record_envelope in self.transform(
@@ -518,8 +530,6 @@ def run(self) -> None:
 
                 self._notify_reporters_on_ingestion_completion()
 
-                self.sink.close()
-
     def transform(self, records: Iterable[RecordEnvelope]) -> Iterable[RecordEnvelope]:
         """
         Transforms the given sequence of records by passing the records through the transformers

diff --git a/metadata-ingestion/src/datahub/ingestion/source/grafana/grafana_source.py b/metadata-ingestion/src/datahub/ingestion/source/grafana/grafana_source.py
@@ -14,6 +14,7 @@
     support_status,
 )
 from datahub.ingestion.api.source import MetadataWorkUnitProcessor
+from datahub.ingestion.api.source_helpers import auto_workunit
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalHandler,
@@ -97,32 +98,30 @@ def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
                 name=_uid,
                 platform_instance=self.source_config.platform_instance,
             )
-            for mcp in MetadataChangeProposalWrapper.construct_many(
-                entityUrn=dashboard_urn,
-                aspects=[
-                    DashboardInfoClass(
-                        description="",
-                        title=_title,
-                        charts=[],
-                        lastModified=ChangeAuditStamps(),
-                        dashboardUrl=full_url,
-                        customProperties={
-                            "displayName": _title,
-                            "id": str(item["id"]),
-                            "uid": _uid,
-                            "title": _title,
-                            "uri": item["uri"],
-                            "type": item["type"],
-                            "folderId": str(item.get("folderId", None)),
-                            "folderUid": item.get("folderUid", None),
-                            "folderTitle": str(item.get("folderTitle", None)),
-                        },
-                    ),
-                    StatusClass(removed=False),
-                ],
-            ):
-                breakpoint()
-                yield MetadataWorkUnit(
-                    id=dashboard_urn,
-                    mcp=mcp,
+
+            yield from auto_workunit(
+                MetadataChangeProposalWrapper.construct_many(
+                    entityUrn=dashboard_urn,
+                    aspects=[
+                        DashboardInfoClass(
+                            description="",
+                            title=_title,
+                            charts=[],
+                            lastModified=ChangeAuditStamps(),
+                            dashboardUrl=full_url,
+                            customProperties={
+                                "displayName": _title,
+                                "id": str(item["id"]),
+                                "uid": _uid,
+                                "title": _title,
+                                "uri": item["uri"],
+                                "type": item["type"],
+                                "folderId": str(item.get("folderId", None)),
+                                "folderUid": item.get("folderUid", None),
+                                "folderTitle": str(item.get("folderTitle", None)),
+                            },
+                        ),
+                        StatusClass(removed=False),
+                    ],
                 )
+            )