common/database/pgdb/pool.go

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"}
}