NianAIGC/backend/internal/postgres/database.go

package postgres

import (
	"context"
	"encoding/json"
	"fmt"
	"time"
)

const ReadinessSQL = `
WITH required_table_privileges(table_name, privilege_name) AS (
  VALUES
    ('assets', 'SELECT'), ('assets', 'INSERT'), ('assets', 'DELETE'),
    ('generation_jobs', 'SELECT'), ('generation_jobs', 'INSERT'), ('generation_jobs', 'UPDATE'), ('generation_jobs', 'DELETE'),
    ('usage_events', 'SELECT'), ('usage_events', 'INSERT'), ('usage_events', 'UPDATE'),
    ('projects', 'SELECT'), ('projects', 'UPDATE'),
    ('image_templates', 'SELECT'), ('image_templates', 'INSERT'), ('image_templates', 'UPDATE'), ('image_templates', 'DELETE'),
    ('platform_organizations', 'SELECT'), ('platform_organizations', 'INSERT'), ('platform_organizations', 'UPDATE'), ('platform_organizations', 'DELETE'),
    ('platform_users', 'SELECT'), ('platform_users', 'INSERT'), ('platform_users', 'UPDATE'), ('platform_users', 'DELETE'),
    ('platform_account_migrations', 'SELECT'), ('platform_account_migrations', 'INSERT'), ('platform_account_migrations', 'UPDATE'),
    ('billing_price_rules', 'SELECT'), ('billing_price_rules', 'INSERT'), ('billing_price_rules', 'UPDATE'),
    ('billing_wallets', 'SELECT'), ('billing_wallets', 'INSERT'), ('billing_wallets', 'UPDATE'),
    ('billing_ledger', 'SELECT'), ('billing_ledger', 'INSERT')
)
SELECT
  NOT EXISTS (
    SELECT 1
    FROM required_table_privileges
    WHERE to_regclass('public.' || table_name) IS NULL
       OR NOT has_table_privilege(current_user, 'public.' || table_name, privilege_name)
  )
  AND has_function_privilege(
    current_user,
    'public.claim_generation_jobs(text,integer,integer)',
    'EXECUTE'
  )
  AND has_function_privilege(
    current_user,
    'public.billing_post_wallet_entry(text,text,text,text,text,bigint,text,text,text,jsonb)',
    'EXECUTE'
  ) AS ready
`

const ClaimGenerationJobsSQL = `SELECT id FROM public.claim_generation_jobs($1::text, $2::integer, $3::integer)`

const PostWalletEntrySQL = `SELECT ledger_id, balance_after_fen, balance_fen, total_recharged_fen, total_charged_fen, created_at, updated_at, delta_fen FROM public.billing_post_wallet_entry($1::text, $2::text, $3::text, $4::text, $5::text, $6::bigint, $7::text, $8::text, $9::text, $10::jsonb)`

type Rows interface {
	Close()
	Err() error
	Next() bool
	Scan(dest ...any) error
}

type Querier interface {
	Query(context.Context, string, ...any) (Rows, error)
}

type Transaction interface {
	Querier
	Exec(context.Context, string, ...any) error
	Commit(context.Context) error
	Rollback(context.Context) error
}

type TransactionBeginner interface {
	Begin(context.Context) (Transaction, error)
}

type Pool interface {
	Querier
	Close()
}

type Database struct {
	config       Config
	querier      Querier
	transactions TransactionBeginner
}

type Store = Database

func NewDatabase(config Config, querier Querier) *Database {
	db := &Database{config: config, querier: querier}
	if transactions, ok := querier.(TransactionBeginner); ok {
		db.transactions = transactions
	}
	return db
}

