From 2fc71da58d2fa32bb2d924227c40fb9576cffdd9 Mon Sep 17 00:00:00 2001
From: Gabe Lyons <itsgabelyons@gmail.com>
Date: Wed, 4 Aug 2021 11:44:22 -0700
Subject: [PATCH] fix(restore-indices): add system metadata restoration to
 restore-indices (#3029)

---
 .../datahub/upgrade/restoreindices/SendMAEStep.java  |  5 ++++-
 .../ElasticSearchSystemMetadataService.java          | 10 ++++++++--
 .../systemmetadata/SystemMetadataService.java        |  3 ++-
 .../ElasticSearchSystemMetadataServiceTest.java      | 12 +++++++++++-
 4 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java
index 6d96c519a1..a55dd1ab31 100644
--- a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java
+++ b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/restoreindices/SendMAEStep.java
@@ -13,6 +13,7 @@ import com.linkedin.metadata.entity.ebean.EbeanUtils;
 import com.linkedin.metadata.models.EntitySpec;
 import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.mxe.MetadataAuditOperation;
+import com.linkedin.mxe.SystemMetadata;
 import io.ebean.EbeanServer;
 import io.ebean.PagedList;
 import java.util.Map;
@@ -101,13 +102,15 @@ public class SendMAEStep implements UpgradeStep {
             return new DefaultUpgradeStepResult(id(), UpgradeStepResult.Result.FAILED);
           }
 
+          SystemMetadata latestSystemMetadata = EbeanUtils.parseSystemMetadata(aspect.getSystemMetadata());
+
           // 5. Produce MAE events for the aspect record
           _entityService.produceMetadataAuditEvent(
               urn,
               null,
               aspectRecord,
               null,
-              null,
+              latestSystemMetadata,
               MetadataAuditOperation.UPDATE
           );
 
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataService.java b/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataService.java
index e87800e5ec..0f29920a83 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataService.java
@@ -17,6 +17,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 import org.elasticsearch.action.delete.DeleteResponse;
@@ -83,9 +84,14 @@ public class ElasticSearchSystemMetadataService implements SystemMetadataService
   }
 
   @Override
-  public void insert(SystemMetadata systemMetadata, String urn, String aspect) {
-    String document = toDocument(systemMetadata, urn, aspect);
+  public void insert(@Nullable SystemMetadata systemMetadata, String urn, String aspect) {
+    if (systemMetadata == null) {
+      return;
+    }
+
     String docId = toDocId(urn, aspect);
+
+    String document = toDocument(systemMetadata, urn, aspect);
     _esDAO.upsertDocument(docId, document);
   }
 
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/SystemMetadataService.java b/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/SystemMetadataService.java
index 75b24b5af9..61c72edcf0 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/SystemMetadataService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/systemmetadata/SystemMetadataService.java
@@ -4,6 +4,7 @@ import com.linkedin.metadata.run.AspectRowSummary;
 import com.linkedin.metadata.run.IngestionRunSummary;
 import com.linkedin.mxe.SystemMetadata;
 import java.util.List;
+import javax.annotation.Nullable;
 
 
 public interface SystemMetadataService {
@@ -11,7 +12,7 @@ public interface SystemMetadataService {
 
   void deleteUrn(String finalOldUrn);
 
-  void insert(SystemMetadata systemMetadata, String urn, String aspect);
+  void insert(@Nullable SystemMetadata systemMetadata, String urn, String aspect);
 
   List<AspectRowSummary> findByRunId(String runId);
 
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataServiceTest.java
index 7a7447fb77..1d9163e5d6 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataServiceTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/systemmetadata/ElasticSearchSystemMetadataServiceTest.java
@@ -22,7 +22,6 @@ import org.testng.annotations.Test;
 
 import static org.testng.Assert.*;
 
-
 public class ElasticSearchSystemMetadataServiceTest {
 
   private ElasticsearchContainer _elasticsearchContainer;
@@ -186,4 +185,15 @@ public class ElasticSearchSystemMetadataServiceTest {
     assertEquals(rows.size(), 2);
     rows.forEach(row -> assertEquals(row.getRunId(), "abc-456"));
   }
+
+  @Test
+  public void testInsertNullData() throws Exception {
+    _client.insert(null, "urn:li:chart:1", "chartKey");
+
+    TimeUnit.SECONDS.sleep(5);
+
+    List<IngestionRunSummary> runs = _client.listRuns(0, 20);
+
+    assertEquals(runs.size(), 0);
+  }
 }