fix: Added caching for key endpoints that lookup keys by id / whoami endpoint #4203
Conversation
Akhileshait
requested review from
Flo4604,
chronark,
imeyer and
perkinsjr
as code owners
|
|
@Akhileshait is attempting to deploy a commit to the Unkey Team on Vercel. A member of the Team first needs to authorize it. |
📝 WalkthroughWalkthroughAdds a LiveKeyByID cache to the caches service and wires a new Changes
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
Possibly related PRs
Suggested labels
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 📒 Files selected for processing (1)
🧰 Additional context used🧠 Learnings (1)📓 Common learnings🧬 Code graph analysis (1)go/apps/api/routes/register.go (8)
🔇 Additional comments (1)
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 |
|
Thank you for following the naming conventions for pull request titles! 🙏 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
go/apps/api/routes/v2_keys_update_key/200_test.go (1)
211-218: Add missing LiveKeyCache field for consistency.The Handler initialization in
TestKeyUpdateCreditsInvalidatesCacheis missing theLiveKeyCachefield that's present in the other two tests in this file (lines 35 and 139). This inconsistency could cause issues if the handler expects this field to be set.Apply this diff to add the missing field:
route := &handler.Handler{ Logger: h.Logger, DB: h.DB, Keys: h.Keys, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, }go/apps/api/routes/v2_keys_update_key/403_test.go (1)
37-44: Critical: Missing LiveKeyCache initialization will cause nil pointer dereference.The handler initialization is missing the
LiveKeyCachefield, but the handler code now requires it (seehandler.goline 68 whereh.LiveKeyCache.SWRis called). This will cause a panic when the test runs.Apply this diff to fix:
route := &handler.Handler{ DB: h.DB, Keys: h.Keys, Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, }
🧹 Nitpick comments (3)
go/internal/services/caches/caches.go (1)
108-118: Right TTLs; consider sizing and memory footprintFresh/stale windows align with SWR goals. MaxSize=1,000,000 with a large value type (db.FindLiveKeyByIDRow) can be memory‑intensive; expose this via config/env and add a rough size budget alert to avoid OOM under churn.
go/apps/api/routes/v2_keys_update_credits/handler.go (2)
280-282: Post‑update: write‑through updated row to cacheEviction is good; also Set the fresh row to minimize stale windows and reduce immediate re‑loads.
h.KeyCache.Remove(ctx, key.Hash) - h.LiveKeyCache.Remove(ctx, key.ID) + // Replace with fresh value to avoid stale re-population and shorten warmup. + h.LiveKeyCache.Set(ctx, key.ID, key)
65-75: Avoid negative caching on write pathsCaching 404s (WriteNull) can cause transient "not found" responses for keys created shortly before an update. Prefer Noop for not‑found to skip negative caching here.
Apply this diff:
- if db.IsNotFound(err) { - return cache.WriteNull - } + if db.IsNotFound(err) { + // Skip negative caching on write path to avoid transient 404s post-create. + return cache.Noop + }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (21)
go/apps/api/routes/register.go(3 hunks)go/apps/api/routes/v2_keys_set_roles/200_test.go(1 hunks)go/apps/api/routes/v2_keys_set_roles/400_test.go(1 hunks)go/apps/api/routes/v2_keys_set_roles/401_test.go(1 hunks)go/apps/api/routes/v2_keys_set_roles/403_test.go(1 hunks)go/apps/api/routes/v2_keys_set_roles/404_test.go(1 hunks)go/apps/api/routes/v2_keys_set_roles/handler.go(3 hunks)go/apps/api/routes/v2_keys_update_credits/200_test.go(1 hunks)go/apps/api/routes/v2_keys_update_credits/400_test.go(1 hunks)go/apps/api/routes/v2_keys_update_credits/401_test.go(1 hunks)go/apps/api/routes/v2_keys_update_credits/403_test.go(1 hunks)go/apps/api/routes/v2_keys_update_credits/404_test.go(1 hunks)go/apps/api/routes/v2_keys_update_credits/handler.go(3 hunks)go/apps/api/routes/v2_keys_update_key/200_test.go(2 hunks)go/apps/api/routes/v2_keys_update_key/400_test.go(1 hunks)go/apps/api/routes/v2_keys_update_key/401_test.go(1 hunks)go/apps/api/routes/v2_keys_update_key/403_test.go(2 hunks)go/apps/api/routes/v2_keys_update_key/404_test.go(1 hunks)go/apps/api/routes/v2_keys_update_key/handler.go(3 hunks)go/apps/api/routes/v2_keys_update_key/three_state_test.go(1 hunks)go/internal/services/caches/caches.go(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (21)
go/apps/api/routes/v2_keys_update_key/handler.go (3)
go/pkg/db/key_find_live_by_id.sql_generated.go (1)
FindLiveKeyByIDRow(105-141)go/pkg/cache/interface.go (3)
Op(51-51)WriteValue(57-57)WriteNull(60-60)go/pkg/db/handle_err_no_rows.go (1)
IsNotFound(8-10)
go/apps/api/routes/v2_keys_update_key/three_state_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_update_key/404_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_update_credits/404_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_update_credits/403_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_set_roles/handler.go (4)
go/pkg/db/cached_key_data.go (1)
CachedKeyData(5-8)go/pkg/db/key_find_live_by_id.sql_generated.go (1)
FindLiveKeyByIDRow(105-141)go/pkg/cache/interface.go (3)
Op(51-51)WriteValue(57-57)WriteNull(60-60)go/pkg/db/handle_err_no_rows.go (1)
IsNotFound(8-10)
go/apps/api/routes/v2_keys_set_roles/403_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_set_roles/401_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_update_key/401_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_update_credits/400_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_update_credits/200_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_set_roles/200_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_update_key/200_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/internal/services/caches/caches.go (5)
go/pkg/cache/interface.go (1)
Cache(7-32)go/pkg/db/key_find_live_by_id.sql_generated.go (1)
FindLiveKeyByIDRow(105-141)go/apps/gw/services/caches/caches.go (3)
New(74-142)Config(38-43)Caches(20-35)go/pkg/cache/cache.go (2)
New(55-99)Config(33-50)go/pkg/cache/middleware/tracing.go (1)
WithTracing(16-18)
go/apps/api/routes/v2_keys_set_roles/404_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_set_roles/400_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_update_credits/401_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_update_key/400_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/register.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_update_credits/handler.go (3)
go/pkg/db/key_find_live_by_id.sql_generated.go (1)
FindLiveKeyByIDRow(105-141)go/pkg/cache/interface.go (3)
Op(51-51)WriteValue(57-57)WriteNull(60-60)go/pkg/db/handle_err_no_rows.go (1)
IsNotFound(8-10)
go/apps/api/routes/v2_keys_update_key/403_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
🔇 Additional comments (30)
go/internal/services/caches/caches.go (2)
28-30: LGTM: new LiveKeyByID cache handle exposedField addition is clear and typed correctly for SWR lookups by key ID.
124-124: LGTM: tracing middleware appliedWiring through middleware.WithTracing mirrors other caches.
go/apps/api/routes/v2_keys_update_credits/handler.go (1)
35-35: LGTM: LiveKeyCache dependency addedKeeps handler testable and consistent with other routes.
go/apps/api/routes/v2_keys_update_credits/404_test.go (1)
24-25: LGTM: tests wire LiveKeyCacheEnsures handler has the new dependency in not‑found scenario.
go/apps/api/routes/v2_keys_set_roles/200_test.go (1)
24-30: LGTM: LiveKeyCache wired into handlerConsistent with service changes; tests remain readable.
go/apps/api/routes/v2_keys_update_credits/400_test.go (1)
25-26: LGTM: LiveKeyCache in bad‑request testsKeeps test setup aligned with handler signature.
go/apps/api/routes/v2_keys_update_key/three_state_test.go (1)
36-36: LGTM! Proper cache dependency wiring.The LiveKeyCache field is correctly initialized with the LiveKeyByID cache from the test harness, consistent with the new caching pattern introduced in this PR.
go/apps/api/routes/v2_keys_update_credits/403_test.go (1)
26-26: LGTM! Cache dependency properly initialized.The LiveKeyCache field is correctly wired to h.Caches.LiveKeyByID, maintaining consistency with other test files in this PR.
go/apps/api/routes/v2_keys_update_credits/200_test.go (1)
31-31: LGTM! Proper cache wiring for credits test.The LiveKeyCache is correctly initialized, and the test at lines 172-218 properly validates cache invalidation after credit updates.
go/apps/api/routes/v2_keys_update_key/200_test.go (2)
35-35: LGTM! Cache dependency properly initialized.The LiveKeyCache field is correctly wired for this test.
139-139: LGTM! Consistent cache initialization.The LiveKeyCache field is properly set up for this test case.
go/apps/api/routes/v2_keys_set_roles/404_test.go (1)
31-31: LGTM! Cache dependency correctly initialized.The LiveKeyCache field is properly wired to h.Caches.LiveKeyByID, maintaining consistency with the caching pattern in this PR.
go/apps/api/routes/v2_keys_update_key/404_test.go (1)
27-27: LGTM! Proper cache wiring.The LiveKeyCache field is correctly initialized for the not-found test scenarios.
go/apps/api/routes/v2_keys_update_key/400_test.go (1)
28-28: LGTM! Cache dependency properly set.The LiveKeyCache is correctly initialized for the invalid refill config test.
go/apps/api/routes/v2_keys_set_roles/403_test.go (1)
24-24: LGTM! Cache dependency correctly wired.The LiveKeyCache field is properly initialized, completing the Handler setup for authorization error tests.
go/apps/api/routes/v2_keys_set_roles/401_test.go (1)
17-23: LGTM!The test setup correctly initializes the new
LiveKeyCachefield, ensuring the handler can use the SWR-based live key lookup pattern introduced in this PR.go/apps/api/routes/v2_keys_update_credits/401_test.go (1)
19-27: LGTM!The handler initialization correctly includes the
LiveKeyCachefield, allowing the test to exercise the new caching behavior.go/apps/api/routes/v2_keys_update_key/403_test.go (2)
90-98: LGTM!The handler correctly initializes
LiveKeyCache, ensuring the test can use the new SWR-based key lookup.
137-145: LGTM!The handler is properly configured with the
LiveKeyCachefield.go/apps/api/routes/v2_keys_update_key/401_test.go (1)
20-28: LGTM!The test setup correctly includes the
LiveKeyCachefield for the handler.go/apps/api/routes/v2_keys_set_roles/400_test.go (1)
17-24: LGTM!The handler is correctly initialized with the new
LiveKeyCachefield.go/apps/api/routes/register.go (3)
420-431: LGTM!The
v2KeysUpdateKeyhandler is correctly wired with the newLiveKeyCache, enabling SWR-based key lookups to reduce database load.
458-469: LGTM!The
v2KeysUpdateCreditshandler is correctly configured with theLiveKeyCache.
472-482: LGTM!The
v2KeysSetRoleshandler is properly configured with theLiveKeyCachefield, completing the caching implementation for all three key endpoints.go/apps/api/routes/v2_keys_set_roles/handler.go (3)
28-35: LGTM!The
LiveKeyCachefield is properly added to the handler, enabling SWR-based caching for live key lookups.
62-84: LGTM!The SWR cache pattern is correctly implemented:
- Caches successful lookups to reduce DB load.
- Caches not-found results to prevent repeated lookups for non-existent keys.
- Doesn't cache transient errors, allowing retry on next request.
- Preserves existing error handling semantics.
278-279: LGTM!Cache invalidation is correctly implemented:
- Both
KeyCache(by hash) andLiveKeyCache(by ID) are invalidated after successful role updates.- Proper placement after transaction success ensures cache consistency.
go/apps/api/routes/v2_keys_update_key/handler.go (3)
33-41: LGTM!The
LiveKeyCachefield is correctly added to enable SWR-based key lookups.
68-94: LGTM!The SWR cache pattern is correctly implemented with appropriate error handling:
- Successful lookups and not-found results are cached.
- Transient errors are not cached.
- Existing error semantics are preserved.
640-649: LGTM!Cache invalidation is correctly implemented:
- Both
KeyCacheandLiveKeyCacheare invalidated after successful key updates.UsageLimiteris additionally invalidated when credits are updated.- Proper placement ensures cache consistency.
Flo4604
left a comment
Flo4604
left a comment
There was a problem hiding this comment.
The tests are not passing, and there are other endpoints that should also use this cache.
12handlers in total that use FindLiveKeyByID
And also please use the caches.DefaultFindFirstOp instead of having to re-use the same cache operation function.
|
@Akhileshait some tests are still not passing, you can run |
|
@Flo4604 Fixed the issue |
There was a problem hiding this comment.
Actionable comments posted: 1
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
go/apps/api/routes/v2_keys_delete_key/handler.go(4 hunks)go/apps/api/routes/v2_keys_get_key/412_test.go(3 hunks)go/apps/api/routes/v2_keys_update_key/200_test.go(3 hunks)go/apps/api/routes/v2_keys_update_key/403_test.go(3 hunks)go/pkg/zen/validation/validator.go(3 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- go/apps/api/routes/v2_keys_update_key/403_test.go
- go/apps/api/routes/v2_keys_update_key/200_test.go
🧰 Additional context used
🧠 Learnings (2)
📓 Common learnings
Learnt from: imeyer
PR: unkeyed/unkey#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
PR: unkeyed/unkey#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
PR: unkeyed/unkey#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
PR: unkeyed/unkey#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.
📚 Learning: 2025-08-27T11:02:17.155Z
Learnt from: chronark
PR: unkeyed/unkey#3860
File: go/pkg/zen/middleware_metrics_test.go:297-313
Timestamp: 2025-08-27T11:02:17.155Z
Learning: The redact function in go/pkg/zen/middleware_metrics.go is not expected to be called concurrently, so concurrency tests are unnecessary for this function.
Applied to files:
go/pkg/zen/validation/validator.go
🧬 Code graph analysis (1)
go/apps/api/routes/v2_keys_delete_key/handler.go (3)
go/pkg/db/cached_key_data.go (1)
CachedKeyData(5-8)go/pkg/db/key_find_live_by_id.sql_generated.go (1)
FindLiveKeyByIDRow(105-141)go/internal/services/caches/op.go (1)
DefaultFindFirstOp(9-22)
🔇 Additional comments (4)
go/apps/api/routes/v2_keys_get_key/412_test.go (1)
19-25: LGTM! Cache dependency properly wired into all test cases.All three Handler initializations consistently include the new
LiveKeyCachefield, correctly sourced fromh.Caches.LiveKeyByID. The tests maintain their focus on verifying precondition error scenarios while properly accommodating the new cache dependency.Also applies to: 90-97, 142-149
go/pkg/zen/validation/validator.go (3)
6-6: LGTM: Import is necessary.The
syncimport is correctly added to support the mutex field.
27-29: Mutex field is properly documented, but appears unrelated to PR objectives.The field and documentation are correct. However, this concurrency fix for the OpenAPI validator seems unrelated to the PR's stated goal of adding caching for key lookup endpoints. Consider whether this should be in a separate PR for better traceability.
66-70: Rewrite review comment with corrected version information and verified findings.The review comment references libopenapi-validator v0.9.0, which does not exist—the latest version is v0.8.0 (released Oct 22, 2025), and your codebase uses v0.4.6. Additionally, I found no public documentation of a race condition in
ValidateHttpRequestacross the library's changelog or issues.More notably, two other validator implementations in your codebase (
go/apps/gw/server/validation/andgo/apps/agent/pkg/openapi/) call the sameValidateHttpRequestmethod on the same library version without mutexes. This inconsistency suggests either the race condition is specific to this context, overly cautious, or outdated.The defer pattern suggestion is still valid for safety:
// Lock to protect concurrent access to the validator // This works around race conditions in libopenapi-validator v0.9.0 and earlier v.mu.Lock() + defer v.mu.Unlock() valid, errors := v.validator.ValidateHttpRequest(r) - v.mu.Unlock()Verify with the team: Is the mutex addressing a known production issue, or can it be removed given that other validators don't use it?
|
Hi @Flo4604 Tests are passing now please review this |
|
@Akhileshait we will review this next week |
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
go/apps/api/routes/v2_keys_delete_key/200_test.go (1)
24-31: Handler initialization looks correct.The LiveKeyCache is properly wired to the handler alongside existing fields. However, consider verifying cache invalidation after the delete operations in your test assertions. Since this PR introduces caching for key lookups, the tests should confirm that both
KeyCacheandLiveKeyCacheare properly invalidated when keys are deleted (soft or hard delete).For example, after the deletion operations, you could add assertions like:
// After soft delete _, found := route.LiveKeyCache.Get(softDeleteKeyID) require.False(t, found, "LiveKeyCache should not contain deleted key") // After hard delete _, found = route.LiveKeyCache.Get(hardDeleteKeyID) require.False(t, found, "LiveKeyCache should not contain permanently deleted key")go/apps/api/routes/v2_keys_set_roles/200_test.go (1)
23-30: LGTM! Handler initialization correctly includes the new cache.The
LiveKeyCachefield is properly initialized withh.Caches.LiveKeyByID, consistent with the other cache dependencies.As an optional enhancement, consider adding assertions to verify the cache is being used and invalidated correctly, since this PR specifically introduces caching to reduce database load. For example:
- Mock or instrument the cache to verify lookups occur before DB queries
- Verify cache invalidation happens after role updates
- Check that subsequent operations use cached data (within the freshness window)
This would provide stronger confidence that the caching feature works as intended, though behavioral correctness is already well-covered by the existing assertions.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (51)
go/apps/api/routes/v2_keys_add_permissions/200_test.go(1 hunks)go/apps/api/routes/v2_keys_add_permissions/400_test.go(1 hunks)go/apps/api/routes/v2_keys_add_permissions/401_test.go(1 hunks)go/apps/api/routes/v2_keys_add_permissions/403_test.go(1 hunks)go/apps/api/routes/v2_keys_add_permissions/404_test.go(1 hunks)go/apps/api/routes/v2_keys_add_roles/200_test.go(1 hunks)go/apps/api/routes/v2_keys_add_roles/400_test.go(1 hunks)go/apps/api/routes/v2_keys_add_roles/401_test.go(1 hunks)go/apps/api/routes/v2_keys_add_roles/403_test.go(1 hunks)go/apps/api/routes/v2_keys_add_roles/404_test.go(1 hunks)go/apps/api/routes/v2_keys_delete_key/200_test.go(1 hunks)go/apps/api/routes/v2_keys_delete_key/401_test.go(1 hunks)go/apps/api/routes/v2_keys_delete_key/403_test.go(1 hunks)go/apps/api/routes/v2_keys_delete_key/404_test.go(1 hunks)go/apps/api/routes/v2_keys_get_key/200_test.go(2 hunks)go/apps/api/routes/v2_keys_get_key/403_test.go(1 hunks)go/apps/api/routes/v2_keys_get_key/404_test.go(1 hunks)go/apps/api/routes/v2_keys_get_key/412_test.go(3 hunks)go/apps/api/routes/v2_keys_remove_permissions/200_test.go(1 hunks)go/apps/api/routes/v2_keys_remove_permissions/400_test.go(1 hunks)go/apps/api/routes/v2_keys_remove_permissions/401_test.go(1 hunks)go/apps/api/routes/v2_keys_remove_permissions/403_test.go(1 hunks)go/apps/api/routes/v2_keys_remove_permissions/404_test.go(1 hunks)go/apps/api/routes/v2_keys_remove_roles/200_test.go(1 hunks)go/apps/api/routes/v2_keys_remove_roles/400_test.go(1 hunks)go/apps/api/routes/v2_keys_remove_roles/401_test.go(1 hunks)go/apps/api/routes/v2_keys_remove_roles/403_test.go(1 hunks)go/apps/api/routes/v2_keys_remove_roles/404_test.go(1 hunks)go/apps/api/routes/v2_keys_reroll_key/200_test.go(1 hunks)go/apps/api/routes/v2_keys_reroll_key/401_test.go(1 hunks)go/apps/api/routes/v2_keys_reroll_key/403_test.go(1 hunks)go/apps/api/routes/v2_keys_reroll_key/handler.go(4 hunks)go/apps/api/routes/v2_keys_set_permissions/200_test.go(1 hunks)go/apps/api/routes/v2_keys_set_permissions/400_test.go(1 hunks)go/apps/api/routes/v2_keys_set_permissions/401_test.go(1 hunks)go/apps/api/routes/v2_keys_set_permissions/403_test.go(1 hunks)go/apps/api/routes/v2_keys_set_permissions/404_test.go(1 hunks)go/apps/api/routes/v2_keys_set_roles/200_test.go(1 hunks)go/apps/api/routes/v2_keys_set_roles/403_test.go(1 hunks)go/apps/api/routes/v2_keys_set_roles/404_test.go(1 hunks)go/apps/api/routes/v2_keys_update_credits/200_test.go(1 hunks)go/apps/api/routes/v2_keys_update_credits/400_test.go(1 hunks)go/apps/api/routes/v2_keys_update_credits/401_test.go(1 hunks)go/apps/api/routes/v2_keys_update_credits/403_test.go(1 hunks)go/apps/api/routes/v2_keys_update_credits/handler.go(4 hunks)go/apps/api/routes/v2_keys_update_key/200_test.go(3 hunks)go/apps/api/routes/v2_keys_update_key/400_test.go(1 hunks)go/apps/api/routes/v2_keys_update_key/401_test.go(1 hunks)go/apps/api/routes/v2_keys_update_key/403_test.go(3 hunks)go/apps/api/routes/v2_keys_update_key/404_test.go(1 hunks)go/apps/api/routes/v2_keys_update_key/three_state_test.go(1 hunks)
🚧 Files skipped from review as they are similar to previous changes (35)
- go/apps/api/routes/v2_keys_remove_permissions/400_test.go
- go/apps/api/routes/v2_keys_update_credits/200_test.go
- go/apps/api/routes/v2_keys_delete_key/403_test.go
- go/apps/api/routes/v2_keys_update_key/404_test.go
- go/apps/api/routes/v2_keys_add_roles/200_test.go
- go/apps/api/routes/v2_keys_update_key/200_test.go
- go/apps/api/routes/v2_keys_remove_permissions/200_test.go
- go/apps/api/routes/v2_keys_set_permissions/200_test.go
- go/apps/api/routes/v2_keys_remove_roles/200_test.go
- go/apps/api/routes/v2_keys_set_permissions/403_test.go
- go/apps/api/routes/v2_keys_update_key/three_state_test.go
- go/apps/api/routes/v2_keys_add_permissions/200_test.go
- go/apps/api/routes/v2_keys_update_credits/handler.go
- go/apps/api/routes/v2_keys_update_key/400_test.go
- go/apps/api/routes/v2_keys_update_key/401_test.go
- go/apps/api/routes/v2_keys_reroll_key/401_test.go
- go/apps/api/routes/v2_keys_remove_roles/401_test.go
- go/apps/api/routes/v2_keys_add_permissions/400_test.go
- go/apps/api/routes/v2_keys_remove_roles/400_test.go
- go/apps/api/routes/v2_keys_add_permissions/404_test.go
- go/apps/api/routes/v2_keys_set_roles/404_test.go
- go/apps/api/routes/v2_keys_set_permissions/400_test.go
- go/apps/api/routes/v2_keys_remove_permissions/403_test.go
- go/apps/api/routes/v2_keys_set_roles/403_test.go
- go/apps/api/routes/v2_keys_add_roles/400_test.go
- go/apps/api/routes/v2_keys_add_roles/403_test.go
- go/apps/api/routes/v2_keys_reroll_key/403_test.go
- go/apps/api/routes/v2_keys_add_roles/404_test.go
- go/apps/api/routes/v2_keys_set_permissions/401_test.go
- go/apps/api/routes/v2_keys_get_key/200_test.go
- go/apps/api/routes/v2_keys_update_credits/400_test.go
- go/apps/api/routes/v2_keys_remove_roles/404_test.go
- go/apps/api/routes/v2_keys_add_permissions/403_test.go
- go/apps/api/routes/v2_keys_delete_key/404_test.go
- go/apps/api/routes/v2_keys_update_credits/401_test.go
🧰 Additional context used
🧠 Learnings (4)
📓 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.710Z
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/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.
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.
📚 Learning: 2025-10-30T15:10:52.710Z
Learnt from: Flo4604
Repo: unkeyed/unkey PR: 4190
File: go/internal/services/keys/verifier.go:51-53
Timestamp: 2025-10-30T15:10:52.710Z
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.
Applied to files:
go/apps/api/routes/v2_keys_update_credits/403_test.gogo/apps/api/routes/v2_keys_add_permissions/401_test.go
📚 Learning: 2025-08-14T16:25:48.167Z
Learnt from: Flo4604
Repo: unkeyed/unkey PR: 3785
File: go/apps/api/routes/v2_keys_reroll_key/401_test.go:52-61
Timestamp: 2025-08-14T16:25:48.167Z
Learning: User Flo4604 requested creation of a GitHub issue to track converting all test files to use table-driven test patterns as a broader codebase improvement, following the suggestion made during review of go/apps/api/routes/v2_keys_reroll_key/401_test.go.
Applied to files:
go/apps/api/routes/v2_keys_update_key/403_test.gogo/apps/api/routes/v2_keys_remove_permissions/401_test.gogo/apps/api/routes/v2_keys_set_permissions/404_test.gogo/apps/api/routes/v2_keys_add_permissions/401_test.gogo/apps/api/routes/v2_keys_add_roles/401_test.gogo/apps/api/routes/v2_keys_get_key/404_test.go
📚 Learning: 2025-08-08T15:09:01.312Z
Learnt from: Flo4604
Repo: unkeyed/unkey PR: 3753
File: go/apps/api/openapi/config.yaml:9-10
Timestamp: 2025-08-08T15:09:01.312Z
Learning: Repo unkeyed/unkey: For go/apps/api/openapi, oapi-codegen is used and does not support OpenAPI 3.1 union types like [T, "null"]; an overlay step is required to downconvert to 3.0-style nullable before code generation.
Applied to files:
go/apps/api/routes/v2_keys_reroll_key/handler.go
🧬 Code graph analysis (4)
go/apps/api/routes/v2_keys_update_credits/403_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_update_key/403_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
go/apps/api/routes/v2_keys_reroll_key/handler.go (3)
go/pkg/db/cached_key_data.go (1)
CachedKeyData(5-8)go/pkg/db/key_find_live_by_id.sql_generated.go (1)
FindLiveKeyByIDRow(105-141)go/internal/services/caches/op.go (1)
DefaultFindFirstOp(9-22)
go/apps/api/routes/v2_keys_set_roles/200_test.go (1)
go/internal/services/caches/caches.go (1)
Caches(15-31)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)
- GitHub Check: Test Go API Local / Test
- GitHub Check: Test API / API Test Local
- GitHub Check: Build / Build
🔇 Additional comments (18)
go/apps/api/routes/v2_keys_update_credits/403_test.go (1)
26-26: LGTM! Cache wiring is correct.The addition of
LiveKeyCachecorrectly wires the new LiveKeyByID cache into the handler for testing. This enables the handler to use SWR-backed key lookups and follows the consistent pattern across the PR.go/apps/api/routes/v2_keys_remove_roles/403_test.go (1)
23-30: LGTM! LiveKeyCache correctly wired into handler initialization.The addition of
LiveKeyCache: h.Caches.LiveKeyByIDfollows the consistent pattern established by the existingKeyCachefield and aligns with the PR's objective to integrate caching for key lookups across handlers.go/apps/api/routes/v2_keys_set_permissions/404_test.go (1)
23-30: LGTM! Cache dependency correctly wired.The addition of
LiveKeyCacheto the handler initialization properly integrates the new LiveKeyByID cache. The test continues to validate 404 scenarios correctly, and the minimal change aligns with the PR's objective to add caching infrastructure for key lookups.go/apps/api/routes/v2_keys_add_permissions/401_test.go (1)
24-29: LGTM! Cache dependency correctly wired.The addition of
LiveKeyCache: h.Caches.LiveKeyByIDproperly wires the new cache dependency into the handler initialization. This is consistent with the PR's objective to add caching for key lookups and follows the same pattern applied across other test files.go/apps/api/routes/v2_keys_add_roles/401_test.go (1)
18-25: LGTM! Handler initialization correctly wired with LiveKeyCache.The addition of
LiveKeyCache: h.Caches.LiveKeyByIDproperly wires the new cache into the test Handler, consistent with the PR's objective to enable key lookups by ID across all route handlers.go/apps/api/routes/v2_keys_delete_key/401_test.go (1)
22-29: Handler initialization looks correct.The LiveKeyCache is properly wired alongside existing fields. All required dependencies are present and correctly initialized for the unauthorized access test scenarios.
go/apps/api/routes/v2_keys_update_key/403_test.go (2)
37-45: LGTM: LiveKeyCache wiring is correct.The addition of
LiveKeyCache: h.Caches.LiveKeyByIDcorrectly wires the new cache dependency into the handler for testing, consistent with the handler's updated structure.
138-146: Cache behavior verification confirmed.Cache invalidation is properly tested in
200_test.go(lines 272-281), which verifies that after a successful key update, the cache is invalidated by retrieving the key again and confirming the updated value is reflected. This is the correct location for such tests since cache invalidation only occurs in success paths. The403_test.gotests authorization failures, which never reach the cache removal code, so it's appropriate that they don't test cache behavior.go/apps/api/routes/v2_keys_reroll_key/handler.go (4)
14-14: LGTM: Cache imports added correctly.The imports for the caching packages are properly added to support the new caching functionality.
Also applies to: 18-18
40-41: LGTM: Cache fields added to Handler.The KeyCache and LiveKeyCache fields are correctly typed and align with the caching strategy described in the PR objectives. KeyCache is used for verification key lookups by hash, while LiveKeyCache is used for live key lookups by ID.
70-72: LGTM: SWR caching pattern correctly implemented.The SWR (Stale-While-Revalidate) pattern is correctly implemented with:
- Cache key based on req.KeyId
- Loader function performing the database query
- DefaultFindFirstOp determining cache behavior based on query result
This reduces database load for repeated key lookups while maintaining consistency.
371-372: LGTM: Cache invalidation properly implemented.Cache invalidation correctly targets the old key after the successful reroll operation:
KeyCache.Remove(ctx, key.Hash): Ensures verification endpoints see the expired keyLiveKeyCache.Remove(ctx, key.ID): Ensures lookup endpoints see the updated key stateThe invalidation order and timing are appropriate. Note that KeyCache is invalidated even though this handler doesn't read from it, which is correct because other handlers (e.g., verification endpoints) may have cached the key by hash.
go/apps/api/routes/v2_keys_reroll_key/200_test.go (1)
24-30: LGTM: Test handler initialization correctly updated.The Handler initialization properly wires the cache instances from the test harness:
KeyCache→h.Caches.VerificationKeyByHashLiveKeyCache→h.Caches.LiveKeyByIDThis ensures the handler can function with the new caching functionality during tests.
go/apps/api/routes/v2_keys_remove_permissions/404_test.go (1)
26-33: LGTM! Correct wiring of the LiveKeyCache dependency.The Handler initialization properly adds the new
LiveKeyCachefield alongside existing dependencies, usingh.Caches.LiveKeyByIDfrom the test harness. This correctly supports the PR's caching objective for key lookups.go/apps/api/routes/v2_keys_remove_permissions/401_test.go (1)
18-25: LGTM! Consistent dependency injection.The Handler initialization correctly adds the
LiveKeyCachefield, matching the pattern used across other test files in this PR. The authentication error tests remain unchanged while properly supporting the new cache dependency.go/apps/api/routes/v2_keys_get_key/404_test.go (1)
20-27: LGTM! Cache dependency correctly wired.The LiveKeyCache field is properly initialized from the test harness, enabling the handler to use the new SWR-backed live key cache.
go/apps/api/routes/v2_keys_get_key/412_test.go (1)
19-25: LGTM! Cache dependency consistently wired across all test cases.All three Handler initializations correctly add LiveKeyCache from the test harness, maintaining consistency across the precondition error test scenarios.
Also applies to: 90-97, 142-149
go/apps/api/routes/v2_keys_get_key/403_test.go (1)
25-32: LGTM! Cache dependency correctly wired.The LiveKeyCache initialization is consistent with the broader pattern across v2_keys_* route tests and properly enables cache usage in the handler.
|
Hello @Akhileshait some merge conflicts popped up, could you please resolve them Thank you |
…caching-for-key-endpoints-unkeyed#4150
…github.com/Akhileshait/unkey into fix-caching-for-key-endpoints-unkeyed#4150
|
Hi @Flo4604 Fixed the merge conflicts |
Flo4604
left a comment
Flo4604
left a comment
There was a problem hiding this comment.
The handlers are all missing their NULL handling
see go/apps/api/routes/v2_apis_get_api/handler.go for an example on how we handle it normally.
| key, _, err := h.LiveKeyCache.SWR(ctx, req.KeyId, func(ctx context.Context) (db.FindLiveKeyByIDRow, error) { | ||
| return db.Query.FindLiveKeyByID(ctx, h.DB.RO(), req.KeyId) | ||
| }, caches.DefaultFindFirstOp) |
There was a problem hiding this comment.
you have to look at the _ which is the hit parameter.
Hit can be NULL which means a key doesn't exist
| key, _, err := h.LiveKeyCache.SWR(ctx, req.KeyId, func(ctx context.Context) (db.FindLiveKeyByIDRow, error) { | ||
| return db.Query.FindLiveKeyByID(ctx, h.DB.RO(), req.KeyId) | ||
| }, caches.DefaultFindFirstOp) |
| key, _, err := h.LiveKeyCache.SWR(ctx, req.KeyId, func(ctx context.Context) (db.FindLiveKeyByIDRow, error) { | ||
| return db.Query.FindLiveKeyByID(ctx, h.DB.RO(), req.KeyId) | ||
| }, caches.DefaultFindFirstOp) |
| key, _, err := h.LiveKeyCache.SWR(ctx, req.KeyId, func(ctx context.Context) (db.FindLiveKeyByIDRow, error) { | ||
| return db.Query.FindLiveKeyByID(ctx, h.DB.RO(), req.KeyId) | ||
| }, caches.DefaultFindFirstOp) |
| key, _, err := h.LiveKeyCache.SWR(ctx, req.KeyId, func(ctx context.Context) (db.FindLiveKeyByIDRow, error) { | ||
| return db.Query.FindLiveKeyByID(ctx, h.DB.RO(), req.KeyId) | ||
| }, caches.DefaultFindFirstOp) |
| key, err := db.Query.FindLiveKeyByID(ctx, h.DB.RO(), req.KeyId) | ||
| key, _, err := h.LiveKeyCache.SWR(ctx, req.KeyId, func(ctx context.Context) (db.FindLiveKeyByIDRow, error) { | ||
| return db.Query.FindLiveKeyByID(ctx, h.DB.RO(), req.KeyId) | ||
| }, caches.DefaultFindFirstOp) |
| } | ||
|
|
||
| key, err := db.Query.FindLiveKeyByID(ctx, h.DB.RO(), req.KeyId) | ||
| key, _, err := h.LiveKeyCache.SWR(ctx, req.KeyId, func(ctx context.Context) (db.FindLiveKeyByIDRow, error) { |
| } | ||
|
|
||
| key, err := db.Query.FindLiveKeyByID(ctx, h.DB.RO(), req.KeyId) | ||
| key, _, err := h.LiveKeyCache.SWR(ctx, req.KeyId, func(ctx context.Context) (db.FindLiveKeyByIDRow, error) { |
| cache.Config[cache.ScopedKey, db.FindKeyAuthsByKeyAuthIdsRow]{ | ||
| Fresh: 10 * time.Minute, | ||
| Fresh: time.Minute, |
There was a problem hiding this comment.
a 10 miniute TTL is fine for keyAuth
| dispatcher, | ||
| cache.Config[cache.ScopedKey, db.FindKeyAuthsByIdsRow]{ | ||
| Fresh: 10 * time.Minute, | ||
| Fresh: time.Minute, |
3 participants
What does this PR do?
This PR adds caching for key lookups in the Go API to reduce database queries when using
updateKey,updateCredits, andsetRolesendpoints.Previously, these handlers were querying the database directly every time they needed to fetch key information. Now they use a new
LiveKeyByIDcache with a Stale-While-Revalidate (SWR) pattern that:This reduces database load and improves API response times for repeated key operations.
Fixes #4150
Type of change
How should this be tested?
Unit Tests
Run the existing unit tests for the modified handlers:
cd go
make test
Checklist
Required
pnpm buildpnpm fmtconsole.logsgit pull origin mainAppreciated