From 4d9a8d4659f2c7dd0a89c08d49518839f095855d Mon Sep 17 00:00:00 2001
From: Harshal Sheth <hsheth2@gmail.com>
Date: Tue, 6 Sep 2022 19:34:01 -0700
Subject: [PATCH] fix(ingest): include urns in rest sink failure logs

---
 .../src/datahub/ingestion/sink/datahub_rest.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py b/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py
index cdfc6e0b920ca..1ea127042a7fc 100644
--- a/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py
+++ b/metadata-ingestion/src/datahub/ingestion/sink/datahub_rest.py
@@ -175,18 +175,18 @@ def _write_done_callback(
                             :200
                         ]
 
+                # Include information about the entity that failed.
+                record = record_envelope.record
+                if isinstance(record, MetadataChangeProposalWrapper):
+                    entity_id = record.entityUrn
+                    e.info["id"] = entity_id
+                elif isinstance(record, MetadataChangeEvent):
+                    entity_id = record.proposedSnapshot.urn
+                    e.info["id"] = entity_id
+
                 if not self.treat_errors_as_warnings:
                     self.report.report_failure({"error": e.message, "info": e.info})
                 else:
-                    record = record_envelope.record
-                    if isinstance(record, MetadataChangeProposalWrapper):
-                        # include information about the entity that failed
-                        entity_id = cast(
-                            MetadataChangeProposalWrapper, record
-                        ).entityUrn
-                        e.info["id"] = entity_id
-                    else:
-                        entity_id = None
                     self.report.report_warning({"warning": e.message, "info": e.info})
                 write_callback.on_failure(record_envelope, e, e.info)
             else: