Use keyring to securely cache the Oauth2 token

Author: mdesmet

README.md

@@ -172,7 +172,7 @@ the [OAuth2 authentication type](https://trino.io/docs/current/security/oauth2.h
 
 A callback to handle the redirect url can be provided via param `redirect_auth_url_handler` of the `trino.auth.OAuth2Authentication` class. By default, it will try to launch a web browser (`trino.auth.WebBrowserRedirectHandler`) to go through the authentication flow and output the redirect url to stdout (`trino.auth.ConsoleRedirectHandler`). Multiple redirect handlers are combined using the `trino.auth.CompositeRedirectHandler` class.
 
-The OAuth2 token will be cached either per `trino.auth.OAuth2Authentication` instance.
+The OAuth2 token will be cached either per `trino.auth.OAuth2Authentication` instance or, when keyring is installed, it will be cached within a secure backend (MacOS keychain, Windows credential locker, etc) under a key including host of the Trino connection. Keyring can be installed using `pip install 'trino[external-authentication-token-cache]'`.
 
 - DBAPI
 

setup.py

@@ -27,7 +27,9 @@
 
 kerberos_require = ["requests_kerberos"]
 sqlalchemy_require = ["sqlalchemy~=1.3"]
+external_authentication_token_cache_require = ["keyring"]
 
+# We don't add localstorage_require to all_require as users must explicitly opt in to use keyring.
 all_require = kerberos_require + sqlalchemy_require
 
 tests_require = all_require + [
@@ -80,6 +82,7 @@
         "kerberos": kerberos_require,
         "sqlalchemy": sqlalchemy_require,
         "tests": tests_require,
+        "external-authentication-token-cache": external_authentication_token_cache_require,
     },
     entry_points={
         "sqlalchemy.dialects": [

trino/auth.py

@@ -22,6 +22,7 @@
 from requests import Request
 from requests.auth import AuthBase, extract_cookies_to_jar
 from requests.utils import parse_dict_header
+import importlib
 
 import trino.logging
 from trino.client import exceptions
@@ -230,6 +231,40 @@ def store_token_to_cache(self, host: str, token: str) -> None:
         self._cache[host] = token
 
 
+class _OAuth2KeyRingTokenCache(_OAuth2TokenCache):
+    """
+    Keyring Token Cache implementation
+    """
+
+    def __init__(self):
+        super().__init__()
+        try:
+            self._keyring = importlib.import_module("keyring")
+        except ImportError:
+            self._keyring = None
+            logger.info("keyring module not found. OAuth2 token will not be stored in keyring.")
+
+    def is_keyring_available(self) -> bool:
+        return self._keyring is not None
+
+    def get_token_from_cache(self, host: str) -> Optional[str]:
+        try:
+            return self._keyring.get_password(host, "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: 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)
+        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
+
+
 class _OAuth2TokenBearer(AuthBase):
     """
     Custom implementation of Trino Oauth2 based authorization to get the token
@@ -239,7 +274,8 @@ class _OAuth2TokenBearer(AuthBase):
 
     def __init__(self, redirect_auth_url_handler: Callable[[str], None]):
         self._redirect_auth_url = redirect_auth_url_handler
-        self._token_cache = _OAuth2TokenInMemoryCache()
+        keyring_cache = _OAuth2KeyRingTokenCache()
+        self._token_cache = keyring_cache if keyring_cache.is_keyring_available() else _OAuth2TokenInMemoryCache()
         self._token_lock = threading.Lock()
         self._inside_oauth_attempt_lock = threading.Lock()
         self._inside_oauth_attempt_blocker = threading.Event()