Worked on the suggested changes

Author: vaibhavatlan

pyatlan/client/asset.py

@@ -584,9 +584,9 @@ def save(
         self,
         entity: Union[Asset, List[Asset]],
         replace_atlan_tags: bool = False,
-        append_atlan_tags: bool = False,
         replace_custom_metadata: bool = False,
         overwrite_custom_metadata: bool = False,
+        append_atlan_tags: bool = False,
     ) -> AssetMutationResponse:
         """
         If an asset with the same qualified_name exists, updates the existing asset. Otherwise, creates the asset.
@@ -597,6 +597,7 @@ def save(
         :param replace_atlan_tags: whether to replace AtlanTags during an update (True) or not (False)
         :param replace_custom_metadata: replaces any custom metadata with non-empty values provided
         :param overwrite_custom_metadata: overwrites any custom metadata, even with empty values
+        :param append_atlan_tags: whether to add/update/remove AtlanTags during an update (True) or not (False)
         :returns: the result of the save
         :raises AtlanError: on any API communication issue
         :raises ApiError: if a connection was created and blocking until policies are synced overruns the retry limit
@@ -875,50 +876,33 @@ def _restore_asset(self, asset: Asset) -> AssetMutationResponse:
         raw_json = self._client._call_api(BULK_UPDATE, query_params, request)
         return AssetMutationResponse(**raw_json)
 
-    @validate_arguments
-    def add_atlan_tags(
+    def _modify_tags(
         self,
+        type_of_modification,
         asset_type: Type[A],
         qualified_name: str,
         atlan_tag_names: List[str],
         propagate: bool = False,
         remove_propagation_on_delete: bool = True,
         restrict_lineage_propagation: bool = False,
         restrict_propagation_through_hierarchy: bool = False,
+        replace_atlan_tags: bool = False,
+        append_atlan_tags: bool = False,
     ) -> A:
-        """
-        Add one or more Atlan tags to the provided asset.
-
-        :param asset_type: type of asset to which to add the Atlan tags
-        :param qualified_name: qualified_name of the asset to which to add the Atlan tags
-        :param atlan_tag_names: human-readable names of the Atlan tags to add to the asset
-        :param propagate: whether to propagate the Atlan tag (True) or not (False)
-        :param remove_propagation_on_delete: whether to remove the propagated Atlan tags
-        when the Atlan tag is removed from this asset (True) or not (False)
-        :param restrict_lineage_propagation: whether to avoid propagating
-        through lineage (True) or do propagate through lineage (False)
-        :param restrict_propagation_through_hierarchy: whether to prevent this Atlan tag from
-        propagating through hierarchy (True) or allow it to propagate through hierarchy (False)
-        :returns: the asset that was updated (note that it will NOT contain details of the added Atlan tags)
-        :raises AtlanError: on any API communication issue
-        """
-        from pyatlan.client.atlan import AtlanClient
-
-        client = AtlanClient.get_current_client()
-        name_asset = client.asset.get_by_qualified_name(
+        reterieved_asset = self.get_by_qualified_name(
             qualified_name=qualified_name,
             asset_type=asset_type,
-            ignore_relationships=False,
+            attributes=["anchor"],
         )
         if asset_type in (AtlasGlossaryTerm, AtlasGlossaryCategory):
             updated_asset = asset_type.updater(
                 qualified_name=qualified_name,
-                name=name_asset.name,
-                glossary_guid=name_asset.anchor.guid,  # type: ignore
+                name=reterieved_asset.name,
+                glossary_guid=reterieved_asset.anchor.guid,  # type: ignore
             )
         else:
             updated_asset = asset_type.updater(
-                qualified_name=qualified_name, name=name_asset.name
+                qualified_name=qualified_name, name=reterieved_asset.name
             )
 
         atlan_tag = [
@@ -932,14 +916,66 @@ def add_atlan_tags(
             for name in atlan_tag_names
         ]
 
-        updated_asset.add_or_update_classifications = atlan_tag
+        if type_of_modification == "add" or "update":
+            updated_asset.add_or_update_classifications = atlan_tag
+        if type_of_modification == "remove":
+            updated_asset.remove_classifications = atlan_tag
+        if type_of_modification == "replace":
+            updated_asset.classifications = atlan_tag
 
-        response = self.save(entity=updated_asset, append_atlan_tags=True)
+        response = self.save(
+            entity=updated_asset,
+            replace_atlan_tags=replace_atlan_tags,
+            append_atlan_tags=append_atlan_tags,
+        )
 
         if assets := response.assets_updated(asset_type=asset_type):
             return assets[0]
         return updated_asset
 
+    @validate_arguments
+    def add_atlan_tags(
+        self,
+        asset_type: Type[A],
+        qualified_name: str,
+        atlan_tag_names: List[str],
+        propagate: bool = False,
+        remove_propagation_on_delete: bool = True,
+        restrict_lineage_propagation: bool = False,
+        restrict_propagation_through_hierarchy: bool = False,
+    ) -> A:
+        """
+        Add one or more Atlan tags to the provided asset.
+
+        :param asset_type: type of asset to which to add the Atlan tags
+        :param qualified_name: qualified_name of the asset to which to add the Atlan tags
+        :param atlan_tag_names: human-readable names of the Atlan tags to add to the asset
+        :param propagate: whether to propagate the Atlan tag (True) or not (False)
+        :param remove_propagation_on_delete: whether to remove the propagated Atlan tags
+        when the Atlan tag is removed from this asset (True) or not (False)
+        :param restrict_lineage_propagation: whether to avoid propagating
+        through lineage (True) or do propagate through lineage (False)
+        :param restrict_propagation_through_hierarchy: whether to prevent this Atlan tag from
+        propagating through hierarchy (True) or allow it to propagate through hierarchy (False)
+        :returns: the asset that was updated (note that it will NOT contain details of the added Atlan tags)
+        :raises AtlanError: on any API communication issue
+        """
+
+        response = self._modify_tags(
+            type_of_modification="add",
+            asset_type=asset_type,
+            qualified_name=qualified_name,
+            atlan_tag_names=atlan_tag_names,
+            propagate=propagate,
+            remove_propagation_on_delete=remove_propagation_on_delete,
+            restrict_lineage_propagation=restrict_lineage_propagation,
+            restrict_propagation_through_hierarchy=restrict_propagation_through_hierarchy,
+            replace_atlan_tags=False,
+            append_atlan_tags=True,
+        )
+
+        return response
+
     @validate_arguments
     def update_atlan_tags(
         self,
@@ -967,40 +1003,21 @@ def update_atlan_tags(
         :returns: the asset that was updated (note that it will NOT contain details of the updated Atlan tags)
         :raises AtlanError: on any API communication issue
         """
