Skip to content

Commit

Permalink
Replace host with key in OAuth utils methods
Browse files Browse the repository at this point in the history
  • Loading branch information
hovaesco authored and hashhar committed Feb 16, 2024
1 parent 5a86e32 commit e90ba3a
Showing 1 changed file with 14 additions and 14 deletions.
28 changes: 14 additions & 14 deletions trino/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,11 +208,11 @@ class _OAuth2TokenCache(metaclass=abc.ABCMeta):
"""

@abc.abstractmethod
def get_token_from_cache(self, host: Optional[str]) -> Optional[str]:
def get_token_from_cache(self, key: Optional[str]) -> Optional[str]:
pass

@abc.abstractmethod
def store_token_to_cache(self, host: Optional[str], token: str) -> None:
def store_token_to_cache(self, key: Optional[str], token: str) -> None:
pass


Expand All @@ -224,11 +224,11 @@ class _OAuth2TokenInMemoryCache(_OAuth2TokenCache):
def __init__(self) -> None:
self._cache: Dict[Optional[str], str] = {}

def get_token_from_cache(self, host: Optional[str]) -> Optional[str]:
return self._cache.get(host)
def get_token_from_cache(self, key: Optional[str]) -> Optional[str]:
return self._cache.get(key)

def store_token_to_cache(self, host: Optional[str], token: str) -> None:
self._cache[host] = token
def store_token_to_cache(self, key: Optional[str], token: str) -> None:
self._cache[key] = token


class _OAuth2KeyRingTokenCache(_OAuth2TokenCache):
Expand All @@ -248,18 +248,18 @@ def is_keyring_available(self) -> bool:
return self._keyring is not None \
and not isinstance(self._keyring.get_keyring(), self._keyring.backends.fail.Keyring)

def get_token_from_cache(self, host: Optional[str]) -> Optional[str]:
def get_token_from_cache(self, key: Optional[str]) -> Optional[str]:
try:
return self._keyring.get_password(host, "token")
return self._keyring.get_password(key, "token")
except self._keyring.errors.NoKeyringError as e:
raise trino.exceptions.NotSupportedError("Although keyring module is installed no backend has been "
"detected, check https://pypi.org/project/keyring/ for more "
"information.") from e

def store_token_to_cache(self, host: Optional[str], token: str) -> None:
def store_token_to_cache(self, key: Optional[str], token: str) -> None:
try:
# keyring is installed, so we can store the token for reuse within multiple threads
self._keyring.set_password(host, "token", token)
self._keyring.set_password(key, "token", token)
except self._keyring.errors.NoKeyringError as e:
raise trino.exceptions.NotSupportedError("Although keyring module is installed no backend has been "
"detected, check https://pypi.org/project/keyring/ for more "
Expand Down Expand Up @@ -382,13 +382,13 @@ def _get_token(self, token_server: str, response: Response, **kwargs: Any) -> st

raise exceptions.TrinoAuthError("Exceeded max attempts while getting the token")

def _get_token_from_cache(self, host: Optional[str]) -> Optional[str]:
def _get_token_from_cache(self, key: Optional[str]) -> Optional[str]:
with self._token_lock:
return self._token_cache.get_token_from_cache(host)
return self._token_cache.get_token_from_cache(key)

def _store_token_to_cache(self, host: Optional[str], token: str) -> None:
def _store_token_to_cache(self, key: Optional[str], token: str) -> None:
with self._token_lock:
self._token_cache.store_token_to_cache(host, token)
self._token_cache.store_token_to_cache(key, token)

@staticmethod
def _determine_host(url: Optional[str]) -> Any:
Expand Down

0 comments on commit e90ba3a

Please sign in to comment.