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

go/apps/api/routes/register.go

@@ -425,6 +425,7 @@ func Register(srv *zen.Server, svc *Services) {
 			Keys:         svc.Keys,
 			Auditlogs:    svc.Auditlogs,
 			KeyCache:     svc.Caches.VerificationKeyByHash,
+			LiveKeyCache: svc.Caches.LiveKeyByID,
 			UsageLimiter: svc.UsageLimiter,
 		},
 	)
@@ -462,6 +463,7 @@ func Register(srv *zen.Server, svc *Services) {
 			Keys:         svc.Keys,
 			Auditlogs:    svc.Auditlogs,
 			KeyCache:     svc.Caches.VerificationKeyByHash,
+			LiveKeyCache: svc.Caches.LiveKeyByID,
 			UsageLimiter: svc.UsageLimiter,
 		},
 	)
@@ -470,11 +472,12 @@ func Register(srv *zen.Server, svc *Services) {
 	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,
 		},
 	)
 

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)

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)

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)

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)

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)

go/apps/api/routes/v2_keys_set_roles/handler.go

@@ -26,11 +26,12 @@ type Response = openapi.V2KeysSetRolesResponseBody
 
 // 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 +59,17 @@ 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)
+	}, func(err error) cache.Op {
+		if err == nil {
+			return cache.WriteValue
+		}
+		if db.IsNotFound(err) {
+			return cache.WriteNull
+		}
+		return cache.Noop
+	})
 	if err != nil {
 		if db.IsNotFound(err) {
 			return fault.New("key not found",
@@ -265,6 +276,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 {

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,
 	}
 

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,
 	}
 

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,
 	}
 

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,
 	}
 

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,
 	}
 

go/apps/api/routes/v2_keys_update_credits/handler.go

@@ -32,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
 }
 
@@ -61,7 +62,17 @@ 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)
+	}, func(err error) cache.Op {
+		if err == nil {
+			return cache.WriteValue
+		}
+		if db.IsNotFound(err) {
+			return cache.WriteNull
+		}
+		return cache.Noop
+	})
 	if err != nil {
 		if db.IsNotFound(err) {
 			return fault.Wrap(
@@ -267,6 +278,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(),

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,
 	}
 
@@ -135,6 +136,7 @@ func TestUpdateKeyUpdateAllFields(t *testing.T) {
 		Logger:       h.Logger,
 		Auditlogs:    h.Auditlogs,
 		KeyCache:     h.Caches.VerificationKeyByHash,
+		LiveKeyCache: h.Caches.LiveKeyByID,
 		UsageLimiter: h.UsageLimiter,
 	}
 

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,
 	}
 

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,
 	}
 

go/apps/api/routes/v2_keys_update_key/403_test.go

@@ -93,6 +93,7 @@ func TestUpdateKeyInsufficientPermissions(t *testing.T) {
 		Logger:       h.Logger,
 		Auditlogs:    h.Auditlogs,
 		KeyCache:     h.Caches.VerificationKeyByHash,
+		LiveKeyCache: h.Caches.LiveKeyByID,
 		UsageLimiter: h.UsageLimiter,
 	}
 
@@ -139,6 +140,7 @@ func TestUpdateKeyCrossWorkspaceIsolation(t *testing.T) {
 		Logger:       h.Logger,
 		Auditlogs:    h.Auditlogs,
 		KeyCache:     h.Caches.VerificationKeyByHash,
+		LiveKeyCache: h.Caches.LiveKeyByID,
 		UsageLimiter: h.UsageLimiter,
 	}
 

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,
 	}
 

go/apps/api/routes/v2_keys_update_key/handler.go

@@ -36,6 +36,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 +65,17 @@ 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)
+	}, func(err error) cache.Op {
+		if err == nil {
+			return cache.WriteValue
+		}
+		if db.IsNotFound(err) {
+			return cache.WriteNull
+		}
+		return cache.Noop
+	})
 	if err != nil {
 		if db.IsNotFound(err) {
 			return fault.Wrap(
@@ -627,6 +638,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",

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,
 	}
 

go/internal/services/caches/caches.go

@@ -24,6 +24,10 @@ type Caches struct {
 	// LiveApiByID caches live API lookups by ID.
 	// Keys are string (ID) and values are db.FindLiveApiByIDRow.
 	LiveApiByID cache.Cache[string, db.FindLiveApiByIDRow]
+
+	// LiveKeyByID caches live key lookups by ID.
+	// Keys are string (ID) and values are db.FindLiveKeyByIDRow.
+	LiveKeyByID cache.Cache[string, db.FindLiveKeyByIDRow]
 }
 
 // Config defines the configuration options for initializing caches.
@@ -101,9 +105,22 @@ func New(config Config) (Caches, error) {
 		return Caches{}, err
 	}
 
+	liveKeyByID, err := cache.New(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,
+	})
+	if err != nil {
+		return Caches{}, err
+	}
+
 	return Caches{
 		RatelimitNamespace:    middleware.WithTracing(ratelimitNamespace),
 		LiveApiByID:           middleware.WithTracing(liveApiByID),
 		VerificationKeyByHash: middleware.WithTracing(verificationKeyByHash),
+		LiveKeyByID:           middleware.WithTracing(liveKeyByID),
 	}, nil
 }