fix(ingest): support multiple entity_registry sections (#7066)

datahub-project · Jan 18, 2023 · f39ed52 · f39ed52
1 parent a3f4c40
commit f39ed52
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/metadata-ingestion/scripts/avro_codegen.py b/metadata-ingestion/scripts/avro_codegen.py
@@ -272,7 +272,13 @@ def generate(entity_registry: str, schemas_path: str, outdir: str) -> None:
         aspect = aspects[entity.keyAspect]
 
         # This requires that entities cannot share a keyAspect.
-        assert "keyForEntity" not in aspect["Aspect"]
+        if (
+            "keyForEntity" in aspect["Aspect"]
+            and aspect["Aspect"]["keyForEntity"] != entity.name
+        ):
+            raise ValueError(
+                f'Entity key {entity.keyAspect} is used by {aspect["Aspect"]["keyForEntity"]} and {entity.name}'
+            )
 
         aspect["Aspect"]["keyForEntity"] = entity.name
         aspect["Aspect"]["entityCategory"] = entity.category