Skip to content

Commit 4ec945c

Browse files
committed
[req-changes] Removed _default_client and cache_key, set _current_client on AtlanClient method calls
- Also fixes mypy violations.
1 parent 4bbb060 commit 4ec945c

30 files changed

+323
-418
lines changed

pyatlan/cache/abstract_asset_cache.py

Lines changed: 11 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -4,30 +4,28 @@
44

55
import threading
66
from abc import ABC, abstractmethod
7-
from typing import Any
7+
from typing import TYPE_CHECKING, Any, Dict
88

99
from pyatlan.errors import ErrorCode
1010
from pyatlan.model.assets import Asset
1111
from pyatlan.model.enums import AtlanConnectorType
1212

13+
if TYPE_CHECKING:
14+
from pyatlan.client.atlan import AtlanClient
15+
1316

1417
class AbstractAssetCache(ABC):
1518
"""
1619
Base class for reusable components that are common
1720
to all caches, where a cache is populated entry-by-entry.
1821
"""
1922

20-
def __init__(self, client):
23+
def __init__(self, client: AtlanClient):
2124
self.client = client
2225
self.lock = threading.Lock()
23-
self.name_to_guid = dict()
24-
self.guid_to_asset = dict()
25-
self.qualified_name_to_guid = dict()
26-
27-
@classmethod
28-
@abstractmethod
29-
def get_cache(cls):
30-
"""Abstract method to retreive cache."""
26+
self.name_to_guid: Dict[str, str] = dict()
27+
self.guid_to_asset: Dict[str, Asset] = dict()
28+
self.qualified_name_to_guid: Dict[str, str] = dict()
3129

3230
@abstractmethod
3331
def lookup_by_guid(self, guid: str):
@@ -84,9 +82,9 @@ def cache(self, asset: Asset):
8482
name = asset and self.get_name(asset)
8583
if not all([name, asset.guid, asset.qualified_name]):
8684
return
87-
self.name_to_guid[name] = asset.guid
88-
self.guid_to_asset[asset.guid] = asset
89-
self.qualified_name_to_guid[asset.qualified_name] = asset.guid
85+
self.name_to_guid[name] = asset.guid # type: ignore[index]
86+
self.guid_to_asset[asset.guid] = asset # type: ignore[index]
87+
self.qualified_name_to_guid[asset.qualified_name] = asset.guid # type: ignore[index]
9088

9189
def _get_by_guid(self, guid: str, allow_refresh: bool = True):
9290
"""

pyatlan/cache/atlan_tag_cache.py

Lines changed: 9 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
# Copyright 2025 Atlan Pte. Ltd.
33
from __future__ import annotations
44

5-
from threading import Lock, local
5+
from threading import Lock
66
from typing import TYPE_CHECKING, Dict, Optional, Set
77

88
from pyatlan.errors import ErrorCode
@@ -13,7 +13,6 @@
1313
from pyatlan.client.atlan import AtlanClient
1414

1515
lock: Lock = Lock()
16-
atlan_tag_cache_tls = local() # Thread-local storage (TLS)
1716

1817

1918
class AtlanTagCache:
@@ -32,61 +31,39 @@ def __init__(self, client: AtlanClient):
3231
self.map_id_to_source_tags_attr_id: Dict[str, str] = {}
3332
self.lock: Lock = Lock()
3433

35-
@classmethod
36-
def get_cache(cls, client: Optional[AtlanClient] = None) -> AtlanTagCache:
37-
from pyatlan.client.atlan import AtlanClient
38-
39-
with lock:
40-
client = client or AtlanClient.get_default_client()
41-
cache_key = client.cache_key
42-
43-
if not hasattr(atlan_tag_cache_tls, "caches"):
44-
atlan_tag_cache_tls.caches = {}
45-
46-
if cache_key not in atlan_tag_cache_tls.caches:
47-
cache_instance = AtlanTagCache(client=client)
48-
cache_instance._refresh_cache() # Refresh on new cache instance
49-
atlan_tag_cache_tls.caches[cache_key] = cache_instance
50-
51-
return atlan_tag_cache_tls.caches[cache_key]
52-
53-
@classmethod
54-
def refresh_cache(cls) -> None:
34+
def refresh_cache(self) -> None:
5535
"""
5636
Refreshes the cache of Atlan tags by requesting the full set of Atlan tags from Atlan.
5737
"""
58-
cls.get_cache()._refresh_cache()
38+
self._refresh_cache()
5939

60-
@classmethod
61-
def get_id_for_name(cls, name: str) -> Optional[str]:
40+
def get_id_for_name(self, name: str) -> Optional[str]:
6241
"""
6342
Translate the provided human-readable Atlan tag name to its Atlan-internal ID string.
6443
6544
:param name: human-readable name of the Atlan tag
6645
:returns: Atlan-internal ID string of the Atlan tag
6746
"""
68-
return cls.get_cache()._get_id_for_name(name=name)
47+
return self._get_id_for_name(name=name)
6948

70-
@classmethod
71-
def get_name_for_id(cls, idstr: str) -> Optional[str]:
49+
def get_name_for_id(self, idstr: str) -> Optional[str]:
7250
"""
7351
Translate the provided Atlan-internal classification ID string to the human-readable Atlan tag name.
7452
7553
:param idstr: Atlan-internal ID string of the Atlan tag
7654
:returns: human-readable name of the Atlan tag
7755
"""
78-
return cls.get_cache()._get_name_for_id(idstr=idstr)
56+
return self._get_name_for_id(idstr=idstr)
7957

80-
@classmethod
81-
def get_source_tags_attr_id(cls, id: str) -> Optional[str]:
58+
def get_source_tags_attr_id(self, id: str) -> Optional[str]:
8259
"""
8360
Translate the provided Atlan-internal Atlan tag ID string to the Atlan-internal name of the attribute that
8461
captures tag attachment details (for source-synced tags).
8562
8663
:param id: Atlan-internal ID string of the Atlan tag
8764
:returns: Atlan-internal ID string of the attribute containing source-synced tag attachment details
8865
"""
89-
return cls.get_cache()._get_source_tags_attr_id(id)
66+
return self._get_source_tags_attr_id(id)
9067

9168
def _refresh_cache(self) -> None:
9269
"""

