Skip to content

feat(auth): add set_token and document push/pull token distribution#661

Merged
LinoGiger merged 2 commits into
mainfrom
feat(auth)/document-push-pull-token-distribution
Jul 9, 2026
Merged

feat(auth): add set_token and document push/pull token distribution#661
LinoGiger merged 2 commits into
mainfrom
feat(auth)/document-push-pull-token-distribution

Conversation

@RapidPoseidon

Copy link
Copy Markdown
Contributor

Why

Follow-up to #655. Review feedback: the "Writing the file yourself" docs section just re-implemented the same file pattern as maintain_token_file. It should instead teach the underlying primitives — how to get a token out of one client and how to ingest it into another — so customers can build any push/pull distribution (Redis, RPC, secret manager, …) instead of being funneled through a file.

There was also an API gap blocking that: you could export a token (get_token()) and bootstrap a client from one (token=), but there was no way to inject a fresh token into a running client — once its token expired you had to reconstruct it, which re-fires the init-time side calls the whole feature exists to avoid.

What

  • RapidataClient.set_token(token) — the counterpart to get_token(): replace the token a live client authenticates with, effective from the next request. Also revives a client whose token already expired. Added on RESTClientObject (both rest.py and rest.mustache).
  • Docs: replaced the "Writing the file yourself" section with "Building your own token distribution" — the three primitives (export → transport of your choice → bootstrap / renew), when to use push vs pull, and the cross-transport invariants (complete token object with absolute expires_at; atomic writes if the transport is a file).

Verification

  • pyright: 0 errors; black clean
  • Smoke test: bootstrap from token object, set_token swap on a live client, expired client revived via set_token without reconstruction
  • mkdocs build clean

🔗 Session: https://session-43887903.poseidon.rapidata.internal/

🤖 Generated with Claude Code

@claude

claude Bot commented Jul 9, 2026

Copy link
Copy Markdown

Review

Overall this is a clean, well-scoped follow-up. set_token is a sensible counterpart to get_token(), the two rest.py/rest.mustache copies are kept byte-for-byte in sync per the repo's CLAUDE.md convention, and the docs rewrite genuinely improves on the old "copy this file-writing loop" section by teaching the underlying primitives instead. A few small things worth a look:

1. Doc example references an undefined variable (docs/distributed_training.md)

```python
token = coordinator.get_token()
...
worker = RapidataClient(token=token)

3. Renew — inject a newer token into a running client at any time;

it is used from the next request on.

worker.set_token(fresh_token)
```

fresh_token is never defined in this snippet — only token is. A reader who copy-pastes this to try it out will hit a NameError. Suggest either worker.set_token(coordinator.get_token()) or adding fresh_token = coordinator.get_token() before the call.

2. No validation of the token shape in set_token

```python
def set_token(self, token: dict):
"""Replace the token the session authenticates with."""
if not self.session:
raise ApiValueError(...)
self.session.token = token
```

The docstring (and the docs page) are explicit that the caller must pass the complete token object with an absolute expires_at. If a caller passes an incomplete dict (e.g. only expires_in, or a raw curl response), the failure won't surface here — it'll show up later as an opaque AttributeError/KeyError from inside authlib's is_expired()/ensure_active_token(), rather than a clear ApiValueError like the sibling guard in this same method. A light shape check (e.g. if "expires_at" not in token: raise ApiValueError(...)) would make the failure mode match the rest of this class.

3. Test coverage

No tests exercise set_token (and none pre-existed for get_token/maintain_token_file either, so this isn't a regression introduced by this PR specifically). Since set_token is now public API with its own branch (uninitialized session), a small unit test mocking rest_client/session would be cheap insurance against a future refactor silently breaking the "revive an expired client" use case called out in the PR description.

Nothing blocking — the core change (a straightforward attribute swap wrapped in the same guard pattern as get_token) is small and low-risk, and it closes the real gap described in the PR (no way to refresh a running client's token without reconstructing it).

@LinoGiger LinoGiger marked this pull request as ready for review July 9, 2026 09:08
@LinoGiger LinoGiger self-requested a review as a code owner July 9, 2026 09:08
@LinoGiger LinoGiger requested a review from LucStr July 9, 2026 09:53
@claude

claude Bot commented Jul 9, 2026

Copy link
Copy Markdown

Review

Overview
Follow-up to #655: adds RapidataClient.set_token() (and the underlying RESTClientObject.set_token() in both rest.py and rest.mustache) so a running client can be handed a fresh token without reconstruction, and rewrites the "Writing the file yourself" docs section into a more general "Building your own token distribution" guide covering the export -> transport -> bootstrap/renew primitives. Good, focused change that closes a real API gap (previously get_token/token= existed but there was no way to push a token into a live client).

Code quality / conventions

  • Correctly duplicates the new method in both rest.py and rest.mustache, matching the mustache-template rule in CLAUDE.md.
  • RapidataClient.set_token uses dict[str, Any] + -> None with from __future__ import annotations already present in the file, consistent with the project's typing convention. The generated RESTClientObject.set_token uses a bare dict, matching the existing (untyped) style of get_token in the same generated file, so no inconsistency there.
  • _logger.warning uses lazy %s formatting correctly, and only logs the names of missing keys, never token contents - good, no secret leakage.
  • Docs rewrite reads well: concrete, not hyperbolic, and reuses the existing expires_at/atomic-write invariants instead of re-explaining them.

Potential issues

  1. Lenient validation may defer failure to a confusing place. set_token only logs a warning when access_token/token_type/expires_at are missing, then assigns the (possibly broken) dict to self.session.token regardless. If access_token is missing, the client will look successfully updated but then fail on the next HTTP request with a less obvious error (or a KeyError deep in authlib) rather than failing fast at the call site. Consider raising ApiValueError (as get_token already does elsewhere in this file) when access_token is missing, and reserve the warning-only behavior for expires_at/token_type where a caller might have a reason to omit them temporarily.
  2. Interaction with token_file-based clients isn't addressed. A client can be constructed with both a token_file (pull mode) and later have set_token called on it manually (push mode). In that combination, get_token()/internal expiry checks still go through _reload_token_from_file_if_expired, which can silently overwrite a manually-pushed token with a stale one from disk the next time the token is near-expired. Not necessarily a bug since the two flows aren't meant to be mixed, but worth a one-line docstring/doc caveat.
  3. No automated test coverage for set_token (mirrors the pre-existing lack of tests for get_token/maintain_token_file, so not a regression introduced by this PR, but worth flagging since this is new public API surface). The repo has no test suite at all currently, so this is more of an observation than a blocker.

Minor

  • Thread-safety: self.session.token = token isn't synchronized against a concurrent in-flight request reading self.session.token. This mirrors the existing _reload_token_from_file_if_expired pattern, so it's not a new risk, but since set_token is explicitly designed to be called from arbitrary user code/threads, it might be worth a short docstring note that swaps should happen between requests, not concurrently with one.

Verification
Per the PR description, pyright/black/mkdocs were run clean and a manual smoke test covered bootstrap-from-token, live swap, and revival-of-expired-client - reasonable given the diff is small and additive. Nothing found that contradicts this.

Overall: solid, well-scoped change with good docs. The main actionable suggestion is #1 (fail fast on a missing access_token rather than warn-and-continue).

@LinoGiger LinoGiger merged commit 344d0bf into main Jul 9, 2026
2 checks passed
@LinoGiger LinoGiger deleted the feat(auth)/document-push-pull-token-distribution branch July 9, 2026 10:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants