Fix refresh tokens flow

Signed-off-by: sinkuladis <[email protected]>

Author: sinkuladis

tests/unit/oauth_test_utils.py

@@ -34,6 +34,14 @@ def __call__(self, url):
         self.redirect_server += url
 
 
+class RedirectHandlerWithException:
+    def __init__(self, exception):
+        self.exception = exception
+
+    def __call__(self, url):
+        raise self.exception
+
+
 class PostStatementCallback:
     def __init__(self, redirect_server, token_server, tokens, sample_post_response_data):
         self.redirect_server = redirect_server
@@ -45,6 +53,9 @@ 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)]
+        elif self.redirect_server is None and self.token_server is not None:
+            return [401, {'Www-Authenticate': f'Bearer x_token_server="{self.token_server}"',
+                          'Basic realm': '"Trino"'}, ""]
         return [401, {'Www-Authenticate': f'Bearer x_redirect_server="{self.redirect_server}", '
                                           f'x_token_server="{self.token_server}"',
                       'Basic realm': '"Trino"'}, ""]

tests/unit/test_client.py

@@ -32,6 +32,7 @@
     MultithreadedTokenServer,
     PostStatementCallback,
     RedirectHandler,
+    RedirectHandlerWithException,
     _get_token_requests,
     _post_statement_requests,
 )
@@ -384,6 +385,48 @@ def test_oauth2_authentication_flow(attempts, sample_post_response_data):
     assert len(_get_token_requests(challenge_id)) == attempts
 
 
+@httprettified
+def test_oauth2_refresh_token_flow(sample_post_response_data):
+    token = str(uuid.uuid4())
+    challenge_id = str(uuid.uuid4())
+
+    token_server = f"{TOKEN_RESOURCE}/{challenge_id}"
+
+    post_statement_callback = PostStatementCallback(None, token_server, [token], sample_post_response_data)
+
+    # bind post statement
+    httpretty.register_uri(
+        method=httpretty.POST,
+        uri=f"{SERVER_ADDRESS}{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 = RedirectHandlerWithException(
+        trino.exceptions.TrinoAuthError(
+            "Do not use redirect handler when there is no redirect_uri in the response"))
+
+    request = TrinoRequest(
+        host="coordinator",
+        port=constants.DEFAULT_TLS_PORT,
+        client_session=ClientSession(
+            user="test",
+        ),
+        http_scheme=constants.HTTPS,
+        auth=trino.auth.OAuth2Authentication(redirect_auth_url_handler=redirect_handler))
+
+    response = request.post("select 1")
+
+    assert response.request.headers['Authorization'] == f"Bearer {token}"
+    assert get_token_callback.attempts == 0
+    assert len(_post_statement_requests()) == 2
+
+
 @pytest.mark.parametrize("attempts", [6, 10])
 @httprettified
 def test_oauth2_exceed_max_attempts(attempts, sample_post_response_data):
@@ -430,10 +473,9 @@ def test_oauth2_exceed_max_attempts(attempts, sample_post_response_data):
 
 @pytest.mark.parametrize("header,error", [
     ("", "Error: header WWW-Authenticate not available in the response."),
-    ('Bearer"', 'Error: header info didn\'t have x_redirect_server'),
+    ('Bearer"', 'Error: header info didn\'t have x_token_server'),
     ('x_redirect_server="redirect_server", x_token_server="token_server"', 'Error: header info didn\'t match x_redirect_server="redirect_server", x_token_server="token_server"'),  # noqa: E501
     ('Bearer x_redirect_server="redirect_server"', 'Error: header info didn\'t have x_token_server'),
-    ('Bearer x_token_server="token_server"', 'Error: header info didn\'t have x_redirect_server'),
 ])
 @httprettified
 def test_oauth2_authentication_missing_headers(header, error):

trino/auth.py

@@ -322,15 +322,13 @@ def _attempt_oauth(self, response, **kwargs):
         auth_info_headers = parse_dict_header(_OAuth2TokenBearer._BEARER_PREFIX.sub("", auth_info, count=1))
 
         auth_server = auth_info_headers.get('x_redirect_server')
-        if auth_server is None:
-            raise exceptions.TrinoAuthError("Error: header info didn't have x_redirect_server")
-
         token_server = auth_info_headers.get('x_token_server')
         if token_server is None:
             raise exceptions.TrinoAuthError("Error: header info didn't have x_token_server")
 
-        # tell app that use this url to proceed with the authentication
-        self._redirect_auth_url(auth_server)
+        if auth_server is not None:
+            # tell app that use this url to proceed with the authentication
+            self._redirect_auth_url(auth_server)
 
         # Consume content and release the original connection
         # to allow our new request to reuse the same one.