EvoBGP/internal/db/open.go

package db

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"
	"time"

	"evobgp/migrations"

	"github.com/jackc/pgx/v5/pgxpool"
	_ "modernc.org/sqlite"
)

// OpenPostgresPool connects, applies embedded migrations, returns a pool.
func OpenPostgresPool(ctx context.Context, dsn string) (*pgxpool.Pool, error) {
	cfg, err := pgxpool.ParseConfig(dsn)
	if err != nil {
		return nil, err
	}
	applyPostgresPoolConfig(cfg)
	pool, err := pgxpool.NewWithConfig(ctx, cfg)
	if err != nil {
		return nil, err
	}
	if err := ApplyPostgresMigrations(ctx, pool, migrations.Postgres, "postgres"); err != nil {
		pool.Close()
		return nil, err
	}
	return pool, nil
}

// applyPostgresPoolConfig sets pool limits. An explicit EVOBGP_DB_MAX_CONNS is kept.
// When it is unset, the pool defaults to 8 connections so it can cover the default
// worker count instead of pgx's default of 4.
func applyPostgresPoolConfig(cfg *pgxpool.Config) {
	if cfg == nil {
		return
	}
	if max := strings.TrimSpace(os.Getenv("EVOBGP_DB_MAX_CONNS")); max != "" {
		if n, err := strconv.Atoi(max); err == nil && n > 0 {
			cfg.MaxConns = int32(n)
		}
	} else {
		cfg.MaxConns = 8
	}
	if min := strings.TrimSpace(os.Getenv("EVOBGP_DB_MIN_CONNS")); min != "" {
		if n, err := strconv.Atoi(min); err == nil && n >= 0 {
			cfg.MinConns = int32(n)
		}
	}
	cfg.MaxConnLifetime = 30 * time.Minute
	cfg.MaxConnIdleTime = 5 * time.Minute
}

// OpenSQLite applies sqlite/*.up.sql then opens the DB (foreign keys on).
func OpenSQLite(ctx context.Context, filePath string) (*sql.DB, error) {
	dsn := fmt.Sprintf("file:%s?_pragma=foreign_keys(1)", strings.TrimPrefix(filePath, "file:"))
	db, err := sql.Open("sqlite", dsn)
	if err != nil {
		return nil, err
	}
	if err := db.PingContext(ctx); err != nil {
		_ = db.Close()
		return nil, err
	}
	if err := applySQLiteMigrations(ctx, db, migrations.SQLite, "sqlite"); err != nil {
		_ = db.Close()
		return nil, err
	}
	return db, nil
}

func applySQLiteMigrations(ctx context.Context, db *sql.DB, fsys fs.FS, subdir string) error {
	if _, err := db.ExecContext(ctx, `CREATE TABLE IF NOT EXISTS schema_migrations (version TEXT NOT NULL PRIMARY KEY)`); err != nil {
		return fmt.Errorf("db: sqlite schema_migrations: %w", err)
	}
	entries, err := fs.ReadDir(fsys, subdir)
	if err != nil {
		return err
	}
	var ups []string
	for _, e := range entries {
		if e.IsDir() || !strings.HasSuffix(e.Name(), ".up.sql") {
			continue
		}
		ups = append(ups, e.Name())
	}
	sort.Strings(ups)
	for _, name := range ups {
		ver := strings.TrimSuffix(name, ".up.sql")
		var dummy int
		err := db.QueryRowContext(ctx, `SELECT 1 FROM schema_migrations WHERE version = ?`, ver).Scan(&dummy)
		if err == nil {
			continue
		}
		if !errors.Is(err, sql.ErrNoRows) {
			return err
		}
		body, err := fs.ReadFile(fsys, path.Join(subdir, name))
		if err != nil {
			return err
		}
		if _, err := db.ExecContext(ctx, string(body)); err != nil {
			// Idempotent ALTER ADD COLUMN for sqlite dev re-runs
			if strings.Contains(err.Error(), "duplicate column") {
				_, _ = db.ExecContext(ctx, `INSERT OR IGNORE INTO schema_migrations (version) VALUES (?)`, ver)
				continue
			}
			return fmt.Errorf("db: sqlite migrate %s: %w", name, err)
		}
		if _, err := db.ExecContext(ctx, `INSERT INTO schema_migrations (version) VALUES (?)`, ver); err != nil {
			return err
		}
	}
	return nil
}