Skip to content

Commit 422f265

Browse files
authored
Allowing for refresh token rotation during token refresh request (#549)
1 parent 1c486fb commit 422f265

File tree

2 files changed

+16
-3
lines changed

2 files changed

+16
-3
lines changed

vertica_python/vertica/connection.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -474,6 +474,14 @@ def closed(self) -> bool:
474474
"""Returns True if the connection is closed."""
475475
return not self.opened()
476476

477+
def get_current_refresh_token(self) -> str:
478+
"""Returns the current refresh token.
479+
480+
This may be different from the user supplied token if token refresh
481+
was required and token rotation is in effect
482+
"""
483+
return self.oauth_refresh_token
484+
477485
def __str__(self) -> str:
478486
safe_options = {key: value for key, value in self.options.items() if key != 'password'}
479487

@@ -920,7 +928,7 @@ def startup_connection(self) -> bool:
920928
# If access token is not set, will attempt to set a new one by using token refresh
921929
if len(self.oauth_access_token) == 0 and self.oauth_manager and not self.oauth_manager.refresh_attempted:
922930
self._logger.info("Issuing an OAuth access token using a refresh token")
923-
self.oauth_access_token = self.oauth_manager.do_token_refresh()
931+
self.oauth_access_token, self.oauth_refresh_token = self.oauth_manager.do_token_refresh()
924932
self.write(messages.Password(self.oauth_access_token, message.code))
925933
else:
926934
self.write(messages.Password(password, message.code,
@@ -940,7 +948,7 @@ def startup_connection(self) -> bool:
940948
raise errors.ConnectionError("Did not receive proper OAuth Authentication response from server. Please upgrade to the latest Vertica server for OAuth Support.")
941949
self.close_socket()
942950
self._logger.info("Issuing a new OAuth access token using a refresh token")
943-
self.oauth_access_token = self.oauth_manager.do_token_refresh()
951+
self.oauth_access_token, self.oauth_refresh_token = self.oauth_manager.do_token_refresh()
944952
return True
945953
raise errors.ConnectionError(message.error_message())
946954
else:

vertica_python/vertica/oauth_manager.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,12 @@ def get_access_token_using_refresh_token(self) -> str:
8080
# TODO handle self.validate_cert_hostname
8181
response = requests.post(self.token_url, headers=headers, data=params, verify=False)
8282
response.raise_for_status()
83-
return response.json()["access_token"]
83+
json_response = response.json()
84+
# If refresh token rotation is used, like in OTDS, we will get both our new valid access token as well as
85+
# a new refresh token to use the next time we need to invoke token refresh.
86+
if 'refresh_token' in json_response:
87+
self.refresh_token = json_response["refresh_token"]
88+
return response.json()["access_token"], self.refresh_token
8489
except requests.exceptions.HTTPError as err:
8590
msg = f'{err_msg}\n{err}\n{response.json()}'
8691
raise OAuthTokenRefreshError(msg)

0 commit comments

Comments
 (0)