Cyrene_v2/backend/internal/db/db.go

// Package db provides database connection management, migration support, and helper utilities.
package db

import (
	"context"
	"database/sql"
	"fmt"
	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
	"strings"
	"time"

	_ "github.com/lib/pq" // PostgreSQL driver
)

// Config holds database connection parameters.
type Config struct {
	Host     string
	Port     int
	User     string
	Password string
	DBName   string
	SSLMode  string
	URL      string // Full DSN takes priority over other fields

	MaxOpenConns    int
	MaxIdleConns    int
	ConnMaxLifetime time.Duration
}

// DefaultConfig returns a Config with sensible defaults.
func DefaultConfig() Config {
	return Config{
		Host:     "localhost",
		Port:     5432,
		User:     "cyrene",
		Password: "cyrene_pass",
		DBName:   "cyrene_ai",
		SSLMode:  "disable",

		MaxOpenConns:    25,
		MaxIdleConns:    5,
		ConnMaxLifetime: 5 * time.Minute,
	}
}

// Connect establishes a PostgreSQL connection pool.
func Connect(cfg Config) (*sql.DB, error) {
	dsn := cfg.URL
	if dsn == "" {
		dsn = fmt.Sprintf("postgres://%s:%s@%s:%d/%s?sslmode=%s",
			cfg.User, cfg.Password, cfg.Host, cfg.Port, cfg.DBName, cfg.SSLMode)
	}

	db, err := sql.Open("postgres", dsn)
	if err != nil {
		return nil, fmt.Errorf("open db: %w", err)
	}

	db.SetMaxOpenConns(cfg.MaxOpenConns)
	db.SetMaxIdleConns(cfg.MaxIdleConns)
	db.SetConnMaxLifetime(cfg.ConnMaxLifetime)

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if err := db.PingContext(ctx); err != nil {
		db.Close()
		return nil, fmt.Errorf("ping db: %w", err)
	}

	logger.Infof("[db] connected to %s:%d/%s (max_open=%d, max_idle=%d)",
		cfg.Host, cfg.Port, cfg.DBName, cfg.MaxOpenConns, cfg.MaxIdleConns)

	return db, nil
}

// MustConnect is like Connect but panics on error.
// Use only in main() where startup failure is unrecoverable.
func MustConnect(cfg Config) *sql.DB {
	db, err := Connect(cfg)
	if err != nil {
		panic(fmt.Sprintf("failed to connect to database: %v", err))
	}
	return db
}

// HealthCheck verifies the database connection is alive.
func HealthCheck(ctx context.Context, db *sql.DB) error {
	return db.PingContext(ctx)
}

// Transaction executes fn within a database transaction.
// If fn returns an error, the transaction is rolled back.
func Transaction(ctx context.Context, db *sql.DB, fn func(tx *sql.Tx) error) error {
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin tx: %w", err)
	}

	defer func() {
		if p := recover(); p != nil {
			tx.Rollback()
			panic(p)
		}
	}()

	if err := fn(tx); err != nil {
		if rbErr := tx.Rollback(); rbErr != nil {
			return fmt.Errorf("rollback failed (original: %w): %v", err, rbErr)
		}
		return err
	}

	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit: %w", err)
	}
	return nil
}

// RunMigration executes a list of SQL migration statements in order.
func RunMigration(ctx context.Context, db *sql.DB, name string, statements []string) error {
	logger.Infof("[db] running migration: %s (%d statements)", name, len(statements))

	for i, stmt := range statements {
		if _, err := db.ExecContext(ctx, stmt); err != nil {
			return fmt.Errorf("migration %s[%d]: %w", name, i, err)
		}
	}
	logger.Infof("[db] migration %s completed", name)
	return nil
}

// SplitSQL splits a multi-statement SQL string on semicolons.
func SplitSQL(sql string) []string {
	parts := strings.Split(sql, ";")
	result := make([]string, 0, len(parts))
	for _, p := range parts {
		p = strings.TrimSpace(p)
		if p != "" && !strings.HasPrefix(p, "--") {
			result = append(result, p)
		}
	}
	return result
}