fix(ingest): lib - revert positional arg change, create new method (d…

…atahub-project#4266) Co-authored-by: Shirshanka Das <[email protected]>
maggiehays · Aug 1, 2022 · 0d7054e · 0d7054e
1 parent 09d27ef
commit 0d7054e
Show file tree

Hide file tree

Showing 8 changed files with 41 additions and 13 deletions.
diff --git a/metadata-ingestion/examples/library/dataset_add_column_tag.py b/metadata-ingestion/examples/library/dataset_add_column_tag.py
@@ -48,7 +48,7 @@ def get_simple_field_path_from_v2_field_path(field_path: str) -> str:
 graph = DataHubGraph(DatahubClientConfig(server=gms_endpoint))
 
 
-current_editable_schema_metadata = graph.get_aspect(
+current_editable_schema_metadata = graph.get_aspect_v2(
     entity_urn=dataset_urn,
     aspect="editableSchemaMetadata",
     aspect_type=EditableSchemaMetadataClass,

diff --git a/metadata-ingestion/examples/library/dataset_add_column_term.py b/metadata-ingestion/examples/library/dataset_add_column_term.py
@@ -48,7 +48,7 @@ def get_simple_field_path_from_v2_field_path(field_path: str) -> str:
 graph = DataHubGraph(DatahubClientConfig(server=gms_endpoint))
 
 
-current_editable_schema_metadata = graph.get_aspect(
+current_editable_schema_metadata = graph.get_aspect_v2(
     entity_urn=dataset_urn,
     aspect="editableSchemaMetadata",
     aspect_type=EditableSchemaMetadataClass,

diff --git a/metadata-ingestion/examples/library/dataset_add_documentation.py b/metadata-ingestion/examples/library/dataset_add_documentation.py
@@ -40,7 +40,7 @@
 gms_endpoint = "http://localhost:8080"
 graph = DataHubGraph(config=DatahubClientConfig(server=gms_endpoint))
 
-current_editable_properties = graph.get_aspect(
+current_editable_properties = graph.get_aspect_v2(
     entity_urn=dataset_urn,
     aspect="editableDatasetProperties",
     aspect_type=EditableDatasetPropertiesClass,
@@ -73,7 +73,7 @@
     log.info("Documentation already exists and is identical, omitting write")
 
 
-current_institutional_memory = graph.get_aspect(
+current_institutional_memory = graph.get_aspect_v2(
     entity_urn=dataset_urn,
     aspect="institutionalMemory",
     aspect_type=InstitutionalMemoryClass,

diff --git a/metadata-ingestion/examples/library/dataset_add_owner.py b/metadata-ingestion/examples/library/dataset_add_owner.py
@@ -34,7 +34,7 @@
 graph = DataHubGraph(DatahubClientConfig(server=gms_endpoint))
 
 
-current_owners: Optional[OwnershipClass] = graph.get_aspect(
+current_owners: Optional[OwnershipClass] = graph.get_aspect_v2(
     entity_urn=dataset_urn,
     aspect="ownership",
     aspect_type=OwnershipClass,

diff --git a/metadata-ingestion/examples/library/dataset_add_tag.py b/metadata-ingestion/examples/library/dataset_add_tag.py
@@ -24,7 +24,7 @@
 
 dataset_urn = make_dataset_urn(platform="hive", name="realestate_db.sales", env="PROD")
 
-current_tags: Optional[GlobalTagsClass] = graph.get_aspect(
+current_tags: Optional[GlobalTagsClass] = graph.get_aspect_v2(
     entity_urn=dataset_urn,
     aspect="globalTags",
     aspect_type=GlobalTagsClass,

diff --git a/metadata-ingestion/examples/library/dataset_add_term.py b/metadata-ingestion/examples/library/dataset_add_term.py
@@ -25,7 +25,7 @@
 
 dataset_urn = make_dataset_urn(platform="hive", name="realestate_db.sales", env="PROD")
 
-current_terms: Optional[GlossaryTermsClass] = graph.get_aspect(
+current_terms: Optional[GlossaryTermsClass] = graph.get_aspect_v2(
     entity_urn=dataset_urn,
     aspect="glossaryTerms",
     aspect_type=GlossaryTermsClass,

diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py
@@ -52,6 +52,8 @@ def get_long_description():
     # Pinning it to a version which works even though we are not using explicitly
     # https://github.com/aws/aws-sam-cli/issues/3661
     "markupsafe==2.0.1",
+    "Deprecated",
+    "types-Deprecated",
 }
 
 kafka_common = {

diff --git a/metadata-ingestion/src/datahub/ingestion/graph/client.py b/metadata-ingestion/src/datahub/ingestion/graph/client.py
@@ -5,6 +5,7 @@
 from typing import Any, Dict, List, Optional, Type
 
 from avro.schema import RecordSchema
+from deprecated import deprecated
 from requests.adapters import Response
 from requests.models import HTTPError
 
@@ -90,13 +91,41 @@ def _guess_entity_type(urn: str) -> str:
         assert urn.startswith("urn:li:"), "urns must start with urn:li:"
         return urn.split(":")[2]
 
+    @deprecated(
+        reason="Use get_aspect_v2 instead which makes aspect_type_name truly optional"
+    )
     def get_aspect(
         self,
         entity_urn: str,
         aspect: str,
+        aspect_type_name: Optional[str],
         aspect_type: Type[Aspect],
+    ) -> Optional[Aspect]:
+        return self.get_aspect_v2(
+            entity_urn=entity_urn,
+            aspect=aspect,
+            aspect_type=aspect_type,
+            aspect_type_name=aspect_type_name,
+        )
+
+    def get_aspect_v2(
+        self,
+        entity_urn: str,
+        aspect_type: Type[Aspect],
+        aspect: str,
         aspect_type_name: Optional[str] = None,
     ) -> Optional[Aspect]:
+        """
+        Get an aspect for an entity.
+
+        :param str entity_urn: The urn of the entity
+        :param Type[Aspect] aspect_type: The type class of the aspect being requested (e.g. datahub.metadata.schema_classes.DatasetProperties)
+        :param str aspect: The name of the aspect being requested (e.g. schemaMetadata, datasetProperties, etc.)
+        :param Optional[str] aspect_type_name: The fully qualified classname of the aspect being requested. Typically not needed and extracted automatically from the class directly. (e.g. com.linkedin.common.DatasetProperties)
+        :return: the Aspect as a dictionary if present, None if no aspect was found (HTTP status 404)
+        :rtype: Optional[Aspect]
+        :raises HttpError: if the HTTP response is not a 200 or a 404
+        """
         url = f"{self._gms_server}/aspects/{urllib.parse.quote(entity_urn)}?aspect={aspect}&version=0"
         response = self._session.get(url)
         if response.status_code == 404:
@@ -126,26 +155,23 @@ def get_config(self) -> Dict[str, Any]:
         return self._get_generic(f"{self.config.server}/config")
 
     def get_ownership(self, entity_urn: str) -> Optional[OwnershipClass]:
-        return self.get_aspect(
+        return self.get_aspect_v2(
             entity_urn=entity_urn,
             aspect="ownership",
-            aspect_type_name="com.linkedin.common.Ownership",
             aspect_type=OwnershipClass,
         )
 
     def get_tags(self, entity_urn: str) -> Optional[GlobalTagsClass]:
-        return self.get_aspect(
+        return self.get_aspect_v2(
             entity_urn=entity_urn,
             aspect="globalTags",
-            aspect_type_name="com.linkedin.common.GlobalTags",
             aspect_type=GlobalTagsClass,
         )
 
     def get_glossary_terms(self, entity_urn: str) -> Optional[GlossaryTermsClass]:
-        return self.get_aspect(
+        return self.get_aspect_v2(
             entity_urn=entity_urn,
             aspect="glossaryTerms",
-            aspect_type_name="com.linkedin.common.GlossaryTerms",
             aspect_type=GlossaryTermsClass,
         )