fix: Added caching for key endpoints that lookup keys by id / whoami endpoint

[Akhileshait]

What does this PR do?

This PR adds caching for key lookups in the Go API to reduce database queries when using updateKey, updateCredits, and setRoles endpoints.

Previously, these handlers were querying the database directly every time they needed to fetch key information. Now they use a new LiveKeyByID cache with a Stale-While-Revalidate (SWR) pattern that:

• Returns fresh data immediately from cache (10-second freshness window)
• Revalidates stale data in the background (10-minute staleness window)
• Properly invalidates cache when keys are updated

This reduces database load and improves API response times for repeated key operations.

Fixes #4150

Type of change

• [x] Bug fix (non-breaking change which fixes an issue)
• [x] Chore (refactoring code, technical debt, workflow improvements)
• [x] Enhancement (small improvements)
• [ ] New feature (non-breaking change which adds functionality)
• [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
• [ ] This change requires a documentation update

How should this be tested?

Unit Tests

Run the existing unit tests for the modified handlers:

cd go make test

Checklist

Required

• [x] Filled out the "How to test" section in this PR
• [x] Read Contributing Guide
• [x] Self-reviewed my own code
• [x] Commented on my code in hard-to-understand areas
• [x] Ran pnpm build
• [x] Ran pnpm fmt
• [x] Checked for warnings, there are none
• [x] Removed all console.logs
• [x] Merged the latest changes from main onto my branch with git pull origin main
• [x] My changes don't cause any responsiveness issues

Appreciated

• [ ] If a UI change was made: Added a screen recording or screenshots to this PR
• [ ] Updated the Unkey Docs if changes were necessary

[changeset-bot[bot]]

⚠️ No Changeset found

Latest commit: a991c684b9606c8d3b6fa1e9fd27a623ea5e3e4a

Merging this PR will not cause a version bump for any packages. If these changes should not result in a new version, you're good to go. If these changes should result in a version bump, you need to add a changeset.

Click here to learn what changesets are, and how to add one.

Click here if you're a maintainer who wants to add a changeset to this PR

[vercel[bot]]

@Akhileshait is attempting to deploy a commit to the Unkey Team on Vercel.

A member of the Team first needs to authorize it.

[coderabbitai[bot]]

📝 Walkthrough

Walkthrough

Adds a LiveKeyByID cache to the caches service and wires a new LiveKeyCache into multiple v2 key handlers and their tests; handlers now use SWR-backed lookups for live-key-by-id and invalidate LiveKeyCache (and KeyCache where applicable) after operations.

Changes

Cohort / File(s)	Summary
Cache Service Initialization go/internal/services/caches/caches.go	Add LiveKeyByID cache.Cache[string, db.FindLiveKeyByIDRow] to Caches; create liveKeyByID with Fresh:10s, Stale:10m, MaxSize:1_000_000; adjust some existing cache sizes and expose via tracing wrapper.
Route Registration go/apps/api/routes/register.go	Wire LiveKeyCache: svc.Caches.LiveKeyByID into v2 keys route handlers during registration.
Handlers — SWR lookup + invalidation go/apps/api/routes/v2_keys_*/handler.go (reroll_key, delete_key, get_key, set_roles, set_permissions, add_permissions, add_roles, remove_permissions, remove_roles, update_key, update_credits)	Add public field LiveKeyCache cache.Cache[string, db.FindLiveKeyByIDRow]; replace direct FindLiveKeyByID DB calls with h.LiveKeyCache.SWR(ctx, keyId, loader, caches.DefaultFindFirstOp) and preserve not-found/error semantics; remove/evict LiveKeyCache.Remove(key.ID) after success; many also continue to evict KeyCache.
v2_keys_reroll_key — extra cache go/apps/api/routes/v2_keys_reroll_key/handler.go	Add KeyCache and LiveKeyCache fields; use SWR for live-key lookup; invalidate KeyCache (by hash) and LiveKeyCache (by id) after success.
Tests — Handler wiring updates go/apps/api/routes/**/_test.go (many files across v2_keys_*)	Update test Handler composite literals to provide LiveKeyCache: h.Caches.LiveKeyByID (and KeyCache where applicable) across updated tests to match new Handler fields.

Sequence Diagram(s)

sequenceDiagram
    rect rgb(250,250,255)
    participant Client
    participant Handler
    participant LiveKeyCache
    participant DB
    participant Tx as Transaction
    end

    Client->>Handler: request (keyId, action)
    Handler->>LiveKeyCache: SWR(ctx, keyId, loader, DefaultFindFirstOp)
    alt Cache hit
        LiveKeyCache-->>Handler: cached key
    else Cache miss / stale
        LiveKeyCache->>DB: FindLiveKeyByID(keyId)
        DB-->>LiveKeyCache: key / not-found / error
        LiveKeyCache-->>Handler: key / result
    end

    Handler->>Tx: perform DB changes
    alt Success
        Tx-->>Handler: commit
        Handler->>LiveKeyCache: Remove(keyId)
        Handler->>Handler: invalidate KeyCache by key.Hash (if present)
        Handler-->>Client: success response
    else Failure
        Tx-->>Handler: rollback
        Handler-->>Client: error response
    end

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

• Pay attention to SWR loader semantics and returning/caching of not-found vs error across handlers.
• Verify correct eviction keys and ordering (pre/post-transaction comments in some handlers).
• Confirm all tests and route registration updated to compile with new LiveKeyCache field.

Possibly related PRs

• unkeyed/unkey#4063 — Modifies route handler structs and cache wiring in the same handler files; likely overlaps with field additions and initialization changes.
• unkeyed/unkey#3638 — Adds/adjusts LiveKeyByID cache and SWR semantics that this PR consumes; strong overlap on cache infra and usage.
• unkeyed/unkey#3636 — Updates SWR usage and cache hit/null handling that handlers now rely on; closely related to the new SWR call patterns.

Suggested labels

Bug, Core Team

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 4.55% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	Title clearly describes adding caching for key endpoints with ID lookup, directly matching the main changeset objective.
Description check	✅ Passed	Description comprehensively covers objectives, testing instructions, and checklist completion following the template structure.
Linked Issues check	✅ Passed	All coding requirements from issue #4150 are met: LiveKeyByID cache created in caches.go, integrated into all v2_keys handlers via SWR pattern, wired in route registration and tests.
Out of Scope Changes check	✅ Passed	All changes directly support LiveKeyByID cache implementation. Minor adjustments to cache configurations (keyAuthToApiRow/apiToKeyAuthRow) relate to cache optimization and are within scope.

✨ Finishing touches

• [ ] 📝 Generate docstrings

🧪 Generate unit tests (beta)

• [ ] Create PR with unit tests
• [ ] Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between a5e915cb7c31ba75e9de2e789a6fa45b98c89cfd and a991c684b9606c8d3b6fa1e9fd27a623ea5e3e4a.

📒 Files selected for processing (1)

• go/apps/api/routes/register.go (5 hunks)

🧰 Additional context used

🧠 Learnings (1)

📓 Common learnings

Learnt from: Flo4604
Repo: unkeyed/unkey PR: 4190
File: go/internal/services/keys/verifier.go:51-53
Timestamp: 2025-10-30T15:10:52.743Z
Learning: PR #4190 for unkeyed/unkey is focused solely on database schema and query changes for identity-based credits. It adds IdentityCredits and KeyCredits fields to structs and queries, but does not implement the priority enforcement logic in the usagelimiter. The logic implementation is intentionally deferred to a later PR in the stack.

Learnt from: imeyer
Repo: unkeyed/unkey PR: 3755
File: .github/workflows/runbook-freshness-check.yaml:157-173
Timestamp: 2025-08-08T14:59:52.283Z
Learning: Repo unkeyed/unkey: When a CI/workflow fix is deferred, imeyer prefers a thorough GitHub issue be opened with sections (Summary, Impact, Where, Repro, Observed vs Expected, Acceptance Criteria, Validation Plan, Out of Scope, References) and assigned to imeyer, including backlinks to the originating PR and comment.

Learnt from: imeyer
Repo: unkeyed/unkey PR: 3755
File: .github/workflows/job_detect_changes.yaml:111-114
Timestamp: 2025-08-08T16:10:00.224Z
Learning: Repo: unkeyed/unkey — When imeyer writes “make me an issue” on a PR thread, automatically open a thorough GitHub issue (Summary, Impact, Where, Repro, Observed vs Expected, Acceptance Criteria, Validation Plan, Out of Scope, References), include backlinks to the PR and comment, and assign to imeyer.

Learnt from: imeyer
Repo: unkeyed/unkey PR: 3755
File: .github/actions/setup-node/action.yaml:0-0
Timestamp: 2025-08-08T15:10:46.436Z
Learning: Repo: unkeyed/unkey — Preference: If imeyer comments “issue” on a PR thread, automatically open a thorough GitHub issue (Summary, Impact, Where, Observed vs Expected, Acceptance Criteria, Validation Plan, Out of Scope, References), include backlinks to the PR and comment, and assign to imeyer.

Learnt from: imeyer
Repo: unkeyed/unkey PR: 3755
File: .github/actions/setup-node/action.yaml:0-0
Timestamp: 2025-08-08T15:10:46.436Z
Learning: For repo unkeyed/unkey and PR review workflows: When imeyer comments "issue" on a thread, automatically create a thorough GitHub issue (sections: Summary, Impact, Where, Repro/Observed vs Expected, Acceptance Criteria, Validation Plan, Out of Scope, References), include backlinks to the PR and the specific comment, and assign the issue to imeyer.

🧬 Code graph analysis (1)

go/apps/api/routes/register.go (8)

go/internal/services/caches/caches.go (1)

• Caches (21-53)

go/apps/api/routes/v2_keys_add_permissions/handler.go (1)

• Handler (29-36)

go/apps/api/routes/v2_keys_get_key/handler.go (1)

• Handler (30-37)

go/apps/api/routes/v2_keys_update_credits/handler.go (1)

• Handler (29-37)

go/apps/api/routes/v2_keys_set_roles/handler.go (1)

• Handler (29-36)

go/apps/api/routes/v2_keys_reroll_key/handler.go (1)

• Handler (34-42)

go/apps/api/routes/v2_keys_update_key/handler.go (1)

• Handler (34-42)

go/apps/api/routes/v2_apis_delete_api/handler.go (1)

• Handler (30-37)

🔇 Additional comments (1)

go/apps/api/routes/register.go (1)

422-429: LGTM: Consistent LiveKeyCache wiring across handlers.

The LiveKeyCache dependency is correctly wired to all modified v2/keys handlers. The dependency injection pattern is consistent and aligns with the existing KeyCache wiring.

Also applies to: 436-442, 454-454, 463-469, 493-493, 502-507, 515-520, 528-533, 541-546, 554-559, 567-572

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[github-actions[bot]]

Thank you for following the naming conventions for pull request titles! 🙏

[Flo4604]

@Akhileshait some tests are still not passing, you can run make test locally to run the test suite on your machine

[Akhileshait]

@Flo4604 Fixed the issue

[Akhileshait]

Hi @Flo4604 Tests are passing now please review this

[Flo4604]

@Akhileshait we will review this next week

[Flo4604]

Hello @Akhileshait some merge conflicts popped up, could you please resolve them

Thank you

[Akhileshait]

Hi @Flo4604 Fixed the merge conflicts