feat(cdk): add REFRESH_TOKEN_THEN_RETRY response action for OAuth token refresh #886
Conversation
…en refresh This introduces a new ResponseAction type called REFRESH_TOKEN_THEN_RETRY that forces an OAuth token refresh before retrying a request. This is useful when APIs return 401 Unauthorized but the stored token expiry hasn't been reached yet. Changes: - Add REFRESH_TOKEN_THEN_RETRY to ResponseAction enum - Update HttpClient._handle_error_resolution to force token refresh when this action is encountered by directly calling refresh_access_token() - Add REFRESH_TOKEN_THEN_RETRY to declarative_component_schema.yaml - Regenerate Pydantic models from schema - Add unit tests for the new functionality This generalizes the pattern from PR #72309 (HubSpot custom component) into the CDK itself. Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>
🤖 Devin AI EngineerI'll be helping with this pull request! Here's what you should know: ✅ I will automatically:
Note: I can only respond to comments from users who have write access to this repository. ⚙️ Control Options:
|
👋 Greetings, Airbyte Team Member!Here are some helpful tips and reminders for your convenience. 💡 Show Tips and TricksTesting This CDK VersionYou can test this version of the CDK using the following: # Run the CLI from this branch:
uvx 'git+https://github.com/airbytehq/airbyte-python-cdk.git@devin/1769697982-refresh-token-then-retry#egg=airbyte-python-cdk[dev]' --help
# Update a connector to use the CDK from this branch ref:
cd airbyte-integrations/connectors/source-example
poe use-cdk-branch devin/1769697982-refresh-token-then-retryPR Slash CommandsAirbyte Maintainers can execute the following slash commands on your PR:
|
Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>
📝 WalkthroughWalkthroughAdds a new response action REFRESH_TOKEN_THEN_RETRY across schema, models, response enums, HTTP client, OAuth authenticators, and unit tests; the HTTP client will attempt to call the authenticator's refresh_and_set_access_token and then raise a backoff to retry when that action is triggered. Changes
Sequence DiagramsequenceDiagram
participant Client
participant HttpClient
participant OAuthAuth as OAuth Authenticator
participant Backoff as Backoff Handler
Client->>HttpClient: Send request
HttpClient->>HttpClient: Execute request
HttpClient-->>HttpClient: Receive 401 -> response filter => REFRESH_TOKEN_THEN_RETRY
alt OAuth authenticator present
HttpClient->>OAuthAuth: refresh_and_set_access_token()
OAuthAuth-->>HttpClient: (updates token & expiry or raises)
HttpClient->>Backoff: raise DefaultBackoffException (retry)
else No oauth authenticator or refresh fails
HttpClient->>HttpClient: log warning
HttpClient->>Backoff: raise DefaultBackoffException (retry)
end
Backoff->>HttpClient: retry request (after backoff)
HttpClient->>HttpClient: Execute request (with updated token if available)
HttpClient-->>Client: Return response
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Would you like me to suggest additional tests for concurrent refresh attempts or token race conditions, wdyt? 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
/autofix
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
…or message Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>
Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Fix all issues with AI agents
In `@airbyte_cdk/sources/streams/http/http_client.py`:
- Around line 469-472: The unpacking of refresh_access_token() assumes a 2-tuple
but Oauth2Authenticator returns 3 values; change the unpacking at the call site
in http_client.py (the try block that calls
self._session.auth.refresh_access_token()) to use extended unpacking (e.g.,
token, expires_in, *_ = self._session.auth.refresh_access_token()) so it safely
handles both 2- and 3-value returns, then use token and expires_in as before
when setting self._session.auth.access_token and calling
self._session.auth.set_token_expiry_date(expires_in).
🧹 Nitpick comments (1)
unit_tests/sources/streams/http/test_http_client.py (1)
842-857: Consider adding a mock for the 3-tuple return variant as well.The
MockOAuthAuthenticatorreturns 2 values fromrefresh_access_token(), which matchesAbstractOauth2Authenticator. Once the http_client code is updated to handle both 2 and 3 value returns, it might be worth adding a test with a mock that returns 3 values (likeOauth2Authenticatordoes) to ensure both variants work correctly, wdyt?
… refresh_access_token Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>
…or to persist new refresh token Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>
| # 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] |
There was a problem hiding this comment.
Should we move this logic into the OAuth authenticator itself as a new method, for example, refresh_and_set_access_token, instead of handling it here?
There was a problem hiding this comment.
That's a good suggestion! Moving this logic into the OAuth authenticator as a new method like refresh_and_set_access_token() would provide better encapsulation and separation of concerns.
The authenticator already has get_access_token() which does similar logic (refresh + set token + set expiry + emit control message for single-use tokens), but it checks token_has_expired() first. A new method could essentially be get_access_token() without the expiry check.
@darynaishchenko - what do you think? Should we refactor this to add a new method to the authenticator interface instead of handling it in HttpClient?
There was a problem hiding this comment.
good point, please proceed with this
There was a problem hiding this comment.
Done! I've refactored the code to add a refresh_and_set_access_token() method to AbstractOauth2Authenticator and overridden it in SingleUseRefreshTokenOauth2Authenticator to handle the additional logic (persisting refresh token + emitting control message).
Now HttpClient simply calls refresh_and_set_access_token() without needing to know about the authenticator internals. This provides better encapsulation and separation of concerns.
Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>
…thod in authenticator Co-Authored-By: Daryna Ishchenko <darina.ishchenko17@gmail.com>
Summary
Introduces a new
REFRESH_TOKEN_THEN_RETRYresponse action type that forces an OAuth token refresh before retrying a request. This is useful when APIs return 401 Unauthorized but the stored token expiry hasn't been reached yet (the token was invalidated server-side before its expected expiration).This generalizes the pattern from PR #72309 (HubSpot custom component) into the CDK itself.
Key changes:
REFRESH_TOKEN_THEN_RETRYtoResponseActionenumHttpClient._handle_error_resolutionto force token refresh by directly callingrefresh_access_token()(bypassing expiry check)declarative_component_schema.yamlfor declarative connectorsUpdates since last revision:
debugtowarningfor the case whenREFRESH_TOKEN_THEN_RETRYis received but the authenticator doesn't support token refresh (per reviewer feedback)warninglog level instead ofdebugReview & Testing Checklist for Human
refresh_access_token()directly (notget_access_token()) to force refresh regardless of stored expiry time - this was the key requirementhasattrchecks athttp_client.py:460-463- ensure they correctly identify OAuth authenticators that support token refreshexcept Exceptionat line 470 - this is intentional for resilience but may mask unexpected errorsRecommended test plan:
Notes
Requested by @darynaishchenko
Link to Devin run: https://app.devin.ai/sessions/3ce943802e5743bead2f54d850ae8455
Summary by CodeRabbit
New Features
Tests
✏️ Tip: You can customize this high-level summary in your review settings.