NianAIGC/backend/internal/postgres/generation_state.go

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:<job-id>.
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() {
		rows.Close()
		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() {
		activation.Close()
		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)
}