From 2978cdb409ca689bf04d3757f9e4935162584f37 Mon Sep 17 00:00:00 2001
From: John Joyce <john@acryl.io>
Date: Mon, 20 Mar 2023 21:14:09 -0700
Subject: [PATCH] fix(entity registry): Fix patching aspects onto existing
 Config based entity  (#7624)

---
 .../metadata/models/registry/MergedEntityRegistry.java    | 8 ++++++++
 metadata-models-custom/registry/entity-registry.yaml      | 3 +++
 2 files changed, 11 insertions(+)

diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/MergedEntityRegistry.java b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/MergedEntityRegistry.java
index 91439f1dca..f0ec57b8d8 100644
--- a/entity-registry/src/main/java/com/linkedin/metadata/models/registry/MergedEntityRegistry.java
+++ b/entity-registry/src/main/java/com/linkedin/metadata/models/registry/MergedEntityRegistry.java
@@ -4,6 +4,7 @@ import com.linkedin.data.schema.compatibility.CompatibilityChecker;
 import com.linkedin.data.schema.compatibility.CompatibilityOptions;
 import com.linkedin.data.schema.compatibility.CompatibilityResult;
 import com.linkedin.metadata.models.AspectSpec;
+import com.linkedin.metadata.models.ConfigEntitySpec;
 import com.linkedin.metadata.models.DefaultEntitySpec;
 import com.linkedin.metadata.models.EntitySpec;
 import com.linkedin.metadata.models.EventSpec;
@@ -108,6 +109,13 @@ public class MergedEntityRegistry implements EntityRegistry {
   private EntitySpec mergeEntitySpecs(EntitySpec existingEntitySpec, EntitySpec newEntitySpec) {
     Map<String, AspectSpec> aspectSpecMap = new HashMap<>(existingEntitySpec.getAspectSpecMap());
     aspectSpecMap.putAll(newEntitySpec.getAspectSpecMap());
+    // If the base is a config spec, always create another config spec.
+    if (existingEntitySpec instanceof ConfigEntitySpec) {
+      return new ConfigEntitySpec(
+          existingEntitySpec.getEntityAnnotation().getName(),
+          existingEntitySpec.getEntityAnnotation().getKeyAspect(),
+          aspectSpecMap.values());
+    }
     return new DefaultEntitySpec(aspectSpecMap.values(), existingEntitySpec.getEntityAnnotation(),
         existingEntitySpec.getSnapshotSchema(), existingEntitySpec.getAspectTyperefSchema());
   }
diff --git a/metadata-models-custom/registry/entity-registry.yaml b/metadata-models-custom/registry/entity-registry.yaml
index 5e281e4af1..2b501946ca 100644
--- a/metadata-models-custom/registry/entity-registry.yaml
+++ b/metadata-models-custom/registry/entity-registry.yaml
@@ -3,3 +3,6 @@ entities:
   - name: dataset
     aspects:
       - customDataQualityRules
+  - name: container
+    aspects:
+      - customDataQualityRules
\ No newline at end of file