fix: use simplified queries for key decrements#3804
Conversation
|
|
The latest updates on your projects. Learn more about Vercel for GitHub. |
|
Thank you for following the naming conventions for pull request titles! 🙏 |
📝 WalkthroughWalkthroughSplits the previous single UpdateKeyCredits DB operation into three explicit DB operations: Set, Increment, and Decrement. Updates the API handler to switch on req.Operation and call the corresponding DB method; updates the usage limiter to call the Decrement method. Replaces combined SQL/sqlc output with three generated query files and updates the Querier interface. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant APIHandler as "API Handler"
participant DB
Client->>APIHandler: PUT /v2/keys/update-credits {operation, value}
APIHandler->>APIHandler: switch(operation)
alt set
APIHandler->>DB: UpdateKeyCreditsSet(id, credits)
else increment
APIHandler->>DB: UpdateKeyCreditsIncrement(id, credits)
else decrement
APIHandler->>DB: UpdateKeyCreditsDecrement(id, credits)
else invalid
APIHandler-->>Client: 400 Invalid operation
end
DB-->>APIHandler: exec result / error
APIHandler-->>Client: response
sequenceDiagram
participant UsageService as "Usage Limiter"
participant DB
UsageService->>DB: UpdateKeyCreditsDecrement(id, credits)
DB-->>UsageService: exec result / error
UsageService-->>UsageService: compute remaining, build response
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (11)
go/apps/api/routes/v2_keys_update_credits/handler.go(1 hunks)go/internal/services/usagelimiter/limit.go(1 hunks)go/pkg/db/key_update_credits.sql_generated.go(0 hunks)go/pkg/db/key_update_credits_decrement.sql_generated.go(1 hunks)go/pkg/db/key_update_credits_increment.sql_generated.go(1 hunks)go/pkg/db/key_update_credits_set.sql_generated.go(1 hunks)go/pkg/db/querier_generated.go(1 hunks)go/pkg/db/queries/key_update_credits.sql(0 hunks)go/pkg/db/queries/key_update_credits_decrement.sql(1 hunks)go/pkg/db/queries/key_update_credits_increment.sql(1 hunks)go/pkg/db/queries/key_update_credits_set.sql(1 hunks)
💤 Files with no reviewable changes (2)
- go/pkg/db/queries/key_update_credits.sql
- go/pkg/db/key_update_credits.sql_generated.go
🧰 Additional context used
📓 Path-based instructions (2)
**/*.go
📄 CodeRabbit Inference Engine (CLAUDE.md)
**/*.go: Follow comprehensive documentation guidelines for Go code as described in go/GO_DOCUMENTATION_GUIDELINES.md
Every public function/type in Go code must be documented
Prefer interfaces for testability in Go code
Use AIDEV-* comments for complex/important code in Go services
Files:
go/pkg/db/key_update_credits_increment.sql_generated.gogo/pkg/db/key_update_credits_set.sql_generated.gogo/pkg/db/key_update_credits_decrement.sql_generated.gogo/pkg/db/querier_generated.gogo/internal/services/usagelimiter/limit.gogo/apps/api/routes/v2_keys_update_credits/handler.go
**/*.{env,js,ts,go}
📄 CodeRabbit Inference Engine (CLAUDE.md)
All environment variables must follow the format: UNKEY_<SERVICE_NAME>_VARNAME
Files:
go/pkg/db/key_update_credits_increment.sql_generated.gogo/pkg/db/key_update_credits_set.sql_generated.gogo/pkg/db/key_update_credits_decrement.sql_generated.gogo/pkg/db/querier_generated.gogo/internal/services/usagelimiter/limit.gogo/apps/api/routes/v2_keys_update_credits/handler.go
🧠 Learnings (1)
📚 Learning: 2025-07-17T14:24:20.403Z
Learnt from: Flo4604
PR: unkeyed/unkey#3631
File: go/pkg/db/bulk_keyring_insert.sql.go:23-25
Timestamp: 2025-07-17T14:24:20.403Z
Learning: In go/pkg/db/bulk_keyring_insert.sql.go and similar bulk insert generated files, hardcoded zero values for fields like size_approx and size_last_updated_at are intentional and reflect the original SQL query structure, not missing parameters.
Applied to files:
go/pkg/db/key_update_credits_increment.sql_generated.gogo/pkg/db/key_update_credits_set.sql_generated.go
🧬 Code Graph Analysis (5)
go/pkg/db/queries/key_update_credits_decrement.sql (1)
go/pkg/db/key_update_credits.sql_generated.go (1)
UpdateKeyCredits(42-56)
go/pkg/db/key_update_credits_set.sql_generated.go (2)
apps/dashboard/lib/trpc/routers/api/keys/query-api-keys/schema.ts (1)
Credits(26-26)apps/dashboard/lib/trpc/context.ts (1)
Context(35-35)
go/pkg/db/querier_generated.go (4)
go/pkg/db/key_update_credits_decrement.sql_generated.go (1)
UpdateKeyCreditsDecrementParams(22-25)go/pkg/db/key_update_credits_increment.sql_generated.go (1)
UpdateKeyCreditsIncrementParams(19-22)go/pkg/db/key_update_credits_refill.sql_generated.go (1)
UpdateKeyCreditsRefillParams(17-21)go/pkg/db/key_update_credits_set.sql_generated.go (1)
UpdateKeyCreditsSetParams(19-22)
go/internal/services/usagelimiter/limit.go (1)
go/pkg/db/key_update_credits_decrement.sql_generated.go (1)
UpdateKeyCreditsDecrementParams(22-25)
go/apps/api/routes/v2_keys_update_credits/handler.go (7)
go/apps/api/openapi/gen.go (3)
Set(24-24)Increment(23-23)Decrement(22-22)go/pkg/db/key_update_credits_set.sql_generated.go (1)
UpdateKeyCreditsSetParams(19-22)apps/dashboard/lib/trpc/routers/api/keys/query-api-keys/schema.ts (1)
Credits(26-26)go/pkg/db/key_update_credits_increment.sql_generated.go (1)
UpdateKeyCreditsIncrementParams(19-22)go/pkg/db/key_update_credits_decrement.sql_generated.go (1)
UpdateKeyCreditsDecrementParams(22-25)go/pkg/codes/unkey_application.go (1)
App(53-71)go/pkg/fault/wrap.go (2)
Internal(75-89)Public(97-111)
⏰ 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 API / API Test Local
- GitHub Check: Test Go API Local / Test
- GitHub Check: Build / Build
🔇 Additional comments (4)
go/pkg/db/queries/key_update_credits_set.sql (1)
1-4: LGTM; using sqlc.narg for nullable sets matches unlimited-semanticsSetting remaining_requests via a nullable named param is correct to support explicitly setting NULL (unlimited). Backticks around table name avoid keyword collisions. No changes needed.
go/pkg/db/key_update_credits_set.sql_generated.go (1)
19-32: Generated code looks correct and matches the SQL; no action neededParam types and Exec usage align with the underlying SQL. Given this is sqlc-generated, documentation omissions are acceptable.
go/pkg/db/key_update_credits_decrement.sql_generated.go (1)
35-37: NULL credits will zero-out remaining_requests on decrement
Verified that all existing call sites passCredits.Valid=true(no paths invokeUpdateKeyCreditsDecrementwith a NULL/invalid credits value):
- go/apps/api/routes/v2_keys_update_credits/handler.go
- go/internal/services/usagelimiter/limit.go
Keep this upstream enforcement in place and avoid adding any new callers that pass
Credits.Valid=false.go/pkg/db/querier_generated.go (1)
1530-1554: Interface surface looks consistent with the new split (Set/Inc/Dec).The new methods and comments match the intended semantics. Once the decrement SQL is simplified to one placeholder, please re-generate to update the doc comment and Exec signature accordingly.
|
|
Graphite Automations"Post a GIF when PR approved" took an action on this PR • (08/19/25)1 gif was posted to this PR based on Andreas Thomas's automation. |
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
go/pkg/db/queries/key_update_credits_decrement.sql (1)
1-7: Make decrement atomic, preserve unlimited keys (NULL), and return affected rowsThe current CASE will:
- Collapse unlimited keys (remaining_requests IS NULL) to 0 via the ELSE branch.
- Allow over-consumption under concurrency (two concurrent decrements can consume more than available because the second update “clamps to 0” instead of rejecting).
Use an atomic conditional decrement with WHERE guards and :execrows so callers can detect “insufficient credits” via 0 affected rows. This also reduces driver args (single credits placeholder), aligning with the PR’s perf goal.
Apply this diff:
--- 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 = ?; +-- name: UpdateKeyCreditsDecrement :execrows +UPDATE `keys` +SET remaining_requests = remaining_requests - sqlc.arg('credits') +WHERE id = ? + AND remaining_requests IS NOT NULL + AND sqlc.arg('credits') > 0 + AND remaining_requests >= sqlc.arg('credits');Follow-ups:
- Update call sites to check rowsAffected; if 0, treat as “insufficient credits.”
- This addresses the earlier nit to use
>=(readable and equivalent).- Keeping CASE/GREATEST is unnecessary here; the WHERE guard is faster and safer. We’ve learned GREATEST is slower in this path; this avoids it entirely while still using a single placeholder.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (3)
go/pkg/db/key_update_credits_decrement.sql_generated.go(1 hunks)go/pkg/db/querier_generated.go(1 hunks)go/pkg/db/queries/key_update_credits_decrement.sql(1 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.go
📄 CodeRabbit Inference Engine (CLAUDE.md)
**/*.go: Follow comprehensive documentation guidelines for Go code as described in go/GO_DOCUMENTATION_GUIDELINES.md
Every public function/type in Go code must be documented
Prefer interfaces for testability in Go code
Use AIDEV-* comments for complex/important code in Go services
Files:
go/pkg/db/key_update_credits_decrement.sql_generated.gogo/pkg/db/querier_generated.go
**/*.{env,js,ts,go}
📄 CodeRabbit Inference Engine (CLAUDE.md)
All environment variables must follow the format: UNKEY_<SERVICE_NAME>_VARNAME
Files:
go/pkg/db/key_update_credits_decrement.sql_generated.gogo/pkg/db/querier_generated.go
🧠 Learnings (2)
📚 Learning: 2025-08-18T08:41:37.670Z
Learnt from: Flo4604
PR: unkeyed/unkey#3804
File: go/pkg/db/key_update_credits_decrement.sql_generated.go:13-20
Timestamp: 2025-08-18T08:41:37.670Z
Learning: In MySQL queries for decrementing key credits, using GREATEST(remaining_requests - ?, 0) is slower than using a CASE expression with duplicate placeholders (CASE WHEN remaining_requests - ? > 0 THEN remaining_requests - ? ELSE 0 END), even though GREATEST reduces parameter count.
Applied to files:
go/pkg/db/key_update_credits_decrement.sql_generated.gogo/pkg/db/queries/key_update_credits_decrement.sql
📚 Learning: 2025-07-17T14:24:20.403Z
Learnt from: Flo4604
PR: unkeyed/unkey#3631
File: go/pkg/db/bulk_keyring_insert.sql.go:23-25
Timestamp: 2025-07-17T14:24:20.403Z
Learning: In go/pkg/db/bulk_keyring_insert.sql.go and similar bulk insert generated files, hardcoded zero values for fields like size_approx and size_last_updated_at are intentional and reflect the original SQL query structure, not missing parameters.
Applied to files:
go/pkg/db/key_update_credits_decrement.sql_generated.go
🧬 Code Graph Analysis (3)
go/pkg/db/key_update_credits_decrement.sql_generated.go (1)
go/pkg/db/key_update_credits.sql_generated.go (1)
UpdateKeyCredits(42-56)
go/pkg/db/queries/key_update_credits_decrement.sql (1)
go/pkg/db/key_update_credits.sql_generated.go (1)
UpdateKeyCredits(42-56)
go/pkg/db/querier_generated.go (4)
go/pkg/db/key_update_credits_decrement.sql_generated.go (1)
UpdateKeyCreditsDecrementParams(22-25)go/pkg/db/key_update_credits_increment.sql_generated.go (1)
UpdateKeyCreditsIncrementParams(19-22)go/pkg/db/key_update_credits_refill.sql_generated.go (1)
UpdateKeyCreditsRefillParams(17-21)go/pkg/db/key_update_credits_set.sql_generated.go (1)
UpdateKeyCreditsSetParams(19-22)
⏰ 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). (4)
- GitHub Check: Test Go API Local / Test
- GitHub Check: Test API / API Test Local
- GitHub Check: Build / Build
- GitHub Check: Analyze (javascript-typescript)
| const updateKeyCreditsDecrement = `-- name: UpdateKeyCreditsDecrement :exec | ||
| UPDATE ` + "`" + `keys` + "`" + ` | ||
| SET remaining_requests = CASE | ||
| WHEN remaining_requests >= ? THEN remaining_requests - ? | ||
| ELSE 0 | ||
| END | ||
| WHERE id = ? | ||
| ` |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Regenerate SQL to remove clamping and add WHERE guards (no NULL collapse, fewer args)
The generated constant mirrors the clamping behavior that zeroes unlimited keys and duplicates the credits arg. After updating the SQL source, re-generate so this block matches the atomic guarded decrement and uses one credits placeholder.
Target shape after re-generation:
-const updateKeyCreditsDecrement = `-- name: UpdateKeyCreditsDecrement :exec
-UPDATE ` + "`" + `keys` + "`" + `
-SET remaining_requests = CASE
- WHEN remaining_requests >= ? THEN remaining_requests - ?
- ELSE 0
-END
-WHERE id = ?
-`
+const updateKeyCreditsDecrement = `-- name: UpdateKeyCreditsDecrement :execrows
+UPDATE ` + "`" + `keys` + "`" + `
+SET remaining_requests = remaining_requests - ?
+WHERE id = ?
+ AND remaining_requests IS NOT NULL
+ AND ? > 0
+ AND remaining_requests >= ?
+`Note: sqlc will duplicate the same credits arg at the call site for the WHERE guards, but the SET uses a single placeholder. This remains performant and prevents incorrect clamping.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| const updateKeyCreditsDecrement = `-- name: UpdateKeyCreditsDecrement :exec | |
| UPDATE ` + "`" + `keys` + "`" + ` | |
| SET remaining_requests = CASE | |
| WHEN remaining_requests >= ? THEN remaining_requests - ? | |
| ELSE 0 | |
| END | |
| WHERE id = ? | |
| ` | |
| const updateKeyCreditsDecrement = `-- name: UpdateKeyCreditsDecrement :execrows | |
| UPDATE ` + "`" + `keys` + "`" + ` | |
| SET remaining_requests = remaining_requests - ? | |
| WHERE id = ? | |
| AND remaining_requests IS NOT NULL | |
| AND ? > 0 | |
| AND remaining_requests >= ? | |
| ` |
🤖 Prompt for AI Agents
In go/pkg/db/key_update_credits_decrement.sql_generated.go around lines 13–20,
the generated SQL still clamps unlimited (NULL) keys to zero and duplicates the
credits placeholder in SET; update the original .sql source and re-generate so
the SQL performs an atomic guarded decrement: use a single credits placeholder
in the SET expression (do not collapse NULLs), and move the guard conditions
into the WHERE clause — e.g. require remaining_requests IS NOT NULL and
remaining_requests >= ? along with id = ? — then run sqlc to regenerate the file
so the SET uses one placeholder while the WHERE duplicates the credits arg as
needed.
| func (q *Queries) UpdateKeyCreditsDecrement(ctx context.Context, db DBTX, arg UpdateKeyCreditsDecrementParams) error { | ||
| _, err := db.ExecContext(ctx, updateKeyCreditsDecrement, arg.Credits, arg.Credits, arg.ID) | ||
| return err | ||
| } |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Return rows affected from decrement to signal “insufficient credits”
Switch this to :execrows semantics so callers can reject consumption when balance is insufficient instead of silently zeroing the counter.
Apply this diff (post re-generation this will be automatic):
-func (q *Queries) UpdateKeyCreditsDecrement(ctx context.Context, db DBTX, arg UpdateKeyCreditsDecrementParams) error {
- _, err := db.ExecContext(ctx, updateKeyCreditsDecrement, arg.Credits, arg.Credits, arg.ID)
- return err
-}
+func (q *Queries) UpdateKeyCreditsDecrement(ctx context.Context, db DBTX, arg UpdateKeyCreditsDecrementParams) (int64, error) {
+ res, err := db.ExecContext(ctx, updateKeyCreditsDecrement,
+ arg.Credits, // SET remaining_requests = remaining_requests - ?
+ arg.ID, // WHERE id = ?
+ arg.Credits, // AND ? > 0
+ arg.Credits, // AND remaining_requests >= ?
+ )
+ if err != nil {
+ return 0, err
+ }
+ return res.RowsAffected()
+}Run this script to find and update call sites to handle the new return type and check rowsAffected == 1:
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate all usages of UpdateKeyCreditsDecrement across the repo with context.
rg -n -C3 --type=go '\bUpdateKeyCreditsDecrement\s*\('
echo
echo "Grepping for the generated interface method signature to spot mismatches:"
rg -n -C1 --type=go 'UpdateKeyCreditsDecrement\(ctx context\.Context, db DBTX, arg UpdateKeyCreditsDecrementParams\)'
echo
echo "Tip: callers should pattern-match like:"
cat <<'EOF'
rows, err := q.UpdateKeyCreditsDecrement(ctx, tx, db.UpdateKeyCreditsDecrementParams{
Credits: sql.NullInt32{Int32: credits, Valid: true},
ID: keyID,
})
if err != nil {
// handle DB error
}
if rows == 0 {
// treat as insufficient credits (no consumption)
}
EOFLength of output: 3684
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the SQL definition for UpdateKeyCreditsDecrement to verify directive
rg -n 'UpdateKeyCreditsDecrement' -g '*.sql'Length of output: 158
Switch UpdateKeyCreditsDecrement to :execrows and handle insufficient credits
To signal “insufficient credits” instead of silently doing nothing, change the SQL tag to :execrows, regenerate, and update all call sites to inspect the returned row count:
• In go/pkg/db/queries/key_update_credits_decrement.sql, change:
-- name: UpdateKeyCreditsDecrement :exec
+-- name: UpdateKeyCreditsDecrement :execrows• Run your SQL code generator (go generate ./go/pkg/db/queries) to update:
– The generated method in key_update_credits_decrement.sql_generated.go to return (int64, error)
– The interface in querier_generated.go accordingly
• Update callers to handle (rows, err):
- go/internal/services/usagelimiter/limit.go (around line 35):
- err = db.Query.UpdateKeyCreditsDecrement(ctx, s.db.RW(), params)
- rows, err := db.Query.UpdateKeyCreditsDecrement(ctx, s.db.RW(), params)
if err != nil {
return UsageResponse{Valid: false, Remaining: 0}, err
} - if rows == 0 {
-
return UsageResponse{Valid: false, Remaining: 0}, nil - }
// …on success, continue with Remaining = previous – cost…
- apps/api/routes/v2_keys_update_credits/handler.go (around line 144):
- err = db.Query.UpdateKeyCreditsDecrement(ctx, tx, db.UpdateKeyCreditsDecrementParams{...})
- if err != nil { … }
- rows, err := db.Query.UpdateKeyCreditsDecrement(ctx, tx, db.UpdateKeyCreditsDecrementParams{...})
- if err != nil {
- return nil, err
- }
- if rows == 0 {
- return &openapi.Error{Code: http.StatusPaymentRequired, Message: "insufficient credits"}, nil
- }
This ensures callers can reject or report when no rows were updated (i.e., balance was too low).
🤖 Prompt for AI Agents
In go/pkg/db/key_update_credits_decrement.sql_generated.go around lines 35–38,
replace the current void-returning Exec usage by switching the SQL tag in
go/pkg/db/queries/key_update_credits_decrement.sql to :execrows, run the SQL
code generator (go generate ./go/pkg/db/queries) so the generated method returns
(int64, error) and the querier interface is updated, then update all call sites
to inspect the returned rows count: in
go/internal/services/usagelimiter/limit.go (around line 35) return rows,err :=
db.Query.UpdateKeyCreditsDecrement(...) and if rows == 0 return an invalid
UsageResponse, and in apps/api/routes/v2_keys_update_credits/handler.go (around
line 144) use rows,err := db.Query.UpdateKeyCreditsDecrement(...), return error
on err and return a PaymentRequired/openapi error when rows == 0; this ensures
insufficient-credit cases are signaled when no rows were updated.
| //UpdateKeyCreditsDecrement | ||
| // | ||
| // UPDATE `keys` | ||
| // SET remaining_requests = | ||
| // CASE | ||
| // WHEN ? = 'set' THEN ? | ||
| // WHEN ? = 'increment' THEN remaining_requests + ? | ||
| // WHEN ? = 'decrement' AND remaining_requests - ? > 0 THEN remaining_requests - ? | ||
| // WHEN ? = 'decrement' AND remaining_requests - ? <= 0 THEN 0 | ||
| // SET remaining_requests = CASE | ||
| // WHEN remaining_requests >= ? THEN remaining_requests - ? | ||
| // ELSE 0 | ||
| // END | ||
| // WHERE id = ? | ||
| UpdateKeyCredits(ctx context.Context, db DBTX, arg UpdateKeyCreditsParams) error | ||
| UpdateKeyCreditsDecrement(ctx context.Context, db DBTX, arg UpdateKeyCreditsDecrementParams) error | ||
| //UpdateKeyCreditsIncrement |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Expose rowsAffected from decrement in the Querier interface
To support atomic “charge-or-reject” semantics, the decrement should return rowsAffected so callers can detect insufficient balance. Update the interface signature accordingly; re-generate after changing the SQL to :execrows.
Apply this diff (will be produced by sqlc after source change):
- //UpdateKeyCreditsDecrement
- //
- // UPDATE `keys`
- // SET remaining_requests = CASE
- // WHEN remaining_requests >= ? THEN remaining_requests - ?
- // ELSE 0
- // END
- // WHERE id = ?
- UpdateKeyCreditsDecrement(ctx context.Context, db DBTX, arg UpdateKeyCreditsDecrementParams) error
+ //UpdateKeyCreditsDecrement
+ //
+ // UPDATE `keys`
+ // SET remaining_requests = remaining_requests - ?
+ // WHERE id = ?
+ // AND remaining_requests IS NOT NULL
+ // AND ? > 0
+ // AND remaining_requests >= ?
+ UpdateKeyCreditsDecrement(ctx context.Context, db DBTX, arg UpdateKeyCreditsDecrementParams) (int64, error)Verification steps:
- Update usages in go/internal/services/usagelimiter/limit.go and go/apps/api/routes/v2_keys_update_credits/handler.go to handle the (int64, error) return and check rows > 0.
- Ensure tests cover:
- Decrement with sufficient credits: rows == 1, value decreases exactly by credits.
- Decrement with insufficient credits: rows == 0, value unchanged.
- Unlimited keys: rows == 0, value (NULL) unchanged.
- Zero/negative credits: rows == 0 (guarded by WHERE).
🏁 Script executed:
#!/bin/bash
set -e
echo "=== Interface signature in querier_generated.go ==="
rg -n "UpdateKeyCreditsDecrement" -C2 go/pkg/db/querier_generated.go
echo
echo "=== Service usage in limit.go ==="
rg -n "UpdateKeyCreditsDecrement" -C2 go/internal/services/usagelimiter/limit.go
echo
echo "=== Handler usage in v2_keys_update_credits handler ==="
rg -n "UpdateKeyCreditsDecrement" -C2 go/apps/api/routes/v2_keys_update_credits/handler.goLength of output: 1302
Expose rowsAffected from decrement in the Querier interface
To support atomic “charge-or-reject” semantics, switch the SQL block for UpdateKeyCreditsDecrement to use :execrows and update its generated signature to return (int64, error). Then update all callers to unpack the rows count and handle rows == 0 as “insufficient balance.”
• go/pkg/db/querier_generated.go
- //UpdateKeyCreditsDecrement
- //
- // UPDATE `keys`
- // SET remaining_requests = CASE
- // WHEN remaining_requests >= ? THEN remaining_requests - ?
- // ELSE 0
- // END
- // WHERE id = ?
- UpdateKeyCreditsDecrement(ctx context.Context, db DBTX, arg UpdateKeyCreditsDecrementParams) error
+ //UpdateKeyCreditsDecrement
+ //
+ // UPDATE `keys`
+ // SET remaining_requests = remaining_requests - ?
+ // WHERE id = ?
+ // AND remaining_requests IS NOT NULL
+ // AND ? > 0
+ // AND remaining_requests >= ?
+ UpdateKeyCreditsDecrement(ctx context.Context, db DBTX, arg UpdateKeyCreditsDecrementParams) (int64, error)• go/internal/services/usagelimiter/limit.go
– Change
err = db.Query.UpdateKeyCreditsDecrement(…)
– To
rows, err := db.Query.UpdateKeyCreditsDecrement(…)
if err != nil { return err }
if rows == 0 { return ErrInsufficientCredits }
• go/apps/api/routes/v2_keys_update_credits/handler.go
– Change
err = db.Query.UpdateKeyCreditsDecrement(…)
– To
rows, err := db.Query.UpdateKeyCreditsDecrement(…)
if err != nil { … }
if rows == 0 { return ctx.JSON(http.StatusPaymentRequired, …) }
Ensure you cover these scenarios in tests:
- Decrement with sufficient credits →
rows == 1, value decreases. - Decrement with insufficient credits →
rows == 0, value unchanged. - Unlimited keys (
remaining_requests IS NULL) →rows == 0, value unchanged. - Zero/negative credits →
rows == 0(guarded byWHERE).
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| //UpdateKeyCreditsDecrement | |
| // | |
| // UPDATE `keys` | |
| // SET remaining_requests = | |
| // CASE | |
| // WHEN ? = 'set' THEN ? | |
| // WHEN ? = 'increment' THEN remaining_requests + ? | |
| // WHEN ? = 'decrement' AND remaining_requests - ? > 0 THEN remaining_requests - ? | |
| // WHEN ? = 'decrement' AND remaining_requests - ? <= 0 THEN 0 | |
| // SET remaining_requests = CASE | |
| // WHEN remaining_requests >= ? THEN remaining_requests - ? | |
| // ELSE 0 | |
| // END | |
| // WHERE id = ? | |
| UpdateKeyCredits(ctx context.Context, db DBTX, arg UpdateKeyCreditsParams) error | |
| UpdateKeyCreditsDecrement(ctx context.Context, db DBTX, arg UpdateKeyCreditsDecrementParams) error | |
| //UpdateKeyCreditsIncrement | |
| //UpdateKeyCreditsDecrement | |
| // | |
| // UPDATE `keys` | |
| // SET remaining_requests = remaining_requests - ? | |
| // WHERE id = ? | |
| // AND remaining_requests IS NOT NULL | |
| // AND ? > 0 | |
| // AND remaining_requests >= ? | |
| UpdateKeyCreditsDecrement(ctx context.Context, db DBTX, arg UpdateKeyCreditsDecrementParams) (int64, error) | |
| //UpdateKeyCreditsIncrement |
🤖 Prompt for AI Agents
In go/pkg/db/querier_generated.go around lines 1530-1539 the generated
UpdateKeyCreditsDecrement currently returns only error and uses a query exec
that hides affected row count; change the SQL block to use :execrows and update
the generated method signature to return (int64, error) so it returns
rowsAffected, then update callers in go/internal/services/usagelimiter/limit.go
and go/apps/api/routes/v2_keys_update_credits/handler.go to unpack rows, treat
rows==0 as insufficient credits (return ErrInsufficientCredits or HTTP 402), and
add tests covering: (1) sufficient credits → rows==1 and value decreased, (2)
insufficient credits → rows==0 and value unchanged, (3) unlimited keys
(remaining_requests IS NULL) → rows==0 and unchanged, and (4) zero/negative
credits → rows==0.
2 participants
What does this PR do?
Uses 3 seperate queries for setting key credits to reduce mysql overhead.
Also needs less parameters which is good for allocations.
Type of change
How should this be tested?
Checklist
Required
pnpm buildpnpm fmtconsole.logsgit pull origin mainAppreciated
Summary by CodeRabbit
New Features
Bug Fixes
Refactor
Chores