From 38ac1007d6a8e1e22291b9f0a1491734dfd44812 Mon Sep 17 00:00:00 2001 From: Shirshanka Das Date: Mon, 14 Oct 2024 01:25:19 -0700 Subject: [PATCH] fix(sdk): platform resource api for non existent resources (#11610) --- .../entities/platformresource/platform_resource.py | 7 +++++++ .../platform_resources/test_platform_resource.py | 13 +++++++++++++ 2 files changed, 20 insertions(+) diff --git a/metadata-ingestion/src/datahub/api/entities/platformresource/platform_resource.py b/metadata-ingestion/src/datahub/api/entities/platformresource/platform_resource.py index 2b730ccb86f513..1556a67a9e5555 100644 --- a/metadata-ingestion/src/datahub/api/entities/platformresource/platform_resource.py +++ b/metadata-ingestion/src/datahub/api/entities/platformresource/platform_resource.py @@ -186,10 +186,17 @@ def to_datahub(self, graph_client: DataHubGraph) -> None: def from_datahub( cls, graph_client: DataHubGraph, key: Union[PlatformResourceKey, str] ) -> Optional["PlatformResource"]: + """ + Fetches a PlatformResource from the graph given a key. + Key can be either a PlatformResourceKey object or an urn string. + Returns None if the resource is not found. + """ if isinstance(key, PlatformResourceKey): urn = PlatformResourceUrn(id=key.id) else: urn = PlatformResourceUrn.from_string(key) + if not graph_client.exists(str(urn)): + return None platform_resource = graph_client.get_entity_semityped(str(urn)) return cls( id=urn.id, diff --git a/smoke-test/tests/platform_resources/test_platform_resource.py b/smoke-test/tests/platform_resources/test_platform_resource.py index 09d25031795728..7c53f72d843c93 100644 --- a/smoke-test/tests/platform_resources/test_platform_resource.py +++ b/smoke-test/tests/platform_resources/test_platform_resource.py @@ -99,3 +99,16 @@ def test_platform_resource_search(graph_client, test_id, cleanup_resources): ] assert len(search_results) == 1 assert search_results[0] == platform_resource + + +def test_platform_resource_non_existent(graph_client, test_id): + key = PlatformResourceKey( + platform=f"test_platform_{test_id}", + resource_type=f"test_resource_type_{test_id}", + primary_key=f"test_primary_key_{test_id}", + ) + platform_resource = PlatformResource.from_datahub( + key=key, + graph_client=graph_client, + ) + assert platform_resource is None