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 }