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
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:
@@ -602,7 +602,7 @@ WITH reset AS (
|
||||
updated_at = NOW()
|
||||
WHERE tenant_id = $1
|
||||
AND quota_type = $2
|
||||
AND status = 'confirmed'
|
||||
AND status IN ('pending', 'confirmed')
|
||||
AND created_at >= $3
|
||||
AND created_at < $4
|
||||
RETURNING reserved_amount
|
||||
@@ -627,7 +627,7 @@ WITH reset AS (
|
||||
updated_at = NOW()
|
||||
WHERE tenant_id = $1
|
||||
AND quota_type = 'ai_points'
|
||||
AND status = 'confirmed'
|
||||
AND status IN ('pending', 'confirmed')
|
||||
AND created_at >= $2
|
||||
AND created_at < $3
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -185,7 +185,7 @@ func CompleteAIPoints(
|
||||
if err != nil || !pending {
|
||||
return err
|
||||
}
|
||||
if err := confirmAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil {
|
||||
if err := confirmAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID); err != nil {
|
||||
return err
|
||||
}
|
||||
var modelPtr *string
|
||||
@@ -226,29 +226,31 @@ func RefundAIPoints(
|
||||
return err
|
||||
}
|
||||
|
||||
status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC())
|
||||
refunded, err := refundAIPointReservationIfPending(ctx, tx, tenantID, reservation.ReservationID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
quotaRepo := repository.NewQuotaRepository(tx)
|
||||
reason := "ai_points_refund"
|
||||
referenceType := "ai_point_usage"
|
||||
usageID := reservation.UsageLogID
|
||||
if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
|
||||
TenantID: tenantID,
|
||||
OperatorID: operatorID,
|
||||
QuotaType: aiPointsQuotaType,
|
||||
Delta: reservation.Points,
|
||||
BalanceAfter: status.Balance + reservation.Points,
|
||||
Reason: &reason,
|
||||
ReferenceType: &referenceType,
|
||||
ReferenceID: &usageID,
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := refundAIPointReservation(ctx, tx, tenantID, reservation.ReservationID); err != nil {
|
||||
return err
|
||||
if refunded {
|
||||
status, err := repository.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, tenantID, time.Now().UTC())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
quotaRepo := repository.NewQuotaRepository(tx)
|
||||
reason := "ai_points_refund"
|
||||
referenceType := "ai_point_usage"
|
||||
usageID := reservation.UsageLogID
|
||||
if _, err := quotaRepo.InsertQuotaLedger(ctx, repository.QuotaLedgerInput{
|
||||
TenantID: tenantID,
|
||||
OperatorID: operatorID,
|
||||
QuotaType: aiPointsQuotaType,
|
||||
Delta: reservation.Points,
|
||||
BalanceAfter: status.Balance + reservation.Points,
|
||||
Reason: &reason,
|
||||
ReferenceType: &referenceType,
|
||||
ReferenceID: &usageID,
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := repository.NewAIPointUsageRepository(tx).MarkRefunded(ctx, tenantID, reservation.UsageLogID, strings.TrimSpace(errorMessage)); err != nil {
|
||||
return err
|
||||
@@ -375,7 +377,7 @@ func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogI
|
||||
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, `
|
||||
UPDATE quota_reservations
|
||||
SET status = 'confirmed',
|
||||
@@ -389,13 +391,10 @@ func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reserva
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return errors.New("ai points reservation is not pending")
|
||||
}
|
||||
return nil
|
||||
return ignoreAIPointReservationReset(ctx, tx, tag.RowsAffected() > 0, tenantID, reservationID)
|
||||
}
|
||||
|
||||
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, `
|
||||
UPDATE quota_reservations
|
||||
SET status = 'refunded',
|
||||
@@ -406,11 +405,32 @@ func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservat
|
||||
AND quota_type = $3
|
||||
AND status = 'pending'
|
||||
`, 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 {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return errors.New("ai points reservation is not pending")
|
||||
if status == "reset" {
|
||||
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 {
|
||||
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
|
||||
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))
|
||||
}
|
||||
}
|
||||
if err := quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID); err != nil {
|
||||
refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID)
|
||||
if err != nil {
|
||||
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 {
|
||||
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
|
||||
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))
|
||||
}
|
||||
var balance int
|
||||
balanceErr := error(nil)
|
||||
if task.OperatorID > 0 {
|
||||
balance, balanceErr = quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
|
||||
}
|
||||
|
||||
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))
|
||||
} 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 {
|
||||
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
|
||||
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))
|
||||
}
|
||||
var balance int
|
||||
balanceErr := error(nil)
|
||||
if task.OperatorID > 0 {
|
||||
balance, balanceErr = quotaRepo.GetCurrentBalance(cleanupCtx, task.TenantID, "article_generation")
|
||||
}
|
||||
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))
|
||||
} 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")
|
||||
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))
|
||||
}
|
||||
}
|
||||
|
||||
if err := quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID); err != nil {
|
||||
refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID)
|
||||
if err != nil {
|
||||
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)
|
||||
if cleanupErr != nil {
|
||||
|
||||
@@ -626,28 +626,29 @@ func (s *TemplateService) failGeneration(ctx context.Context, job generationJob,
|
||||
}
|
||||
|
||||
balance, balanceErr := quotaRepo.GetCurrentBalance(cleanupCtx, job.TenantID, "article_generation")
|
||||
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))
|
||||
}
|
||||
}
|
||||
|
||||
if err := quotaRepo.RefundReservation(cleanupCtx, job.ReservationID, job.TenantID); err != nil {
|
||||
refunded, err := quotaRepo.RefundPendingReservation(cleanupCtx, job.ReservationID, job.TenantID)
|
||||
if err != nil {
|
||||
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)
|
||||
if cleanupErr != nil {
|
||||
|
||||
@@ -1417,7 +1417,7 @@ func (s *Service) completeComplianceLLMJudgePoints(ctx context.Context, job revi
|
||||
if err != nil || !pending {
|
||||
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
|
||||
}
|
||||
var modelPtr *string
|
||||
@@ -1447,28 +1447,31 @@ func (s *Service) refundComplianceLLMJudgePoints(ctx context.Context, job review
|
||||
if err != nil || !pending {
|
||||
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 {
|
||||
return err
|
||||
}
|
||||
quotaRepo := tenantrepo.NewQuotaRepository(tx)
|
||||
reason := "ai_points_refund"
|
||||
referenceType := "ai_point_usage"
|
||||
usageID := reservation.UsageLogID
|
||||
if _, err := quotaRepo.InsertQuotaLedger(ctx, tenantrepo.QuotaLedgerInput{
|
||||
TenantID: job.TenantID,
|
||||
OperatorID: job.CheckedBy,
|
||||
QuotaType: aiPointsQuotaType,
|
||||
Delta: reservation.Points,
|
||||
BalanceAfter: status.Balance + reservation.Points,
|
||||
Reason: &reason,
|
||||
ReferenceType: &referenceType,
|
||||
ReferenceID: &usageID,
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := refundAIPointReservation(ctx, tx, job.TenantID, reservation.ReservationID); err != nil {
|
||||
return err
|
||||
if refunded {
|
||||
status, err := tenantrepo.NewQuotaRepository(tx).GetAIQuotaStatus(ctx, job.TenantID, time.Now().UTC())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
quotaRepo := tenantrepo.NewQuotaRepository(tx)
|
||||
reason := "ai_points_refund"
|
||||
referenceType := "ai_point_usage"
|
||||
usageID := reservation.UsageLogID
|
||||
if _, err := quotaRepo.InsertQuotaLedger(ctx, tenantrepo.QuotaLedgerInput{
|
||||
TenantID: job.TenantID,
|
||||
OperatorID: job.CheckedBy,
|
||||
QuotaType: aiPointsQuotaType,
|
||||
Delta: reservation.Points,
|
||||
BalanceAfter: status.Balance + reservation.Points,
|
||||
Reason: &reason,
|
||||
ReferenceType: &referenceType,
|
||||
ReferenceID: &usageID,
|
||||
}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if err := tenantrepo.NewAIPointUsageRepository(tx).MarkRefunded(ctx, job.TenantID, reservation.UsageLogID, truncateRunes(errorString(jobErr), 2000)); err != nil {
|
||||
return err
|
||||
@@ -1512,7 +1515,7 @@ func lockPendingAIPointUsage(ctx context.Context, tx pgx.Tx, tenantID, usageLogI
|
||||
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, `
|
||||
UPDATE quota_reservations
|
||||
SET status = 'confirmed',
|
||||
@@ -1526,13 +1529,10 @@ func confirmAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reserva
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return errors.New("ai points reservation is not pending")
|
||||
}
|
||||
return nil
|
||||
return ignoreAIPointReservationReset(ctx, tx, tag.RowsAffected() > 0, tenantID, reservationID)
|
||||
}
|
||||
|
||||
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, `
|
||||
UPDATE quota_reservations
|
||||
SET status = 'refunded',
|
||||
@@ -1543,13 +1543,34 @@ func refundAIPointReservation(ctx context.Context, tx pgx.Tx, tenantID, reservat
|
||||
AND quota_type = $3
|
||||
AND status = 'pending'
|
||||
`, 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 {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return errors.New("ai points reservation is not pending")
|
||||
if status == "reset" {
|
||||
return nil
|
||||
}
|
||||
return nil
|
||||
return errors.New("ai points reservation is not pending")
|
||||
}
|
||||
|
||||
func (s *Service) markReviewJobDone(ctx context.Context, jobID int64) error {
|
||||
|
||||
@@ -13,7 +13,7 @@ import (
|
||||
|
||||
const confirmReservation = `-- name: ConfirmReservation :exec
|
||||
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 {
|
||||
@@ -114,7 +114,7 @@ func (q *Queries) InsertQuotaLedger(ctx context.Context, arg InsertQuotaLedgerPa
|
||||
|
||||
const refundReservation = `-- name: RefundReservation :exec
|
||||
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 {
|
||||
|
||||
@@ -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
|
||||
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
|
||||
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
|
||||
ConfirmReservation(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 {
|
||||
@@ -256,3 +257,19 @@ func (r *quotaRepository) RefundReservation(ctx context.Context, reservationID,
|
||||
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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user