package pgdb import ( "context" "errors" "fmt" "os" "time" "github.com/jackc/pgx/v5/pgxpool" "go.ntppool.org/common/database" "gopkg.in/yaml.v3" ) // PoolOptions configures pgxpool connection behavior. // // Default values match pgxpool defaults from github.com/jackc/pgx/v5/pgxpool. // See: https://pkg.go.dev/github.com/jackc/pgx/v5/pgxpool#Config // // To customize pool settings, either: // - Modify PoolOptions before calling OpenPool (for config file mode) // - Use URI query parameters like ?pool_max_conns=25 (for DATABASE_URI mode) type PoolOptions struct { // ConfigFiles is a list of config file paths to search for database configuration ConfigFiles []string // MinConns is the minimum number of connections in the pool. // Default: 0 (matches pgxpool default) MinConns int32 // MaxConns is the maximum number of connections in the pool. // Default: 4 (matches pgxpool default) // For higher concurrency, increase via PoolOptions or URI ?pool_max_conns=N MaxConns int32 // MaxConnLifetime is the maximum lifetime of a connection. // Default: 1 hour (matches pgxpool default) MaxConnLifetime time.Duration // MaxConnIdleTime is the maximum idle time of a connection. // Default: 30 minutes (matches pgxpool default) MaxConnIdleTime time.Duration // HealthCheckPeriod is how often to check connection health. // Default: 1 minute (matches pgxpool default) HealthCheckPeriod time.Duration } // DefaultPoolOptions returns defaults matching pgxpool. // See https://pkg.go.dev/github.com/jackc/pgx/v5/pgxpool#Config for pgxpool defaults. func DefaultPoolOptions() PoolOptions { return PoolOptions{ ConfigFiles: GetConfigFiles(), MinConns: 0, MaxConns: 4, MaxConnLifetime: time.Hour, MaxConnIdleTime: 30 * time.Minute, HealthCheckPeriod: time.Minute, } } // OpenPool opens a native pgx connection pool with the specified configuration // This is the primary and recommended way to connect to PostgreSQL // // Configuration precedence (highest to lowest): // 1. DATABASE_URI environment variable (pool settings can be included in URI) // 2. DATABASE_CONFIG_FILE environment variable (YAML) // 3. Default config files (database.yaml, /vault/secrets/database.yaml) // // When using DATABASE_URI, pool settings (pool_max_conns, pool_min_conns, etc.) // can be specified in the URI query string and PoolOptions are ignored. // When using config files, PoolOptions are applied. func OpenPool(ctx context.Context, options PoolOptions) (*pgxpool.Pool, error) { // Validate PoolOptions if options.MaxConns <= 0 { return nil, fmt.Errorf("pgdb: MaxConns must be positive, got: %d", options.MaxConns) } if options.MinConns < 0 { return nil, fmt.Errorf("pgdb: MinConns must be non-negative, got: %d", options.MinConns) } if options.MinConns > options.MaxConns { return nil, fmt.Errorf("pgdb: MinConns (%d) cannot exceed MaxConns (%d)", options.MinConns, options.MaxConns) } var poolConfig *pgxpool.Config var err error // Check DATABASE_URI environment variable first (highest priority) // When using DATABASE_URI, pool settings come from URI query parameters // (e.g., ?pool_max_conns=25). PoolOptions are not applied since our defaults // match pgxpool defaults. if uri := os.Getenv("DATABASE_URI"); uri != "" { poolConfig, err = pgxpool.ParseConfig(uri) if err != nil { return nil, fmt.Errorf("failed to parse DATABASE_URI: %w", err) } } else { // Fall back to config file approach pgCfg, _, err := FindConfig(options.ConfigFiles) if err != nil { return nil, err } poolConfig, err = CreatePoolConfig(pgCfg) if err != nil { return nil, err } // Apply pool-specific settings from PoolOptions (config files don't support these) poolConfig.MinConns = options.MinConns poolConfig.MaxConns = options.MaxConns poolConfig.MaxConnLifetime = options.MaxConnLifetime poolConfig.MaxConnIdleTime = options.MaxConnIdleTime poolConfig.HealthCheckPeriod = options.HealthCheckPeriod } // Create the pool pool, err := pgxpool.NewWithConfig(ctx, poolConfig) if err != nil { return nil, fmt.Errorf("failed to create connection pool: %w", err) } // Test the connection if err := pool.Ping(ctx); err != nil { pool.Close() // pgxpool.Pool.Close() doesn't return an error return nil, fmt.Errorf("failed to ping database: %w", err) } return pool, nil } // OpenPoolWithConfigFile opens a connection pool using an explicit config file path // This is a convenience function for when you have a specific config file func OpenPoolWithConfigFile(ctx context.Context, configFile string) (*pgxpool.Pool, error) { options := DefaultPoolOptions() options.ConfigFiles = []string{configFile} return OpenPool(ctx, options) } // FindConfig searches for and parses the first existing config file. // Returns the PostgresConfig, the path to the config file used, and any error. // If DATABASE_URI env var is set, returns nil config with empty path (use ParseURIConfig instead). func FindConfig(configFiles []string) (*database.PostgresConfig, string, error) { // Check if DATABASE_URI takes precedence if os.Getenv("DATABASE_URI") != "" { return nil, "", nil } var errs []error var triedFiles []string for _, configFile := range configFiles { if configFile == "" { continue } triedFiles = append(triedFiles, configFile) // Check if file exists if _, err := os.Stat(configFile); err != nil { errs = append(errs, fmt.Errorf("%s: %w", configFile, err)) continue } // Try to read and parse the file pgCfg, err := parseConfigFile(configFile) if err != nil { errs = append(errs, fmt.Errorf("%s: %w", configFile, err)) continue } return pgCfg, configFile, nil } if len(errs) > 0 { return nil, "", fmt.Errorf("no valid config file found (tried: %v): %w", triedFiles, errors.Join(errs...)) } return nil, "", fmt.Errorf("no valid config files provided") } // ParseURIConfig extracts connection info from DATABASE_URI environment variable. // Returns nil if DATABASE_URI is not set. func ParseURIConfig() (*pgxpool.Config, error) { uri := os.Getenv("DATABASE_URI") if uri == "" { return nil, nil } return pgxpool.ParseConfig(uri) } // parseConfigFile reads and parses a YAML config file func parseConfigFile(configFile string) (*database.PostgresConfig, error) { file, err := os.Open(configFile) if err != nil { return nil, fmt.Errorf("failed to open config file: %w", err) } defer func() { _ = file.Close() }() dec := yaml.NewDecoder(file) cfg := database.Config{} if err := dec.Decode(&cfg); err != nil { return nil, fmt.Errorf("failed to decode config: %w", err) } // Extract PostgreSQL config if cfg.Postgres != nil { return cfg.Postgres, nil } // Check for legacy flat format if cfg.User != "" && cfg.Name != "" { return &database.PostgresConfig{ User: cfg.User, Pass: cfg.Pass, Host: cfg.Host, Port: cfg.Port, Name: cfg.Name, SSLMode: cfg.SSLMode, }, nil } return nil, fmt.Errorf("no PostgreSQL configuration found in %s", configFile) } // GetConfigFiles returns the list of config files to search for database configuration. // Checks DATABASE_CONFIG_FILE env var first, otherwise returns default paths. func GetConfigFiles() []string { if configFile := os.Getenv("DATABASE_CONFIG_FILE"); configFile != "" { return []string{configFile} } return []string{"database.yaml", "/vault/secrets/database.yaml"} }