Files
NianAIGC/backend/internal/postgres/open.go

109 lines
2.9 KiB
Go

package postgres
import (
"context"
"fmt"
"strconv"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
)
type Module struct {
pool Pool
Store *Store
}
func (module *Module) Close() {
if module != nil && module.pool != nil {
module.pool.Close()
}
}
func Open(ctx context.Context, config Config) (*Module, error) {
if config.Backend == BackendLocal {
return &Module{Store: NewDatabase(config, nil)}, nil
}
if config.Backend != BackendPostgres {
return nil, fmt.Errorf("unsupported data backend %q", config.Backend)
}
poolConfig, err := pgxpool.ParseConfig(config.DatabaseURL)
if err != nil {
return nil, fmt.Errorf("parse DATABASE_URL: %w", err)
}
poolConfig.MaxConns = config.PoolMax
poolConfig.MaxConnIdleTime = config.IdleTimeout
poolConfig.ConnConfig.ConnectTimeout = config.ConnectionTimeout
poolConfig.ConnConfig.RuntimeParams["statement_timeout"] = strconv.FormatInt(config.StatementTimeout.Milliseconds(), 10)
poolConfig.ConnConfig.RuntimeParams["application_name"] = config.ApplicationName
switch config.SSLMode {
case SSLDisable:
poolConfig.ConnConfig.TLSConfig = nil
poolConfig.ConnConfig.Fallbacks = nil
case SSLVerifyFull:
if config.TLSConfig == nil {
return nil, fmt.Errorf("TLS configuration is required when DATABASE_SSL_MODE=verify-full")
}
tlsConfig := config.TLSConfig.Clone()
if tlsConfig.ServerName == "" {
tlsConfig.ServerName = poolConfig.ConnConfig.Host
}
poolConfig.ConnConfig.TLSConfig = tlsConfig
poolConfig.ConnConfig.Fallbacks = nil
default:
return nil, fmt.Errorf("unsupported DATABASE_SSL_MODE %q", config.SSLMode)
}
pool, err := pgxpool.NewWithConfig(ctx, poolConfig)
if err != nil {
return nil, fmt.Errorf("open PostgreSQL pool: %w", err)
}
adapter := &pgxPoolAdapter{pool: pool}
return &Module{pool: adapter, Store: NewDatabase(config, adapter)}, nil
}
type pgxPoolAdapter struct {
pool *pgxpool.Pool
}
func (p *pgxPoolAdapter) Query(ctx context.Context, sql string, args ...any) (Rows, error) {
return p.pool.Query(ctx, sql, args...)
}
func (p *pgxPoolAdapter) Begin(ctx context.Context) (Transaction, error) {
tx, err := p.pool.Begin(ctx)
if err != nil {
return nil, err
}
return &pgxTransactionAdapter{tx: tx}, nil
}
func (p *pgxPoolAdapter) Close() {
p.pool.Close()
}
var _ Pool = (*pgxPoolAdapter)(nil)
var _ TransactionBeginner = (*pgxPoolAdapter)(nil)
type pgxTransactionAdapter struct {
tx pgx.Tx
}
func (t *pgxTransactionAdapter) Query(ctx context.Context, sql string, args ...any) (Rows, error) {
return t.tx.Query(ctx, sql, args...)
}
func (t *pgxTransactionAdapter) Exec(ctx context.Context, sql string, args ...any) error {
_, err := t.tx.Exec(ctx, sql, args...)
return err
}
func (t *pgxTransactionAdapter) Commit(ctx context.Context) error {
return t.tx.Commit(ctx)
}
func (t *pgxTransactionAdapter) Rollback(ctx context.Context) error {
return t.tx.Rollback(ctx)
}
var _ Transaction = (*pgxTransactionAdapter)(nil)