feat(ingest): adds get_entities_v3 method to DataHubGraph (#13045)

Author: sgomezvillamor

metadata-ingestion/examples/datahubgraph/batch_get_entities.py

@@ -0,0 +1,75 @@
+import argparse
+
+from datahub.ingestion.graph.client import DataHubGraph
+from datahub.ingestion.graph.config import DatahubClientConfig
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Fetch entities from DataHub using get_entities"
+    )
+    parser.add_argument("--token", required=False, help="DataHub access token")
+    parser.add_argument(
+        "--server_url",
+        required=False,
+        default="http://localhost:8080",
+        help="DataHub server URL (defaults to http://localhost:8080)",
+    )
+    parser.add_argument(
+        "--entity_name",
+        required=True,
+        help="Entity type name (e.g., dataset, dashboard, chart)",
+    )
+    parser.add_argument(
+        "--urn",
+        required=True,
+        action="append",
+        dest="urns",
+        help="Entity URN(s) to fetch. Can specify multiple times.",
+    )
+    parser.add_argument(
+        "--aspect",
+        action="append",
+        dest="aspects",
+        help="Aspect name(s) to fetch. Can specify multiple times. If none provided, all aspects will be fetched.",
+    )
+    parser.add_argument(
+        "--with-system-metadata",
+        action="store_true",
+        help="Include system metadata in the response.",
+    )
+    args = parser.parse_args()
+
+    # Validate that at least one URN is provided
+    if not args.urns:
+        parser.error("At least one --urn argument is required")
+
+    client = DataHubGraph(
+        config=DatahubClientConfig(
+            server=args.server_url,
+            token=args.token,
+        )
+    )
+
+    response = client.get_entities(
+        entity_name=args.entity_name,
+        urns=args.urns,
+        aspects=args.aspects,
+        with_system_metadata=args.with_system_metadata,
+    )
+
+    print(f"Received {len(response)} entities")
+
+    for urn, entity in response.items():
+        print(f"Entity: {urn}")
+
+        if not entity:
+            print("\tNo aspects found for this entity")
+            continue
+
+        for aspect_name, (aspect, system_metadata) in entity.items():
+            print(f"\tAspect: {aspect_name} Type: {type(aspect).__name__}")
+            print(f"\t\t{aspect}")
+            if system_metadata:
+                print(f"\tSystem Metadata: {system_metadata}")
+
+        print()

metadata-ingestion/src/datahub/ingestion/graph/client.py

@@ -27,6 +27,7 @@
 from requests.models import HTTPError
 from typing_extensions import deprecated
 
+from datahub._codegen.aspect import _Aspect
 from datahub.cli import config_utils
 from datahub.configuration.common import ConfigModel, GraphError, OperationalError
 from datahub.emitter.aspect import TIMESERIES_ASPECT_MAP
