package postgres import ( "context" "encoding/json" "errors" "fmt" "time" "git.nianxx.cn/wangxuming/NianAIGC/backend/internal/billing" "github.com/jackc/pgx/v5/pgconn" ) const activateChargedCreationSQL = `UPDATE public.generation_jobs SET billing = $2::jsonb, dispatch_ready_at = COALESCE(dispatch_ready_at, now()), updated_at = now() WHERE id = $1::text AND (dispatch_ready_at IS NULL OR billing = $2::jsonb) RETURNING id` const reconcileChargedCreationSQL = `SELECT billing, dispatch_ready_at::text FROM public.generation_jobs WHERE id = $1::text LIMIT 1` // ChargeAndActivateCreation closes the only non-idempotent creation window: // the wallet function and job dispatch gate commit in one database transaction. // A deterministic ledger ID is safe because billing_post_wallet_entry already // owns payload-idempotency validation for job-charge:. func (db *Database) ChargeAndActivateCreation(ctx context.Context, request billing.ChargeRequest, pending json.RawMessage) (json.RawMessage, error) { if db == nil || db.config.Backend != BackendPostgres || db.transactions == nil { return nil, fmt.Errorf("charge and activate generation creation: PostgreSQL transaction is unavailable") } if request.JobID == "" || request.OrganizationID == "" || request.AmountFen <= 0 || len(pending) == 0 || !json.Valid(pending) { return nil, fmt.Errorf("charge and activate generation creation: invalid request") } metadata, err := json.Marshal(request.Metadata) if err != nil { return nil, fmt.Errorf("charge and activate generation creation: encode metadata: %w", err) } tx, err := db.transactions.Begin(ctx) if err != nil { return nil, fmt.Errorf("charge and activate generation creation: begin: %w", err) } committed := false defer func() { if !committed { _ = tx.Rollback(ctx) } }() rows, err := tx.Query(ctx, PostWalletEntrySQL, "job-charge-entry:"+request.JobID, request.OrganizationID, optionalDatabaseText(request.AccountID), request.JobID, "charge", -request.AmountFen, billing.CurrencyCNY, "job-charge:"+request.JobID, request.Description, json.RawMessage(metadata), ) if err != nil { return nil, fmt.Errorf("charge and activate generation creation: post wallet: %w", err) } var posting WalletEntry if !rows.Next() { rowsErr := rows.Err() rows.Close() if rowsErr != nil { return nil, fmt.Errorf("charge and activate generation creation: read wallet: %w", rowsErr) } return nil, fmt.Errorf("charge and activate generation creation: wallet returned no row") } if err := rows.Scan(&posting.LedgerID, &posting.BalanceAfterFen, &posting.BalanceFen, &posting.TotalRechargedFen, &posting.TotalChargedFen, &posting.CreatedAt, &posting.UpdatedAt, &posting.DeltaFen); err != nil { rows.Close() return nil, fmt.Errorf("charge and activate generation creation: scan wallet: %w", err) } rows.Close() var state map[string]any if err := json.Unmarshal(pending, &state); err != nil { return nil, fmt.Errorf("charge and activate generation creation: decode billing: %w", err) } state["status"] = "charged" state["ledgerEntryId"] = posting.LedgerID chargedAt := posting.CreatedAt if chargedAt.IsZero() { chargedAt = time.Now().UTC() } state["chargedAt"] = chargedAt.UTC().Format(time.RFC3339Nano) charged, err := json.Marshal(state) if err != nil { return nil, fmt.Errorf("charge and activate generation creation: encode billing: %w", err) } activation, err := tx.Query(ctx, activateChargedCreationSQL, request.JobID, json.RawMessage(charged)) if err != nil { return nil, fmt.Errorf("charge and activate generation creation: activate job: %w", err) } if !activation.Next() { rowsErr := activation.Err() activation.Close() if rowsErr != nil { return nil, fmt.Errorf("charge and activate generation creation: read activation: %w", rowsErr) } return nil, fmt.Errorf("charge and activate generation creation: job is not pending") } var id string if err := activation.Scan(&id); err != nil { activation.Close() return nil, fmt.Errorf("charge and activate generation creation: scan job: %w", err) } activation.Close() commitErr := tx.Commit(ctx) committed = true if commitErr != nil { return db.reconcileChargedCreation(request.JobID, commitErr) } return charged, nil } func (db *Database) reconcileChargedCreation(jobID string, commitErr error) (json.RawMessage, error) { queryCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() rows, err := db.querier.Query(queryCtx, reconcileChargedCreationSQL, jobID) if err != nil { return nil, fmt.Errorf("%w: reconcile charge and activate generation creation after commit error: %v (commit: %v)", billing.ErrCommitOutcomeUnknown, err, commitErr) } defer rows.Close() if !rows.Next() { if err := rows.Err(); err != nil { return nil, fmt.Errorf("%w: reconcile charge and activate generation creation after commit error: %v (commit: %v)", billing.ErrCommitOutcomeUnknown, err, commitErr) } return unresolvedCommitOutcome(commitErr, "job not found during reconciliation") } var persistedBytes []byte var dispatchReadyAt string if err := rows.Scan(&persistedBytes, &dispatchReadyAt); err != nil { return nil, fmt.Errorf("%w: scan charge and activate reconciliation: %v (commit: %v)", billing.ErrCommitOutcomeUnknown, err, commitErr) } persisted := json.RawMessage(persistedBytes) var state struct { Status string `json:"status"` } if dispatchReadyAt == "" || json.Unmarshal(persisted, &state) != nil || state.Status != "charged" { return unresolvedCommitOutcome(commitErr, "job was not activated during reconciliation") } return persisted, nil } func unresolvedCommitOutcome(commitErr error, reconciliation string) (json.RawMessage, error) { var serverError *pgconn.PgError if errors.As(commitErr, &serverError) { return nil, fmt.Errorf("charge and activate generation creation: PostgreSQL rejected commit (%s): %w", reconciliation, commitErr) } return nil, fmt.Errorf("%w: %s after commit error: %v", billing.ErrCommitOutcomeUnknown, reconciliation, commitErr) }