pyatlan/cache/connection_cache.py

Lines changed: 8 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44

55
import logging
66
import threading
7-
from threading import local
87
from typing import TYPE_CHECKING, Optional, Union
98

109
from pyatlan.cache.abstract_asset_cache import AbstractAssetCache, AbstractAssetName
@@ -17,7 +16,6 @@
1716
from pyatlan.client.atlan import AtlanClient
1817

1918
lock = threading.Lock()
20-
connection_cache_tls = local() # Thread-local storage (TLS)
2119
LOGGER = logging.getLogger(__name__)
2220

2321

@@ -44,25 +42,7 @@ class ConnectionCache(AbstractAssetCache):
4442
def __init__(self, client: AtlanClient):
4543
super().__init__(client)
4644

47-
@classmethod
48-
def get_cache(cls, client: Optional[AtlanClient] = None) -> ConnectionCache:
49-
from pyatlan.client.atlan import AtlanClient
50-
51-
with lock:
52-
client = client or AtlanClient.get_default_client()
53-
cache_key = client.cache_key
54-
55-
if not hasattr(connection_cache_tls, "caches"):
56-
connection_cache_tls.caches = {}
57-
58-
if cache_key not in connection_cache_tls.caches:
59-
cache_instance = ConnectionCache(client=client)
60-
connection_cache_tls.caches[cache_key] = cache_instance
61-
62-
return connection_cache_tls.caches[cache_key]
63-
64-
@classmethod
65-
def get_by_guid(cls, guid: str, allow_refresh: bool = True) -> Connection:
45+
def get_by_guid(self, guid: str, allow_refresh: bool = True) -> Connection:
6646
"""
6747
Retrieve a connection from the cache by its UUID.
6848
If the asset is not found, it will be looked up and added to the cache.
@@ -74,11 +54,10 @@ def get_by_guid(cls, guid: str, allow_refresh: bool = True) -> Connection:
7454
:raises NotFoundError: if the connection cannot be found (does not exist) in Atlan
7555
:raises InvalidRequestError: if no UUID was provided for the connection to retrieve
7656
"""
77-
return cls.get_cache()._get_by_guid(guid=guid, allow_refresh=allow_refresh)
57+
return self._get_by_guid(guid=guid, allow_refresh=allow_refresh)
7858

79-
@classmethod
8059
def get_by_qualified_name(
81-
cls, qualified_name: str, allow_refresh: bool = True
60+
self, qualified_name: str, allow_refresh: bool = True
8261
) -> Connection:
8362
"""
8463
Retrieve a connection from the cache by its unique Atlan-internal name.
@@ -92,13 +71,12 @@ def get_by_qualified_name(
9271
:raises NotFoundError: if the connection cannot be found (does not exist) in Atlan
9372
:raises InvalidRequestError: if no qualified_name was provided for the connection to retrieve
9473
"""
95-
return cls.get_cache()._get_by_qualified_name(
74+
return self._get_by_qualified_name(
9675
qualified_name=qualified_name, allow_refresh=allow_refresh
9776
)
9877

99-
@classmethod
10078
def get_by_name(
101-
cls, name: ConnectionName, allow_refresh: bool = True
79+
self, name: ConnectionName, allow_refresh: bool = True
10280
) -> Connection:
10381
"""
10482
Retrieve an connection from the cache by its uniquely identifiable name.
@@ -112,7 +90,7 @@ def get_by_name(
11290
:raises NotFoundError: if the connection cannot be found (does not exist) in Atlan
11391
:raises InvalidRequestError: if no name was provided for the connection to retrieve
11492
"""
115-
return cls.get_cache()._get_by_name(name=name, allow_refresh=allow_refresh)
93+
return self._get_by_name(name=name, allow_refresh=allow_refresh)
11694

11795
def lookup_by_guid(self, guid: str) -> None:
11896
if not guid:
@@ -149,8 +127,8 @@ def lookup_by_name(self, name: ConnectionName) -> None:
149127
return
150128
with self.lock:
151129
results = self.client.asset.find_connections_by_name(
152-
name=name.name,
153-
connector_type=name.type,
130+
name=name.name, # type: ignore[arg-type]
131+
connector_type=name.type, # type: ignore[arg-type]
154132
attributes=self.SEARCH_ATTRIBUTES,
155133
)
156134
if not results:

0 commit comments

Comments
 (0)