Skip to content

Commit d0cc2e8

Browse files
vaibhavatlanAryamanz29
authored andcommitted
Added the logic and test cases
1 parent 6cbd422 commit d0cc2e8

File tree

9 files changed

+156
-116
lines changed

9 files changed

+156
-116
lines changed

pyatlan/client/asset.py

Lines changed: 83 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -40,13 +40,11 @@
4040
ADD_BUSINESS_ATTRIBUTE_BY_ID,
4141
BULK_UPDATE,
4242
DELETE_ENTITIES_BY_GUIDS,
43-
DELETE_ENTITY_BY_ATTRIBUTE,
4443
GET_ENTITY_BY_GUID,
4544
GET_ENTITY_BY_UNIQUE_ATTRIBUTE,
4645
GET_LINEAGE_LIST,
4746
INDEX_SEARCH,
4847
PARTIAL_UPDATE_ENTITY_BY_ATTRIBUTE,
49-
UPDATE_ENTITY_BY_ATTRIBUTE,
5048
)
5149
from pyatlan.errors import AtlanError, ErrorCode
5250
from pyatlan.model.aggregation import Aggregations
@@ -74,7 +72,6 @@
7472
AtlanObject,
7573
AtlanTag,
7674
AtlanTagName,
77-
AtlanTags,
7875
BulkRequest,
7976
SearchRequest,
8077
)
@@ -587,6 +584,7 @@ def save(
587584
self,
588585
entity: Union[Asset, List[Asset]],
589586
replace_atlan_tags: bool = False,
587+
append_atlan_tags: bool = False,
590588
replace_custom_metadata: bool = False,
591589
overwrite_custom_metadata: bool = False,
592590
) -> AssetMutationResponse:
@@ -604,7 +602,8 @@ def save(
604602
:raises ApiError: if a connection was created and blocking until policies are synced overruns the retry limit
605603
"""
606604
query_params = {
607-
"replaceClassifications": replace_atlan_tags,
605+
"replaceTags": replace_atlan_tags,
606+
"appendTags": append_atlan_tags,
608607
"replaceBusinessAttributes": replace_custom_metadata,
609608
"overwriteBusinessAttributes": overwrite_custom_metadata,
610609
}
@@ -876,36 +875,6 @@ def _restore_asset(self, asset: Asset) -> AssetMutationResponse:
876875
raw_json = self._client._call_api(BULK_UPDATE, query_params, request)
877876
return AssetMutationResponse(**raw_json)
878877

879-
def _modify_tags(
880-
self,
881-
api: API,
882-
asset_type: Type[A],
883-
qualified_name: str,
884-
atlan_tag_names: List[str],
885-
propagate: bool = False,
886-
remove_propagation_on_delete: bool = True,
887-
restrict_lineage_propagation: bool = False,
888-
restrict_propagation_through_hierarchy: bool = False,
889-
) -> None:
890-
atlan_tags = AtlanTags(
891-
__root__=[
892-
AtlanTag(
893-
type_name=AtlanTagName(display_text=name),
894-
propagate=propagate,
895-
remove_propagations_on_entity_delete=remove_propagation_on_delete,
896-
restrict_propagation_through_lineage=restrict_lineage_propagation,
897-
restrict_propagation_through_hierarchy=restrict_propagation_through_hierarchy,
898-
)
899-
for name in atlan_tag_names
900-
]
901-
)
902-
query_params = {"attr:qualifiedName": qualified_name}
903-
self._client._call_api(
904-
api.format_path_with_params(asset_type.__name__, "classifications"),
905-
query_params,
906-
atlan_tags,
907-
)
908-
909878
@validate_arguments
910879
def add_atlan_tags(
911880
self,
@@ -916,11 +885,9 @@ def add_atlan_tags(
916885
remove_propagation_on_delete: bool = True,
917886
restrict_lineage_propagation: bool = False,
918887
restrict_propagation_through_hierarchy: bool = False,
919-
) -> None:
888+
) -> A:
920889
"""
921890
Add one or more Atlan tags to the provided asset.
922-
Note: if one or more of the provided Atlan tags already exist on the asset, an error
923-
will be raised. (In other words, this operation is NOT idempotent.)
924891
925892
:param asset_type: type of asset to which to add the Atlan tags
926893
:param qualified_name: qualified_name of the asset to which to add the Atlan tags
@@ -932,19 +899,40 @@ def add_atlan_tags(
932899
through lineage (True) or do propagate through lineage (False)
933900
:param restrict_propagation_through_hierarchy: whether to prevent this Atlan tag from
934901
propagating through hierarchy (True) or allow it to propagate through hierarchy (False)
902+
:returns: the asset that was updated (note that it will NOT contain details of the added Atlan tags)
935903
:raises AtlanError: on any API communication issue
936904
"""
937-
self._modify_tags(
938-
UPDATE_ENTITY_BY_ATTRIBUTE,
939-
asset_type,
940-
qualified_name,
941-
atlan_tag_names,
942-
propagate,
943-
remove_propagation_on_delete,
944-
restrict_lineage_propagation,
945-
restrict_propagation_through_hierarchy,
905+
from pyatlan.client.atlan import AtlanClient
906+
907+
client = AtlanClient.get_current_client()
908+
909+
name_asset = client.asset.get_by_qualified_name(
910+
qualified_name=qualified_name, asset_type=asset_type
911+
)
912+
913+
updated_asset = asset_type.updater(
914+
qualified_name=qualified_name, name=name_asset.name
946915
)
947916

917+
atlan_tag = [
918+
AtlanTag(
919+
type_name=AtlanTagName(display_text=name),
920+
propagate=propagate,
921+
remove_propagations_on_entity_delete=remove_propagation_on_delete,
922+
restrict_propagation_through_lineage=restrict_lineage_propagation,
923+
restrict_propagation_through_hierarchy=restrict_propagation_through_hierarchy,
924+
)
925+
for name in atlan_tag_names
926+
]
927+
928+
updated_asset.add_or_update_classifications = atlan_tag
929+
930+
response = self.save(entity=updated_asset, append_atlan_tags=True)
931+
932+
if assets := response.assets_updated(asset_type=asset_type):
933+
return assets[0]
934+
return updated_asset
935+
948936
@validate_arguments
949937
def update_atlan_tags(
950938
self,
@@ -955,11 +943,9 @@ def update_atlan_tags(
955943
remove_propagation_on_delete: bool = True,
956944
restrict_lineage_propagation: bool = True,
957945
restrict_propagation_through_hierarchy: bool = False,
958-
) -> None:
946+
) -> A:
959947
"""
960948
Update one or more Atlan tags to the provided asset.
961-
Note: if one or more of the provided Atlan tags already exist on the asset, an error
962-
will be raised. (In other words, this operation is NOT idempotent.)
963949
964950
:param asset_type: type of asset to which to update the Atlan tags
965951
:param qualified_name: qualified_name of the asset to which to update the Atlan tags
@@ -971,52 +957,70 @@ def update_atlan_tags(
971957
through lineage (True) or do propagate through lineage (False)
972958
:param restrict_propagation_through_hierarchy: whether to prevent this Atlan tag from
973959
propagating through hierarchy (True) or allow it to propagate through hierarchy (False)
960+
:returns: the asset that was updated (note that it will NOT contain details of the updated Atlan tags)
974961
:raises AtlanError: on any API communication issue
975962
"""
976-
self._modify_tags(
977-
PARTIAL_UPDATE_ENTITY_BY_ATTRIBUTE,
978-
asset_type,
979-
qualified_name,
980-
atlan_tag_names,
981-
propagate,
982-
remove_propagation_on_delete,
983-
restrict_lineage_propagation,
984-
restrict_propagation_through_hierarchy,
963+
from pyatlan.client.atlan import AtlanClient
964+
965+
client = AtlanClient.get_current_client()
966+
name_asset = client.asset.get_by_qualified_name(
967+
qualified_name=qualified_name, asset_type=asset_type
968+
)
969+
updated_asset = asset_type.updater(
970+
qualified_name=qualified_name, name=name_asset.name
985971
)
986972

973+
atlan_tag = [
974+
AtlanTag(
975+
type_name=AtlanTagName(display_text=name),
976+
propagate=propagate,
977+
remove_propagations_on_entity_delete=remove_propagation_on_delete,
978+
restrict_propagation_through_lineage=restrict_lineage_propagation,
979+
restrict_propagation_through_hierarchy=restrict_propagation_through_hierarchy,
980+
)
981+
for name in atlan_tag_names
982+
]
983+
updated_asset.add_or_update_classifications = atlan_tag
984+
response = self.save(entity=updated_asset, append_atlan_tags=True)
985+
if assets := response.assets_updated(asset_type=asset_type):
986+
return assets[0]
987+
return updated_asset
988+
987989
@validate_arguments
988-
def remove_atlan_tag(
989-
self, asset_type: Type[A], qualified_name: str, atlan_tag_name: str
990-
) -> None:
990+
def remove_atlan_tags(
991+
self,
992+
asset_type: Type[A],
993+
qualified_name: str,
994+
atlan_tag_names: List[str],
995+
) -> A:
991996
"""
992-
Removes a single Atlan tag from the provided asset.
993-
Note: if the provided Atlan tag does not exist on the asset, an error will be raised.
994-
(In other words, this operation is NOT idempotent.)
997+
Removes one or more Atlan tag from the provided asset.
995998
996999
:param asset_type: type of asset to which to add the Atlan tags
9971000
:param qualified_name: qualified_name of the asset to which to add the Atlan tags
998-
:param atlan_tag_name: human-readable name of the Atlan tag to remove from the asset
1001+
:param atlan_tag_names: human-readable name of the Atlan tag to remove from the asset
9991002
:raises AtlanError: on any API communication issue
10001003
"""
10011004
from pyatlan.client.atlan import AtlanClient
10021005

1003-
classification_id = (
1004-
AtlanClient.get_current_client().atlan_tag_cache.get_id_for_name(
1005-
atlan_tag_name
1006-
)
1006+
client = AtlanClient.get_current_client()
1007+
name_asset = client.asset.get_by_qualified_name(
1008+
qualified_name=qualified_name, asset_type=asset_type
10071009
)
1008-
if not classification_id:
1009-
raise ErrorCode.ATLAN_TAG_NOT_FOUND_BY_NAME.exception_with_parameters(
1010-
atlan_tag_name
1011-
)
1012-
query_params = {"attr:qualifiedName": qualified_name}
1013-
self._client._call_api(
1014-
DELETE_ENTITY_BY_ATTRIBUTE.format_path_with_params(
1015-
asset_type.__name__, "classification", classification_id
1016-
),
1017-
query_params,
1010+
updated_asset = asset_type.updater(
1011+
qualified_name=qualified_name, name=name_asset.name
10181012
)
10191013

1014+
atlan_tag = [
1015+
AtlanTag(type_name=AtlanTagName(display_text=name))
1016+
for name in atlan_tag_names
1017+
]
1018+
updated_asset.remove_classifications = atlan_tag
1019+
response = self.save(entity=updated_asset, append_atlan_tags=True)
1020+
if assets := response.assets_updated(asset_type=asset_type):
1021+
return assets[0]
1022+
return updated_asset
1023+
10201024
def _update_asset_by_attribute(
10211025
self, asset: A, asset_type: Type[A], qualified_name: str
10221026
):

pyatlan/client/atlan.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1339,19 +1339,19 @@ def add_atlan_tags(
13391339

13401340
@validate_arguments
13411341
def remove_atlan_tag(
1342-
self, asset_type: Type[A], qualified_name: str, atlan_tag_name: str
1342+
self, asset_type: Type[A], qualified_name: str, atlan_tag_names: List[str]
13431343
) -> None:
1344-
"""Deprecated - use asset.remove_atlan_tag() instead."""
1344+
"""Deprecated - use asset.remove_atlan_tags() instead."""
13451345
warn(
1346-
"This method is deprecated, please use 'asset.remove_atlan_tag' instead, which offers identical "
1346+
"This method is deprecated, please use 'asset.remove_atlan_tags' instead, which offers identical "
13471347
"functionality.",
13481348
DeprecationWarning,
13491349
stacklevel=2,
13501350
)
1351-
self.asset.remove_atlan_tag(
1351+
self.asset.remove_atlan_tags(
13521352
asset_type=asset_type,
13531353
qualified_name=qualified_name,
1354-
atlan_tag_name=atlan_tag_name,
1354+
atlan_tag_names=atlan_tag_names
13551355
)
13561356

13571357
@validate_arguments

pyatlan/model/assets/core/referenceable.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -362,6 +362,14 @@ def validate_required(self):
362362
default=None,
363363
description="Map of remove relationship attributes.",
364364
)
365+
add_or_update_classifications: Optional[List[AtlanTag]] = Field(
366+
default=None,
367+
description="Map of add/update classifcations of atlan tag.",
368+
)
369+
remove_classifications: Optional[List[AtlanTag]] = Field(
370+
default=None,
371+
description="Map of remove classifcations of atlan tag.",
372+
)
365373
semantic: Optional[SaveSemantic] = Field(
366374
default=None,
367375
exclude=True,

testing.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
import logging
2+
3+
from pyatlan.client.atlan import AtlanClient
4+
from pyatlan.model.assets import Table
5+
6+
logging.basicConfig(level=logging.DEBUG)
7+
client = AtlanClient()
8+
# a = client.asset.get_by_guid(guid="21757812-b08c-49f7-b723-7d8dbb2e2dd1",ignore_relationships=False, attributes=[Asset.ATLAN_TAGS])
9+
# column = client.asset.append_terms( #
10+
# asset_type=Table, #
11+
# qualified_name="default/snowflake/1746022526/WIDE_WORLD_IMPORTERS/BRONZE_SALES/ORDERS", #
12+
# terms=[AtlasGlossaryTerm.ref_by_guid(guid="e46b0ac0-5451-4161-b162-5e98aafede7f"),
13+
# AtlasGlossaryTerm.ref_by_guid(guid="e46b0ac0-5451-4161-b162-5e98aafede7f") ]
14+
# ) #
15+
16+
client = AtlanClient()
17+
client.asset.add_atlan_tags(
18+
asset_type=Table,
19+
qualified_name="default/snowflake/1746022526/WIDE_WORLD_IMPORTERS/BRONZE_WAREHOUSE/FIVETRAN_AUDIT",
20+
atlan_tag_names=["Alert: DQ-F2Cfn", "Alert: DQ-q4lhS"], # n
21+
)
22+
# column = client.asset.add_atlan_tags( #
23+
# asset_type=Table, #
24+
# qualified_name="default/snowflake/1746022526/WIDE_WORLD_IMPORTERS/BRONZE_SALES/ORDERS", #
25+
# terms=[AtlasGlossaryTerm.ref_by_guid(guid="e46b0ac0-5451-4161-b162-5e98aafede7f") ]
26+
# ) #
27+
28+
# print(column)

0 commit comments

Comments
 (0)