fix: use simplified queries for key decrements

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

@@ -129,11 +129,29 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error {
 	}
 
 	key, err = db.TxWithResult(ctx, h.DB.RW(), func(ctx context.Context, tx db.DBTX) (db.FindLiveKeyByIDRow, error) {
-		err = db.Query.UpdateKeyCredits(ctx, tx, db.UpdateKeyCreditsParams{
-			ID:        key.ID,
-			Operation: string(req.Operation),
-			Credits:   credits,
-		})
+		switch req.Operation {
+		case openapi.Set:
+			err = db.Query.UpdateKeyCreditsSet(ctx, tx, db.UpdateKeyCreditsSetParams{
+				ID:      key.ID,
+				Credits: credits,
+			})
+		case openapi.Increment:
+			err = db.Query.UpdateKeyCreditsIncrement(ctx, tx, db.UpdateKeyCreditsIncrementParams{
+				ID:      key.ID,
+				Credits: credits,
+			})
+		case openapi.Decrement:
+			err = db.Query.UpdateKeyCreditsDecrement(ctx, tx, db.UpdateKeyCreditsDecrementParams{
+				ID:      key.ID,
+				Credits: credits,
+			})
+		default:
+			return db.FindLiveKeyByIDRow{}, fault.New("invalid operation",
+				fault.Code(codes.App.Validation.InvalidInput.URN()),
+				fault.Internal(fmt.Sprintf("invalid operation: %s", req.Operation)),
+				fault.Public("Invalid operation specified."),
+			)
+		}
 		if err != nil {
 			return db.FindLiveKeyByIDRow{}, fault.Wrap(err,
 				fault.Code(codes.App.Internal.ServiceUnavailable.URN()),

go/internal/services/usagelimiter/limit.go

@@ -32,10 +32,9 @@ func (s *service) Limit(ctx context.Context, req UsageRequest) (UsageResponse, e
 		return UsageResponse{Valid: false, Remaining: 0}, nil
 	}
 
-	err = db.Query.UpdateKeyCredits(ctx, s.db.RW(), db.UpdateKeyCreditsParams{
-		ID:        req.KeyId,
-		Operation: "decrement",
-		Credits:   sql.NullInt32{Int32: req.Cost, Valid: true},
+	err = db.Query.UpdateKeyCreditsDecrement(ctx, s.db.RW(), db.UpdateKeyCreditsDecrementParams{
+		ID:      req.KeyId,
+		Credits: sql.NullInt32{Int32: req.Cost, Valid: true},
 	})
 	if err != nil {
 		return UsageResponse{}, err

go/pkg/db/queries/key_update_credits_decrement.sql

@@ -0,0 +1,7 @@
+-- name: UpdateKeyCreditsDecrement :exec
+UPDATE `keys`
+SET remaining_requests = CASE
+    WHEN remaining_requests >= sqlc.arg('credits') THEN remaining_requests - sqlc.arg('credits')
+    ELSE 0
+END
+WHERE id = ?;

go/pkg/db/queries/key_update_credits_increment.sql

@@ -0,0 +1,4 @@
+-- name: UpdateKeyCreditsIncrement :exec
+UPDATE `keys`
+SET remaining_requests = remaining_requests + sqlc.arg('credits')
+WHERE id = ?;

go/pkg/db/queries/key_update_credits_set.sql

@@ -0,0 +1,4 @@
+-- name: UpdateKeyCreditsSet :exec
+UPDATE `keys`
+SET remaining_requests = sqlc.narg('credits')
+WHERE id = ?;