-        from pyatlan.client.atlan import AtlanClient
 
-        client = AtlanClient.get_current_client()
-        name_asset = client.asset.get_by_qualified_name(
-            qualified_name=qualified_name,
+        response = self._modify_tags(
+            type_of_modification="update",
             asset_type=asset_type,
-            ignore_relationships=False,
+            qualified_name=qualified_name,
+            atlan_tag_names=atlan_tag_names,
+            propagate=propagate,
+            remove_propagation_on_delete=remove_propagation_on_delete,
+            restrict_lineage_propagation=restrict_lineage_propagation,
+            restrict_propagation_through_hierarchy=restrict_propagation_through_hierarchy,
+            replace_atlan_tags=False,
+            append_atlan_tags=True,
         )
-        if asset_type in (AtlasGlossaryTerm, AtlasGlossaryCategory):
-            updated_asset = asset_type.updater(
-                qualified_name=qualified_name,
-                name=name_asset.name,
-                glossary_guid=name_asset.anchor.guid,  # type: ignore
-            )
-        else:
-            updated_asset = asset_type.updater(
-                qualified_name=qualified_name, name=name_asset.name
-            )
 
-        atlan_tag = [
-            AtlanTag(
-                type_name=AtlanTagName(display_text=name),
-                propagate=propagate,
-                remove_propagations_on_entity_delete=remove_propagation_on_delete,
-                restrict_propagation_through_lineage=restrict_lineage_propagation,
-                restrict_propagation_through_hierarchy=restrict_propagation_through_hierarchy,
-            )
-            for name in atlan_tag_names
-        ]
-        updated_asset.add_or_update_classifications = atlan_tag
-        response = self.save(entity=updated_asset, append_atlan_tags=True)
-        if assets := response.assets_updated(asset_type=asset_type):
-            return assets[0]
-        return updated_asset
+        return response
 
     @validate_arguments
     def remove_atlan_tag(
@@ -1018,32 +1035,17 @@ def remove_atlan_tag(
         :returns: the asset that was updated (note that it will NOT contain details of the deleted Atlan tag)
         :raises AtlanError: on any API communication issue
         """
-        from pyatlan.client.atlan import AtlanClient
 
-        client = AtlanClient.get_current_client()
-        name_asset = client.asset.get_by_qualified_name(
-            qualified_name=qualified_name,
+        response = self._modify_tags(
+            type_of_modification="remove",
             asset_type=asset_type,
-            ignore_relationships=False,
+            qualified_name=qualified_name,
+            atlan_tag_names=[atlan_tag_name],
+            replace_atlan_tags=False,
+            append_atlan_tags=True,
         )
 
-        if asset_type in (AtlasGlossaryTerm, AtlasGlossaryCategory):
-            updated_asset = asset_type.updater(
-                qualified_name=qualified_name,
-                name=name_asset.name,
-                glossary_guid=name_asset.anchor.guid,  # type: ignore
-            )
-        else:
-            updated_asset = asset_type.updater(
-                qualified_name=qualified_name, name=name_asset.name
-            )
-
-        atlan_tag = [AtlanTag(type_name=AtlanTagName(display_text=atlan_tag_name))]
-        updated_asset.remove_classifications = atlan_tag
-        response = self.save(entity=updated_asset, append_atlan_tags=True)
-        if assets := response.assets_updated(asset_type=asset_type):
-            return assets[0]
-        return updated_asset
+        return response
 
     @validate_arguments
     def remove_atlan_tags(
@@ -1061,35 +1063,16 @@ def remove_atlan_tags(
         :returns: the asset that was updated (note that it will NOT contain details of the deleted Atlan tags)
         :raises AtlanError: on any API communication issue
         """
-        from pyatlan.client.atlan import AtlanClient
-
-        client = AtlanClient.get_current_client()
-        name_asset = client.asset.get_by_qualified_name(
-            qualified_name=qualified_name,
+        response = self._modify_tags(
+            type_of_modification="remove",
             asset_type=asset_type,
-            ignore_relationships=False,
+            qualified_name=qualified_name,
+            atlan_tag_names=atlan_tag_names,
+            replace_atlan_tags=False,
+            append_atlan_tags=True,
         )
 
-        if asset_type in (AtlasGlossaryTerm, AtlasGlossaryCategory):
-            updated_asset = asset_type.updater(
-                qualified_name=qualified_name,
-                name=name_asset.name,
-                glossary_guid=name_asset.anchor.guid,  # type: ignore
-            )
-        else:
-            updated_asset = asset_type.updater(
-                qualified_name=qualified_name, name=name_asset.name
-            )
-
-        atlan_tag = [
-            AtlanTag(type_name=AtlanTagName(display_text=name))
-            for name in atlan_tag_names
-        ]
-        updated_asset.remove_classifications = atlan_tag
-        response = self.save(entity=updated_asset, append_atlan_tags=True)
-        if assets := response.assets_updated(asset_type=asset_type):
-            return assets[0]
-        return updated_asset
+        return response
 
     def _update_asset_by_attribute(
         self, asset: A, asset_type: Type[A], qualified_name: str

tests/integration/test_client.py

@@ -44,7 +44,7 @@
 from tests.integration.requests_test import delete_token
 
 CLASSIFICATION_NAME = "Issue"
-CLASSIFICATION_NAME_2 = "Confidential"
+CLASSIFICATION_NAME2 = "Confidential"
 SL_SORT_BY_TIMESTAMP = SortItem(field="timestamp", order=SortOrder.ASCENDING)
 SL_SORT_BY_GUID = SortItem(field="entityGuidsAll", order=SortOrder.ASCENDING)
 SL_SORT_BY_QUALIFIED_NAME = SortItem(
@@ -751,7 +751,7 @@ def test_add_classification(client: AtlanClient, term1: AtlasGlossaryTerm):
         AtlasGlossaryTerm, term1.qualified_name, [CLASSIFICATION_NAME]
     )
     glossary_term = client.asset.get_by_guid(
-        term1.guid, asset_type=AtlasGlossaryTerm, ignore_relationships=False
+        term1.guid, asset_type=AtlasGlossaryTerm, attributes=["atlanTags"]
     )
     assert glossary_term.atlan_tags
     assert len(glossary_term.atlan_tags) == 1
@@ -798,7 +798,7 @@ def test_update_classification(client: AtlanClient, term1: AtlasGlossaryTerm):
         remove_propagation_on_delete=False,
     )
     glossary_term = client.asset.get_by_guid(
-        term1.guid, asset_type=AtlasGlossaryTerm, ignore_relationships=False
+        term1.guid, asset_type=AtlasGlossaryTerm, attributes=["atlanTags"]
     )
     assert glossary_term.atlan_tags
     assert len(glossary_term.atlan_tags) == 1
@@ -817,7 +817,7 @@ def test_remove_classification(client: AtlanClient, term1: AtlasGlossaryTerm):
         AtlasGlossaryTerm, term1.qualified_name, CLASSIFICATION_NAME
     )
     glossary_term = client.asset.get_by_guid(
-        term1.guid, asset_type=AtlasGlossaryTerm, ignore_relationships=False
+        term1.guid, asset_type=AtlasGlossaryTerm, attributes=["atlanTags"]
     )
     assert not glossary_term.atlan_tags
 
@@ -827,19 +827,45 @@ def test_multiple_add_classification(client: AtlanClient, term1: AtlasGlossaryTe
     client.asset.add_atlan_tags(
         AtlasGlossaryTerm,
         term1.qualified_name,
-        [CLASSIFICATION_NAME, CLASSIFICATION_NAME_2],
+        [CLASSIFICATION_NAME, CLASSIFICATION_NAME2],
     )
     glossary_term = client.asset.get_by_guid(
-        term1.guid, asset_type=AtlasGlossaryTerm, ignore_relationships=False
+        term1.guid, asset_type=AtlasGlossaryTerm, attributes=["atlanTags"]
     )
     assert glossary_term.atlan_tags
     assert len(glossary_term.atlan_tags) == 2
     classification = glossary_term.atlan_tags[0]
-    assert str(classification.type_name) == CLASSIFICATION_NAME or CLASSIFICATION_NAME_2
-    classification_2 = glossary_term.atlan_tags[1]
-    assert (
-        str(classification_2.type_name) == CLASSIFICATION_NAME or CLASSIFICATION_NAME_2
+    assert str(classification.type_name) == CLASSIFICATION_NAME or CLASSIFICATION_NAME2
+    classification2 = glossary_term.atlan_tags[1]
+    assert str(classification2.type_name) == CLASSIFICATION_NAME or CLASSIFICATION_NAME2
+
+
+def test_multiple_update_classification(client: AtlanClient, term1: AtlasGlossaryTerm):
+    assert term1.qualified_name
+    client.asset.update_atlan_tags(
+        AtlasGlossaryTerm,
+        term1.qualified_name,
+        [CLASSIFICATION_NAME, CLASSIFICATION_NAME2],
+        propagate=True,
+        remove_propagation_on_delete=False,
     )
+    glossary_term = client.asset.get_by_guid(
+        term1.guid, asset_type=AtlasGlossaryTerm, attributes=["atlanTags"]
+    )
+    assert glossary_term.atlan_tags
+    assert len(glossary_term.atlan_tags) == 2
+    classification = glossary_term.atlan_tags[0]
+    assert str(classification.type_name) == CLASSIFICATION_NAME or CLASSIFICATION_NAME2
+    assert classification.propagate
+    assert not classification.remove_propagations_on_entity_delete
+    assert classification.restrict_propagation_through_lineage
+    assert not classification.restrict_propagation_through_hierarchy
+    classification2 = glossary_term.atlan_tags[1]
+    assert str(classification2.type_name) == CLASSIFICATION_NAME or CLASSIFICATION_NAME2
+    assert classification2.propagate
+    assert not classification2.remove_propagations_on_entity_delete
+    assert classification2.restrict_propagation_through_lineage
+    assert not classification2.restrict_propagation_through_hierarchy
 
 
 @pytest.mark.order(after="test_multiple_add_classification")
@@ -848,10 +874,10 @@ def test_multiple_remove_classification(client: AtlanClient, term1: AtlasGlossar
     client.asset.remove_atlan_tags(
         AtlasGlossaryTerm,
         term1.qualified_name,
-        [CLASSIFICATION_NAME, CLASSIFICATION_NAME_2],
+        [CLASSIFICATION_NAME, CLASSIFICATION_NAME2],
     )
     glossary_term = client.asset.get_by_guid(
-        term1.guid, asset_type=AtlasGlossaryTerm, ignore_relationships=False
+        term1.guid, asset_type=AtlasGlossaryTerm, attributes=["atlanTags"]
     )
     assert not glossary_term.atlan_tags