Merge pull request #43 from atlanhq/link-terms

Update code to link terms to assets

HISTORY.md

@@ -1,3 +1,9 @@
+## 0.0.26 (May 4, 2023)
+
+* Add remove_terms method to AtlanClient
+* Add append_terms method to AtlanClient
+* Add replace_terms method to AtlanClient
+
 ## 0.0.25 (May 2, 2023)
 
 * Update create method for Readme asset

pyatlan/client/atlan.py

@@ -313,8 +313,7 @@ def get_asset_by_qualified_name(
             raw_json["entity"]["attributes"].update(
                 raw_json["entity"]["relationshipAttributes"]
             )
-            asset = AssetResponse[A](**raw_json).entity
-            asset.is_incomplete = False
+            asset = self.handle_relationships(raw_json)
             if not isinstance(asset, asset_type):
                 raise NotFoundError(
                     message=f"Asset with qualifiedName {qualified_name} "
@@ -345,16 +344,7 @@ def get_asset_by_guid(
                 GET_ENTITY_BY_GUID.format_path_with_params(guid),
                 query_params,
             )
-            if (
-                "relationshipAttributes" in raw_json["entity"]
-                and raw_json["entity"]["relationshipAttributes"]
-            ):
-                raw_json["entity"]["attributes"].update(
-                    raw_json["entity"]["relationshipAttributes"]
-                )
-            raw_json["entity"]["relationshipAttributes"] = {}
-            asset = AssetResponse[A](**raw_json).entity
-            asset.is_incomplete = False
+            asset = self.handle_relationships(raw_json)
             if not isinstance(asset, asset_type):
                 raise NotFoundError(
                     message=f"Asset with GUID {guid} is not of the type requested: {asset_type.__name__}.",
@@ -366,6 +356,19 @@ def get_asset_by_guid(
                 raise NotFoundError(message=ae.user_message, code=ae.code) from ae
             raise ae
 
+    def handle_relationships(self, raw_json):
+        if (
+            "relationshipAttributes" in raw_json["entity"]
+            and raw_json["entity"]["relationshipAttributes"]
+        ):
+            raw_json["entity"]["attributes"].update(
+                raw_json["entity"]["relationshipAttributes"]
+            )
+        raw_json["entity"]["relationshipAttributes"] = {}
+        asset = AssetResponse[A](**raw_json).entity
+        asset.is_incomplete = False
+        return asset
+
     @validate_arguments()
     def retrieve_minimal(self, guid: str, asset_type: Type[A]) -> A:
         return self.get_asset_by_guid(
@@ -602,3 +605,100 @@ def replace_custom_metadata(self, guid: str, custom_metadata: CustomMetadata):
             None,
             custom_metadata_request,
         )
+
+    @validate_arguments()
+    def append_terms(
+        self,
+        asset_type: Type[A],
+        terms: list[AtlasGlossaryTerm],
+        guid: Optional[str] = None,
+        qualified_name: Optional[str] = None,
+    ) -> A:
+        if guid:
+            if qualified_name:
+                raise ValueError(
+                    "Either guid or qualified_name can be be specified not both"
+                )
+            asset = self.get_asset_by_guid(guid=guid, asset_type=asset_type)
+        elif qualified_name:
+            asset = self.get_asset_by_qualified_name(
+                qualified_name=qualified_name, asset_type=asset_type
+            )
+        else:
+            raise ValueError("Either guid or qualified name must be specified")
+        if not terms:
+            return asset
+        replacement_terms: list[AtlasGlossaryTerm] = []
+        if existing_terms := asset.terms:
+            replacement_terms.extend(
+                term for term in existing_terms if term.relationship_status != "DELETED"
+            )
+        replacement_terms.extend(terms)
+        asset.terms = replacement_terms
+        response = self.upsert(entity=asset)
+        if assets := response.assets_updated(asset_type=asset_type):
+            return assets[0]
+        return asset
+
+    @validate_arguments()
+    def replace_terms(
+        self,
+        asset_type: Type[A],
+        terms: list[AtlasGlossaryTerm],
+        guid: Optional[str] = None,
+        qualified_name: Optional[str] = None,
+    ) -> A:
+        if guid:
+            if qualified_name:
+                raise ValueError(
+                    "Either guid or qualified_name can be be specified not both"
+                )
+            asset = self.get_asset_by_guid(guid=guid, asset_type=asset_type)
+        elif qualified_name:
+            asset = self.get_asset_by_qualified_name(
+                qualified_name=qualified_name, asset_type=asset_type
+            )
+        else:
+            raise ValueError("Either guid or qualified name must be specified")
+        asset.terms = terms
+        response = self.upsert(entity=asset)
+        if assets := response.assets_updated(asset_type=asset_type):
+            return assets[0]
+        return asset
+
+    @validate_arguments()
+    def remove_terms(
+        self,
+        asset_type: Type[A],
+        terms: list[AtlasGlossaryTerm],
+        guid: Optional[str] = None,
+        qualified_name: Optional[str] = None,
+    ) -> A:
+        if not terms:
+            raise ValueError("A list of terms to remove must be specified")
+        if guid:
+            if qualified_name:
+                raise ValueError(
+                    "Either guid or qualified_name can be be specified not both"
+                )
+            asset = self.get_asset_by_guid(guid=guid, asset_type=asset_type)
+        elif qualified_name:
+            asset = self.get_asset_by_qualified_name(
+                qualified_name=qualified_name, asset_type=asset_type
+            )
+        else:
+            raise ValueError("Either guid or qualified name must be specified")
+        replacement_terms: list[AtlasGlossaryTerm] = []
+        guids_to_be_removed = {t.guid for t in terms}
+        if existing_terms := asset.terms:
+            replacement_terms.extend(
+                term
+                for term in existing_terms
+                if term.relationship_status != "DELETED"
+                and term.guid not in guids_to_be_removed
+            )
+        asset.terms = replacement_terms
+        response = self.upsert(entity=asset)
+        if assets := response.assets_updated(asset_type=asset_type):
+            return assets[0]
+        return asset

pyatlan/generator/templates/entity.jinja2

@@ -71,11 +71,15 @@ def validate_required_fields(field_names:list[str], values:list[Any]):
             raise ValueError(f"{field_name} is required")
         if isinstance(value, str) and not value.strip():
             raise ValueError(f"{field_name} cannot be blank")
-{%- macro gen_properties(attribute_defs) %}
+{%- macro gen_properties(attribute_defs, additional_names=[]) %}
     _convience_properties: ClassVar[list[str]] = [
         {%- for attribute_def in attribute_defs %}
         "{{ attribute_def.name | to_snake_case }}",
-        {%- endfor %}]
+        {%- endfor %}
+        {%- for name in additional_names  %}
+        "{{ name }}",
+         {%- endfor %}]
+
     {%- for attribute_def in attribute_defs %}
     {%- set type = attribute_def.typeName | get_type %}
     {%- set property_type %}{% if attribute_def.isOptional %}Optional[{% endif %}{{type}}{% if attribute_def.isOptional %}]{% endif %}{% endset %}
@@ -107,7 +111,19 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
                 return object.__setattr__(self, name, value)
             super().__setattr__( name, value)
 
-    {{ gen_properties(entity_def.attribute_defs) }}
+    {{ gen_properties(entity_def.attribute_defs, ["terms"]) }}
+
+    @property
+    def terms(self) -> list[AtlasGlossaryTerm]:
+        if self.attributes is None:
+            self.attributes = self.Attributes()
+        return [] if self.attributes.meanings is None else self.attributes.meanings
+
+    @terms.setter
+    def terms(self, terms: list[AtlasGlossaryTerm]):
+        if self.attributes is None:
+            self.attributes = self.Attributes()
+        self.attributes.meanings = terms
 
 {%- if entity_def.name == "Referenceable"  %}