refactor terms patch

datahub-project · Dec 31, 2024 · f7dfa64 · f7dfa64
1 parent 68573cb
commit f7dfa64
Show file tree

Hide file tree

Showing 6 changed files with 70 additions and 139 deletions.
diff --git a/metadata-ingestion/src/datahub/specific/aspect_helpers/terms.py b/metadata-ingestion/src/datahub/specific/aspect_helpers/terms.py
@@ -0,0 +1,43 @@
+from typing import Union
+
+from typing_extensions import Self
+
+from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
+from datahub.metadata.schema_classes import (
+    GlossaryTermAssociationClass as Term,
+    GlossaryTermsClass,
+)
+from datahub.metadata.urns import GlossaryTermUrn, Urn
+
+
+class HasTermsPatch(MetadataPatchProposal):
+    def add_term(self, term: Term) -> Self:
+        """Adds a glossary term to the entity.
+
+        Args:
+            term: The Term object representing the glossary term to be added.
+
+        Returns:
+            The patch builder instance.
+        """
+        # TODO: Make this support raw strings, in addition to Term objects.
+        self._add_patch(
+            GlossaryTermsClass.ASPECT_NAME, "add", path=("terms", term.urn), value=term
+        )
+        return self
+
+    def remove_term(self, term: Union[str, Urn]) -> Self:
+        """Removes a glossary term from the entity.
+
+        Args:
+            term: The term to remove, specified as a string or Urn object.
+
+        Returns:
+            The patch builder instance.
+        """
+        if isinstance(term, str) and not term.startswith("urn:li:glossaryTerm:"):
+            term = GlossaryTermUrn(term)
+        self._add_patch(
+            GlossaryTermsClass.ASPECT_NAME, "remove", path=("terms", term), value={}
+        )
+        return self
diff --git a/metadata-ingestion/src/datahub/specific/chart.py b/metadata-ingestion/src/datahub/specific/chart.py
@@ -7,19 +7,22 @@
     ChartInfoClass as ChartInfo,
     ChartTypeClass,
     EdgeClass as Edge,
-    GlossaryTermAssociationClass as Term,
-    GlossaryTermsClass as GlossaryTerms,
     KafkaAuditHeaderClass,
     SystemMetadataClass,
 )
 from datahub.specific.aspect_helpers.custom_properties import HasCustomPropertiesPatch
 from datahub.specific.aspect_helpers.ownership import HasOwnershipPatch
 from datahub.specific.aspect_helpers.tags import HasTagsPatch
+from datahub.specific.aspect_helpers.terms import HasTermsPatch
 from datahub.utilities.urns.urn import Urn
 
 
 class ChartPatchBuilder(
-    HasOwnershipPatch, HasCustomPropertiesPatch, HasTagsPatch, MetadataPatchProposal
+    HasOwnershipPatch,
+    HasCustomPropertiesPatch,
+    HasTagsPatch,
+    HasTermsPatch,
+    MetadataPatchProposal,
 ):
     def __init__(
         self,
@@ -117,38 +120,6 @@ def set_input_edges(self, inputs: List[Edge]) -> "ChartPatchBuilder":
         )
         return self
 
-    def add_term(self, term: Term) -> "ChartPatchBuilder":
-        """
-        Adds a glossary term to the ChartPatchBuilder.
-
-        Args:
-            term: The Term object representing the glossary term to be added.
-
-        Returns:
-            The ChartPatchBuilder instance.
-        """
-        self._add_patch(
-            GlossaryTerms.ASPECT_NAME, "add", path=("terms", term.urn), value=term
-        )
-        return self
-
-    def remove_term(self, term: Union[str, Urn]) -> "ChartPatchBuilder":
-        """
-        Removes a glossary term from the ChartPatchBuilder.
-
-        Args:
-            term: The term to remove, specified as a string or Urn object.
-
-        Returns:
-            The ChartPatchBuilder instance.
-        """
-        if isinstance(term, str) and not term.startswith("urn:li:glossaryTerm:"):
-            term = "urn:li:glossaryTerm:" + term
-        self._add_patch(
-            GlossaryTerms.ASPECT_NAME, "remove", path=("terms", term), value={}
-        )
-        return self
-
     def set_title(self, title: str) -> "ChartPatchBuilder":
         assert title, "ChartInfo title should not be None"
         self._add_patch(

diff --git a/metadata-ingestion/src/datahub/specific/dashboard.py b/metadata-ingestion/src/datahub/specific/dashboard.py
@@ -6,19 +6,22 @@
     ChangeAuditStampsClass,
     DashboardInfoClass as DashboardInfo,
     EdgeClass as Edge,
-    GlossaryTermAssociationClass as Term,
-    GlossaryTermsClass as GlossaryTerms,
     KafkaAuditHeaderClass,
     SystemMetadataClass,
 )
 from datahub.specific.aspect_helpers.custom_properties import HasCustomPropertiesPatch
 from datahub.specific.aspect_helpers.ownership import HasOwnershipPatch
 from datahub.specific.aspect_helpers.tags import HasTagsPatch
