fix: use the refresh code for oauth token#88
fix: use the refresh code for oauth token#88vmaurin wants to merge 1 commit intocrispthinking:mainfrom
Conversation
In order for a channel to keep going, it needs to get a fresh token when needed at individual request level. The token fetched only once at channel creation where it should be fetched for each request. The CredentialHelper has to be moved to non asyncio as it seems to be full sync code in underlying grpc.
iwillspeak
requested review from
Copilot and
snus-kin
and removed request for
snus-kin
iwillspeak
left a comment
iwillspeak
left a comment
There was a problem hiding this comment.
Nice. This definitely seems like a sensible way of doing it.
There was a problem hiding this comment.
Pull request overview
Updates OAuth handling so tokens are refreshed per-request (via gRPC metadata credentials) rather than only at channel creation time, and converts the credential helper implementation to synchronous to match underlying (threaded) gRPC auth plugin execution.
Changes:
- Convert
CredentialHelpertoken fetching/refresh from async (httpx.AsyncClient,asyncio.Lock) to sync (httpx.Client,threading.Lock). - Replace fixed access-token call credentials with an auth metadata plugin that calls
CredentialHelper.get_token()for each request. - Update tests/docs/examples to reflect the sync token API and removal of “test token acquisition” steps.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/functional/conftest.py | Stops eagerly acquiring a token during fixture creation; returns helper directly. |
| tests/client/test_channel.py | Updates tests to the sync token API; removes tests tied to the previous token plugin / creation-time token acquisition. |
| src/resolver_athena_client/client/channel.py | Implements per-request token refresh via metadata plugin; rewrites credential fetching to synchronous with thread locks. |
| examples/example.py | Removes proactive token acquisition step before channel creation. |
| examples/classify_single_example.py | Removes proactive token acquisition step before channel creation. |
| docs/authentication.rst | Updates examples to use sync get_token() and removes token check step in setup flow. |
| docs/api/exceptions.rst | Updates example to use sync get_token(). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| self._lock = threading.Lock() | ||
|
|
||
| async def get_token(self) -> str: | ||
| def get_token(self) -> str: |
There was a problem hiding this comment.
get_token() is annotated to return str but can return None (e.g., if _is_token_valid() is true while _token is None, or if _token is cleared between the validity check and the return). Make the return value guaranteed by (a) performing the validity check + read under the same lock, and (b) raising a clear error if the token is unexpectedly missing after a 'valid' check / refresh.
| if self._is_token_valid(): | ||
| return self._token | ||
|
|
||
| with self._lock: | ||
| if not self._is_token_valid(): | ||
| self._refresh_token() | ||
| return self._token |
There was a problem hiding this comment.
get_token() is annotated to return str but can return None (e.g., if _is_token_valid() is true while _token is None, or if _token is cleared between the validity check and the return). Make the return value guaranteed by (a) performing the validity check + read under the same lock, and (b) raising a clear error if the token is unexpectedly missing after a 'valid' check / refresh.
| if self._is_token_valid(): | ||
| return self._token | ||
|
|
||
| with self._lock: | ||
| if not self._is_token_valid(): | ||
| self._refresh_token() | ||
| return self._token |
There was a problem hiding this comment.
There’s a race between the unlocked fast-path in get_token() and invalidate_token(): another thread can invalidate between the _is_token_valid() check and return self._token, causing a None return. Remove the unlocked fast-path and do the check/refresh/return entirely inside the lock (or copy the token to a local variable while holding the lock).
| with self._lock: | ||
| self._token = None | ||
| self._token_expires_at = None |
There was a problem hiding this comment.
There’s a race between the unlocked fast-path in get_token() and invalidate_token(): another thread can invalidate between the _is_token_valid() check and return self._token, causing a None return. Remove the unlocked fast-path and do the check/refresh/return entirely inside the lock (or copy the token to a local variable while holding the lock).
| except OAuthError as err: | ||
| callback(None, err) | ||
|
|
||
| async def create_channel_with_credentials( |
There was a problem hiding this comment.
create_channel_with_credentials() no longer acquires a token during channel creation, so it likely won’t raise OAuthError at creation time anymore (OAuth failures will occur during RPC auth metadata injection). Update the function docstring/raises section to reflect when/where OAuth errors are expected now.
| self._token_expires_at = None | ||
|
|
||
|
|
||
| class _AutoRefreshTokenAuthMetadataPlugin(grpc.AuthMetadataPlugin): |
There was a problem hiding this comment.
New per-request auth behavior is introduced via _AutoRefreshTokenAuthMetadataPlugin, but the PR removes the previous unit test coverage around auth metadata injection. Add focused tests that (1) the plugin calls credential_helper.get_token() and passes the expected authorization metadata to the callback, and (2) an OAuthError results in callback(None, err).
|
Code needs formatting, probably needs linting too. There's a pre-commit hook that I think will cover everything but otherwise instructions are in the readme. Do agree this cred helper stuff was bad though, bit of a rush over the other bits in the repo. |
Yes sorry, I am not so familiar with the project and what I need to run here. I also not very familiar with the grpc stack (so not sure what should be locked or not here). Maybe you can use this PR more as a detailed bug report, more than the actual way it should be fixed. Meanwhile, we use a workaround by duck typing the token/credential helper so it does actually the refresh from our code base. |
All good, I'm not actively working on this project right now. Will try and get something done about it if possible though 😄 |
I've raised this internally and we're tracking it. |
3 participants
In order for a channel to keep going, it needs to get a fresh token when needed at individual request level.
The token fetched only once at channel creation where it should be fetched for each request.
The CredentialHelper has to be moved to non asyncio as it seems to be full sync code in underlying grpc.