refactor: move token refresh logic to refresh_and_set_access_token method in authenticator

Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>

Author: devin-ai-integration[bot]

airbyte_cdk/sources/streams/http/http_client.py

@@ -50,9 +50,6 @@
     rate_limit_default_backoff_handler,
     user_defined_backoff_handler,
 )
-from airbyte_cdk.sources.streams.http.requests_native_auth import (
-    SingleUseRefreshTokenOauth2Authenticator,
-)
 from airbyte_cdk.sources.utils.types import JsonType
 from airbyte_cdk.utils.airbyte_secrets_utils import filter_secrets
 from airbyte_cdk.utils.constants import ENV_REQUEST_CACHE_PATH
@@ -461,31 +458,16 @@ def _handle_error_resolution(
 
         # Handle REFRESH_TOKEN_THEN_RETRY: Force refresh the OAuth token before retry
         # This is useful when the API returns 401 but the stored token expiry hasn't been reached yet
-        # Only OAuth authenticators have refresh_access_token method
+        # Only OAuth authenticators have refresh_and_set_access_token method
         # Non-OAuth auth types (e.g., BearerAuthenticator) will fall through to normal retry
         if error_resolution.response_action == ResponseAction.REFRESH_TOKEN_THEN_RETRY:
             if (
                 hasattr(self._session, "auth")
                 and self._session.auth is not None
-                and hasattr(self._session.auth, "refresh_access_token")
+                and hasattr(self._session.auth, "refresh_and_set_access_token")
             ):
                 try:
-                    if isinstance(self._session.auth, SingleUseRefreshTokenOauth2Authenticator):
-                        # For single-use refresh tokens, we must persist the new refresh token
-                        # and emit a control message to update the connector config
-                        token, expires_in, new_refresh_token = (
-                            self._session.auth.refresh_access_token()
-                        )
-                        self._session.auth.access_token = token
-                        self._session.auth.set_refresh_token(new_refresh_token)
-                        self._session.auth.set_token_expiry_date(expires_in)
-                        self._session.auth._emit_control_message()
-                    else:
-                        # Use extended unpacking to handle both 2-tuple (AbstractOauth2Authenticator)
-                        # and 3-tuple (Oauth2Authenticator which also returns refresh_token) returns
-                        token, expires_in, *_ = self._session.auth.refresh_access_token()  # type: ignore[union-attr]
-                        self._session.auth.access_token = token  # type: ignore[union-attr]
-                        self._session.auth.set_token_expiry_date(expires_in)  # type: ignore[union-attr]
+                    self._session.auth.refresh_and_set_access_token()  # type: ignore[union-attr]
                     self._logger.info(
                         "Refreshed OAuth token due to REFRESH_TOKEN_THEN_RETRY response action"
                     )

airbyte_cdk/sources/streams/http/requests_native_auth/abstract_oauth.py

@@ -88,12 +88,21 @@ def get_auth_header(self) -> Mapping[str, Any]:
     def get_access_token(self) -> str:
         """Returns the access token"""
         if self.token_has_expired():
-            token, expires_in = self.refresh_access_token()
-            self.access_token = token
-            self.set_token_expiry_date(expires_in)
+            self.refresh_and_set_access_token()
 
         return self.access_token
 
+    def refresh_and_set_access_token(self) -> None:
+        """Force refresh the access token and update internal state.
+
+        This method refreshes the access token regardless of whether it has expired,
+        and updates the internal token and expiry date. Subclasses may override this
+        to handle additional state updates (e.g., persisting new refresh tokens).
+        """
+        token, expires_in = self.refresh_access_token()
+        self.access_token = token
+        self.set_token_expiry_date(expires_in)
+
     def token_has_expired(self) -> bool:
         """Returns True if the token is expired"""
         return ab_datetime_now() > self.get_token_expiry_date()

airbyte_cdk/sources/streams/http/requests_native_auth/oauth.py

@@ -318,20 +318,28 @@ def token_has_expired(self) -> bool:
 
     def get_access_token(self) -> str:
         """Retrieve new access and refresh token if the access token has expired.
-        The new refresh token is persisted with the set_refresh_token function
+
+        The new refresh token is persisted with the set_refresh_token function.
+
         Returns:
             str: The current access_token, updated if it was previously expired.
         """
         if self.token_has_expired():
-            new_access_token, access_token_expires_in, new_refresh_token = (
-                self.refresh_access_token()
-            )
-            self.access_token = new_access_token
-            self.set_refresh_token(new_refresh_token)
-            self.set_token_expiry_date(access_token_expires_in)
-            self._emit_control_message()
+            self.refresh_and_set_access_token()
         return self.access_token
 
+    def refresh_and_set_access_token(self) -> None:
+        """Force refresh the access token and update internal state.
+
+        For single-use refresh tokens, this also persists the new refresh token
+        and emits a control message to update the connector config.
+        """
+        new_access_token, access_token_expires_in, new_refresh_token = self.refresh_access_token()
+        self.access_token = new_access_token
+        self.set_refresh_token(new_refresh_token)
+        self.set_token_expiry_date(access_token_expires_in)
+        self._emit_control_message()
+
     def refresh_access_token(self) -> Tuple[str, AirbyteDateTime, str]:  # type: ignore[override]
         """
         Refreshes the access token by making a handled request and extracting the necessary token information.

unit_tests/sources/streams/http/test_http_client.py

@@ -845,12 +845,10 @@ def __init__(self):
         self._token_expiry_date = None
         self.refresh_called = False
 
-    def refresh_access_token(self):
+    def refresh_and_set_access_token(self):
         self.refresh_called = True
-        return ("new_refreshed_token", "2099-01-01T00:00:00Z")
-
-    def set_token_expiry_date(self, value):
-        self._token_expiry_date = value
+        self.access_token = "new_refreshed_token"
+        self._token_expiry_date = "2099-01-01T00:00:00Z"
 
     def __call__(self, request):
         request.headers["Authorization"] = f"Bearer {self.access_token}"
@@ -932,7 +930,7 @@ class FailingOAuthAuthenticator:
         def __init__(self):
             self.access_token = "old_token"
 
-        def refresh_access_token(self):
+        def refresh_and_set_access_token(self):
             raise Exception("Token refresh failed")
 
         def __call__(self, request):
@@ -978,11 +976,6 @@ def test_refresh_token_then_retry_action_with_single_use_refresh_token_authentic
     )
 
     mock_authenticator = MagicMock(spec=SingleUseRefreshTokenOauth2Authenticator)
-    mock_authenticator.refresh_access_token.return_value = (
-        "new_access_token",
-        "2099-01-01T00:00:00Z",
-        "new_refresh_token",
-    )
 
     mocked_session = MagicMock(spec=requests.Session)
     mocked_session.auth = mock_authenticator
@@ -1013,10 +1006,7 @@ def test_refresh_token_then_retry_action_with_single_use_refresh_token_authentic
     with pytest.raises(DefaultBackoffException):
         http_client._send(prepared_request, {})
 
-    mock_authenticator.refresh_access_token.assert_called_once()
-    mock_authenticator.set_refresh_token.assert_called_once_with("new_refresh_token")
-    mock_authenticator.set_token_expiry_date.assert_called_once_with("2099-01-01T00:00:00Z")
-    mock_authenticator._emit_control_message.assert_called_once()
+    mock_authenticator.refresh_and_set_access_token.assert_called_once()
 
 
 @pytest.mark.usefixtures("mock_sleep")