+from datahub.specific.aspect_helpers.terms import HasTermsPatch
 from datahub.utilities.urns.urn import Urn
 
 
 class DashboardPatchBuilder(
-    HasOwnershipPatch, HasCustomPropertiesPatch, HasTagsPatch, MetadataPatchProposal
+    HasOwnershipPatch,
+    HasCustomPropertiesPatch,
+    HasTagsPatch,
+    HasTermsPatch,
+    MetadataPatchProposal,
 ):
     def __init__(
         self,
@@ -210,38 +213,6 @@ def set_chart_edges(self, charts: List[Edge]) -> "DashboardPatchBuilder":
         )
         return self
 
-    def add_term(self, term: Term) -> "DashboardPatchBuilder":
-        """
-        Adds a glossary term to the DashboardPatchBuilder.
-
-        Args:
-            term: The Term object representing the glossary term to be added.
-
-        Returns:
-            The DashboardPatchBuilder instance.
-        """
-        self._add_patch(
-            GlossaryTerms.ASPECT_NAME, "add", path=("terms", term.urn), value=term
-        )
-        return self
-
-    def remove_term(self, term: Union[str, Urn]) -> "DashboardPatchBuilder":
-        """
-        Removes a glossary term from the DashboardPatchBuilder.
-
-        Args:
-            term: The term to remove, specified as a string or Urn object.
-
-        Returns:
-            The DashboardPatchBuilder instance.
-        """
-        if isinstance(term, str) and not term.startswith("urn:li:glossaryTerm:"):
-            term = "urn:li:glossaryTerm:" + term
-        self._add_patch(
-            GlossaryTerms.ASPECT_NAME, "remove", path=("terms", term), value={}
-        )
-        return self
-
     def set_title(self, title: str) -> "DashboardPatchBuilder":
         assert title, "DashboardInfo title should not be None"
         self._add_patch(

diff --git a/metadata-ingestion/src/datahub/specific/datajob.py b/metadata-ingestion/src/datahub/specific/datajob.py
@@ -5,19 +5,22 @@
     DataJobInfoClass as DataJobInfo,
     DataJobInputOutputClass as DataJobInputOutput,
     EdgeClass as Edge,
-    GlossaryTermAssociationClass as Term,
-    GlossaryTermsClass as GlossaryTerms,
     KafkaAuditHeaderClass,
     SystemMetadataClass,
 )
 from datahub.metadata.urns import SchemaFieldUrn, Urn
 from datahub.specific.aspect_helpers.custom_properties import HasCustomPropertiesPatch
 from datahub.specific.aspect_helpers.ownership import HasOwnershipPatch
 from datahub.specific.aspect_helpers.tags import HasTagsPatch
+from datahub.specific.aspect_helpers.terms import HasTermsPatch
 
 
 class DataJobPatchBuilder(
-    HasOwnershipPatch, HasCustomPropertiesPatch, HasTagsPatch, MetadataPatchProposal
+    HasOwnershipPatch,
+    HasCustomPropertiesPatch,
+    HasTagsPatch,
+    HasTermsPatch,
+    MetadataPatchProposal,
 ):
     def __init__(
         self,
@@ -427,35 +430,3 @@ def set_output_dataset_fields(self, outputs: List[Edge]) -> "DataJobPatchBuilder
             value=outputs,
         )
         return self
