diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/PatchEntityRegistry.java b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/PatchEntityRegistry.java index 382df1084aa5a..ad32b530b784e 100644 --- a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/PatchEntityRegistry.java +++ b/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 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 = diff --git a/entity-registry/src/test/java/com/linkedin/metadata/models/registry/PatchEntityRegistryTest.java b/entity-registry/src/test/java/com/linkedin/metadata/models/registry/PatchEntityRegistryTest.java index 58b322704677a..38664fedb1570 100644 --- a/entity-registry/src/test/java/com/linkedin/metadata/models/registry/PatchEntityRegistryTest.java +++ b/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 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 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)); } } diff --git a/entity-registry/src/test_plugins/mycompany-full-model/0.0.1/entity-registry.yaml b/entity-registry/src/test_plugins/mycompany-full-model/0.0.1/entity-registry.yaml index 2b46962f28ba7..b49e206cdbb21 100644 --- a/entity-registry/src/test_plugins/mycompany-full-model/0.0.1/entity-registry.yaml +++ b/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 \ No newline at end of file + - testDataQualityRules \ No newline at end of file