fix(patch-entity-registry): Remove exception for entities with key aspects. (datahub-project#5831)

* fix(patch-entity-registry): Remove exception for entities with key aspects.

* test(patch-entity-registry): Fix failing unit test for entities with key aspect

* test(patch-entity-registry): Fix checkstyle for PatchEntityRegistryTest.java

* fix(patch-entity-registry): Remove duplicate semicolon.

Co-authored-by: Parham Ghazanfari <[email protected]>

Author: pghazanfari

entity-registry/src/main/java/com/linkedin/metadata/models/registry/PatchEntityRegistry.java

@@ -138,14 +138,16 @@ private PatchEntityRegistry(DataSchemaFactory dataSchemaFactory, InputStream con
           entity.getAspects().stream().collect(Collectors.joining()));
       List<AspectSpec> aspectSpecs = new ArrayList<>();
       if (entity.getKeyAspect() != null) {
-        throw new EntityRegistryException(
-            "Patch Entities cannot define entities yet. They can only enhance an existing entity with additional (non-key) aspects");
-        // aspectSpecs.add(getAspectSpec(entity.getKeyAspect(), entitySpecBuilder));
+        AspectSpec keyAspectSpec = buildAspectSpec(entity.getKeyAspect(), entitySpecBuilder);
+        log.info("Adding key aspect {} with spec {}", entity.getKeyAspect(), keyAspectSpec);
+        aspectSpecs.add(keyAspectSpec);
       }
       entity.getAspects().forEach(aspect -> {
-        AspectSpec aspectSpec = buildAspectSpec(aspect, entitySpecBuilder);
-        log.info("Adding aspect {} with spec {}", aspect, aspectSpec);
-        aspectSpecs.add(aspectSpec);
+        if (!aspect.equals(entity.getKeyAspect())) {
+          AspectSpec aspectSpec = buildAspectSpec(aspect, entitySpecBuilder);
+          log.info("Adding aspect {} with spec {}", aspect, aspectSpec);
+          aspectSpecs.add(aspectSpec);
+        }
       });
 
       EntitySpec entitySpec =

entity-registry/src/test/java/com/linkedin/metadata/models/registry/PatchEntityRegistryTest.java

@@ -1,9 +1,10 @@
 package com.linkedin.metadata.models.registry;
 
+import com.linkedin.metadata.models.DataSchemaFactory;
 import com.linkedin.metadata.models.EntitySpec;
 import com.linkedin.metadata.models.EventSpec;
+import java.nio.file.Paths;
 import java.util.Map;
-import org.apache.maven.artifact.versioning.ComparableVersion;
 import org.testng.annotations.Test;
 
 import static org.testng.Assert.*;
@@ -14,7 +15,11 @@ public class PatchEntityRegistryTest {
   @Test
   public void testEntityRegistryLoad() throws Exception, EntityRegistryException {
     PatchEntityRegistry patchEntityRegistry = new PatchEntityRegistry(
-        TestConstants.BASE_DIRECTORY + "/" + TestConstants.TEST_REGISTRY + "/" + TestConstants.TEST_VERSION.toString(),
+        TestConstants.BASE_DIRECTORY 
+        + "/" 
+        + TestConstants.TEST_REGISTRY 
+        + "/" 
+        + TestConstants.TEST_VERSION.toString(),
         TestConstants.TEST_REGISTRY, TestConstants.TEST_VERSION);
 
     Map<String, EntitySpec> entitySpecs = patchEntityRegistry.getEntitySpecs();
@@ -34,14 +39,28 @@ public void testEntityRegistryLoad() throws Exception, EntityRegistryException {
   }
 
   /**
-   * Validate that patch entity registries cannot have key aspects
+   * Validate that patch entity registries can have key aspects
    * @throws Exception
    * @throws EntityRegistryException
    */
   @Test
-  public void testEntityRegistryWithKeyLoad() {
-    assertThrows(EntityRegistryException.class,
-        () -> new PatchEntityRegistry("src/test_plugins/mycompany-full-model/0.0.1", "mycompany-full-model",
-            new ComparableVersion("0.0.1")));
+  public void testEntityRegistryWithKeyLoad() throws Exception, EntityRegistryException {
+    DataSchemaFactory dataSchemaFactory = DataSchemaFactory.withCustomClasspath(
+      Paths.get(TestConstants.BASE_DIRECTORY 
+        + "/" 
+        + TestConstants.TEST_REGISTRY 
+        + "/" 
+        + TestConstants.TEST_VERSION.toString()));
+
+    PatchEntityRegistry patchEntityRegistry = new PatchEntityRegistry(
+      dataSchemaFactory, Paths.get("src/test_plugins/mycompany-full-model/0.0.1/entity-registry.yaml"),
+            TestConstants.TEST_REGISTRY, TestConstants.TEST_VERSION);
+
+    Map<String, EntitySpec> entitySpecs = patchEntityRegistry.getEntitySpecs();
+    assertEquals(entitySpecs.values().size(), 1);
+    EntitySpec newThingSpec = patchEntityRegistry.getEntitySpec("newThing");
+    assertNotNull(newThingSpec);
+    assertNotNull(newThingSpec.getKeyAspectSpec());
+    assertNotNull(newThingSpec.getAspectSpec(TestConstants.TEST_ASPECT_NAME));
   }
 }

entity-registry/src/test_plugins/mycompany-full-model/0.0.1/entity-registry.yaml

@@ -3,6 +3,6 @@
 id: mycompany-full-model
 entities:
   - name: newThing
-    keyAspect: ownership
+    keyAspect: testDataQualityRules
     aspects:
-      - ownership
+      - testDataQualityRules