fix(aspects): fix default aspect generation for non-restli locations

datahub-project · Jan 31, 2024 · 2c65863 · 2c65863
1 parent f378fb6
commit 2c65863
Show file tree

Hide file tree

Showing 61 changed files with 1,004 additions and 764 deletions.
diff --git a/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/DataMigrationStep.java b/datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/nocode/DataMigrationStep.java
@@ -10,6 +10,7 @@
 import com.linkedin.datahub.upgrade.UpgradeStepResult;
 import com.linkedin.datahub.upgrade.impl.DefaultUpgradeStepResult;
 import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.aspect.utils.DefaultAspectsUtil;
 import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.entity.ebean.EbeanAspectV1;
 import com.linkedin.metadata.entity.ebean.EbeanAspectV2;
@@ -170,7 +171,7 @@ public Function<UpgradeContext, UpgradeStepResult> executable() {
             // Emit a browse path aspect.
             final BrowsePaths browsePaths;
             try {
-              browsePaths = _entityService.buildDefaultBrowsePath(urn);
+              browsePaths = DefaultAspectsUtil.buildDefaultBrowsePath(urn, _entityService);
 
               final AuditStamp browsePathsStamp = new AuditStamp();
               browsePathsStamp.setActor(Urn.createFromString(Constants.SYSTEM_ACTOR));

diff --git a/...main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillBrowsePathsV2Step.java b/...main/java/com/linkedin/datahub/upgrade/system/entity/steps/BackfillBrowsePathsV2Step.java
@@ -15,6 +15,7 @@
 import com.linkedin.datahub.upgrade.impl.DefaultUpgradeStepResult;
 import com.linkedin.events.metadata.ChangeType;
 import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.aspect.utils.DefaultAspectsUtil;
 import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.query.SearchFlags;
 import com.linkedin.metadata.query.filter.Condition;
@@ -181,7 +182,8 @@ private Filter backfillDefaultBrowsePathsV2Filter() {
   }
 
   private void ingestBrowsePathsV2(Urn urn, AuditStamp auditStamp) throws Exception {
-    BrowsePathsV2 browsePathsV2 = _entityService.buildDefaultBrowsePathV2(urn, true);
+    BrowsePathsV2 browsePathsV2 =
+        DefaultAspectsUtil.buildDefaultBrowsePathV2(urn, true, _entityService);
     log.debug(String.format("Adding browse path v2 for urn %s with value %s", urn, browsePathsV2));
     MetadataChangeProposal proposal = new MetadataChangeProposal();
     proposal.setEntityUrn(urn);

diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/AspectsBatch.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/AspectsBatch.java
@@ -3,6 +3,7 @@
 import com.linkedin.metadata.aspect.plugins.validation.AspectRetriever;
 import com.linkedin.mxe.SystemMetadata;
 import com.linkedin.util.Pair;
+import java.util.Collection;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -17,14 +18,14 @@
  * SystemMetadata} and record/message created time
  */
 public interface AspectsBatch {
-  List<? extends BatchItem> getItems();
+  Collection<? extends BatchItem> getItems();
 
   /**
    * Returns MCP items. Can be patch, upsert, etc.
    *
    * @return batch items
    */
-  default List<? extends MCPBatchItem> getMCPItems() {
+  default Collection<? extends MCPBatchItem> getMCPItems() {
     return getItems().stream()
         .filter(item -> item instanceof MCPBatchItem)
         .map(item -> (MCPBatchItem) item)

diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/BatchItem.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/BatchItem.java
@@ -2,11 +2,13 @@
 
 import com.linkedin.common.AuditStamp;
 import com.linkedin.common.urn.Urn;
+import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.events.metadata.ChangeType;
 import com.linkedin.metadata.models.AspectSpec;
 import com.linkedin.metadata.models.EntitySpec;
 import com.linkedin.mxe.SystemMetadata;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
 public interface BatchItem {
   /**
@@ -63,4 +65,12 @@ default String getAspectName() {
    */
   @Nonnull
   AspectSpec getAspectSpec();
+
+  /**
+   * The aspect's record template. Null when patch
+   *
+   * @return record template if it exists
+   */
+  @Nullable
+  RecordTemplate getRecordTemplate();
 }
diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCLBatchItem.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCLBatchItem.java
@@ -26,7 +26,7 @@ default String getAspectName() {
     if (getMetadataChangeLog().getAspectName() != null) {
       return getMetadataChangeLog().getAspectName();
     } else {
-      return getAspect().schema().getName();
+      return getRecordTemplate().schema().getName();
     }
   }
 
@@ -40,10 +40,7 @@ default SystemMetadata getPreviousSystemMetadata() {
   }
 
   @Nullable
-  RecordTemplate getPreviousAspect();
-
-  @Nonnull
-  RecordTemplate getAspect();
+  RecordTemplate getPreviousRecordTemplate();
 
   @Override
   @Nonnull

diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/UpsertItem.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/UpsertItem.java
@@ -11,8 +11,6 @@
  * related data stored along with the aspect
  */
 public abstract class UpsertItem extends MCPBatchItem {
-  public abstract RecordTemplate getAspect();
-
   public abstract SystemAspect toLatestEntityAspect();
 
   public abstract void validatePreCommit(

diff --git a/...ta-integration/java/spark-lineage/src/test/java/datahub/spark/TestCoalesceJobLineage.java b/...ta-integration/java/spark-lineage/src/test/java/datahub/spark/TestCoalesceJobLineage.java
@@ -99,7 +99,9 @@ public static void resetBaseExpectations() {
 
   @BeforeClass
   public static void initMockServer() {
-    mockServer = startClientAndServer(GMS_PORT);
+    if (mockServer == null) {
+      mockServer = startClientAndServer(GMS_PORT);
+    }
     resetBaseExpectations();
   }
 

diff --git a/...data-integration/java/spark-lineage/src/test/java/datahub/spark/TestSparkJobsLineage.java b/...data-integration/java/spark-lineage/src/test/java/datahub/spark/TestSparkJobsLineage.java
@@ -138,7 +138,9 @@ public static void resetBaseExpectations() {
 
   @BeforeClass
   public static void init() {
-    mockServer = startClientAndServer(GMS_PORT);
+    if (mockServer == null) {
+      mockServer = startClientAndServer(GMS_PORT);
+    }
     resetBaseExpectations();
   }