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