-
-    def add_term(self, term: Term) -> "DataJobPatchBuilder":
-        """
-        Adds a glossary term to the DataJobPatchBuilder.
-
-        Args:
-            term: The Term object representing the glossary term to be added.
-
-        Returns:
-            The DataJobPatchBuilder instance.
-        """
-        self._add_patch(
-            GlossaryTerms.ASPECT_NAME, "add", path=("terms", term.urn), value=term
-        )
-        return self
-
-    def remove_term(self, term: Union[str, Urn]) -> "DataJobPatchBuilder":
-        """
-        Removes a glossary term from the DataJobPatchBuilder.
-
-        Args:
-            term: The term to remove, specified as a string or Urn object.
-
-        Returns:
-            The DataJobPatchBuilder instance.
-        """
-        if isinstance(term, str) and not term.startswith("urn:li:glossaryTerm:"):
-            term = "urn:li:glossaryTerm:" + term
-        self._add_patch(
-            GlossaryTerms.ASPECT_NAME, "remove", path=("terms", term), value={}
-        )
-        return self
diff --git a/metadata-ingestion/src/datahub/specific/dataproduct.py b/metadata-ingestion/src/datahub/specific/dataproduct.py
@@ -1,22 +1,24 @@
-from typing import List, Optional, Tuple, Union
+from typing import List, Optional, Tuple
 
 from datahub.emitter.mcp_patch_builder import MetadataPatchProposal, PatchPath
 from datahub.metadata.schema_classes import (
     DataProductAssociationClass as DataProductAssociation,
     DataProductPropertiesClass as DataProductProperties,
-    GlossaryTermAssociationClass as Term,
-    GlossaryTermsClass as GlossaryTerms,
     KafkaAuditHeaderClass,
     SystemMetadataClass,
 )
 from datahub.specific.aspect_helpers.custom_properties import HasCustomPropertiesPatch
 from datahub.specific.aspect_helpers.ownership import HasOwnershipPatch
 from datahub.specific.aspect_helpers.tags import HasTagsPatch
-from datahub.utilities.urns.urn import Urn
+from datahub.specific.aspect_helpers.terms import HasTermsPatch
 
 
 class DataProductPatchBuilder(
-    HasOwnershipPatch, HasCustomPropertiesPatch, HasTagsPatch, MetadataPatchProposal
+    HasOwnershipPatch,
+    HasCustomPropertiesPatch,
+    HasTagsPatch,
+    HasTermsPatch,
+    MetadataPatchProposal,
 ):
     def __init__(
         self,
@@ -34,20 +36,6 @@ def __init__(
     def _custom_properties_location(cls) -> Tuple[str, PatchPath]:
         return DataProductProperties.ASPECT_NAME, ("customProperties",)
 
-    def add_term(self, term: Term) -> "DataProductPatchBuilder":
-        self._add_patch(
-            GlossaryTerms.ASPECT_NAME, "add", path=("terms", term.urn), value=term
-        )
-        return self
-
-    def remove_term(self, term: Union[str, Urn]) -> "DataProductPatchBuilder":
-        if isinstance(term, str) and not term.startswith("urn:li:glossaryTerm:"):
-            term = "urn:li:glossaryTerm:" + term
-        self._add_patch(
-            GlossaryTerms.ASPECT_NAME, "remove", path=("terms", term), value={}
-        )
-        return self
-
     def set_name(self, name: str) -> "DataProductPatchBuilder":
         self._add_patch(
             DataProductProperties.ASPECT_NAME,

diff --git a/metadata-ingestion/src/datahub/specific/dataset.py b/metadata-ingestion/src/datahub/specific/dataset.py
@@ -8,7 +8,6 @@
     EditableSchemaMetadataClass as EditableSchemaMetadata,
     FineGrainedLineageClass as FineGrainedLineage,
     GlossaryTermAssociationClass as Term,
-    GlossaryTermsClass as GlossaryTerms,
     KafkaAuditHeaderClass,
     SchemaMetadataClass,
     SystemMetadataClass,
@@ -22,6 +21,7 @@
     HasStructuredPropertiesPatch,
 )
 from datahub.specific.aspect_helpers.tags import HasTagsPatch
+from datahub.specific.aspect_helpers.terms import HasTermsPatch
 from datahub.utilities.urns.tag_urn import TagUrn
 from datahub.utilities.urns.urn import Urn
 
@@ -99,6 +99,7 @@ class DatasetPatchBuilder(
     HasCustomPropertiesPatch,
     HasStructuredPropertiesPatch,
     HasTagsPatch,
+    HasTermsPatch,
     MetadataPatchProposal,
 ):
     def __init__(
@@ -214,20 +215,6 @@ def set_fine_grained_upstream_lineages(
         )
         return self
 
-    def add_term(self, term: Term) -> "DatasetPatchBuilder":
-        self._add_patch(
-            GlossaryTerms.ASPECT_NAME, "add", path=("terms", term.urn), value=term
-        )
-        return self
-
-    def remove_term(self, term: Union[str, Urn]) -> "DatasetPatchBuilder":
-        if isinstance(term, str) and not term.startswith("urn:li:glossaryTerm:"):
-            term = "urn:li:glossaryTerm:" + term
-        self._add_patch(
-            GlossaryTerms.ASPECT_NAME, "remove", path=("terms", term), value={}
-        )
-        return self
-
     def for_field(
         self, field_path: str, editable: bool = True
     ) -> FieldPatchHelper["DatasetPatchBuilder"]: