fix(ingest): better handling around sink errors (#9003)

datahub-project · Oct 25, 2023 · 32f5dcb · 32f5dcb
1 parent 1c77bca
commit 32f5dcb
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 21 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/run/pipeline.py b/metadata-ingestion/src/datahub/ingestion/run/pipeline.py
@@ -390,7 +390,15 @@ def run(self) -> None:
                         record_envelopes = self.extractor.get_records(wu)
                         for record_envelope in self.transform(record_envelopes):
                             if not self.dry_run:
-                                self.sink.write_record_async(record_envelope, callback)
+                                try:
+                                    self.sink.write_record_async(
+                                        record_envelope, callback
+                                    )
+                                except Exception as e:
+                                    # In case the sink's error handling is bad, we still want to report the error.
+                                    self.sink.report.report_failure(
+                                        f"Failed to write record: {e}"
+                                    )
 
                     except RuntimeError:
                         raise

diff --git a/metadata-ingestion/src/datahub/ingestion/sink/datahub_kafka.py b/metadata-ingestion/src/datahub/ingestion/sink/datahub_kafka.py
@@ -9,7 +9,6 @@
     MetadataChangeEvent,
     MetadataChangeProposal,
 )
-from datahub.metadata.schema_classes import MetadataChangeProposalClass
 
 
 class KafkaSinkConfig(KafkaEmitterConfig):
@@ -58,27 +57,21 @@ def write_record_async(
         ],
         write_callback: WriteCallback,
     ) -> None:
-        record = record_envelope.record
-        if isinstance(record, MetadataChangeEvent):
-            self.emitter.emit_mce_async(
+        callback = _KafkaCallback(
+            self.report, record_envelope, write_callback
+        ).kafka_callback
+        try:
+            record = record_envelope.record
+            self.emitter.emit(
                 record,
-                callback=_KafkaCallback(
-                    self.report, record_envelope, write_callback
-                ).kafka_callback,
-            )
-        elif isinstance(
-            record, (MetadataChangeProposalWrapper, MetadataChangeProposalClass)
-        ):
-            self.emitter.emit_mcp_async(
-                record,
-                callback=_KafkaCallback(
-                    self.report, record_envelope, write_callback
-                ).kafka_callback,
-            )
-        else:
-            raise ValueError(
-                f"The datahub-kafka sink only supports MetadataChangeEvent/MetadataChangeProposal[Wrapper] classes, not {type(record)}"
+                callback=callback,
             )
+        except Exception as err:
+            # In case we throw an exception while trying to emit the record,
+            # catch it and report the failure. This might happen if the schema
+            # registry is down or otherwise misconfigured, in which case we'd
+            # fail when serializing the record.
+            callback(err, f"Failed to write record: {err}")
 
     def close(self) -> None:
         self.emitter.flush()