func (db *Database) Readiness(ctx context.Context) error {
	if db.config.Backend == BackendLocal {
		return nil
	}
	if db.querier == nil {
		return fmt.Errorf("PostgreSQL pool is not open")
	}
	rows, err := db.querier.Query(ctx, ReadinessSQL)
	if err != nil {
		return fmt.Errorf("query PostgreSQL readiness: %w", err)
	}
	defer rows.Close()
	if !rows.Next() {
		if err := rows.Err(); err != nil {
			return fmt.Errorf("read PostgreSQL readiness: %w", err)
		}
		return fmt.Errorf("PostgreSQL readiness query returned no row")
	}
	var ready bool
	if err := rows.Scan(&ready); err != nil {
		return fmt.Errorf("scan PostgreSQL readiness: %w", err)
	}
	if !ready {
		return fmt.Errorf("PostgreSQL schema or application privileges are not ready")
	}
	return rows.Err()
}

type GenerationJob struct {
	ID string
}

func (db *Database) ClaimGenerationJobs(ctx context.Context, workerID string, limit, lockTimeoutSeconds int) ([]GenerationJob, error) {
	if db.config.Backend != BackendPostgres || db.querier == nil {
		return nil, fmt.Errorf("PostgreSQL is unavailable when ZHINIAN_DATA_BACKEND=%s", db.config.Backend)
	}
	limit = max(1, min(limit, 20))
	rows, err := db.querier.Query(ctx, ClaimGenerationJobsSQL, workerID, limit, lockTimeoutSeconds)
	if err != nil {
		return nil, fmt.Errorf("claim generation jobs: %w", err)
	}
	defer rows.Close()
	var jobs []GenerationJob
	for rows.Next() {
		var job GenerationJob
		if err := rows.Scan(&job.ID); err != nil {
			return nil, fmt.Errorf("scan claimed generation job: %w", err)
		}
		jobs = append(jobs, job)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("read claimed generation jobs: %w", err)
	}
	return jobs, nil
}

type WalletEntryParams struct {
	LedgerID       string
	OrganizationID string
	AccountID      string
	JobID          string
	Kind           string
	DeltaFen       int64
	Currency       string
	IdempotencyKey string
	Description    string
	Metadata       json.RawMessage
}

type WalletEntry struct {
	LedgerID          string
	BalanceAfterFen   int64
	BalanceFen        int64
	TotalRechargedFen int64
	TotalChargedFen   int64
	CreatedAt         time.Time
	UpdatedAt         time.Time
	DeltaFen          int64
}

func (db *Database) PostWalletEntry(ctx context.Context, params WalletEntryParams) (WalletEntry, error) {
	if db.config.Backend != BackendPostgres || db.querier == nil {
		return WalletEntry{}, fmt.Errorf("PostgreSQL is unavailable when ZHINIAN_DATA_BACKEND=%s", db.config.Backend)
	}
	accountID := optionalDatabaseText(params.AccountID)
	if params.Kind == "recharge" || params.Kind == "adjustment" {
		accountID = nil
	}
	jobID := optionalDatabaseText(params.JobID)
	currency := params.Currency
	if currency == "" {
		currency = "CNY"
	}
	rows, err := db.querier.Query(ctx, PostWalletEntrySQL,
		params.LedgerID, params.OrganizationID, accountID, jobID, params.Kind,
		params.DeltaFen, currency, params.IdempotencyKey, params.Description, params.Metadata,
	)
	if err != nil {
		return WalletEntry{}, fmt.Errorf("post wallet entry: %w", err)
	}
	defer rows.Close()
	if !rows.Next() {
		if err := rows.Err(); err != nil {
			return WalletEntry{}, fmt.Errorf("read wallet entry: %w", err)
		}
		return WalletEntry{}, fmt.Errorf("billing_post_wallet_entry returned no row")
	}
	var entry WalletEntry
	if err := rows.Scan(
		&entry.LedgerID, &entry.BalanceAfterFen, &entry.BalanceFen,
		&entry.TotalRechargedFen, &entry.TotalChargedFen, &entry.CreatedAt,
		&entry.UpdatedAt, &entry.DeltaFen,
	); err != nil {
		return WalletEntry{}, fmt.Errorf("scan wallet entry: %w", err)
	}
	return entry, rows.Err()
}

func optionalDatabaseText(value string) any {
	if value == "" {
		return nil
	}
	return value
}