package pgdb import ( "context" "fmt" "os" "time" "github.com/jackc/pgx/v5/pgxpool" "go.ntppool.org/common/database" "gopkg.in/yaml.v3" ) // PoolOptions configures pgxpool connection behavior 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 (no minimum) MinConns int32 // MaxConns is the maximum number of connections in the pool // Default: 25 MaxConns int32 // MaxConnLifetime is the maximum lifetime of a connection // Default: 1 hour MaxConnLifetime time.Duration // MaxConnIdleTime is the maximum idle time of a connection // Default: 30 minutes MaxConnIdleTime time.Duration // HealthCheckPeriod is how often to check connection health // Default: 1 minute HealthCheckPeriod time.Duration } // DefaultPoolOptions returns sensible defaults for pgxpool func DefaultPoolOptions() PoolOptions { return PoolOptions{ ConfigFiles: getConfigFiles(), MinConns: 0, MaxConns: 25, 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) { var poolConfig *pgxpool.Config var err error // Check DATABASE_URI environment variable first (highest priority) 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) } // Pool settings from URI are used; PoolOptions ignored } else { // Fall back to config file approach pgCfg, err := findAndParseConfig(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() 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) } // findAndParseConfig searches for and parses the first existing config file func findAndParseConfig(configFiles []string) (*database.PostgresConfig, error) { var firstErr error for _, configFile := range configFiles { if configFile == "" { continue } // Check if file exists if _, err := os.Stat(configFile); err != nil { if firstErr == nil { firstErr = err } continue } // Try to read and parse the file pgCfg, err := parseConfigFile(configFile) if err != nil { if firstErr == nil { firstErr = err } continue } return pgCfg, nil } if firstErr != nil { return nil, fmt.Errorf("no config file found: %w", firstErr) } return nil, fmt.Errorf("no valid config files provided") } // 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 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 func getConfigFiles() []string { if configFile := os.Getenv("DATABASE_CONFIG_FILE"); configFile != "" { return []string{configFile} } return []string{"database.yaml", "/vault/secrets/database.yaml"} }