diff --git a/go/apps/api/routes/register.go b/go/apps/api/routes/register.go index 4e05b81fa2..fe1e41756a 100644 --- a/go/apps/api/routes/register.go +++ b/go/apps/api/routes/register.go @@ -419,11 +419,13 @@ func Register(srv *zen.Server, svc *Services, info zen.InstanceInfo) { srv.RegisterRoute( defaultMiddlewares, &v2KeysRerollKey.Handler{ - Logger: svc.Logger, - DB: svc.Database, - Keys: svc.Keys, - Auditlogs: svc.Auditlogs, - Vault: svc.Vault, + Logger: svc.Logger, + DB: svc.Database, + Keys: svc.Keys, + Auditlogs: svc.Auditlogs, + Vault: svc.Vault, + KeyCache: svc.Caches.VerificationKeyByHash, + LiveKeyCache: svc.Caches.LiveKeyByID, }, ) @@ -431,11 +433,12 @@ func Register(srv *zen.Server, svc *Services, info zen.InstanceInfo) { srv.RegisterRoute( defaultMiddlewares, &v2KeysDeleteKey.Handler{ - KeyCache: svc.Caches.VerificationKeyByHash, - Logger: svc.Logger, - DB: svc.Database, - Keys: svc.Keys, - Auditlogs: svc.Auditlogs, + KeyCache: svc.Caches.VerificationKeyByHash, + LiveKeyCache: svc.Caches.LiveKeyByID, + Logger: svc.Logger, + DB: svc.Database, + Keys: svc.Keys, + Auditlogs: svc.Auditlogs, }, ) @@ -448,6 +451,7 @@ func Register(srv *zen.Server, svc *Services, info zen.InstanceInfo) { Keys: svc.Keys, Auditlogs: svc.Auditlogs, KeyCache: svc.Caches.VerificationKeyByHash, + LiveKeyCache: svc.Caches.LiveKeyByID, UsageLimiter: svc.UsageLimiter, }, ) @@ -456,11 +460,12 @@ func Register(srv *zen.Server, svc *Services, info zen.InstanceInfo) { srv.RegisterRoute( defaultMiddlewares, &v2KeysGetKey.Handler{ - Logger: svc.Logger, - DB: svc.Database, - Keys: svc.Keys, - Auditlogs: svc.Auditlogs, - Vault: svc.Vault, + Logger: svc.Logger, + DB: svc.Database, + Keys: svc.Keys, + Auditlogs: svc.Auditlogs, + Vault: svc.Vault, + LiveKeyCache: svc.Caches.LiveKeyByID, }, ) @@ -485,6 +490,7 @@ func Register(srv *zen.Server, svc *Services, info zen.InstanceInfo) { Keys: svc.Keys, Auditlogs: svc.Auditlogs, KeyCache: svc.Caches.VerificationKeyByHash, + LiveKeyCache: svc.Caches.LiveKeyByID, UsageLimiter: svc.UsageLimiter, }, ) @@ -493,11 +499,12 @@ func Register(srv *zen.Server, svc *Services, info zen.InstanceInfo) { srv.RegisterRoute( defaultMiddlewares, &v2KeysSetRoles.Handler{ - Logger: svc.Logger, - DB: svc.Database, - Keys: svc.Keys, - Auditlogs: svc.Auditlogs, - KeyCache: svc.Caches.VerificationKeyByHash, + Logger: svc.Logger, + DB: svc.Database, + Keys: svc.Keys, + Auditlogs: svc.Auditlogs, + KeyCache: svc.Caches.VerificationKeyByHash, + LiveKeyCache: svc.Caches.LiveKeyByID, }, ) @@ -505,11 +512,12 @@ func Register(srv *zen.Server, svc *Services, info zen.InstanceInfo) { srv.RegisterRoute( defaultMiddlewares, &v2KeysSetPermissions.Handler{ - Logger: svc.Logger, - DB: svc.Database, - Keys: svc.Keys, - Auditlogs: svc.Auditlogs, - KeyCache: svc.Caches.VerificationKeyByHash, + Logger: svc.Logger, + DB: svc.Database, + Keys: svc.Keys, + Auditlogs: svc.Auditlogs, + KeyCache: svc.Caches.VerificationKeyByHash, + LiveKeyCache: svc.Caches.LiveKeyByID, }, ) @@ -517,11 +525,12 @@ func Register(srv *zen.Server, svc *Services, info zen.InstanceInfo) { srv.RegisterRoute( defaultMiddlewares, &v2KeysAddPermissions.Handler{ - Logger: svc.Logger, - DB: svc.Database, - Keys: svc.Keys, - Auditlogs: svc.Auditlogs, - KeyCache: svc.Caches.VerificationKeyByHash, + Logger: svc.Logger, + DB: svc.Database, + Keys: svc.Keys, + Auditlogs: svc.Auditlogs, + KeyCache: svc.Caches.VerificationKeyByHash, + LiveKeyCache: svc.Caches.LiveKeyByID, }, ) @@ -529,11 +538,12 @@ func Register(srv *zen.Server, svc *Services, info zen.InstanceInfo) { srv.RegisterRoute( defaultMiddlewares, &v2KeysAddRoles.Handler{ - Logger: svc.Logger, - DB: svc.Database, - Keys: svc.Keys, - Auditlogs: svc.Auditlogs, - KeyCache: svc.Caches.VerificationKeyByHash, + Logger: svc.Logger, + DB: svc.Database, + Keys: svc.Keys, + Auditlogs: svc.Auditlogs, + KeyCache: svc.Caches.VerificationKeyByHash, + LiveKeyCache: svc.Caches.LiveKeyByID, }, ) @@ -541,11 +551,12 @@ func Register(srv *zen.Server, svc *Services, info zen.InstanceInfo) { srv.RegisterRoute( defaultMiddlewares, &v2KeysRemovePermissions.Handler{ - Logger: svc.Logger, - DB: svc.Database, - Keys: svc.Keys, - Auditlogs: svc.Auditlogs, - KeyCache: svc.Caches.VerificationKeyByHash, + Logger: svc.Logger, + DB: svc.Database, + Keys: svc.Keys, + Auditlogs: svc.Auditlogs, + KeyCache: svc.Caches.VerificationKeyByHash, + LiveKeyCache: svc.Caches.LiveKeyByID, }, ) @@ -553,11 +564,12 @@ func Register(srv *zen.Server, svc *Services, info zen.InstanceInfo) { srv.RegisterRoute( defaultMiddlewares, &v2KeysRemoveRoles.Handler{ - Logger: svc.Logger, - DB: svc.Database, - Keys: svc.Keys, - Auditlogs: svc.Auditlogs, - KeyCache: svc.Caches.VerificationKeyByHash, + Logger: svc.Logger, + DB: svc.Database, + Keys: svc.Keys, + Auditlogs: svc.Auditlogs, + KeyCache: svc.Caches.VerificationKeyByHash, + LiveKeyCache: svc.Caches.LiveKeyByID, }, ) diff --git a/go/apps/api/routes/v2_keys_add_permissions/200_test.go b/go/apps/api/routes/v2_keys_add_permissions/200_test.go index 71b064fac4..6c72ae9e86 100644 --- a/go/apps/api/routes/v2_keys_add_permissions/200_test.go +++ b/go/apps/api/routes/v2_keys_add_permissions/200_test.go @@ -19,11 +19,12 @@ func TestSuccess(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_add_permissions/400_test.go b/go/apps/api/routes/v2_keys_add_permissions/400_test.go index e5558da8df..ca1fdac7ec 100644 --- a/go/apps/api/routes/v2_keys_add_permissions/400_test.go +++ b/go/apps/api/routes/v2_keys_add_permissions/400_test.go @@ -23,11 +23,12 @@ func TestValidationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_add_permissions/401_test.go b/go/apps/api/routes/v2_keys_add_permissions/401_test.go index 202e178d7a..4e649a9684 100644 --- a/go/apps/api/routes/v2_keys_add_permissions/401_test.go +++ b/go/apps/api/routes/v2_keys_add_permissions/401_test.go @@ -21,11 +21,12 @@ func TestAuthenticationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_add_permissions/403_test.go b/go/apps/api/routes/v2_keys_add_permissions/403_test.go index 07c0cfa693..73d38b84d6 100644 --- a/go/apps/api/routes/v2_keys_add_permissions/403_test.go +++ b/go/apps/api/routes/v2_keys_add_permissions/403_test.go @@ -21,11 +21,12 @@ func TestAuthorizationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_add_permissions/404_test.go b/go/apps/api/routes/v2_keys_add_permissions/404_test.go index 4ef5971b0b..0bb9e88a4b 100644 --- a/go/apps/api/routes/v2_keys_add_permissions/404_test.go +++ b/go/apps/api/routes/v2_keys_add_permissions/404_test.go @@ -22,11 +22,12 @@ func TestNotFoundErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_add_permissions/handler.go b/go/apps/api/routes/v2_keys_add_permissions/handler.go index 781a7be78b..e169da4e65 100644 --- a/go/apps/api/routes/v2_keys_add_permissions/handler.go +++ b/go/apps/api/routes/v2_keys_add_permissions/handler.go @@ -9,6 +9,7 @@ import ( "github.com/unkeyed/unkey/go/apps/api/openapi" "github.com/unkeyed/unkey/go/internal/services/auditlogs" + "github.com/unkeyed/unkey/go/internal/services/caches" "github.com/unkeyed/unkey/go/internal/services/keys" "github.com/unkeyed/unkey/go/pkg/auditlog" "github.com/unkeyed/unkey/go/pkg/cache" @@ -26,11 +27,12 @@ type Request = openapi.V2KeysAddPermissionsRequestBody type Response = openapi.V2KeysAddPermissionsResponseBody type Handler struct { - Logger logging.Logger - DB db.Database - Keys keys.KeyService - Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.CachedKeyData] + Logger logging.Logger + DB db.Database + Keys keys.KeyService + Auditlogs auditlogs.AuditLogService + KeyCache cache.Cache[string, db.CachedKeyData] + LiveKeyCache cache.Cache[string, db.FindLiveKeyByIDRow] } // Method returns the HTTP method this route responds to @@ -56,7 +58,9 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } - 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) if err != nil { if db.IsNotFound(err) { return fault.New("key not found", @@ -292,6 +296,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { } h.KeyCache.Remove(ctx, key.Hash) + h.LiveKeyCache.Remove(ctx, key.ID) responseData := make(openapi.V2KeysAddPermissionsResponseData, 0) diff --git a/go/apps/api/routes/v2_keys_add_roles/200_test.go b/go/apps/api/routes/v2_keys_add_roles/200_test.go index cb4f705de3..0985b168be 100644 --- a/go/apps/api/routes/v2_keys_add_roles/200_test.go +++ b/go/apps/api/routes/v2_keys_add_roles/200_test.go @@ -20,11 +20,12 @@ func TestSuccess(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - Logger: h.Logger, - DB: h.DB, - Keys: h.Keys, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + Logger: h.Logger, + DB: h.DB, + Keys: h.Keys, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_add_roles/400_test.go b/go/apps/api/routes/v2_keys_add_roles/400_test.go index 00cfdec065..1e0276c075 100644 --- a/go/apps/api/routes/v2_keys_add_roles/400_test.go +++ b/go/apps/api/routes/v2_keys_add_roles/400_test.go @@ -16,11 +16,12 @@ func TestValidationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - Logger: h.Logger, - DB: h.DB, - Keys: h.Keys, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + Logger: h.Logger, + DB: h.DB, + Keys: h.Keys, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_add_roles/401_test.go b/go/apps/api/routes/v2_keys_add_roles/401_test.go index 6846f9435a..ecad708ea2 100644 --- a/go/apps/api/routes/v2_keys_add_roles/401_test.go +++ b/go/apps/api/routes/v2_keys_add_roles/401_test.go @@ -16,11 +16,12 @@ func TestAuthenticationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - Logger: h.Logger, - DB: h.DB, - Keys: h.Keys, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + Logger: h.Logger, + DB: h.DB, + Keys: h.Keys, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_add_roles/403_test.go b/go/apps/api/routes/v2_keys_add_roles/403_test.go index 6e85047994..9fcc2774a0 100644 --- a/go/apps/api/routes/v2_keys_add_roles/403_test.go +++ b/go/apps/api/routes/v2_keys_add_roles/403_test.go @@ -23,11 +23,12 @@ func TestAuthorizationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - Logger: h.Logger, - DB: h.DB, - Keys: h.Keys, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + Logger: h.Logger, + DB: h.DB, + Keys: h.Keys, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_add_roles/404_test.go b/go/apps/api/routes/v2_keys_add_roles/404_test.go index 5cdc0e0443..abb729b545 100644 --- a/go/apps/api/routes/v2_keys_add_roles/404_test.go +++ b/go/apps/api/routes/v2_keys_add_roles/404_test.go @@ -22,11 +22,12 @@ func TestNotFoundErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - Logger: h.Logger, - DB: h.DB, - Keys: h.Keys, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + Logger: h.Logger, + DB: h.DB, + Keys: h.Keys, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_add_roles/handler.go b/go/apps/api/routes/v2_keys_add_roles/handler.go index 7a51da3670..aa01c92bd2 100644 --- a/go/apps/api/routes/v2_keys_add_roles/handler.go +++ b/go/apps/api/routes/v2_keys_add_roles/handler.go @@ -8,6 +8,7 @@ import ( "github.com/unkeyed/unkey/go/apps/api/openapi" "github.com/unkeyed/unkey/go/internal/services/auditlogs" + "github.com/unkeyed/unkey/go/internal/services/caches" "github.com/unkeyed/unkey/go/internal/services/keys" "github.com/unkeyed/unkey/go/pkg/auditlog" "github.com/unkeyed/unkey/go/pkg/cache" @@ -25,11 +26,12 @@ type ( ) type Handler struct { - Logger logging.Logger - DB db.Database - Keys keys.KeyService - Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.CachedKeyData] + Logger logging.Logger + DB db.Database + Keys keys.KeyService + Auditlogs auditlogs.AuditLogService + KeyCache cache.Cache[string, db.CachedKeyData] + LiveKeyCache cache.Cache[string, db.FindLiveKeyByIDRow] } // Method returns the HTTP method this route responds to @@ -57,7 +59,9 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } - 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) if err != nil { if db.IsNotFound(err) { return fault.New("key not found", @@ -217,6 +221,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { } h.KeyCache.Remove(ctx, key.Hash) + h.LiveKeyCache.Remove(ctx, key.ID) } responseData := make(openapi.V2KeysAddRolesResponseData, 0) diff --git a/go/apps/api/routes/v2_keys_delete_key/200_test.go b/go/apps/api/routes/v2_keys_delete_key/200_test.go index f52e344e4f..9d47af12ea 100644 --- a/go/apps/api/routes/v2_keys_delete_key/200_test.go +++ b/go/apps/api/routes/v2_keys_delete_key/200_test.go @@ -22,11 +22,12 @@ func TestKeyDeleteSuccess(t *testing.T) { ctx := context.Background() route := &handler.Handler{ - Logger: h.Logger, - DB: h.DB, - Keys: h.Keys, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + Logger: h.Logger, + DB: h.DB, + Keys: h.Keys, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_delete_key/400_test.go b/go/apps/api/routes/v2_keys_delete_key/400_test.go index 147f38cfc9..ee2e20e9cf 100644 --- a/go/apps/api/routes/v2_keys_delete_key/400_test.go +++ b/go/apps/api/routes/v2_keys_delete_key/400_test.go @@ -15,11 +15,12 @@ func TestKeyDeleteBadRequest(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_delete_key/401_test.go b/go/apps/api/routes/v2_keys_delete_key/401_test.go index 3313b9582a..20dd8ecfd7 100644 --- a/go/apps/api/routes/v2_keys_delete_key/401_test.go +++ b/go/apps/api/routes/v2_keys_delete_key/401_test.go @@ -20,11 +20,12 @@ func TestKeyDeleteUnauthorized(t *testing.T) { ctx := t.Context() route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_delete_key/403_test.go b/go/apps/api/routes/v2_keys_delete_key/403_test.go index 88db514cd9..e6a27d523c 100644 --- a/go/apps/api/routes/v2_keys_delete_key/403_test.go +++ b/go/apps/api/routes/v2_keys_delete_key/403_test.go @@ -23,11 +23,12 @@ func TestKeyDeleteForbidden(t *testing.T) { ctx := context.Background() route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_delete_key/404_test.go b/go/apps/api/routes/v2_keys_delete_key/404_test.go index 6329fc387f..6fe9a254be 100644 --- a/go/apps/api/routes/v2_keys_delete_key/404_test.go +++ b/go/apps/api/routes/v2_keys_delete_key/404_test.go @@ -21,11 +21,12 @@ func TestKeyDeleteNotFound(t *testing.T) { ctx := t.Context() route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_delete_key/handler.go b/go/apps/api/routes/v2_keys_delete_key/handler.go index 98edf482b9..5139b8adfb 100644 --- a/go/apps/api/routes/v2_keys_delete_key/handler.go +++ b/go/apps/api/routes/v2_keys_delete_key/handler.go @@ -9,6 +9,7 @@ import ( "github.com/unkeyed/unkey/go/apps/api/openapi" "github.com/unkeyed/unkey/go/internal/services/auditlogs" + "github.com/unkeyed/unkey/go/internal/services/caches" "github.com/unkeyed/unkey/go/internal/services/keys" "github.com/unkeyed/unkey/go/pkg/auditlog" "github.com/unkeyed/unkey/go/pkg/cache" @@ -26,11 +27,12 @@ type Response = openapi.V2KeysDeleteKeyResponseBody // Handler implements zen.Route interface for the v2 keys.deleteKey endpoint type Handler struct { - Logger logging.Logger - DB db.Database - Keys keys.KeyService - Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.CachedKeyData] + Logger logging.Logger + DB db.Database + Keys keys.KeyService + Auditlogs auditlogs.AuditLogService + KeyCache cache.Cache[string, db.CachedKeyData] + LiveKeyCache cache.Cache[string, db.FindLiveKeyByIDRow] } // Method returns the HTTP method this route responds to @@ -59,7 +61,9 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } - 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) if err != nil { if db.IsNotFound(err) { return fault.Wrap( @@ -103,6 +107,11 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } + // Evict caches BEFORE transaction to minimize the window where stale data is served. + // We evict again after the transaction to handle repopulation during the transaction. + h.KeyCache.Remove(ctx, key.Hash) + h.LiveKeyCache.Remove(ctx, key.ID) + err = db.Tx(ctx, h.DB.RW(), func(ctx context.Context, tx db.DBTX) (err error) { description := "Deleted" if ptr.SafeDeref(req.Permanent) { @@ -152,7 +161,11 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } + // Re-evict caches after transaction commits to handle the race where another + // goroutine repopulated the cache while the transaction was running. + // This ensures the deleted key cannot remain cached during the SWR staleness window. h.KeyCache.Remove(ctx, key.Hash) + h.LiveKeyCache.Remove(ctx, key.ID) return s.JSON(http.StatusOK, Response{ Meta: openapi.Meta{ diff --git a/go/apps/api/routes/v2_keys_get_key/200_test.go b/go/apps/api/routes/v2_keys_get_key/200_test.go index 324fb0860a..a97b9ae95d 100644 --- a/go/apps/api/routes/v2_keys_get_key/200_test.go +++ b/go/apps/api/routes/v2_keys_get_key/200_test.go @@ -25,11 +25,12 @@ func TestGetKeyByKeyID(t *testing.T) { ctx := context.Background() route := &handler.Handler{ - Logger: h.Logger, - DB: h.DB, - Keys: h.Keys, - Auditlogs: h.Auditlogs, - Vault: h.Vault, + Logger: h.Logger, + DB: h.DB, + Keys: h.Keys, + Auditlogs: h.Auditlogs, + Vault: h.Vault, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) @@ -124,11 +125,12 @@ func TestGetKeyByKeyID(t *testing.T) { func TestGetKey_AdditionalScenarios(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - Logger: h.Logger, - DB: h.DB, - Keys: h.Keys, - Auditlogs: h.Auditlogs, - Vault: h.Vault, + Logger: h.Logger, + DB: h.DB, + Keys: h.Keys, + Auditlogs: h.Auditlogs, + Vault: h.Vault, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_get_key/400_test.go b/go/apps/api/routes/v2_keys_get_key/400_test.go index 773aa67727..63802f1836 100644 --- a/go/apps/api/routes/v2_keys_get_key/400_test.go +++ b/go/apps/api/routes/v2_keys_get_key/400_test.go @@ -16,11 +16,12 @@ func TestGetKeyBadRequest(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - Vault: h.Vault, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + Vault: h.Vault, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_get_key/401_test.go b/go/apps/api/routes/v2_keys_get_key/401_test.go index b3a339abed..b2500538e0 100644 --- a/go/apps/api/routes/v2_keys_get_key/401_test.go +++ b/go/apps/api/routes/v2_keys_get_key/401_test.go @@ -16,11 +16,12 @@ func TestGetKeyUnauthorized(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - Vault: h.Vault, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + Vault: h.Vault, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_get_key/403_test.go b/go/apps/api/routes/v2_keys_get_key/403_test.go index 79f1a496ee..cdb8f38c0a 100644 --- a/go/apps/api/routes/v2_keys_get_key/403_test.go +++ b/go/apps/api/routes/v2_keys_get_key/403_test.go @@ -23,11 +23,12 @@ func TestGetKeyForbidden(t *testing.T) { ctx := context.Background() route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - Vault: h.Vault, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + Vault: h.Vault, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_get_key/404_test.go b/go/apps/api/routes/v2_keys_get_key/404_test.go index a1c8b37f48..c5daf31da6 100644 --- a/go/apps/api/routes/v2_keys_get_key/404_test.go +++ b/go/apps/api/routes/v2_keys_get_key/404_test.go @@ -18,11 +18,12 @@ func TestGetKeyNotFound(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - Vault: h.Vault, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + Vault: h.Vault, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_get_key/412_test.go b/go/apps/api/routes/v2_keys_get_key/412_test.go index acd12e1c3e..5af66279ca 100644 --- a/go/apps/api/routes/v2_keys_get_key/412_test.go +++ b/go/apps/api/routes/v2_keys_get_key/412_test.go @@ -17,10 +17,11 @@ func TestPreconditionError(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - Logger: h.Logger, - DB: h.DB, - Keys: h.Keys, - Vault: h.Vault, + Logger: h.Logger, + DB: h.DB, + Keys: h.Keys, + Vault: h.Vault, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) @@ -87,11 +88,12 @@ func TestPreconditionError(t *testing.T) { }) route := &handler.Handler{ - Logger: h.Logger, - DB: h.DB, - Keys: h.Keys, - Auditlogs: h.Auditlogs, - Vault: h.Vault, + Logger: h.Logger, + DB: h.DB, + Keys: h.Keys, + Auditlogs: h.Auditlogs, + Vault: h.Vault, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) @@ -138,11 +140,12 @@ func TestPreconditionError(t *testing.T) { // Create route with nil vault routeNoVault := &handler.Handler{ - Logger: h.Logger, - DB: h.DB, - Keys: h.Keys, - Auditlogs: h.Auditlogs, - Vault: nil, // No vault + Logger: h.Logger, + DB: h.DB, + Keys: h.Keys, + Auditlogs: h.Auditlogs, + Vault: nil, // No vault + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(routeNoVault) diff --git a/go/apps/api/routes/v2_keys_get_key/500_test.go b/go/apps/api/routes/v2_keys_get_key/500_test.go index f65b3d1fd3..8ffcb3ca0e 100644 --- a/go/apps/api/routes/v2_keys_get_key/500_test.go +++ b/go/apps/api/routes/v2_keys_get_key/500_test.go @@ -16,11 +16,12 @@ import ( func TestInternalError(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - Logger: h.Logger, - DB: h.DB, - Keys: h.Keys, - Auditlogs: h.Auditlogs, - Vault: h.Vault, + Logger: h.Logger, + DB: h.DB, + Keys: h.Keys, + Auditlogs: h.Auditlogs, + Vault: h.Vault, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) rootKey := h.CreateRootKey(h.Resources().UserWorkspace.ID, "api.*.read_key") diff --git a/go/apps/api/routes/v2_keys_get_key/handler.go b/go/apps/api/routes/v2_keys_get_key/handler.go index 58a70ffa27..b66445247c 100644 --- a/go/apps/api/routes/v2_keys_get_key/handler.go +++ b/go/apps/api/routes/v2_keys_get_key/handler.go @@ -8,7 +8,9 @@ import ( "github.com/unkeyed/unkey/go/apps/api/openapi" vaultv1 "github.com/unkeyed/unkey/go/gen/proto/vault/v1" "github.com/unkeyed/unkey/go/internal/services/auditlogs" + "github.com/unkeyed/unkey/go/internal/services/caches" "github.com/unkeyed/unkey/go/internal/services/keys" + "github.com/unkeyed/unkey/go/pkg/cache" "github.com/unkeyed/unkey/go/pkg/codes" "github.com/unkeyed/unkey/go/pkg/db" "github.com/unkeyed/unkey/go/pkg/fault" @@ -26,11 +28,12 @@ type ( // Handler implements zen.Route interface for the v2 keys.getKey endpoint type Handler struct { - Logger logging.Logger - DB db.Database - Keys keys.KeyService - Auditlogs auditlogs.AuditLogService - Vault *vault.Service + Logger logging.Logger + DB db.Database + Keys keys.KeyService + Auditlogs auditlogs.AuditLogService + Vault *vault.Service + LiveKeyCache cache.Cache[string, db.FindLiveKeyByIDRow] } func (h *Handler) Method() string { @@ -57,7 +60,9 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } - 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) if err != nil { if db.IsNotFound(err) { return fault.Wrap( diff --git a/go/apps/api/routes/v2_keys_remove_permissions/200_test.go b/go/apps/api/routes/v2_keys_remove_permissions/200_test.go index 9789bddd7a..3aac48653f 100644 --- a/go/apps/api/routes/v2_keys_remove_permissions/200_test.go +++ b/go/apps/api/routes/v2_keys_remove_permissions/200_test.go @@ -21,11 +21,12 @@ func TestSuccess(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_remove_permissions/400_test.go b/go/apps/api/routes/v2_keys_remove_permissions/400_test.go index 29dd622d9e..44ce9e4029 100644 --- a/go/apps/api/routes/v2_keys_remove_permissions/400_test.go +++ b/go/apps/api/routes/v2_keys_remove_permissions/400_test.go @@ -21,11 +21,12 @@ func TestValidationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_remove_permissions/401_test.go b/go/apps/api/routes/v2_keys_remove_permissions/401_test.go index 2d04f97e4c..2aa16a5d30 100644 --- a/go/apps/api/routes/v2_keys_remove_permissions/401_test.go +++ b/go/apps/api/routes/v2_keys_remove_permissions/401_test.go @@ -16,11 +16,12 @@ func TestAuthenticationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_remove_permissions/403_test.go b/go/apps/api/routes/v2_keys_remove_permissions/403_test.go index af945304f6..d55f85128b 100644 --- a/go/apps/api/routes/v2_keys_remove_permissions/403_test.go +++ b/go/apps/api/routes/v2_keys_remove_permissions/403_test.go @@ -21,11 +21,12 @@ func TestAuthorizationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_remove_permissions/404_test.go b/go/apps/api/routes/v2_keys_remove_permissions/404_test.go index a94582d813..33dced309e 100644 --- a/go/apps/api/routes/v2_keys_remove_permissions/404_test.go +++ b/go/apps/api/routes/v2_keys_remove_permissions/404_test.go @@ -24,11 +24,12 @@ func TestNotFoundErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_remove_permissions/handler.go b/go/apps/api/routes/v2_keys_remove_permissions/handler.go index 4854eac038..572ed1b572 100644 --- a/go/apps/api/routes/v2_keys_remove_permissions/handler.go +++ b/go/apps/api/routes/v2_keys_remove_permissions/handler.go @@ -7,6 +7,7 @@ import ( "github.com/unkeyed/unkey/go/apps/api/openapi" "github.com/unkeyed/unkey/go/internal/services/auditlogs" + "github.com/unkeyed/unkey/go/internal/services/caches" "github.com/unkeyed/unkey/go/internal/services/keys" "github.com/unkeyed/unkey/go/pkg/auditlog" "github.com/unkeyed/unkey/go/pkg/cache" @@ -24,11 +25,12 @@ type Response = openapi.V2KeysRemovePermissionsResponseBody // Handler implements zen.Route interface for the v2 keys remove permissions endpoint type Handler struct { // Services as public fields - Logger logging.Logger - DB db.Database - Keys keys.KeyService - Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.CachedKeyData] + Logger logging.Logger + DB db.Database + Keys keys.KeyService + Auditlogs auditlogs.AuditLogService + KeyCache cache.Cache[string, db.CachedKeyData] + LiveKeyCache cache.Cache[string, db.FindLiveKeyByIDRow] } // Method returns the HTTP method this route responds to @@ -57,7 +59,9 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } - 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) if err != nil { if db.IsNotFound(err) { return fault.New("key not found", @@ -220,6 +224,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { } h.KeyCache.Remove(ctx, key.Hash) + h.LiveKeyCache.Remove(ctx, key.ID) } responseData := make(openapi.V2KeysRemovePermissionsResponseData, 0) diff --git a/go/apps/api/routes/v2_keys_remove_roles/200_test.go b/go/apps/api/routes/v2_keys_remove_roles/200_test.go index c6871a6d45..9135f260ac 100644 --- a/go/apps/api/routes/v2_keys_remove_roles/200_test.go +++ b/go/apps/api/routes/v2_keys_remove_roles/200_test.go @@ -21,11 +21,12 @@ func TestSuccess(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_remove_roles/400_test.go b/go/apps/api/routes/v2_keys_remove_roles/400_test.go index 19319d64cd..487d525539 100644 --- a/go/apps/api/routes/v2_keys_remove_roles/400_test.go +++ b/go/apps/api/routes/v2_keys_remove_roles/400_test.go @@ -21,11 +21,12 @@ func TestValidationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_remove_roles/401_test.go b/go/apps/api/routes/v2_keys_remove_roles/401_test.go index 60ecd4ce21..fc9ed0b425 100644 --- a/go/apps/api/routes/v2_keys_remove_roles/401_test.go +++ b/go/apps/api/routes/v2_keys_remove_roles/401_test.go @@ -21,11 +21,12 @@ func TestAuthenticationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_remove_roles/403_test.go b/go/apps/api/routes/v2_keys_remove_roles/403_test.go index 0482246c56..67d06b7293 100644 --- a/go/apps/api/routes/v2_keys_remove_roles/403_test.go +++ b/go/apps/api/routes/v2_keys_remove_roles/403_test.go @@ -21,11 +21,12 @@ func TestAuthorizationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_remove_roles/404_test.go b/go/apps/api/routes/v2_keys_remove_roles/404_test.go index 4b68607b83..42739cbdd8 100644 --- a/go/apps/api/routes/v2_keys_remove_roles/404_test.go +++ b/go/apps/api/routes/v2_keys_remove_roles/404_test.go @@ -21,11 +21,12 @@ func TestNotFoundErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_remove_roles/handler.go b/go/apps/api/routes/v2_keys_remove_roles/handler.go index 9c6e1b102b..3479c9b632 100644 --- a/go/apps/api/routes/v2_keys_remove_roles/handler.go +++ b/go/apps/api/routes/v2_keys_remove_roles/handler.go @@ -7,6 +7,7 @@ import ( "github.com/unkeyed/unkey/go/apps/api/openapi" "github.com/unkeyed/unkey/go/internal/services/auditlogs" + "github.com/unkeyed/unkey/go/internal/services/caches" "github.com/unkeyed/unkey/go/internal/services/keys" "github.com/unkeyed/unkey/go/pkg/auditlog" "github.com/unkeyed/unkey/go/pkg/cache" @@ -25,11 +26,12 @@ type ( // Handler implements zen.Route interface for the v2 keys remove roles endpoint type Handler struct { - Logger logging.Logger - DB db.Database - Keys keys.KeyService - Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.CachedKeyData] + Logger logging.Logger + DB db.Database + Keys keys.KeyService + Auditlogs auditlogs.AuditLogService + KeyCache cache.Cache[string, db.CachedKeyData] + LiveKeyCache cache.Cache[string, db.FindLiveKeyByIDRow] } // Method returns the HTTP method this route responds to @@ -58,7 +60,9 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } - 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) if err != nil { if db.IsNotFound(err) { return fault.New("key not found", @@ -206,6 +210,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { } h.KeyCache.Remove(ctx, key.Hash) + h.LiveKeyCache.Remove(ctx, key.ID) } responseData := make(openapi.V2KeysRemoveRolesResponseData, 0) diff --git a/go/apps/api/routes/v2_keys_reroll_key/200_test.go b/go/apps/api/routes/v2_keys_reroll_key/200_test.go index 18709f1235..2af67e4bcf 100644 --- a/go/apps/api/routes/v2_keys_reroll_key/200_test.go +++ b/go/apps/api/routes/v2_keys_reroll_key/200_test.go @@ -21,11 +21,13 @@ func TestRerollKeySuccess(t *testing.T) { ctx := t.Context() route := &handler.Handler{ - Logger: h.Logger, - DB: h.DB, - Keys: h.Keys, - Auditlogs: h.Auditlogs, - Vault: h.Vault, + Logger: h.Logger, + DB: h.DB, + Keys: h.Keys, + Auditlogs: h.Auditlogs, + Vault: h.Vault, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_reroll_key/400_test.go b/go/apps/api/routes/v2_keys_reroll_key/400_test.go index eb7b407219..93829b6e34 100644 --- a/go/apps/api/routes/v2_keys_reroll_key/400_test.go +++ b/go/apps/api/routes/v2_keys_reroll_key/400_test.go @@ -17,11 +17,13 @@ func TestRerollKeyBadRequest(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - Vault: h.Vault, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + Vault: h.Vault, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_reroll_key/401_test.go b/go/apps/api/routes/v2_keys_reroll_key/401_test.go index 2ac5ff6178..e946702c67 100644 --- a/go/apps/api/routes/v2_keys_reroll_key/401_test.go +++ b/go/apps/api/routes/v2_keys_reroll_key/401_test.go @@ -18,11 +18,13 @@ func TestRerollKeyUnauthorized(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - Vault: h.Vault, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + Vault: h.Vault, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_reroll_key/403_test.go b/go/apps/api/routes/v2_keys_reroll_key/403_test.go index a0686712ca..0bad6b9137 100644 --- a/go/apps/api/routes/v2_keys_reroll_key/403_test.go +++ b/go/apps/api/routes/v2_keys_reroll_key/403_test.go @@ -16,11 +16,13 @@ func TestRerollKeyForbidden(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - Vault: h.Vault, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + Vault: h.Vault, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_reroll_key/404_test.go b/go/apps/api/routes/v2_keys_reroll_key/404_test.go index 100e3a3b0e..452b4ace24 100644 --- a/go/apps/api/routes/v2_keys_reroll_key/404_test.go +++ b/go/apps/api/routes/v2_keys_reroll_key/404_test.go @@ -16,11 +16,13 @@ func TestRerollKeyNotFound(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - Vault: h.Vault, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + Vault: h.Vault, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_reroll_key/handler.go b/go/apps/api/routes/v2_keys_reroll_key/handler.go index 85d5019b11..f00cf3c9fd 100644 --- a/go/apps/api/routes/v2_keys_reroll_key/handler.go +++ b/go/apps/api/routes/v2_keys_reroll_key/handler.go @@ -11,9 +11,11 @@ import ( "github.com/unkeyed/unkey/go/apps/api/openapi" vaultv1 "github.com/unkeyed/unkey/go/gen/proto/vault/v1" "github.com/unkeyed/unkey/go/internal/services/auditlogs" + "github.com/unkeyed/unkey/go/internal/services/caches" "github.com/unkeyed/unkey/go/internal/services/keys" "github.com/unkeyed/unkey/go/pkg/auditlog" + "github.com/unkeyed/unkey/go/pkg/cache" "github.com/unkeyed/unkey/go/pkg/codes" "github.com/unkeyed/unkey/go/pkg/db" "github.com/unkeyed/unkey/go/pkg/fault" @@ -30,11 +32,13 @@ type ( ) type Handler struct { - Logger logging.Logger - DB db.Database - Keys keys.KeyService - Auditlogs auditlogs.AuditLogService - Vault *vault.Service + Logger logging.Logger + DB db.Database + Keys keys.KeyService + Auditlogs auditlogs.AuditLogService + Vault *vault.Service + KeyCache cache.Cache[string, db.CachedKeyData] + LiveKeyCache cache.Cache[string, db.FindLiveKeyByIDRow] } // Method returns the HTTP method this route responds to @@ -63,7 +67,9 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } - 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) if err != nil { if db.IsNotFound(err) { return fault.New("key not found", @@ -361,6 +367,9 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } + h.KeyCache.Remove(ctx, key.Hash) + h.LiveKeyCache.Remove(ctx, key.ID) + return s.JSON(http.StatusOK, Response{ Meta: openapi.Meta{ RequestId: s.RequestID(), diff --git a/go/apps/api/routes/v2_keys_set_permissions/200_test.go b/go/apps/api/routes/v2_keys_set_permissions/200_test.go index 4b2414d3f3..ba0407b404 100644 --- a/go/apps/api/routes/v2_keys_set_permissions/200_test.go +++ b/go/apps/api/routes/v2_keys_set_permissions/200_test.go @@ -21,11 +21,12 @@ func TestSuccess(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_set_permissions/400_test.go b/go/apps/api/routes/v2_keys_set_permissions/400_test.go index 16db4312be..b1cffd6157 100644 --- a/go/apps/api/routes/v2_keys_set_permissions/400_test.go +++ b/go/apps/api/routes/v2_keys_set_permissions/400_test.go @@ -16,11 +16,12 @@ func TestBadRequest(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_set_permissions/401_test.go b/go/apps/api/routes/v2_keys_set_permissions/401_test.go index 8fad806f2c..7932cc722a 100644 --- a/go/apps/api/routes/v2_keys_set_permissions/401_test.go +++ b/go/apps/api/routes/v2_keys_set_permissions/401_test.go @@ -23,11 +23,12 @@ func TestAuthenticationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_set_permissions/403_test.go b/go/apps/api/routes/v2_keys_set_permissions/403_test.go index f0dd4879bb..c6c6f1cb59 100644 --- a/go/apps/api/routes/v2_keys_set_permissions/403_test.go +++ b/go/apps/api/routes/v2_keys_set_permissions/403_test.go @@ -21,11 +21,12 @@ func TestForbidden(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_set_permissions/404_test.go b/go/apps/api/routes/v2_keys_set_permissions/404_test.go index c8725b5ecf..eab6477c6b 100644 --- a/go/apps/api/routes/v2_keys_set_permissions/404_test.go +++ b/go/apps/api/routes/v2_keys_set_permissions/404_test.go @@ -21,11 +21,12 @@ func TestNotFound(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_set_permissions/handler.go b/go/apps/api/routes/v2_keys_set_permissions/handler.go index 6098effffe..a818e97b13 100644 --- a/go/apps/api/routes/v2_keys_set_permissions/handler.go +++ b/go/apps/api/routes/v2_keys_set_permissions/handler.go @@ -9,6 +9,7 @@ import ( "github.com/unkeyed/unkey/go/apps/api/openapi" "github.com/unkeyed/unkey/go/internal/services/auditlogs" + "github.com/unkeyed/unkey/go/internal/services/caches" "github.com/unkeyed/unkey/go/internal/services/keys" "github.com/unkeyed/unkey/go/pkg/auditlog" "github.com/unkeyed/unkey/go/pkg/cache" @@ -29,11 +30,12 @@ type ( // Handler implements zen.Route interface for the v2 keys set permissions endpoint type Handler struct { - Logger logging.Logger - DB db.Database - Keys keys.KeyService - Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.CachedKeyData] + Logger logging.Logger + DB db.Database + Keys keys.KeyService + Auditlogs auditlogs.AuditLogService + KeyCache cache.Cache[string, db.CachedKeyData] + LiveKeyCache cache.Cache[string, db.FindLiveKeyByIDRow] } // Method returns the HTTP method this route responds to @@ -62,7 +64,9 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } - 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) if err != nil { if db.IsNotFound(err) { return fault.New("key not found", @@ -361,6 +365,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { } h.KeyCache.Remove(ctx, key.Hash) + h.LiveKeyCache.Remove(ctx, key.ID) responseData := make(openapi.V2KeysSetPermissionsResponseData, 0) for _, permission := range permissionsToSet { diff --git a/go/apps/api/routes/v2_keys_set_roles/200_test.go b/go/apps/api/routes/v2_keys_set_roles/200_test.go index 38e2ccf1c5..8ba3b47f67 100644 --- a/go/apps/api/routes/v2_keys_set_roles/200_test.go +++ b/go/apps/api/routes/v2_keys_set_roles/200_test.go @@ -21,11 +21,12 @@ func TestSuccess(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_set_roles/400_test.go b/go/apps/api/routes/v2_keys_set_roles/400_test.go index 30d15700ba..e8bb2e41cc 100644 --- a/go/apps/api/routes/v2_keys_set_roles/400_test.go +++ b/go/apps/api/routes/v2_keys_set_roles/400_test.go @@ -15,11 +15,12 @@ func TestValidationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_set_roles/401_test.go b/go/apps/api/routes/v2_keys_set_roles/401_test.go index 1d714a7255..e56f96e773 100644 --- a/go/apps/api/routes/v2_keys_set_roles/401_test.go +++ b/go/apps/api/routes/v2_keys_set_roles/401_test.go @@ -14,11 +14,12 @@ func TestAuthenticationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_set_roles/403_test.go b/go/apps/api/routes/v2_keys_set_roles/403_test.go index 579a8a8280..df091b9dc3 100644 --- a/go/apps/api/routes/v2_keys_set_roles/403_test.go +++ b/go/apps/api/routes/v2_keys_set_roles/403_test.go @@ -16,11 +16,12 @@ func TestAuthorizationErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_set_roles/404_test.go b/go/apps/api/routes/v2_keys_set_roles/404_test.go index 2156e1c279..9a76f4f4da 100644 --- a/go/apps/api/routes/v2_keys_set_roles/404_test.go +++ b/go/apps/api/routes/v2_keys_set_roles/404_test.go @@ -23,11 +23,12 @@ func TestNotFoundErrors(t *testing.T) { h := testutil.NewHarness(t) route := &handler.Handler{ - DB: h.DB, - Keys: h.Keys, - Logger: h.Logger, - Auditlogs: h.Auditlogs, - KeyCache: h.Caches.VerificationKeyByHash, + DB: h.DB, + Keys: h.Keys, + Logger: h.Logger, + Auditlogs: h.Auditlogs, + KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, } h.Register(route) diff --git a/go/apps/api/routes/v2_keys_set_roles/handler.go b/go/apps/api/routes/v2_keys_set_roles/handler.go index 6177d0a178..8587d635f4 100644 --- a/go/apps/api/routes/v2_keys_set_roles/handler.go +++ b/go/apps/api/routes/v2_keys_set_roles/handler.go @@ -8,6 +8,7 @@ import ( "github.com/unkeyed/unkey/go/apps/api/openapi" "github.com/unkeyed/unkey/go/internal/services/auditlogs" + "github.com/unkeyed/unkey/go/internal/services/caches" "github.com/unkeyed/unkey/go/internal/services/keys" "github.com/unkeyed/unkey/go/pkg/auditlog" "github.com/unkeyed/unkey/go/pkg/cache" @@ -26,11 +27,12 @@ type ( // Handler implements zen.Route interface for the v2 keys set roles endpoint type Handler struct { - Logger logging.Logger - DB db.Database - Keys keys.KeyService - Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.CachedKeyData] + Logger logging.Logger + DB db.Database + Keys keys.KeyService + Auditlogs auditlogs.AuditLogService + KeyCache cache.Cache[string, db.CachedKeyData] + LiveKeyCache cache.Cache[string, db.FindLiveKeyByIDRow] } // Method returns the HTTP method this route responds to @@ -58,7 +60,9 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } - 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) if err != nil { if db.IsNotFound(err) { return fault.New("key not found", @@ -265,6 +269,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { } h.KeyCache.Remove(ctx, key.Hash) + h.LiveKeyCache.Remove(ctx, key.ID) responseData := make(openapi.V2KeysSetRolesResponseData, 0) for _, role := range foundRoles { diff --git a/go/apps/api/routes/v2_keys_update_credits/200_test.go b/go/apps/api/routes/v2_keys_update_credits/200_test.go index 7acfd62f93..9c9526cb35 100644 --- a/go/apps/api/routes/v2_keys_update_credits/200_test.go +++ b/go/apps/api/routes/v2_keys_update_credits/200_test.go @@ -28,6 +28,7 @@ func TestKeyUpdateCreditsSuccess(t *testing.T) { Keys: h.Keys, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } diff --git a/go/apps/api/routes/v2_keys_update_credits/400_test.go b/go/apps/api/routes/v2_keys_update_credits/400_test.go index 8ad7d1eab8..4e1be477d0 100644 --- a/go/apps/api/routes/v2_keys_update_credits/400_test.go +++ b/go/apps/api/routes/v2_keys_update_credits/400_test.go @@ -22,6 +22,7 @@ func TestKeyUpdateCreditsBadRequest(t *testing.T) { Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } diff --git a/go/apps/api/routes/v2_keys_update_credits/401_test.go b/go/apps/api/routes/v2_keys_update_credits/401_test.go index fa92247569..0580ef4529 100644 --- a/go/apps/api/routes/v2_keys_update_credits/401_test.go +++ b/go/apps/api/routes/v2_keys_update_credits/401_test.go @@ -22,6 +22,7 @@ func TestKeyUpdateCreditsUnauthorized(t *testing.T) { Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } diff --git a/go/apps/api/routes/v2_keys_update_credits/403_test.go b/go/apps/api/routes/v2_keys_update_credits/403_test.go index 2aa9fad85d..d446b3d233 100644 --- a/go/apps/api/routes/v2_keys_update_credits/403_test.go +++ b/go/apps/api/routes/v2_keys_update_credits/403_test.go @@ -23,6 +23,7 @@ func TestKeyUpdateCreditsForbidden(t *testing.T) { Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } diff --git a/go/apps/api/routes/v2_keys_update_credits/404_test.go b/go/apps/api/routes/v2_keys_update_credits/404_test.go index 7167106998..452e2cc387 100644 --- a/go/apps/api/routes/v2_keys_update_credits/404_test.go +++ b/go/apps/api/routes/v2_keys_update_credits/404_test.go @@ -21,6 +21,7 @@ func TestUpdateKeyCreditsNotFound(t *testing.T) { Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } diff --git a/go/apps/api/routes/v2_keys_update_credits/handler.go b/go/apps/api/routes/v2_keys_update_credits/handler.go index d4c7a7516b..f1d25af6ad 100644 --- a/go/apps/api/routes/v2_keys_update_credits/handler.go +++ b/go/apps/api/routes/v2_keys_update_credits/handler.go @@ -9,6 +9,7 @@ import ( "github.com/oapi-codegen/nullable" "github.com/unkeyed/unkey/go/apps/api/openapi" "github.com/unkeyed/unkey/go/internal/services/auditlogs" + "github.com/unkeyed/unkey/go/internal/services/caches" "github.com/unkeyed/unkey/go/internal/services/keys" "github.com/unkeyed/unkey/go/internal/services/usagelimiter" "github.com/unkeyed/unkey/go/pkg/auditlog" @@ -31,6 +32,7 @@ type Handler struct { Keys keys.KeyService Auditlogs auditlogs.AuditLogService KeyCache cache.Cache[string, db.CachedKeyData] + LiveKeyCache cache.Cache[string, db.FindLiveKeyByIDRow] UsageLimiter usagelimiter.Service } @@ -60,7 +62,9 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } - 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) if err != nil { if db.IsNotFound(err) { return fault.Wrap( @@ -266,6 +270,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { } h.KeyCache.Remove(ctx, key.Hash) + h.LiveKeyCache.Remove(ctx, key.ID) if err := h.UsageLimiter.Invalidate(ctx, key.ID); err != nil { h.Logger.Error("Failed to invalidate usage limit", "error", err.Error(), diff --git a/go/apps/api/routes/v2_keys_update_key/200_test.go b/go/apps/api/routes/v2_keys_update_key/200_test.go index 6c00af99fe..3129e1f982 100644 --- a/go/apps/api/routes/v2_keys_update_key/200_test.go +++ b/go/apps/api/routes/v2_keys_update_key/200_test.go @@ -32,6 +32,7 @@ func TestUpdateKeySuccess(t *testing.T) { Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } @@ -136,6 +137,7 @@ func TestUpdateKeyUpdateAllFields(t *testing.T) { Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } @@ -213,6 +215,7 @@ func TestKeyUpdateCreditsInvalidatesCache(t *testing.T) { Keys: h.Keys, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } diff --git a/go/apps/api/routes/v2_keys_update_key/400_test.go b/go/apps/api/routes/v2_keys_update_key/400_test.go index 01eb894ce0..4d9d10400f 100644 --- a/go/apps/api/routes/v2_keys_update_key/400_test.go +++ b/go/apps/api/routes/v2_keys_update_key/400_test.go @@ -25,6 +25,7 @@ func TestUpdateKeyInvalidRefillConfig(t *testing.T) { Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } diff --git a/go/apps/api/routes/v2_keys_update_key/401_test.go b/go/apps/api/routes/v2_keys_update_key/401_test.go index 53b3d84b85..888f3d40da 100644 --- a/go/apps/api/routes/v2_keys_update_key/401_test.go +++ b/go/apps/api/routes/v2_keys_update_key/401_test.go @@ -23,6 +23,7 @@ func TestUpdateKeyUnauthorized(t *testing.T) { Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } diff --git a/go/apps/api/routes/v2_keys_update_key/403_test.go b/go/apps/api/routes/v2_keys_update_key/403_test.go index 1cc00337d9..af00e8602b 100644 --- a/go/apps/api/routes/v2_keys_update_key/403_test.go +++ b/go/apps/api/routes/v2_keys_update_key/403_test.go @@ -41,6 +41,7 @@ func TestUpdateKeyCorrectPermissions(t *testing.T) { Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } @@ -94,6 +95,7 @@ func TestUpdateKeyInsufficientPermissions(t *testing.T) { Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } @@ -140,6 +142,7 @@ func TestUpdateKeyCrossWorkspaceIsolation(t *testing.T) { Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } diff --git a/go/apps/api/routes/v2_keys_update_key/404_test.go b/go/apps/api/routes/v2_keys_update_key/404_test.go index 65ecbaea04..48b21665e4 100644 --- a/go/apps/api/routes/v2_keys_update_key/404_test.go +++ b/go/apps/api/routes/v2_keys_update_key/404_test.go @@ -24,6 +24,7 @@ func TestUpdateKeyNotFound(t *testing.T) { Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } diff --git a/go/apps/api/routes/v2_keys_update_key/handler.go b/go/apps/api/routes/v2_keys_update_key/handler.go index dd58ab8393..8acb3adb7e 100644 --- a/go/apps/api/routes/v2_keys_update_key/handler.go +++ b/go/apps/api/routes/v2_keys_update_key/handler.go @@ -10,6 +10,7 @@ import ( "github.com/unkeyed/unkey/go/apps/api/openapi" "github.com/unkeyed/unkey/go/internal/services/auditlogs" + "github.com/unkeyed/unkey/go/internal/services/caches" "github.com/unkeyed/unkey/go/internal/services/keys" "github.com/unkeyed/unkey/go/internal/services/usagelimiter" @@ -36,6 +37,7 @@ type Handler struct { Keys keys.KeyService Auditlogs auditlogs.AuditLogService KeyCache cache.Cache[string, db.CachedKeyData] + LiveKeyCache cache.Cache[string, db.FindLiveKeyByIDRow] UsageLimiter usagelimiter.Service } @@ -64,7 +66,9 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { return err } - 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) if err != nil { if db.IsNotFound(err) { return fault.Wrap( @@ -633,6 +637,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { } h.KeyCache.Remove(ctx, key.Hash) + h.LiveKeyCache.Remove(ctx, key.ID) if req.Credits.IsSpecified() { if err := h.UsageLimiter.Invalidate(ctx, key.ID); err != nil { h.Logger.Error("Failed to invalidate usage limit", diff --git a/go/apps/api/routes/v2_keys_update_key/three_state_test.go b/go/apps/api/routes/v2_keys_update_key/three_state_test.go index f699538df0..0160fb8663 100644 --- a/go/apps/api/routes/v2_keys_update_key/three_state_test.go +++ b/go/apps/api/routes/v2_keys_update_key/three_state_test.go @@ -33,6 +33,7 @@ func TestThreeStateUpdateLogic(t *testing.T) { Logger: h.Logger, Auditlogs: h.Auditlogs, KeyCache: h.Caches.VerificationKeyByHash, + LiveKeyCache: h.Caches.LiveKeyByID, UsageLimiter: h.UsageLimiter, } diff --git a/go/internal/services/caches/caches.go b/go/internal/services/caches/caches.go index 85c7d4ee5e..857bdb7fc9 100644 --- a/go/internal/services/caches/caches.go +++ b/go/internal/services/caches/caches.go @@ -43,6 +43,10 @@ type Caches struct { // Keys are string (api_id) and values are db.FindKeyAuthsByIdsRow (has both KeyAuthID and ApiID). ApiToKeyAuthRow cache.Cache[cache.ScopedKey, db.FindKeyAuthsByIdsRow] + // LiveKeyByID caches live key lookups by ID. + // Keys are string (ID) and values are db.FindLiveKeyByIDRow. + LiveKeyByID cache.Cache[string, db.FindLiveKeyByIDRow] + // dispatcher handles routing of invalidation events to all caches in this process. // This is not exported as it's an internal implementation detail. dispatcher *clustering.InvalidationDispatcher @@ -226,7 +230,7 @@ func New(config Config) (Caches, error) { return Caches{}, err } - // Create API cache (uses ScopedKey) + // Create live API cache (uses ScopedKey) liveApiByID, err := createCache( config, dispatcher, @@ -245,6 +249,26 @@ func New(config Config) (Caches, error) { return Caches{}, err } + // Create live key cache (uses string keys) + liveKeyByID, err := createCache( + config, + dispatcher, + cache.Config[string, db.FindLiveKeyByIDRow]{ + Fresh: 10 * time.Second, + Stale: 10 * time.Minute, + Logger: config.Logger, + MaxSize: 1_000_000, + Resource: "live_key_by_id", + Clock: config.Clock, + }, + nil, + nil, + ) + if err != nil { + return Caches{}, err + } + + // Create ClickHouse settings cache clickhouseSetting, err := createCache( config, dispatcher, @@ -252,7 +276,7 @@ func New(config Config) (Caches, error) { Fresh: time.Minute, Stale: 24 * time.Hour, Logger: config.Logger, - MaxSize: 1_000_000, + MaxSize: 10_000, Resource: "clickhouse_setting", Clock: config.Clock, }, @@ -263,15 +287,15 @@ func New(config Config) (Caches, error) { return Caches{}, err } - // Create key_auth_id -> api row cache + // Create KeyAuth to API row cache keyAuthToApiRow, err := createCache( config, dispatcher, cache.Config[cache.ScopedKey, db.FindKeyAuthsByKeyAuthIdsRow]{ - Fresh: 10 * time.Minute, + Fresh: time.Minute, Stale: 24 * time.Hour, Logger: config.Logger, - MaxSize: 1_000_000, + MaxSize: 100_000, Resource: "key_auth_to_api_row", Clock: config.Clock, }, @@ -282,15 +306,15 @@ func New(config Config) (Caches, error) { return Caches{}, err } - // Create api_id -> key_auth row cache + // Create API to KeyAuth row cache apiToKeyAuthRow, err := createCache( config, dispatcher, cache.Config[cache.ScopedKey, db.FindKeyAuthsByIdsRow]{ - Fresh: 10 * time.Minute, + Fresh: time.Minute, Stale: 24 * time.Hour, Logger: config.Logger, - MaxSize: 1_000_000, + MaxSize: 100_000, Resource: "api_to_key_auth_row", Clock: config.Clock, }, @@ -305,6 +329,7 @@ func New(config Config) (Caches, error) { RatelimitNamespace: middleware.WithTracing(ratelimitNamespace), LiveApiByID: middleware.WithTracing(liveApiByID), VerificationKeyByHash: middleware.WithTracing(verificationKeyByHash), + LiveKeyByID: middleware.WithTracing(liveKeyByID), ClickhouseSetting: middleware.WithTracing(clickhouseSetting), KeyAuthToApiRow: middleware.WithTracing(keyAuthToApiRow), ApiToKeyAuthRow: middleware.WithTracing(apiToKeyAuthRow),