Skip to content

Commit 579278e

Browse files
committed
[changes] Use instance bound client and simplified usage of AtlanFields in search attributes
1 parent a00828a commit 579278e

File tree

1 file changed

+23
-25
lines changed

1 file changed

+23
-25
lines changed

pyatlan/client/asset.py

Lines changed: 23 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -385,6 +385,14 @@ def find_purposes_by_name(
385385
allow_multiple=True,
386386
)
387387

388+
def _normalize_search_fields(
389+
self,
390+
fields: Optional[Union[List[str], List[AtlanField]]],
391+
) -> List[str]:
392+
if not fields:
393+
return []
394+
return [f.atlan_field_name if isinstance(f, AtlanField) else f for f in fields]
395+
388396
@validate_arguments(config=dict(arbitrary_types_allowed=True))
389397
def get_by_qualified_name(
390398
self,
@@ -408,29 +416,27 @@ def get_by_qualified_name(
408416
:raises NotFoundError: if the asset does not exist
409417
:raises AtlanError: on any API communication issue
410418
"""
411-
from pyatlan.client.atlan import AtlanClient
412419
from pyatlan.model.fluent_search import FluentSearch
413420

414421
query_params = {
415422
"attr:qualifiedName": qualified_name,
416423
"minExtInfo": min_ext_info,
417424
"ignoreRelationships": ignore_relationships,
418425
}
419-
attributes = attributes or []
420-
related_attributes = related_attributes or []
426+
attributes = self._normalize_search_fields(attributes)
427+
related_attributes = self._normalize_search_fields(related_attributes)
421428

422429
if (attributes and len(attributes)) or (
423430
related_attributes and len(related_attributes)
424431
):
425-
client = AtlanClient.get_current_client()
426432
search = (
427433
FluentSearch().select().where(Asset.QUALIFIED_NAME.eq(qualified_name))
428434
)
429435
for attribute in attributes:
430436
search = search.include_on_results(attribute)
431437
for relation_attribute in related_attributes:
432438
search = search.include_on_relations(relation_attribute)
433-
results = search.execute(client=client)
439+
results = search.execute(client=self._client) # type: ignore[arg-type]
434440
if results and results.current_page():
435441
first_result = results.current_page()[0]
436442
if isinstance(first_result, asset_type):
@@ -482,26 +488,24 @@ def get_by_guid(
482488
:raises NotFoundError: if the asset does not exist, or is not of the type requested
483489
:raises AtlanError: on any API communication issue
484490
"""
485-
from pyatlan.client.atlan import AtlanClient
486491
from pyatlan.model.fluent_search import FluentSearch
487492

488493
query_params = {
489494
"minExtInfo": min_ext_info,
490495
"ignoreRelationships": ignore_relationships,
491496
}
492-
attributes = attributes or []
493-
related_attributes = related_attributes or []
497+
attributes = self._normalize_search_fields(attributes)
498+
related_attributes = self._normalize_search_fields(related_attributes)
494499

495500
if (attributes and len(attributes)) or (
496501
related_attributes and len(related_attributes)
497502
):
498-
client = AtlanClient.get_current_client()
499503
search = FluentSearch().select().where(Asset.GUID.eq(guid))
500504
for attribute in attributes:
501505
search = search.include_on_results(attribute)
502506
for relation_attribute in related_attributes:
503507
search = search.include_on_relations(relation_attribute)
504-
results = search.execute(client=client)
508+
results = search.execute(client=self._client) # type: ignore[arg-type]
505509
if results and results.current_page():
506510
first_result = results.current_page()[0]
507511
if isinstance(first_result, asset_type):
@@ -892,13 +896,13 @@ def _modify_tags(
892896
reterieved_asset = self.get_by_qualified_name(
893897
qualified_name=qualified_name,
894898
asset_type=asset_type,
895-
attributes=[AtlasGlossaryTerm.ANCHOR.atlan_field_name],
899+
attributes=[AtlasGlossaryTerm.ANCHOR], # type: ignore[arg-type]
896900
)
897901
if asset_type in (AtlasGlossaryTerm, AtlasGlossaryCategory):
898902
updated_asset = asset_type.updater(
899903
qualified_name=qualified_name,
900904
name=reterieved_asset.name,
901-
glossary_guid=reterieved_asset.anchor.guid, # type: ignore
905+
glossary_guid=reterieved_asset.anchor.guid, # type: ignore[attr-defined]
902906
)
903907
else:
904908
updated_asset = asset_type.updater(
@@ -1428,10 +1432,8 @@ def append_terms(
14281432
:param qualified_name: the qualified_name of the asset to which to link the terms
14291433
:returns: the asset that was updated (note that it will NOT contain details of the appended terms)
14301434
"""
1431-
from pyatlan.client.atlan import AtlanClient
14321435
from pyatlan.model.fluent_search import FluentSearch
14331436

1434-
client = AtlanClient.get_current_client()
14351437
if guid:
14361438
if qualified_name:
14371439
raise ErrorCode.QN_OR_GUID_NOT_BOTH.exception_with_parameters()
@@ -1440,15 +1442,15 @@ def append_terms(
14401442
.select()
14411443
.where(Asset.TYPE_NAME.eq(asset_type.__name__))
14421444
.where(asset_type.GUID.eq(guid))
1443-
.execute(client=client)
1445+
.execute(client=self._client) # type: ignore[arg-type]
14441446
)
14451447
elif qualified_name:
14461448
results = (
14471449
FluentSearch()
14481450
.select()
14491451
.where(Asset.TYPE_NAME.eq(asset_type.__name__))
14501452
.where(asset_type.QUALIFIED_NAME.eq(qualified_name))
1451-
.execute(client=client)
1453+
.execute(client=self._client) # type: ignore[arg-type]
14521454
)
14531455
else:
14541456
raise ErrorCode.QN_OR_GUID.exception_with_parameters()
@@ -1507,10 +1509,8 @@ def replace_terms(
15071509
:param qualified_name: the qualified_name of the asset to which to replace the terms
15081510
:returns: the asset that was updated (note that it will NOT contain details of the replaced terms)
15091511
"""
1510-
from pyatlan.client.atlan import AtlanClient
15111512
from pyatlan.model.fluent_search import FluentSearch
15121513

1513-
client = AtlanClient.get_current_client()
15141514
if guid:
15151515
if qualified_name:
15161516
raise ErrorCode.QN_OR_GUID_NOT_BOTH.exception_with_parameters()
@@ -1519,15 +1519,15 @@ def replace_terms(
15191519
.select()
15201520
.where(Asset.TYPE_NAME.eq(asset_type.__name__))
15211521
.where(asset_type.GUID.eq(guid))
1522-
.execute(client=client)
1522+
.execute(client=self._client) # type: ignore[arg-type]
15231523
)
15241524
elif qualified_name:
15251525
results = (
15261526
FluentSearch()
15271527
.select()
15281528
.where(Asset.TYPE_NAME.eq(asset_type.__name__))
15291529
.where(asset_type.QUALIFIED_NAME.eq(qualified_name))
1530-
.execute(client=client)
1530+
.execute(client=self._client) # type: ignore[arg-type]
15311531
)
15321532
else:
15331533
raise ErrorCode.QN_OR_GUID.exception_with_parameters()
@@ -1588,10 +1588,8 @@ def remove_terms(
15881588
:param qualified_name: the qualified_name of the asset from which to remove the terms
15891589
:returns: the asset that was updated (note that it will NOT contain details of the resulting terms)
15901590
"""
1591-
from pyatlan.client.atlan import AtlanClient
15921591
from pyatlan.model.fluent_search import FluentSearch
15931592

1594-
client = AtlanClient.get_current_client()
15951593
if guid:
15961594
if qualified_name:
15971595
raise ErrorCode.QN_OR_GUID_NOT_BOTH.exception_with_parameters()
@@ -1600,15 +1598,15 @@ def remove_terms(
16001598
.select()
16011599
.where(Asset.TYPE_NAME.eq(asset_type.__name__))
16021600
.where(asset_type.GUID.eq(guid))
1603-
.execute(client=client)
1601+
.execute(client=self._client) # type: ignore[arg-type]
16041602
)
16051603
elif qualified_name:
16061604
results = (
16071605
FluentSearch()
16081606
.select()
16091607
.where(Asset.TYPE_NAME.eq(asset_type.__name__))
16101608
.where(asset_type.QUALIFIED_NAME.eq(qualified_name))
1611-
.execute(client=client)
1609+
.execute(client=self._client) # type: ignore[arg-type]
16121610
)
16131611
else:
16141612
raise ErrorCode.QN_OR_GUID.exception_with_parameters()
@@ -2567,7 +2565,7 @@ def flush(self) -> Optional[AssetMutationResponse]:
25672565
)
25682566
results = search.page_size(
25692567
max(self._max_size * 2, DSL.__fields__.get("size").default) # type: ignore[union-attr]
2570-
).execute(client=self._client)
2568+
).execute(client=self._client) # type: ignore[arg-type]
25712569

25722570
for asset in results:
25732571
asset_id = AssetIdentity(

0 commit comments

Comments
 (0)