feat(sdk): add ContainerKey.parent_key method (datahub-project#12423)

metadata-ingestion/src/datahub/emitter/mcp_builder.py

@@ -90,6 +90,25 @@ def property_dict(self) -> Dict[str, str]:
     def as_urn(self) -> str:
         return make_container_urn(guid=self.guid())
 
+    def parent_key(self) -> Optional["ContainerKey"]:
+        # Find the immediate base class of self.
+        # This is a bit of a hack, but it works.
+        base_classes = self.__class__.__bases__
+        if len(base_classes) != 1:
+            # TODO: Raise a more specific error.
+            raise ValueError(
+                f"Unable to determine parent key for {self.__class__}: {self}"
+            )
+        base_class = base_classes[0]
+        if base_class is DatahubKey or base_class is ContainerKey:
+            return None
+
+        # We need to use `__dict__` instead of `pydantic.BaseModel.dict()`
+        # in order to include "excluded" fields e.g. `backcompat_env_as_instance`.
+        # Tricky: this only works because DatahubKey is a BaseModel and hence
+        # allows extra fields.
+        return base_class(**self.__dict__)
+
 
 # DEPRECATION: Keeping the `PlatformKey` name around for backwards compatibility.
 PlatformKey = ContainerKey

metadata-ingestion/tests/unit/sdk/test_mcp_builder.py

@@ -86,6 +86,70 @@ def test_guid_generators():
     assert guid == guid_datahub
 
 
+def _assert_eq(a: builder.ContainerKey, b: builder.ContainerKey) -> None:
+    assert a == b
+    assert a.__class__ is b.__class__
+    assert a.guid() == b.guid()
+    assert a.property_dict() == b.property_dict()
+
+
+def test_parent_key() -> None:
+    schema_key = builder.SchemaKey(
+        database="test",
+        schema="Test",
+        platform="mysql",
+        instance="TestInstance",
+        env="DEV",
+    )
+    db_key = schema_key.parent_key()
+    assert db_key is not None
+    _assert_eq(
+        db_key,
+        builder.DatabaseKey(
+            database="test",
+            platform="mysql",
+            instance="TestInstance",
+            env="DEV",
+        ),
+    )
+
+    assert db_key.parent_key() is None
+
+
+def test_parent_key_with_backcompat_env_as_instance() -> None:
+    schema_key = builder.SchemaKey(
+        database="test",
+        schema="Test",
+        platform="mysql",
+        instance=None,
+        env="PROD",
+        backcompat_env_as_instance=True,
+    )
+    db_key = schema_key.parent_key()
+    assert db_key is not None
+    _assert_eq(
+        db_key,
+        builder.DatabaseKey(
+            database="test",
+            platform="mysql",
+            instance=None,
+            env="PROD",
+            backcompat_env_as_instance=True,
+        ),
+    )
+
+
+def test_parent_key_on_container_key() -> None:
+    # In general, people shouldn't be calling parent_key() on ContainerKey directly.
+    # But just in case, we should make sure it works.
+    container_key = builder.ContainerKey(
+        platform="bigquery",
+        name="test",
+        env="DEV",
+    )
+    assert container_key.parent_key() is None
+
+
 def test_entity_supports_aspect():
     assert builder.entity_supports_aspect("dataset", StatusClass)
     assert not builder.entity_supports_aspect("telemetry", StatusClass)