Repository: trinodb/trino-python-client

Issue:
Improve Oauth authentication flow to launch webbrowser and cache token
The Trino cli client and JDBC driver allow the Oauth2 token to be cached, this prevents prompting the user to go through authentication flow for each separate query. 

**Current behavior**

Currently the python trino client prints out urls in the console for the user. 

![image](https://user-images.githubusercontent.com/9265503/155877654-52f33993-0c72-4675-8003-a4082190bd4e.png)

**Improved behavior**

The first query per thread shall launch a webbrowser to go through the authentication flow with the external authenticator. On success it will cache the token in the user keychain (secure local storage) if available (or in a per connection cache), similar to the trino CLI and JDBC driver.

Ideally the token can be stored to be reused over multiple threads (eg. dbt launches multiple threads with each their own designated database connection).




Discussion:
mdesmet: Adding the keyring feature instead of the threadlocal made the test fail. IMHO the threadlocal seems not so useful, this works different from the Trino CLI and JDBC driver. However maybe from security perspective it's a good default to not cache the token over multiple threads.

In trino CLI:
See https://trino.io/docs/current/installation/cli.html#external-authentication-sso

Further queries in the CLI session do not require additional logins while the authentication token remains valid. Token expiration depends on the external authentication type configuration.

Expired tokens force you to log in again.

in trino JDBC:
See https://trino.io/docs/current/installation/jdbc.html#parameter-reference
externalAuthenticationTokenCache | Allows the sharing of external authentication tokens between different connections for the same authenticated user until the cache is invalidated, such as when a client is restarted or when the classloader reloads the JDBC driver. This is disabled by default, with a value of NONE. To enable, set the value to MEMORY. If the JDBC driver is used in a shared mode by different users, the first registered token is stored and authenticates all users.

@lukasz-walkiewicz : what do you think?


mdesmet: Open items: 

- Thread safety impact: writing and reading token is not thread safe. Might need locking to ensure only one token is obtained simultaneously.
- Hardcoded name of token, should be based on url and username.
lukasz-walkiewicz: First of all, thanks for the contribution @mdesmet.
I've tried running this locally in a new virtual environment however it didn't work:
```
(keyring)  ✘  ~/dev/prestosql/trino-python-client   feature/webbrowser  python main.py                                                                                      ~/dev/prestosql/trino-python-client
error could be tests
Exception in thread 1: :
Traceback (most recent call last):
  File "/nix/store/lrw7ifwnpz7rskfkphq9mj9fwm5i9ldc-python3-3.9.6/lib/python3.9/threading.py", line 973, in _bootstrap_inner
Exception in thread 2: :
Traceback (most recent call last):
  File "/nix/store/lrw7ifwnpz7rskfkphq9mj9fwm5i9ldc-python3-3.9.6/lib/python3.9/threading.py", line 973, in _bootstrap_inner
    self.run()
  File "/Users/lukasz/dev/prestosql/trino-python-client/main.py", line 30, in run
    self.run()
  File "/Users/lukasz/dev/prestosql/trino-python-client/main.py", line 30, in run
    cur.execute('SELECT * FROM system.runtime.nodes')
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/dbapi.py", line 441, in execute
    cur.execute('SELECT * FROM system.runtime.nodes')
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/dbapi.py", line 441, in execute
    result = self._query.execute()
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/client.py", line 561, in execute
    result = self._query.execute()
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/client.py", line 561, in execute
    response = self._request.post(self._sql, additional_http_headers)
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/client.py", line 361, in post
    response = self._request.post(self._sql, additional_http_headers)
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/client.py", line 361, in post
    http_response = self._post(
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/exceptions.py", line 136, in decorated
    http_response = self._post(
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/exceptions.py", line 136, in decorated
    raise error
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/exceptions.py", line 123, in decorated
    raise error
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/exceptions.py", line 123, in decorated
    result = func(*args, **kwargs)
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/requests/sessions.py", line 577, in post
    result = func(*args, **kwargs)
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/requests/sessions.py", line 577, in post
    return self.request('POST', url, data=data, json=json, **kwargs)
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/requests/sessions.py", line 515, in request
    return self.request('POST', url, data=data, json=json, **kwargs)
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/requests/sessions.py", line 515, in request
    prep = self.prepare_request(req)
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/requests/sessions.py", line 443, in prepare_request
    prep = self.prepare_request(req)
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/requests/sessions.py", line 443, in prepare_request
    p.prepare(
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/requests/models.py", line 322, in prepare
    p.prepare(
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/requests/models.py", line 322, in prepare
    self.prepare_auth(auth, url)
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/requests/models.py", line 558, in prepare_auth
    self.prepare_auth(auth, url)
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/requests/models.py", line 558, in prepare_auth
    r = auth(self)
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/auth.py", line 192, in __call__
    r = auth(self)
    token = self._get_token_from_cache(*user_host)
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/auth.py", line 287, in _get_token_from_cache
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/auth.py", line 192, in __call__
    return self._keyring.get_password(
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/keyring/core.py", line 55, in get_password
    token = self._get_token_from_cache(*user_host)
  File "/Users/lukasz/dev/prestosql/trino-python-client/trino/auth.py", line 287, in _get_token_from_cache
    return get_keyring().get_password(service_name, username)
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/keyring/backends/fail.py", line 25, in get_password
    return self._keyring.get_password(
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/keyring/core.py", line 55, in get_password
    raise NoKeyringError(msg)
keyring.errors.NoKeyringError: No recommended backend was available. Install a recommended 3rd party backend package; or, install the keyrings.alt package if you want to use the non-recommended backends. See https://pypi.org/project/keyring for details.
    return get_keyring().get_password(service_name, username)
  File "/Users/lukasz/venvs/keyring/lib/python3.9/site-packages/keyring/backends/fail.py", line 25, in get_password
    raise NoKeyringError(msg)
keyring.errors.NoKeyringError: No recommended backend was available. Install a recommended 3rd party backend package; or, install the keyrings.alt package if you want to use the non-recommended backends. See https://pypi.org/project/keyring for details.
```
Environment:
```
(keyring)  ~/dev/prestosql/trino-python-client   feature/webbrowser  python                                                                                                  ~/dev/prestosql/trino-python-client
Python 3.9.6 (default, Oct 31 2021, 15:37:00)
[Clang 7.1.0 (tags/RELEASE_710/final)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import platform
>>> platform.platform()
'macOS-10.16-x86_64-i386-64bit'
```
Which platform you tested it on?
mdesmet: Tested on following platform:
```
>>> platform.platform()
'macOS-12.2.1-arm64-arm-64bit'
```

Checking on https://pypi.org/project/keyring/

Compatibility - macOS
macOS keychain support **macOS 11** (Big Sur) and later requires Python 3.8.7 or later with the “universal2” binary. See [#525](https://github.com/jaraco/keyring/issues/525) for details.

mdesmet: > First of all, thanks for the contribution @mdesmet. I've tried running this locally in a new virtual environment however it didn't work:
> 

So it's possible that even if you have installed keyring, it doesn't work as there is no suitable backend. I adapted the code that it will fallback to just store the token as if keyring is not installed.
mdesmet: > > I think i can add another test to proof that the caching of the token is per connection and not per cursor as before.
> 
> That would be great.

I added the test: test_token_only_retrieved_once_per_connection
mdesmet: @lukasz-walkiewicz, @hovaesco  : Any other remarks or concerns? 

FYI launching the webbrowser and storing credentials in keyring is also what Snowflake is doing in their python connector, I think it will greatly improve the user experience of dbt in combination with oauth authentication. Am I missing any usage patterns that this behavior for OAuth authentication won't be compatible with?
mdesmet: > Looks good. One question, if keyring is installed then a token is cached by default? Basically, having keyring installed is like an enabler to cache a token?

Even without keyring installed the token is always cached per connection (actually per instance of OAuth2Authentication), compared to the old code where the token was cached per cursor. When you install keyring, the token will be reused even when you restart or rerun your application (eg multiple dbt runs won't prompt authentication again if your token is still valid).
hovaesco: Awesome job 🎉 One general question. Have you fully tested it with dbt-trino? 
mdesmet: > Awesome job 🎉 One general question. Have you fully tested it with dbt-trino?

Haven't tried that. Will try it during next week. On first sight we might have to ensure that only one instance of `trino.auth.OAuth2Authentication()` is created.
mdesmet: > Awesome job 🎉 One general question. Have you fully tested it with dbt-trino?

I have succesfully tested the feature with dbt and also added the changes in the [dbt PR](https://github.com/starburstdata/dbt-trino/pull/41). We still will need to update the trino-python-client version there, after release.

If you want to try it yourself, you can run following commands

```
pip install git+https://github.com/mdesmet/dbt-trino.git@feature/oidc 
pip install git+https://github.com/mdesmet/trino-python-client.git@feature/webbrowser
```

mdesmet: > Impl looks good to me. I'd love to be able to integration test these instead of doing manually but that's separate effort.
> 

There was still some issue with the locking that i fixed yesterday. The integration tests actually lock on each query, which means that every request including retry is queued instead of executed in parallel. The locking is required to make httpretty threadsafe. So the existing tests are not completely covering the real nature of a multiple requests blocking until one finishes the oauth authentication process.

https://github.com/trinodb/trino-python-client/blob/0430df37d8074442daa22ec6c35a742efd75e363/tests/unit/test_client.py#L589-L590



hashhar: I don't think `cursor` is expected to be thread-safe. What does DB-API say about it?

Or do you mean we lock per-instance of connection?
mdesmet: > I don't think `cursor` is expected to be thread-safe. What does DB-API say about it?
> 
> Or do you mean we lock per-instance of connection?

It's not that we want to share a cursor over multiple threads. It's about the requests that happen within the cursor execution. The cursor is the first place where we actually talk to the coordinator, as the connection is just a cursor factory. So within a mulithreaded execution (each thread with their own cursor), multiple requests receive a 401 error and should wait until the first lock goes through the OAuth authentication flow. When that one thread completes and stores the token, the other threads should wake up and retry the request.

Within the integration tests, as we put a lock on every cursor execution, this effectively lets only one thread go through the request, 401, retrieve token and retry request flow. while all other threads are waiting and then execute with the token provided by the first thread. This doesn't match the real behaviour where actually multiple threads fail on 401 (as mentioned above).

I tried without the lock but then the httpretty request counts not match up anymore. So the issue lies within the test framework as mentioned in the comment in the code.

https://github.com/trinodb/trino-python-client/blob/0430df37d8074442daa22ec6c35a742efd75e363/tests/unit/test_client.py#L587-L591
mdesmet: ` keyring.errors.NoKeyringError: No recommended backend was available. Install a recommended 3rd party backend package; or, install the keyrings.alt package if you want to use the non-recommended backends. See https://pypi.org/project/keyring for details.`

We can't add add `localstorage_require` to `all_require` as within the tests there is no keyring backend. I will revert that change. I also don't think we should install keyring without explicit approval from the user.
hashhar: Thanks, please add that as a comment to `localstorage_require` about why it's missing from `all_require`.

Patch:
diff --git a/README.md b/README.md
index 09a4c5e..12641a7 100644
--- a/README.md
+++ b/README.md
@@ -167,11 +167,14 @@ the [`JWT` authentication type](https://trino.io/docs/current/security/jwt.html)
 
 ### OAuth2 Authentication
 
-- `OAuth2Authentication` class can be used to connect to a Trino cluster configured with
+The `OAuth2Authentication` class can be used to connect to a Trino cluster configured with
 the [OAuth2 authentication type](https://trino.io/docs/current/security/oauth2.html).
-- A callback to handle the redirect url can be provided via param `redirect_auth_url_handler`, by default it just outputs the redirect url to stdout.
 
-* DBAPI
+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 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
 
     ```python
     from trino.dbapi import connect
@@ -185,7 +188,7 @@ the [OAuth2 authentication type](https://trino.io/docs/current/security/oauth2.h
     )
     ```
 
-* SQLAlchemy
+- SQLAlchemy
 
     ```python
     from sqlalchemy import create_engine
diff --git a/setup.py b/setup.py
index c5f4ef2..2695a53 100755
--- a/setup.py
+++ b/setup.py
@@ -27,7 +27,9 @@ with open("trino/__init__.py", "rb") as f:
 
 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 @@ setup(
         "kerberos": kerberos_require,
         "sqlalchemy": sqlalchemy_require,
         "tests": tests_require,
+        "external-authentication-token-cache": external_authentication_token_cache_require,
     },
     entry_points={
         "sqlalchemy.dialects": [
diff --git a/trino/auth.py b/trino/auth.py
index 236794e..e6b4f04 100644
--- a/trino/auth.py
+++ b/trino/auth.py
@@ -15,11 +15,14 @@ import json
 import os
 import re
 import threading
-from typing import Optional
+import webbrowser
+from typing import Optional, List, Callable
+from urllib.parse import urlparse
 
 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
@@ -130,6 +133,7 @@ class _BearerAuth(AuthBase):
     """
     Custom implementation of Authentication class for bearer token
     """
+
     def __init__(self, token):
         self.token = token
 
@@ -156,9 +160,109 @@ class JWTAuthentication(Authentication):
         return self.token == other.token
 
 
-def handle_redirect_auth_url(auth_url):
-    print("Open the following URL in browser for the external authentication:")
-    print(auth_url)
+class RedirectHandler(metaclass=abc.ABCMeta):
+    """
+    Abstract class for OAuth redirect handlers, inherit from this class to implement your own redirect handler.
+    """
+
+    @abc.abstractmethod
+    def __call__(self, url: str) -> None:
+        raise NotImplementedError()
+
+
+class ConsoleRedirectHandler(RedirectHandler):
+    """
+    Handler for OAuth redirections to log to console.
+    """
+
+    def __call__(self, url: str) -> None:
+        print("Open the following URL in browser for the external authentication:")
+        print(url)
+
+
+class WebBrowserRedirectHandler(RedirectHandler):
+    """
+    Handler for OAuth redirections to open in web browser.
+    """
+
+    def __call__(self, url: str) -> None:
+        webbrowser.open_new(url)
+
+
+class CompositeRedirectHandler(RedirectHandler):
+    """
+    Composite handler for OAuth redirect handlers.
+    """
+
+    def __init__(self, handlers: List[Callable[[str], None]]):
+        self.handlers = handlers
+
+    def __call__(self, url: str):
+        for handler in self.handlers:
+            handler(url)
+
+
+class _OAuth2TokenCache(metaclass=abc.ABCMeta):
+    """
+    Abstract class for OAuth token cache, inherit from this class to implement your own token cache.
+    """
+
+    @abc.abstractmethod
+    def get_token_from_cache(self, host: str) -> Optional[str]:
+        pass
+
+    @abc.abstractmethod
+    def store_token_to_cache(self, host: str, token: str) -> None:
+        pass
+
+
+class _OAuth2TokenInMemoryCache(_OAuth2TokenCache):
+    """
+    In-memory token cache implementation. The token is stored per host, so multiple clients can share the same cache.
+    """
+
+    def __init__(self):
+        self._cache = {}
+
+    def get_token_from_cache(self, host: str) -> Optional[str]:
+        return self._cache.get(host)
+
+    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):
@@ -168,14 +272,20 @@ class _OAuth2TokenBearer(AuthBase):
     MAX_OAUTH_ATTEMPTS = 5
     _BEARER_PREFIX = re.compile(r"bearer", flags=re.IGNORECASE)
 
-    def __init__(self, http_session, redirect_auth_url_handler=handle_redirect_auth_url):
+    def __init__(self, redirect_auth_url_handler: Callable[[str], None]):
         self._redirect_auth_url = redirect_auth_url_handler
-        self._thread_local = threading.local()
-        http_session.hooks['response'].append(self._authenticate)
+        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()
 
     def __call__(self, r):
-        if hasattr(self._thread_local, 'token') and self._thread_local.token:
-            r.headers['Authorization'] = "Bearer " + self._thread_local.token
+        host = self._determine_host(r.url)
+        token = self._get_token_from_cache(host)
+
+        if token is not None:
+            r.headers['Authorization'] = "Bearer " + token
 
         r.register_hook('response', self._authenticate)
 
@@ -185,7 +295,23 @@ class _OAuth2TokenBearer(AuthBase):
         if not 400 <= response.status_code < 500:
             return response
 
-        # we have to handle the authentication, may be token the token expired or it wasn't there at all
+        acquired = self._inside_oauth_attempt_lock.acquire(blocking=False)
+        if acquired:
+            try:
+                # Lock is acquired, attempt the OAuth2 flow
+                self._attempt_oauth(response, **kwargs)
+                self._inside_oauth_attempt_blocker.set()
+            finally:
+                self._inside_oauth_attempt_lock.release()
+                self._inside_oauth_attempt_blocker.clear()
+        else:
+            # Lock is not acquired, we are already in the OAuth2 flow, so we block until OAuth2 flow is finished.
+            self._inside_oauth_attempt_blocker.wait()
+
+        return self._retry_request(response, **kwargs)
+
+    def _attempt_oauth(self, response, **kwargs):
+        # we have to handle the authentication, may be token the token expired, or it wasn't there at all
         auth_info = response.headers.get('WWW-Authenticate')
         if not auth_info:
             raise exceptions.TrinoAuthError("Error: header WWW-Authenticate not available in the response.")
@@ -203,8 +329,6 @@ class _OAuth2TokenBearer(AuthBase):
         if token_server is None:
             raise exceptions.TrinoAuthError("Error: header info didn't have x_token_server")
 
-        self._thread_local.token_server = token_server
-
         # tell app that use this url to proceed with the authentication
         self._redirect_auth_url(auth_server)
 
@@ -213,15 +337,19 @@ class _OAuth2TokenBearer(AuthBase):
         response.content
         response.close()
 
-        self._thread_local.token = self._get_token(token_server, response, **kwargs)
-        return self._retry_request(response, **kwargs)
+        token = self._get_token(token_server, response, **kwargs)
+
+        request = response.request
+        host = self._determine_host(request.url)
+        self._store_token_to_cache(host, token)
 
     def _retry_request(self, response, **kwargs):
         request = response.request.copy()
         extract_cookies_to_jar(request._cookies, response.request, response.raw)
         request.prepare_cookies(request._cookies)
 
-        request.headers['Authorization'] = "Bearer " + self._thread_local.token
+        host = self._determine_host(response.request.url)
+        request.headers['Authorization'] = "Bearer " + self._get_token_from_cache(host)
         retry_response = response.connection.send(request, **kwargs)
         retry_response.history.append(response)
         retry_response.request = request
@@ -251,13 +379,29 @@ class _OAuth2TokenBearer(AuthBase):
 
         raise exceptions.TrinoAuthError("Exceeded max attempts while getting the token")
 
+    def _get_token_from_cache(self, host: str) -> Optional[str]:
+        with self._token_lock:
+            return self._token_cache.get_token_from_cache(host)
+
+    def _store_token_to_cache(self, host: str, token: str) -> None:
+        with self._token_lock:
+            self._token_cache.store_token_to_cache(host, token)
+
+    @staticmethod
+    def _determine_host(url) -> Optional[str]:
+        return urlparse(url).hostname
+
 
 class OAuth2Authentication(Authentication):
-    def __init__(self, redirect_auth_url_handler=handle_redirect_auth_url):
+    def __init__(self, redirect_auth_url_handler=CompositeRedirectHandler([
+        WebBrowserRedirectHandler(),
+        ConsoleRedirectHandler()
+    ])):
         self._redirect_auth_url = redirect_auth_url_handler
+        self._bearer = _OAuth2TokenBearer(self._redirect_auth_url)
 
     def set_http_session(self, http_session):
-        http_session.auth = _OAuth2TokenBearer(http_session, self._redirect_auth_url)
+        http_session.auth = self._bearer
         return http_session
 
     def get_exceptions(self):


Test patch:
diff --git a/tests/unit/oauth_test_utils.py b/tests/unit/oauth_test_utils.py
new file mode 100644
index 0000000..77b891c
--- /dev/null
+++ b/tests/unit/oauth_test_utils.py
@@ -0,0 +1,128 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import re
+import uuid
+from collections import namedtuple
+
+import httpretty
+
+from trino import constants
+
+SERVER_ADDRESS = "https://coordinator"
+REDIRECT_PATH = "oauth2/initiate"
+TOKEN_PATH = "oauth2/token"
+REDIRECT_RESOURCE = f"{SERVER_ADDRESS}/{REDIRECT_PATH}"
+TOKEN_RESOURCE = f"{SERVER_ADDRESS}/{TOKEN_PATH}"
+
+
+class RedirectHandler:
+    def __init__(self):
+        self.redirect_server = ""
+
+    def __call__(self, url):
+        self.redirect_server += url
+
+
+class PostStatementCallback:
+    def __init__(self, redirect_server, token_server, tokens, sample_post_response_data):
+        self.redirect_server = redirect_server
+        self.token_server = token_server
+        self.tokens = tokens
+        self.sample_post_response_data = sample_post_response_data
+
+    def __call__(self, request, uri, response_headers):
+        authorization = request.headers.get("Authorization")
+        if authorization and authorization.replace("Bearer ", "") in self.tokens:
+            return [200, response_headers, json.dumps(self.sample_post_response_data)]
+        return [401, {'Www-Authenticate': f'Bearer x_redirect_server="{self.redirect_server}", '
+                                          f'x_token_server="{self.token_server}"',
+                      'Basic realm': '"Trino"'}, ""]
+
+
+class GetTokenCallback:
+    def __init__(self, token_server, token, attempts=1):
+        self.token_server = token_server
+        self.token = token
+        self.attempts = attempts
+
+    def __call__(self, request, uri, response_headers):
+        self.attempts -= 1
+        if self.attempts < 0:
+            return [404, response_headers, "{}"]
+        if self.attempts == 0:
+            return [200, response_headers, f'{{"token": "{self.token}"}}']
+        return [200, response_headers, f'{{"nextUri": "{self.token_server}"}}']
+
+
+def _get_token_requests(challenge_id):
+    return list(filter(
+        lambda r: r.method == "GET" and r.path == f"/{TOKEN_PATH}/{challenge_id}",
+        httpretty.latest_requests()))
+
+
+def _post_statement_requests():
+    return list(filter(
+        lambda r: r.method == "POST" and r.path == constants.URL_STATEMENT_PATH,
+        httpretty.latest_requests()))
+
+
+class MultithreadedTokenServer:
+    Challenge = namedtuple('Challenge', ['token', 'attempts'])
+
+    def __init__(self, sample_post_response_data, attempts=1):
+        self.tokens = set()
+        self.challenges = {}
+        self.sample_post_response_data = sample_post_response_data
+        self.attempts = attempts
+
+        # bind post statement
+        httpretty.register_uri(
+            method=httpretty.POST,
+            uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
+            body=self.post_statement_callback)
+
+        # bind get token
+        httpretty.register_uri(
+            method=httpretty.GET,
+            uri=re.compile(rf"{TOKEN_RESOURCE}/.*"),
+            body=self.get_token_callback)
+
+    # noinspection PyUnusedLocal
+    def post_statement_callback(self, request, uri, response_headers):
+        authorization = request.headers.get("Authorization")
+
+        if authorization and authorization.replace("Bearer ", "") in self.tokens:
+            return [200, response_headers, json.dumps(self.sample_post_response_data)]
+
+        challenge_id = str(uuid.uuid4())
+        token = str(uuid.uuid4())
+        self.tokens.add(token)
+        self.challenges[challenge_id] = MultithreadedTokenServer.Challenge(token, self.attempts)
+        redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}"
+        token_server = f"{TOKEN_RESOURCE}/{challenge_id}"
+        return [401, {'Www-Authenticate': f'Bearer x_redirect_server="{redirect_server}", '
+                                          f'x_token_server="{token_server}"',
+                      'Basic realm': '"Trino"'}, ""]
+
+    # noinspection PyUnusedLocal
+    def get_token_callback(self, request, uri, response_headers):
+        challenge_id = uri.replace(f"{TOKEN_RESOURCE}/", "")
+        challenge = self.challenges[challenge_id]
+        challenge = challenge._replace(attempts=challenge.attempts - 1)
+        self.challenges[challenge_id] = challenge
+        if challenge.attempts < 0:
+            return [404, response_headers, "{}"]
+        if challenge.attempts == 0:
+            return [200, response_headers, f'{{"token": "{challenge.token}"}}']
+        return [200, response_headers, f'{{"nextUri": "{uri}"}}']
diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py
index 3c26795..f03b5f7 100644
--- a/tests/unit/test_client.py
+++ b/tests/unit/test_client.py
@@ -10,11 +10,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import json
-import re
 import threading
 import time
 import uuid
-from collections import namedtuple
 from unittest import mock
 from urllib.parse import urlparse
 
@@ -25,6 +23,9 @@ from httpretty import httprettified
 from requests_kerberos.exceptions import KerberosExchangeError
 
 import trino.exceptions
+from tests.unit.oauth_test_utils import RedirectHandler, GetTokenCallback, PostStatementCallback, \
+    MultithreadedTokenServer, _post_statement_requests, _get_token_requests, REDIRECT_RESOURCE, TOKEN_RESOURCE, \
+    SERVER_ADDRESS
 from trino import constants
 from trino.auth import KerberosAuthentication, _OAuth2TokenBearer
 from trino.client import TrinoQuery, TrinoRequest, TrinoResult
@@ -259,52 +260,6 @@ def test_request_timeout():
     httpretty.reset()
 
 
-SERVER_ADDRESS = "https://coordinator"
-REDIRECT_PATH = "oauth2/initiate"
-TOKEN_PATH = "oauth2/token"
-REDIRECT_RESOURCE = f"{SERVER_ADDRESS}/{REDIRECT_PATH}"
-TOKEN_RESOURCE = f"{SERVER_ADDRESS}/{TOKEN_PATH}"
-
-
-class RedirectHandler:
-    def __init__(self):
-        self.redirect_server = ""
-
-    def __call__(self, url):
-        self.redirect_server += url
-
-
-class PostStatementCallback:
-    def __init__(self, redirect_server, token_server, tokens, sample_post_response_data):
-        self.redirect_server = redirect_server
-        self.token_server = token_server
-        self.tokens = tokens
-        self.sample_post_response_data = sample_post_response_data
-
-    def __call__(self, request, uri, response_headers):
-        authorization = request.headers.get("Authorization")
-        if authorization and authorization.replace("Bearer ", "") in self.tokens:
-            return [200, response_headers, json.dumps(self.sample_post_response_data)]
-        return [401, {'Www-Authenticate': f'Bearer x_redirect_server="{self.redirect_server}", '
-                                          f'x_token_server="{self.token_server}"',
-                      'Basic realm': '"Trino"'}, ""]
-
-
-class GetTokenCallback:
-    def __init__(self, token_server, token, attempts=1):
-        self.token_server = token_server
-        self.token = token
-        self.attempts = attempts
-
-    def __call__(self, request, uri, response_headers):
-        self.attempts -= 1
-        if self.attempts < 0:
-            return [404, response_headers, "{}"]
-        if self.attempts == 0:
-            return [200, response_headers, f'{{"token": "{self.token}"}}']
-        return [200, response_headers, f'{{"nextUri": "{self.token_server}"}}']
-
-
 @pytest.mark.parametrize("attempts", [1, 3, 5])
 @httprettified
 def test_oauth2_authentication_flow(attempts, sample_post_response_data):
@@ -511,57 +466,6 @@ def test_oauth2_authentication_fail_token_server(http_status, sample_post_respon
     assert len(_get_token_requests(challenge_id)) == 1
 
 
-class MultithreadedTokenServer:
-    Challenge = namedtuple('Challenge', ['token', 'attempts'])
-
-    def __init__(self, sample_post_response_data, attempts=1):
-        self.tokens = set()
-        self.challenges = {}
-        self.sample_post_response_data = sample_post_response_data
-        self.attempts = attempts
-
-        # bind post statement
-        httpretty.register_uri(
-            method=httpretty.POST,
-            uri=f"{SERVER_ADDRESS}{constants.URL_STATEMENT_PATH}",
-            body=self.post_statement_callback)
-
-        # bind get token
-        httpretty.register_uri(
-            method=httpretty.GET,
-            uri=re.compile(rf"{TOKEN_RESOURCE}/.*"),
-            body=self.get_token_callback)
-
-    # noinspection PyUnusedLocal
-    def post_statement_callback(self, request, uri, response_headers):
-        authorization = request.headers.get("Authorization")
-
-        if authorization and authorization.replace("Bearer ", "") in self.tokens:
-            return [200, response_headers, json.dumps(self.sample_post_response_data)]
-
-        challenge_id = str(uuid.uuid4())
-        token = str(uuid.uuid4())
-        self.tokens.add(token)
-        self.challenges[challenge_id] = MultithreadedTokenServer.Challenge(token, self.attempts)
-        redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}"
-        token_server = f"{TOKEN_RESOURCE}/{challenge_id}"
-        return [401, {'Www-Authenticate': f'Bearer x_redirect_server="{redirect_server}", '
-                                          f'x_token_server="{token_server}"',
-                      'Basic realm': '"Trino"'}, ""]
-
-    # noinspection PyUnusedLocal
-    def get_token_callback(self, request, uri, response_headers):
-        challenge_id = uri.replace(f"{TOKEN_RESOURCE}/", "")
-        challenge = self.challenges[challenge_id]
-        challenge = challenge._replace(attempts=challenge.attempts - 1)
-        self.challenges[challenge_id] = challenge
-        if challenge.attempts < 0:
-            return [404, response_headers, "{}"]
-        if challenge.attempts == 0:
-            return [200, response_headers, f'{{"token": "{challenge.token}"}}']
-        return [200, response_headers, f'{{"nextUri": "{uri}"}}']
-
-
 @httprettified
 def test_multithreaded_oauth2_authentication_flow(sample_post_response_data):
     redirect_handler = RedirectHandler()
@@ -598,31 +502,19 @@ def test_multithreaded_oauth2_authentication_flow(sample_post_response_data):
     for thread in threads:
         thread.join()
 
-    # should issue only 3 tokens and each thread should get one
-    assert len(token_server.tokens) == 3
+    # should issue only 1 token and each thread should reuse it
+    assert len(token_server.tokens) == 1
     for thread in threads:
         assert thread.token in token_server.tokens
 
-    # should start only 3 challenges and every token should be obtained
-    assert len(token_server.challenges.keys()) == 3
+    # should start only 1 challenge
+    assert len(token_server.challenges.keys()) == 1
     for challenge_id, challenge in token_server.challenges.items():
         assert f"{REDIRECT_RESOURCE}/{challenge_id}" in redirect_handler.redirect_server
         assert challenge.attempts == 0
         assert len(_get_token_requests(challenge_id)) == 1
     # 3 threads * (10 POST /statement each + 1 replied request by authentication)
-    assert len(_post_statement_requests()) == 33
-
-
-def _get_token_requests(challenge_id):
-    return list(filter(
-        lambda r: r.method == "GET" and r.path == f"/{TOKEN_PATH}/{challenge_id}",
-        httpretty.latest_requests()))
-
-
-def _post_statement_requests():
-    return list(filter(
-        lambda r: r.method == "POST" and r.path == constants.URL_STATEMENT_PATH,
-        httpretty.latest_requests()))
+    assert len(_post_statement_requests()) == 31
 
 
 @mock.patch("trino.client.TrinoRequest.http")
diff --git a/tests/unit/test_dbapi.py b/tests/unit/test_dbapi.py
index 4e6dbbb..690aa50 100644
--- a/tests/unit/test_dbapi.py
+++ b/tests/unit/test_dbapi.py
@@ -9,9 +9,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import threading
+import uuid
+from unittest.mock import patch
 
+import httpretty
+from httpretty import httprettified
 from requests import Session
-from unittest.mock import patch
+
+from tests.unit.oauth_test_utils import _post_statement_requests, _get_token_requests, RedirectHandler, \
+    GetTokenCallback, REDIRECT_RESOURCE, TOKEN_RESOURCE, PostStatementCallback, SERVER_ADDRESS
+from trino import constants
+from trino.auth import OAuth2Authentication
 from trino.dbapi import connect
 
 
@@ -39,3 +48,177 @@ def test_http_session_is_defaulted_when_not_specified(mock_client):
     # THEN
     request_args, _ = mock_client.TrinoRequest.call_args
     assert mock_client.TrinoRequest.http.Session.return_value in request_args
+
+
+@httprettified
+def test_token_retrieved_once_per_auth_instance(sample_post_response_data):
+    token = str(uuid.uuid4())
+    challenge_id = str(uuid.uuid4())
+
+    redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}"
+    token_server = f"{TOKEN_RESOURCE}/{challenge_id}"
+
+    post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data)
+
+    # bind post statement
+    httpretty.register_uri(
+        method=httpretty.POST,
+        uri=f"{SERVER_ADDRESS}:8080{constants.URL_STATEMENT_PATH}",
+        body=post_statement_callback)
+
+    # bind get token
+    get_token_callback = GetTokenCallback(token_server, token)
+    httpretty.register_uri(
+        method=httpretty.GET,
+        uri=token_server,
+        body=get_token_callback)
+
+    redirect_handler = RedirectHandler()
+
+    with connect(
+            "coordinator",
+            user="test",
+            auth=OAuth2Authentication(redirect_auth_url_handler=redirect_handler),
+            http_scheme=constants.HTTPS
+    ) as conn:
+        conn.cursor().execute("SELECT 1")
+        conn.cursor().execute("SELECT 2")
+        conn.cursor().execute("SELECT 3")
+
+    # bind get token
+    get_token_callback = GetTokenCallback(token_server, token)
+    httpretty.register_uri(
+        method=httpretty.GET,
+        uri=token_server,
+        body=get_token_callback)
+
+    redirect_handler = RedirectHandler()
+
+    with connect(
+            "coordinator",
+            user="test",
+            auth=OAuth2Authentication(redirect_auth_url_handler=redirect_handler),
+            http_scheme=constants.HTTPS
+    ) as conn2:
+        conn2.cursor().execute("SELECT 1")
+        conn2.cursor().execute("SELECT 2")
+        conn2.cursor().execute("SELECT 3")
+
+    assert len(_get_token_requests(challenge_id)) == 2
+
+
+@httprettified
+def test_token_retrieved_once_when_authentication_instance_is_shared(sample_post_response_data):
+    token = str(uuid.uuid4())
+    challenge_id = str(uuid.uuid4())
+
+    redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}"
+    token_server = f"{TOKEN_RESOURCE}/{challenge_id}"
+
+    post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data)
+
+    # bind post statement
+    httpretty.register_uri(
+        method=httpretty.POST,
+        uri=f"{SERVER_ADDRESS}:8080{constants.URL_STATEMENT_PATH}",
+        body=post_statement_callback)
+
+    # bind get token
+    get_token_callback = GetTokenCallback(token_server, token)
+    httpretty.register_uri(
+        method=httpretty.GET,
+        uri=token_server,
+        body=get_token_callback)
+
+    redirect_handler = RedirectHandler()
+
+    authentication = OAuth2Authentication(redirect_auth_url_handler=redirect_handler)
+
+    with connect(
+            "coordinator",
+            user="test",
+            auth=authentication,
+            http_scheme=constants.HTTPS
+    ) as conn:
+        conn.cursor().execute("SELECT 1")
+        conn.cursor().execute("SELECT 2")
+        conn.cursor().execute("SELECT 3")
+
+    # bind get token
+    get_token_callback = GetTokenCallback(token_server, token)
+    httpretty.register_uri(
+        method=httpretty.GET,
+        uri=token_server,
+        body=get_token_callback)
+
+    with connect(
+            "coordinator",
+            user="test",
+            auth=authentication,
+            http_scheme=constants.HTTPS
+    ) as conn2:
+        conn2.cursor().execute("SELECT 1")
+        conn2.cursor().execute("SELECT 2")
+        conn2.cursor().execute("SELECT 3")
+
+    assert len(_post_statement_requests()) == 7
+    assert len(_get_token_requests(challenge_id)) == 1
+
+
+@httprettified
+def test_token_retrieved_once_when_multithreaded(sample_post_response_data):
+    token = str(uuid.uuid4())
+    challenge_id = str(uuid.uuid4())
+
+    redirect_server = f"{REDIRECT_RESOURCE}/{challenge_id}"
+    token_server = f"{TOKEN_RESOURCE}/{challenge_id}"
+
+    post_statement_callback = PostStatementCallback(redirect_server, token_server, [token], sample_post_response_data)
+
+    # bind post statement
+    httpretty.register_uri(
+        method=httpretty.POST,
+        uri=f"{SERVER_ADDRESS}:8080{constants.URL_STATEMENT_PATH}",
+        body=post_statement_callback)
+
+    # bind get token
+    get_token_callback = GetTokenCallback(token_server, token)
+    httpretty.register_uri(
+        method=httpretty.GET,
+        uri=token_server,
+        body=get_token_callback)
+
+    redirect_handler = RedirectHandler()
+
+    authentication = OAuth2Authentication(redirect_auth_url_handler=redirect_handler)
+
+    conn = connect(
+        "coordinator",
+        user="test",
+        auth=authentication,
+        http_scheme=constants.HTTPS
+    )
+
+    class RunningThread(threading.Thread):
+        lock = threading.Lock()
+
+        def __init__(self):
+            super().__init__()
+
+        def run(self) -> None:
+            with RunningThread.lock:
+                conn.cursor().execute("SELECT 1")
+
+    threads = [
+        RunningThread(),
+        RunningThread(),
+        RunningThread()
+    ]
+
+    # run and join all threads
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+
+    assert len(_get_token_requests(challenge_id)) == 1