@@ -1697,6 +1698,7 @@ def run_assertions_for_asset(
 
         return res["runAssertionsForAsset"]
 
+    @deprecated("Use get_entities instead which returns typed aspects")
     def get_entities_v2(
         self,
         entity_name: str,
@@ -1736,6 +1738,108 @@ def get_entities_v2(
                     retval[entity_urn][aspect_key] = aspect_value
         return retval
 
+    def get_entities(
+        self,
+        entity_name: str,
+        urns: List[str],
+        aspects: Optional[List[str]] = None,
+        with_system_metadata: bool = False,
+    ) -> Dict[str, Dict[str, Tuple[_Aspect, Optional[SystemMetadataClass]]]]:
+        """
+        Get entities using the OpenAPI v3 endpoint, deserializing aspects into typed objects.
+
+        Args:
+            entity_name: The entity type name
+            urns: List of entity URNs to fetch
+            aspects: Optional list of aspect names to fetch. If None, all aspects will be fetched.
+            with_system_metadata: If True, return system metadata along with each aspect.
+
+        Returns:
+            A dictionary mapping URNs to a dictionary of aspect name to tuples of
+            (typed aspect object, system metadata). If with_system_metadata is False,
+            the system metadata in the tuple will be None.
+        """
+        aspects = aspects or []
+
+        request_payload = []
+        for urn in urns:
+            entity_request: Dict[str, Any] = {"urn": urn}
+            for aspect_name in aspects:
+                entity_request[aspect_name] = {}
+            request_payload.append(entity_request)
+
+        headers: Dict[str, Any] = {
+            "Accept": "application/json",
+            "Content-Type": "application/json",
+        }
+
+        url = f"{self.config.server}/openapi/v3/entity/{entity_name}/batchGet"
+        if with_system_metadata:
+            url += "?systemMetadata=true"
+
+        response = self._session.post(
+            url, data=json.dumps(request_payload), headers=headers
+        )
+        response.raise_for_status()
+        entities = response.json()
+
+        result: Dict[str, Dict[str, Tuple[_Aspect, Optional[SystemMetadataClass]]]] = {}
+
+        for entity in entities:
+            entity_urn = entity.get("urn")
+            if entity_urn is None:
+                logger.warning(
+                    f"Missing URN in entity response: {entity}, skipping deserialization"
+                )
+                continue
+
+            entity_aspects: Dict[
+                str, Tuple[_Aspect, Optional[SystemMetadataClass]]
+            ] = {}
+
+            for aspect_name, aspect_obj in entity.items():
+                if aspect_name == "urn":
+                    continue
+
+                aspect_class = ASPECT_NAME_MAP.get(aspect_name)
+                if aspect_class is None:
+                    logger.warning(
+                        f"Unknown aspect type {aspect_name}, skipping deserialization"
+                    )
+                    continue
+
+                aspect_value = aspect_obj.get("value")
+                if aspect_value is None:
+                    logger.warning(
+                        f"Unknown aspect value for aspect {aspect_name}, skipping deserialization"
+                    )
+                    continue
+
+                try:
+                    post_json_obj = post_json_transform(aspect_value)
+                    typed_aspect = aspect_class.from_obj(post_json_obj)
+                    assert isinstance(typed_aspect, aspect_class) and isinstance(
+                        typed_aspect, _Aspect
+                    )
+
+                    system_metadata = None
+                    if with_system_metadata:
+                        system_metadata_obj = aspect_obj.get("systemMetadata")
+                        if system_metadata_obj:
+                            system_metadata = SystemMetadataClass.from_obj(
+                                system_metadata_obj
+                            )
+
+                    entity_aspects[aspect_name] = (typed_aspect, system_metadata)
+                except Exception as e:
+                    logger.error(f"Error deserializing aspect {aspect_name}: {e}")
+                    raise
+
+            if entity_aspects:
+                result[entity_urn] = entity_aspects
+
+        return result
+
     def upsert_custom_assertion(
         self,
         urn: Optional[str],

metadata-ingestion/tests/test_helpers/graph_helpers.py

@@ -1,7 +1,8 @@
 from datetime import timedelta
 from pathlib import Path
-from typing import Any, Callable, Dict, Iterable, List, Optional, Type, Union
+from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Type, Union
 
+from datahub._codegen.aspect import _Aspect
 from datahub.emitter.mce_builder import Aspect
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import mcps_from_mce
@@ -17,6 +18,7 @@
 from datahub.metadata.schema_classes import (
     ASPECT_NAME_MAP,
     DomainPropertiesClass,
+    SystemMetadataClass,
     UsageAggregationClass,
 )
 
@@ -128,6 +130,27 @@ def emit_mcp(
     ) -> None:
         self.emitted.append(mcp)
 
+    def get_entities(
+        self,
+        entity_name: str,
+        urns: List[str],
+        aspects: Optional[List[str]] = None,
+        with_system_metadata: bool = False,
+    ) -> Dict[str, Dict[str, Tuple[_Aspect, Optional[SystemMetadataClass]]]]:
+        result: Dict[str, Dict[str, Tuple[_Aspect, Optional[SystemMetadataClass]]]] = {}
+        for urn, entity in self.entity_graph.items():
+            if urn not in urns:
+                continue
+            if urn not in result:
+                result[urn] = {}
+            for aspect_name, aspect in entity.items():
+                if aspects and aspect_name not in aspects:
+                    continue
+                # Mock implementation always returns None for system metadata
+                system_metadata = None
+                result[urn][aspect_name] = (aspect, system_metadata)
+        return result
+
     def get_emitted(
         self,
     ) -> List[

metadata-ingestion/tests/test_helpers/test_graph_helpers.py

@@ -0,0 +1,57 @@
+import pytest
+
+from datahub.metadata.schema_classes import DatasetPropertiesClass
+from tests.test_helpers.graph_helpers import MockDataHubGraph
+
+
+@pytest.fixture
+def mock_datahub_graph(pytestconfig):
+    graph = MockDataHubGraph()
+    demo_data_path = pytestconfig.rootpath / "examples" / "demo_data" / "demo_data.json"
+    graph.import_file(demo_data_path)
+    return graph
+
+
+def test_get_entities(mock_datahub_graph):
+    aspect_name = "datasetProperties"
+    urn_1 = "urn:li:dataset:(urn:li:dataPlatform:bigquery,bigquery-public-data.covid19_aha.hospital_beds,PROD)"
+    urn_2 = "urn:li:dataset:(urn:li:dataPlatform:bigquery,bigquery-public-data.covid19_geotab_mobility_impact.city_congestion,PROD)"
+
+    entities = mock_datahub_graph.get_entities(
+        entity_name="dataset",
+        urns=[
+            urn_1,  # misses datasetProperties aspect
+            urn_2,  # has datasetProperties aspect
+        ],
+        aspects=[aspect_name],
+    )
+
+    # Verify results
+    assert entities
+    assert (
+        urn_1 in entities
+        and aspect_name not in entities[urn_1]
+        and entities[urn_1] == {}
+    )
+    assert (
+        urn_2 in entities
+        and aspect_name in entities[urn_2]
+        and isinstance(entities[urn_2][aspect_name][0], DatasetPropertiesClass)
+        and entities[urn_2][aspect_name][1] is None
+    )
+
+    # Test system_metadata is always None regardless of flag
+    entities_with_metadata = mock_datahub_graph.get_entities(
+        entity_name="dataset",
+        urns=[urn_2],
+        aspects=[aspect_name],
+        with_system_metadata=True,
+    )
+    assert (
+        urn_2 in entities_with_metadata
+        and aspect_name in entities_with_metadata[urn_2]
+        and isinstance(
+            entities_with_metadata[urn_2][aspect_name][0], DatasetPropertiesClass
+        )
+        and entities_with_metadata[urn_2][aspect_name][1] is None
+    )

smoke-test/tests/cli/datahub_graph_test.py

@@ -4,7 +4,13 @@
 import tenacity
 
 from datahub.ingestion.graph.client import DataHubGraph
-from datahub.metadata.schema_classes import KafkaSchemaClass, SchemaMetadataClass
+from datahub.metadata.schema_classes import (
+    DatasetPropertiesClass,
+    KafkaSchemaClass,
+    OwnershipClass,
+    SchemaMetadataClass,
+    SystemMetadataClass,
+)
 from tests.utils import delete_urns_from_file, get_sleep_info, ingest_file_via_rest
 
 sleep_sec, sleep_times = get_sleep_info()
@@ -41,6 +47,64 @@ def test_get_aspect_v2(graph_client, ingest_cleanup_data):
     )
 
 
+def test_get_entities_v3(graph_client, ingest_cleanup_data):
+    ownership_aspect_name = "ownership"
+    dataset_properties_aspect_name = "datasetProperties"
+    urn = "urn:li:dataset:(urn:li:dataPlatform:kafka,test-rollback,PROD)"
+    entities = graph_client.get_entities(
+        entity_name="dataset",
+        urns=[urn],
+        aspects=[ownership_aspect_name, dataset_properties_aspect_name],
+    )
+
+    assert entities
+    assert len(entities) == 1 and urn in entities
+    assert (
+        len(entities[urn]) == 2
+        and ownership_aspect_name in entities[urn]
+        and dataset_properties_aspect_name in entities[urn]
+        and isinstance(entities[urn][ownership_aspect_name][0], OwnershipClass)
+        and isinstance(
+            entities[urn][dataset_properties_aspect_name][0], DatasetPropertiesClass
+        )
+        and entities[urn][ownership_aspect_name][1] is None
+        and entities[urn][dataset_properties_aspect_name][1] is None
+    )
+    assert {
+        owner.owner for owner in entities[urn][ownership_aspect_name][0].owners
+    } == {
+        "urn:li:corpuser:datahub",
+        "urn:li:corpuser:jdoe",
+    }
+    assert not entities[urn][dataset_properties_aspect_name][0].description
+    assert entities[urn][dataset_properties_aspect_name][0].customProperties == {
+        "prop1": "fakeprop",
+        "prop2": "pikachu",
+    }
+
+    # Test with system metadata
+    entities_with_metadata = graph_client.get_entities(
+        entity_name="dataset",
+        urns=[urn],
+        aspects=[ownership_aspect_name],
+        with_system_metadata=True,
+    )
+
+    assert entities_with_metadata
+    assert len(entities_with_metadata) == 1 and urn in entities_with_metadata
+    assert (
+        ownership_aspect_name in entities_with_metadata[urn]
+        and entities_with_metadata[urn][ownership_aspect_name][0]
+        and isinstance(
+            entities_with_metadata[urn][ownership_aspect_name][0], OwnershipClass
+        )
+        and entities_with_metadata[urn][ownership_aspect_name][1]
+        and isinstance(
+            entities_with_metadata[urn][ownership_aspect_name][1], SystemMetadataClass
+        )
+    )
+
+
 @tenacity.retry(
     stop=tenacity.stop_after_attempt(sleep_times), wait=tenacity.wait_fixed(sleep_sec)
 )