feat: add Go migration compatibility foundation
This commit is contained in:
zn-admin committed
1 parent
7de3300034
commit
716a8031b1
27 files changed
+2582
No files matched your search
@@ -0,0 +1,198 @@
|
||||
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 Pool interface {
|
||||
Querier
|
||||
Close()
|
||||
}
|
||||
|
||||
type Database struct {
|
||||
config Config
|
||||
querier Querier
|
||||
}
|
||||
|
||||
type Store = Database
|
||||
|
||||
func NewDatabase(config Config, querier Querier) *Database {
|
||||
return &Database{config: config, querier: querier}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
Reference in new issue
Block a user