fix quota reset reservation race
Desktop Client Build / Resolve Build Metadata (push) Successful in 28s
Frontend CI / Frontend (push) Successful in 3m57s
Backend CI / Backend (push) Successful in 17m3s
Desktop Client Build / Build Desktop Client (push) Successful in 26m8s
Desktop Client Build / Publish Client Artifacts to NAS (push) Successful in 51s

This commit is contained in:
2026-05-12 02:24:57 +08:00
parent 1aad002a3e
commit 01fa2b8309
13 changed files with 337 additions and 168 deletions
+2 -2
View File
@@ -602,7 +602,7 @@ WITH reset AS (
updated_at = NOW() updated_at = NOW()
WHERE tenant_id = $1 WHERE tenant_id = $1
AND quota_type = $2 AND quota_type = $2
AND status = 'confirmed' AND status IN ('pending', 'confirmed')
AND created_at >= $3 AND created_at >= $3
AND created_at < $4 AND created_at < $4
RETURNING reserved_amount RETURNING reserved_amount
@@ -627,7 +627,7 @@ WITH reset AS (
updated_at = NOW() updated_at = NOW()
WHERE tenant_id = $1 WHERE tenant_id = $1
AND quota_type = 'ai_points' AND quota_type = 'ai_points'
AND status = 'confirmed' AND status IN ('pending', 'confirmed')
AND created_at >= $2 AND created_at >= $2
AND created_at < $3 AND created_at < $3
RETURNING reserved_amount RETURNING reserved_amount
@@ -0,0 +1,77 @@
package repository
import (
"context"
"errors"
"strings"
"testing"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
)
var (
testStart = time.Unix(0, 0).UTC()
testEnd = time.Unix(3600, 0).UTC()
)
type captureResetSQLTx struct {
sql string
}
func (tx *captureResetSQLTx) QueryRow(_ context.Context, sql string, _ ...any) pgx.Row {
tx.sql = sql
return failingRow{}
}
func (tx *captureResetSQLTx) Begin(context.Context) (pgx.Tx, error) { return nil, errors.New("unused") }
func (tx *captureResetSQLTx) Commit(context.Context) error { return nil }
func (tx *captureResetSQLTx) Rollback(context.Context) error { return nil }
func (tx *captureResetSQLTx) CopyFrom(context.Context, pgx.Identifier, []string, pgx.CopyFromSource) (int64, error) {
return 0, errors.New("unused")
}
func (tx *captureResetSQLTx) SendBatch(context.Context, *pgx.Batch) pgx.BatchResults {
return nil
}
func (tx *captureResetSQLTx) LargeObjects() pgx.LargeObjects { return pgx.LargeObjects{} }
func (tx *captureResetSQLTx) Prepare(context.Context, string, string) (*pgconn.StatementDescription, error) {
return nil, errors.New("unused")
}
func (tx *captureResetSQLTx) Exec(context.Context, string, ...any) (pgconn.CommandTag, error) {
return pgconn.CommandTag{}, errors.New("unused")
}
func (tx *captureResetSQLTx) Query(context.Context, string, ...any) (pgx.Rows, error) {
return nil, errors.New("unused")
}
func (tx *captureResetSQLTx) Conn() *pgx.Conn { return nil }
type failingRow struct{}
func (failingRow) Scan(...any) error { return errors.New("stop after capturing sql") }
func TestResetQuotaReservationsSQLMatchesEffectiveUsage(t *testing.T) {
tx := &captureResetSQLTx{}
_, _, _ = resetQuotaReservations(context.Background(), tx, 1, "article_generation", testStart, testEnd)
if !strings.Contains(tx.sql, "status IN ('pending', 'confirmed')") {
t.Fatalf("reset must clear every status counted as used; query:\n%s", tx.sql)
}
if strings.Contains(tx.sql, "ai_point_usage_logs") {
t.Fatalf("reset must not mutate AI point usage logs; query:\n%s", tx.sql)
}
}
func TestResetAIPointReservationsSQLDoesNotTouchUsageLogs(t *testing.T) {
tx := &captureResetSQLTx{}
_, _, _ = resetAIPointReservations(context.Background(), tx, 1, testStart, testEnd)
if !strings.Contains(tx.sql, "status IN ('pending', 'confirmed')") {
t.Fatalf("reset must clear every status counted as used; query:\n%s", tx.sql)
}
if strings.Contains(tx.sql, "ai_point_usage_logs") {
t.Fatalf("reset must not mutate AI point usage logs; query:\n%s", tx.sql)
}
}
+50 -30
View File
@@ -185,7 +185,7 @@ func CompleteAIPoints(
if err != nil || !pending { if err != nil || !pending {
return err return err
} }
if err := confirmAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil { if err := confirmAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID); err != nil {
return err return err
} }
var modelPtr *string var modelPtr *string
@@ -226,29 +226,31 @@ func RefundAIPoints(
return err return err
} }
status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC()) refunded, err := refundAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID)
if err != nil { if err != nil {
return err return err
} }
if refunded {
quotaRepo := repository.NewQuotaRepository(tx) status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC())
reason := "ai_points_refund" if err != nil {
referenceType := "ai_point_usage" return err
usageID := reservation.UsageLogID }
if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{ quotaRepo := repository.NewQuotaRepository(tx)
TenantID: tenantID, reason := "ai_points_refund"
OperatorID: operatorID, referenceType := "ai_point_usage"
QuotaType: aiPointsQuotaType, usageID := reservation.UsageLogID
Delta: reservation.Points, if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
BalanceAfter: status.Balance + reservation.Points, TenantID: tenantID,
Reason: &reason, OperatorID: operatorID,
ReferenceType: &referenceType, QuotaType: aiPointsQuotaType,
ReferenceID: &usageID, Delta: reservation.Points,
}); err != nil { BalanceAfter: status.Balance + reservation.Points,
return err Reason: &reason,
} ReferenceType: &referenceType,
if err := refundAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil { ReferenceID: &usageID,
return err }); err != nil {
return err
}
} }
if err := repository.NewAIPointUsageRepository(tx).MarkRefunded(ctx, tenantID, reservation.UsageLogID, strings.TrimSpace(errorMessage)); err != nil { if err := repository.NewAIPointUsageRepository(tx).MarkRefunded(ctx, tenantID, reservation.UsageLogID, strings.TrimSpace(errorMessage)); err != nil {
return err return err
@@ -375,7 +377,7 @@ func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogI
return status == "pending", nil return status == "pending", nil
} }
func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { func confirmAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error {
tag, err := tx.Exec(ctx, ` tag, err := tx.Exec(ctx, `
UPDATE quota_reservations UPDATE quota_reservations
SET status = 'confirmed', SET status = 'confirmed',
@@ -389,13 +391,10 @@ func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reserva
if err != nil { if err != nil {
return err return err
} }
if tag.RowsAffected() == 0 { return ignoreAIPointReservationReset(ctx, tx, tag.RowsAffected() > 0, tenantID, reservationID)
return errors.New("ai points reservation is not pending")
}
return nil
} }
func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { func refundAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) (bool, error) {
tag, err := tx.Exec(ctx, ` tag, err := tx.Exec(ctx, `
UPDATE quota_reservations UPDATE quota_reservations
SET status = 'refunded', SET status = 'refunded',
@@ -406,11 +405,32 @@ func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservat
AND quota_type = $3 AND quota_type = $3
AND status = 'pending' AND status = 'pending'
`, reservationID, tenantID, aiPointsQuotaType) `, reservationID, tenantID, aiPointsQuotaType)
if err != nil {
return false, err
}
if tag.RowsAffected() > 0 {
return true, nil
}
return false, ignoreAIPointReservationReset(ctx, tx, false, tenantID, reservationID)
}
func ignoreAIPointReservationReset(ctx context.Context, tx pgx.Tx, updated bool, tenantID, reservationID int64) error {
if updated {
return nil
}
var status string
err := tx.QueryRow(ctx, `
SELECT status
FROM quota_reservations
WHERE id = $1
AND tenant_id = $2
AND quota_type = $3
`, reservationID, tenantID, aiPointsQuotaType).Scan(&status)
if err != nil { if err != nil {
return err return err
} }
if tag.RowsAffected() == 0 { if status == "reset" {
return errors.New("ai points reservation is not pending") return nil
} }
return nil return errors.New("ai points reservation is not pending")
} }
@@ -512,27 +512,29 @@ func (s *ArticleImitationService) failGeneration(ctx context.Context, job articl
if job.ReservationID > 0 { if job.ReservationID > 0 {
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation") balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
if balanceErr != nil { refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID)
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) if err != nil {
} else {
reason := "generation_refund"
referenceType := "generation_task"
taskReferenceID := job.TaskID
if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
TenantID: job.TenantID,
OperatorID: job.OperatorID,
QuotaType: "article_generation",
Delta: 1,
BalanceAfter: balance + 1,
Reason: &reason,
ReferenceType: &referenceType,
ReferenceID: &taskReferenceID,
}); err != nil {
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err))
}
}
if err := quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID); err != nil {
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err)) cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err))
} else if refunded {
if balanceErr != nil {
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
} else {
reason := "generation_refund"
referenceType := "generation_task"
taskReferenceID := job.TaskID
if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
TenantID: job.TenantID,
OperatorID: job.OperatorID,
QuotaType: "article_generation",
Delta: 1,
BalanceAfter: balance + 1,
Reason: &reason,
ReferenceType: &referenceType,
ReferenceID: &taskReferenceID,
}); err != nil {
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err))
}
}
} }
} }
@@ -79,29 +79,34 @@ func HandleClaimedGenerationTaskFailure(
} }
if task.QuotaReservationID > 0 { if task.QuotaReservationID > 0 {
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation") var balance int
if balanceErr != nil { balanceErr := error(nil)
failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) if task.OperatorID > 0 {
} else if task.OperatorID > 0 { balance, balanceErr = quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
reason := "generation_refund"
referenceType := "generation_task"
taskReferenceID := task.ID
if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
TenantID: task.TenantID,
OperatorID: task.OperatorID,
QuotaType: "article_generation",
Delta: 1,
BalanceAfter: balance + 1,
Reason: &reason,
ReferenceType: &referenceType,
ReferenceID: &taskReferenceID,
}); err != nil {
failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err))
}
} }
refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, task.QuotaReservationID, task.TenantID)
if err := quotaRepo.RefundReservation(cleanupCtx, task.QuotaReservationID, task.TenantID); err != nil { if err != nil {
failureErr = errors.Join(failureErr, fmt.Errorf("refund generation reservation: %w", err)) failureErr = errors.Join(failureErr, fmt.Errorf("refund generation reservation: %w", err))
} else if refunded {
if balanceErr != nil {
failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
} else if task.OperatorID > 0 {
reason := "generation_refund"
referenceType := "generation_task"
taskReferenceID := task.ID
if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
TenantID: task.TenantID,
OperatorID: task.OperatorID,
QuotaType: "article_generation",
Delta: 1,
BalanceAfter: balance + 1,
Reason: &reason,
ReferenceType: &referenceType,
ReferenceID: &taskReferenceID,
}); err != nil {
failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err))
}
}
} }
} }
@@ -445,28 +445,34 @@ func HandleKolGenerationTaskFailure(
} }
if task.QuotaReservationID > 0 { if task.QuotaReservationID > 0 {
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation") var balance int
if balanceErr != nil { balanceErr := error(nil)
failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) if task.OperatorID > 0 {
} else if task.OperatorID > 0 { balance, balanceErr = quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
reason := "generation_refund"
referenceType := "generation_task"
taskReferenceID := task.ID
if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
TenantID: task.TenantID,
OperatorID: task.OperatorID,
QuotaType: "article_generation",
Delta: 1,
BalanceAfter: balance + 1,
Reason: &reason,
ReferenceType: &referenceType,
ReferenceID: &taskReferenceID,
}); err != nil {
failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err))
}
} }
if err := quotaRepo.RefundReservation(cleanupCtx, task.QuotaReservationID, task.TenantID); err != nil { refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, task.QuotaReservationID, task.TenantID)
if err != nil {
failureErr = errors.Join(failureErr, fmt.Errorf("refund generation reservation: %w", err)) failureErr = errors.Join(failureErr, fmt.Errorf("refund generation reservation: %w", err))
} else if refunded {
if balanceErr != nil {
failureErr = errors.Join(failureErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
} else if task.OperatorID > 0 {
reason := "generation_refund"
referenceType := "generation_task"
taskReferenceID := task.ID
if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
TenantID: task.TenantID,
OperatorID: task.OperatorID,
QuotaType: "article_generation",
Delta: 1,
BalanceAfter: balance + 1,
Reason: &reason,
ReferenceType: &referenceType,
ReferenceID: &taskReferenceID,
}); err != nil {
failureErr = errors.Join(failureErr, fmt.Errorf("insert generation refund ledger: %w", err))
}
}
} }
} }
@@ -739,28 +739,29 @@ func (s *PromptRuleGenerationService) failGeneration(ctx context.Context, job pr
} }
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation") balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
if balanceErr != nil { refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID)
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) if err != nil {
} else {
reason := "generation_refund"
referenceType := "generation_task"
taskReferenceID := job.TaskID
if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
TenantID: job.TenantID,
OperatorID: job.OperatorID,
QuotaType: "article_generation",
Delta: 1,
BalanceAfter: balance + 1,
Reason: &reason,
ReferenceType: &referenceType,
ReferenceID: &taskReferenceID,
}); err != nil {
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err))
}
}
if err := quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID); err != nil {
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err)) cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err))
} else if refunded {
if balanceErr != nil {
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
} else {
reason := "generation_refund"
referenceType := "generation_task"
taskReferenceID := job.TaskID
if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
TenantID: job.TenantID,
OperatorID: job.OperatorID,
QuotaType: "article_generation",
Delta: 1,
BalanceAfter: balance + 1,
Reason: &reason,
ReferenceType: &referenceType,
ReferenceID: &taskReferenceID,
}); err != nil {
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err))
}
}
} }
invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID) invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID)
if cleanupErr != nil { if cleanupErr != nil {
+22 -21
View File
@@ -626,28 +626,29 @@ func (s *TemplateService) failGeneration(ctx context.Context, job generationJob,
} }
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation") balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
if balanceErr != nil { refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID)
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr)) if err != nil {
} else {
reason := "generation_refund"
referenceType := "generation_task"
taskReferenceID := job.TaskID
if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
TenantID: job.TenantID,
OperatorID: job.OperatorID,
QuotaType: "article_generation",
Delta: 1,
BalanceAfter: balance + 1,
Reason: &reason,
ReferenceType: &referenceType,
ReferenceID: &taskReferenceID,
}); err != nil {
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err))
}
}
if err := quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID); err != nil {
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err)) cleanupErr = errors.Join(cleanupErr, fmt.Errorf("refund generation reservation: %w", err))
} else if refunded {
if balanceErr != nil {
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("load quota balance for refund: %w", balanceErr))
} else {
reason := "generation_refund"
referenceType := "generation_task"
taskReferenceID := job.TaskID
if _, err := quotaRepo.InsertQuotaLedger(cleanupCtx, repository.QuotaLedgerInput{
TenantID: job.TenantID,
OperatorID: job.OperatorID,
QuotaType: "article_generation",
Delta: 1,
BalanceAfter: balance + 1,
Reason: &reason,
ReferenceType: &referenceType,
ReferenceID: &taskReferenceID,
}); err != nil {
cleanupErr = errors.Join(cleanupErr, fmt.Errorf("insert generation refund ledger: %w", err))
}
}
} }
invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID) invalidateArticleCaches(cleanupCtx, s.cache, job.TenantID, &job.ArticleID)
if cleanupErr != nil { if cleanupErr != nil {
+50 -29
View File
@@ -1417,7 +1417,7 @@ func (s *Service) completeComplianceLLMJudgePoints(ctx context.Context, job revi
if err != nil || !pending { if err != nil || !pending {
return err return err
} }
if err := confirmAIPointReservation(ctx, tx, job.TenantID, reservation.ReservationID); err != nil { if err := confirmAIPointReservationIfPending(ctx, tx, job.TenantID, reservation.ReservationID); err != nil {
return err return err
} }
var modelPtr *string var modelPtr *string
@@ -1447,28 +1447,31 @@ func (s *Service) refundComplianceLLMJudgePoints(ctx context.Context, job review
if err != nil || !pending { if err != nil || !pending {
return err return err
} }
status, err := tenantrepo.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, job.TenantID, time.Now().UTC()) refunded, err := refundAIPointReservationIfPending(ctx, tx, job.TenantID, reservation.ReservationID)
if err != nil { if err != nil {
return err return err
} }
quotaRepo := tenantrepo.NewQuotaRepository(tx) if refunded {
reason := "ai_points_refund" status, err := tenantrepo.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, job.TenantID, time.Now().UTC())
referenceType := "ai_point_usage" if err != nil {
usageID := reservation.UsageLogID return err
if _, err := quotaRepo.InsertQuotaLedger(ctx, tenantrepo.QuotaLedgerInput{ }
TenantID: job.TenantID, quotaRepo := tenantrepo.NewQuotaRepository(tx)
OperatorID: job.CheckedBy, reason := "ai_points_refund"
QuotaType: aiPointsQuotaType, referenceType := "ai_point_usage"
Delta: reservation.Points, usageID := reservation.UsageLogID
BalanceAfter: status.Balance + reservation.Points, if _, err := quotaRepo.InsertQuotaLedger(ctx, tenantrepo.QuotaLedgerInput{
Reason: &reason, TenantID: job.TenantID,
ReferenceType: &referenceType, OperatorID: job.CheckedBy,
ReferenceID: &usageID, QuotaType: aiPointsQuotaType,
}); err != nil { Delta: reservation.Points,
return err BalanceAfter: status.Balance + reservation.Points,
} Reason: &reason,
if err := refundAIPointReservation(ctx, tx, job.TenantID, reservation.ReservationID); err != nil { ReferenceType: &referenceType,
return err ReferenceID: &usageID,
}); err != nil {
return err
}
} }
if err := tenantrepo.NewAIPointUsageRepository(tx).MarkRefunded(ctx, job.TenantID, reservation.UsageLogID, truncateRunes(errorString(jobErr), 2000)); err != nil { if err := tenantrepo.NewAIPointUsageRepository(tx).MarkRefunded(ctx, job.TenantID, reservation.UsageLogID, truncateRunes(errorString(jobErr), 2000)); err != nil {
return err return err
@@ -1512,7 +1515,7 @@ func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogI
return status == "pending", nil return status == "pending", nil
} }
func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { func confirmAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error {
tag, err := tx.Exec(ctx, ` tag, err := tx.Exec(ctx, `
UPDATE quota_reservations UPDATE quota_reservations
SET status = 'confirmed', SET status = 'confirmed',
@@ -1526,13 +1529,10 @@ func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reserva
if err != nil { if err != nil {
return err return err
} }
if tag.RowsAffected() == 0 { return ignoreAIPointReservationReset(ctx, tx, tag.RowsAffected() > 0, tenantID, reservationID)
return errors.New("ai points reservation is not pending")
}
return nil
} }
func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) error { func refundAIPointReservationIfPending(ctx context.Context, tx pgx.Tx, tenantID, reservationID int64) (bool, error) {
tag, err := tx.Exec(ctx, ` tag, err := tx.Exec(ctx, `
UPDATE quota_reservations UPDATE quota_reservations
SET status = 'refunded', SET status = 'refunded',
@@ -1543,13 +1543,34 @@ func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservat
AND quota_type = $3 AND quota_type = $3
AND status = 'pending' AND status = 'pending'
`, reservationID, tenantID, aiPointsQuotaType) `, reservationID, tenantID, aiPointsQuotaType)
if err != nil {
return false, err
}
if tag.RowsAffected() > 0 {
return true, nil
}
return false, ignoreAIPointReservationReset(ctx, tx, false, tenantID, reservationID)
}
func ignoreAIPointReservationReset(ctx context.Context, tx pgx.Tx, updated bool, tenantID, reservationID int64) error {
if updated {
return nil
}
var status string
err := tx.QueryRow(ctx, `
SELECT status
FROM quota_reservations
WHERE id = $1
AND tenant_id = $2
AND quota_type = $3
`, reservationID, tenantID, aiPointsQuotaType).Scan(&status)
if err != nil { if err != nil {
return err return err
} }
if tag.RowsAffected() == 0 { if status == "reset" {
return errors.New("ai points reservation is not pending") return nil
} }
return nil return errors.New("ai points reservation is not pending")
} }
func (s *Service) markReviewJobDone(ctx context.Context, jobID int64) error { func (s *Service) markReviewJobDone(ctx context.Context, jobID int64) error {
@@ -13,7 +13,7 @@ import (
const confirmReservation = `-- name: ConfirmReservation :exec const confirmReservation = `-- name: ConfirmReservation :exec
UPDATE quota_reservations SET status = 'confirmed', consumed_amount = reserved_amount, updated_at = NOW() UPDATE quota_reservations SET status = 'confirmed', consumed_amount = reserved_amount, updated_at = NOW()
WHERE id = $1 AND tenant_id = $2 WHERE id = $1 AND tenant_id = $2 AND status = 'pending'
` `
type ConfirmReservationParams struct { type ConfirmReservationParams struct {
@@ -114,7 +114,7 @@ func (q *Queries) InsertQuotaLedger(ctx context.Context, arg InsertQuotaLedgerPa
const refundReservation = `-- name: RefundReservation :exec const refundReservation = `-- name: RefundReservation :exec
UPDATE quota_reservations SET status = 'refunded', refunded_amount = reserved_amount, updated_at = NOW() UPDATE quota_reservations SET status = 'refunded', refunded_amount = reserved_amount, updated_at = NOW()
WHERE id = $1 AND tenant_id = $2 WHERE id = $1 AND tenant_id = $2 AND status = 'pending'
` `
type RefundReservationParams struct { type RefundReservationParams struct {
@@ -0,0 +1,19 @@
package generated
import (
"strings"
"testing"
)
func TestQuotaReservationTransitionsOnlyUpdatePending(t *testing.T) {
for name, query := range map[string]string{
"confirm": confirmReservation,
"refund": refundReservation,
} {
t.Run(name, func(t *testing.T) {
if !strings.Contains(query, "status = 'pending'") {
t.Fatalf("%s reservation query must only update pending reservations; query:\n%s", name, query)
}
})
}
}
@@ -23,8 +23,8 @@ WHERE id = sqlc.arg(id) AND tenant_id = sqlc.arg(tenant_id);
-- name: ConfirmReservation :exec -- name: ConfirmReservation :exec
UPDATE quota_reservations SET status = 'confirmed', consumed_amount = reserved_amount, updated_at = NOW() UPDATE quota_reservations SET status = 'confirmed', consumed_amount = reserved_amount, updated_at = NOW()
WHERE id = sqlc.arg(id) AND tenant_id = sqlc.arg(tenant_id); WHERE id = sqlc.arg(id) AND tenant_id = sqlc.arg(tenant_id) AND status = 'pending';
-- name: RefundReservation :exec -- name: RefundReservation :exec
UPDATE quota_reservations SET status = 'refunded', refunded_amount = reserved_amount, updated_at = NOW() UPDATE quota_reservations SET status = 'refunded', refunded_amount = reserved_amount, updated_at = NOW()
WHERE id = sqlc.arg(id) AND tenant_id = sqlc.arg(tenant_id); WHERE id = sqlc.arg(id) AND tenant_id = sqlc.arg(tenant_id) AND status = 'pending';
@@ -59,6 +59,7 @@ type QuotaRepository interface {
UpdateQuotaReservationResource(ctx context.Context, reservationID, tenantID, resourceID int64) error UpdateQuotaReservationResource(ctx context.Context, reservationID, tenantID, resourceID int64) error
ConfirmReservation(ctx context.Context, reservationID, tenantID int64) error ConfirmReservation(ctx context.Context, reservationID, tenantID int64) error
RefundReservation(ctx context.Context, reservationID, tenantID int64) error RefundReservation(ctx context.Context, reservationID, tenantID int64) error
RefundPendingReservation(ctx context.Context, reservationID, tenantID int64) (bool, error)
} }
type quotaRepository struct { type quotaRepository struct {
@@ -256,3 +257,19 @@ func (r *quotaRepository) RefundReservation(ctx context.Context, reservationID,
TenantID: tenantID, TenantID: tenantID,
}) })
} }
func (r *quotaRepository) RefundPendingReservation(ctx context.Context, reservationID, tenantID int64) (bool, error) {
tag, err := r.db.Exec(ctx, `
UPDATE quota_reservations
SET status = 'refunded',
refunded_amount = reserved_amount,
updated_at = NOW()
WHERE id = $1
AND tenant_id = $2
AND status = 'pending'
`, reservationID, tenantID)
if err != nil {
return false, err
}
return tag.RowsAffected() > 0, nil
}