fix: add retries to getKey operation

go/apps/api/routes/v2_ratelimit_limit/handler.go

@@ -74,11 +74,14 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error {
 	namespace, hit, err := h.RatelimitNamespaceCache.SWR(ctx,
 		cache.ScopedKey{WorkspaceID: auth.AuthorizedWorkspaceID, Key: namespaceKey},
 		func(ctx context.Context) (db.FindRatelimitNamespace, error) {
-			response, err := db.Query.FindRatelimitNamespace(ctx, h.DB.RO(), db.FindRatelimitNamespaceParams{
-				WorkspaceID: auth.AuthorizedWorkspaceID,
-				Namespace:   namespaceKey,
-			})
 			result := db.FindRatelimitNamespace{} // nolint:exhaustruct
+
+			response, err := db.WithRetry(func() (db.FindRatelimitNamespaceRow, error) {
+				return db.Query.FindRatelimitNamespace(ctx, h.DB.RO(), db.FindRatelimitNamespaceParams{
+					WorkspaceID: auth.AuthorizedWorkspaceID,
+					Namespace:   namespaceKey,
+				})
+			})
 			if err != nil {
 				return result, err
 			}
@@ -186,6 +189,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error {
 		Cost:       cost,
 		Time:       time.Time{},
 	}
+
 	if h.TestMode {
 		header := s.Request().Header.Get("X-Test-Time")
 		if header != "" {
@@ -216,6 +220,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error {
 			Passed:      result.Success,
 		})
 	}
+
 	res := Response{
 		Meta: openapi.Meta{
 			RequestId: s.RequestID(),

go/internal/services/caches/op.go

@@ -1,10 +1,8 @@
 package caches
 
 import (
-	"database/sql"
-	"errors"
-
 	"github.com/unkeyed/unkey/go/pkg/cache"
+	"github.com/unkeyed/unkey/go/pkg/db"
 )
 
 // DefaultFindFirstOp returns the appropriate cache operation based on the sql error
@@ -14,7 +12,7 @@ func DefaultFindFirstOp(err error) cache.Op {
 		return cache.WriteValue
 	}
 
-	if errors.Is(err, sql.ErrNoRows) {
+	if db.IsNotFound(err) {
 		// the response is empty, we need to store that the row does not exist
 		return cache.WriteNull
 	}

go/internal/services/keys/get.go

@@ -68,8 +68,12 @@ func (s *service) Get(ctx context.Context, sess *zen.Session, rawKey string) (*K
 
 	h := hash.Sha256(rawKey)
 	key, hit, err := s.keyCache.SWR(ctx, h, func(ctx context.Context) (db.FindKeyForVerificationRow, error) {
-		return db.Query.FindKeyForVerification(ctx, s.db.RO(), h)
+		// Use database retry with exponential backoff, skipping non-transient errors
+		return db.WithRetry(func() (db.FindKeyForVerificationRow, error) {
+			return db.Query.FindKeyForVerification(ctx, s.db.RO(), h)
+		})
 	}, caches.DefaultFindFirstOp)
+
 	if err != nil {
 		if db.IsNotFound(err) {
 			// nolint:exhaustruct

go/internal/services/usagelimiter/limit.go

@@ -13,7 +13,9 @@ func (s *service) Limit(ctx context.Context, req UsageRequest) (UsageResponse, e
 	ctx, span := tracing.Start(ctx, "usagelimiter.Limit")
 	defer span.End()
 
-	limit, err := db.Query.FindKeyCredits(ctx, s.db.RW(), req.KeyId)
+	limit, err := db.WithRetry(func() (sql.NullInt32, error) {
+		return db.Query.FindKeyCredits(ctx, s.db.RO(), req.KeyId)
+	})
 	if err != nil {
 		if db.IsNotFound(err) {
 			return UsageResponse{Valid: false, Remaining: 0}, nil

go/internal/services/usagelimiter/redis.go

@@ -232,7 +232,10 @@ func (s *counterService) initializeFromDatabase(ctx context.Context, req UsageRe
 	ctx, span := tracing.Start(ctx, "usagelimiter.counter.initializeFromDatabase")
 	defer span.End()
 
-	limit, err := db.Query.FindKeyCredits(ctx, s.db.RO(), req.KeyId)
+	limit, err := db.WithRetry(func() (sql.NullInt32, error) {
+		return db.Query.FindKeyCredits(ctx, s.db.RO(), req.KeyId)
+	})
+
 	if err != nil {
 		if db.IsNotFound(err) {
 			return UsageResponse{Valid: false, Remaining: 0}, nil

go/pkg/db/retry.go

@@ -0,0 +1,67 @@
+package db
+
+import (
+	"time"
+
+	"github.com/unkeyed/unkey/go/pkg/retry"
+)
+
+const (
+	// DefaultBackoff is the base duration for exponential backoff in database retries
+	DefaultBackoff = 50 * time.Millisecond
+)
+
+// WithRetry executes a database operation with optimized retry configuration.
+// It retries transient errors with exponential backoff but skips non-retryable errors
+// like "not found" or "duplicate key" to avoid unnecessary delays.
+//
+// Configuration:
+//   - 3 attempts maximum
+//   - Exponential backoff: 50ms, 100ms, 200ms
+//   - Skips retries for "not found" and "duplicate key" errors
+//
+// Usage:
+//
+//	result, err := db.WithRetry(func() (SomeType, error) {
+//		return db.Query.SomeOperation(ctx, db.RO(), params)
+//	})
+func WithRetry[T any](fn func() (T, error)) (T, error) {
+	retrier := retry.New(
+		retry.Attempts(3),
+		retry.Backoff(func(n int) time.Duration {
+			// Predefined backoff delays: 50ms, 100ms, 200ms
+			delays := []time.Duration{
+				DefaultBackoff,     // 50ms for attempt 1
+				DefaultBackoff * 2, // 100ms for attempt 2  
+				DefaultBackoff * 4, // 200ms for attempt 3
+			}
+			if n <= 0 || n > len(delays) {
+				return DefaultBackoff // fallback to base delay
+			}
+			return delays[n-1]
+		}),
+		retry.ShouldRetry(func(err error) bool {
+			// Don't retry if resource is not found - this is a valid response
+			if IsNotFound(err) {
+				return false
+			}
+
+			// Don't retry duplicate key errors - these won't succeed on retry
+			if IsDuplicateKeyError(err) {
+				return false
+			}
+
+			// Retry all other errors (network issues, timeouts, deadlocks, etc.)
+			return true
+		}),
+	)
+
+	var result T
+	err := retrier.Do(func() error {
+		var retryErr error
+		result, retryErr = fn()
+		return retryErr
+	})
+
+	return result, err
+}