sql.go

package prism

import (
	"database/sql"
	"fmt"
	"net/url"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	_ "github.com/go-sql-driver/mysql"
	_ "github.com/jackc/pgx/v5/stdlib"
	_ "modernc.org/sqlite"
)

// NewSQL creates and configures a new SQL database connection based on the provided DSN.
func NewSQL(dsn string) (*sql.DB, error) {
	var driver, addr string
	var isSQLite bool

	switch {
	case strings.HasPrefix(dsn, "mysql://"), strings.HasPrefix(dsn, "mariadb://"):
		driver = "mysql"
		addr = strings.TrimPrefix(strings.TrimPrefix(dsn, "mysql://"), "mariadb://")
	case strings.HasPrefix(dsn, "postgres://"), strings.HasPrefix(dsn, "postgresql://"):
		driver = "pgx"
		addr = dsn
	case strings.HasPrefix(dsn, "sqlite://"):
		driver = "sqlite"
		isSQLite = true
		addr = strings.TrimPrefix(dsn, "sqlite://")
	default:
		return nil, fmt.Errorf("prism: unsupported database protocol")
	}

	// Handle SQLite Directory Creation
	if isSQLite {
		if err := ensureDir(addr); err != nil {
			return nil, fmt.Errorf("prism: failed to create sqlite directory: %w", err)
		}
	}

	db, err := sql.Open(driver, addr)
	if err != nil {
		return nil, err
	}

	maxOpen := 25
	maxIdle := 10
	lifetime := 5 * time.Minute

	if u, err := url.Parse(dsn); err == nil {
		q := u.Query()

		if v, err := strconv.Atoi(q.Get("max_open")); err == nil && v > 0 {
			maxOpen = v
		}

		if v, err := strconv.Atoi(q.Get("max_idle")); err == nil && v > 0 {
			maxIdle = v
		} else if q.Get("max_open") != "" {
			maxIdle = maxOpen / 2
		}

		if v, err := time.ParseDuration(q.Get("conn_lifetime")); err == nil {
			lifetime = v
		}
	}

	db.SetMaxOpenConns(maxOpen)
	db.SetMaxIdleConns(maxIdle)
	db.SetConnMaxLifetime(lifetime)

	if isSQLite {
		_, _ = db.Exec("PRAGMA journal_mode=WAL;")
		_, _ = db.Exec("PRAGMA foreign_keys=ON;")

		if !strings.Contains(dsn, "max_open") {
			db.SetMaxOpenConns(1)
		}
	}

	if err := db.Ping(); err != nil {
		_ = db.Close()
		return nil, err
	}

	return db, nil
}

// EnsureDir ensures that the directory for the given file-based DSN exists.
func ensureDir(addr string) error {
	path := strings.TrimPrefix(addr, "file:")
	if idx := strings.Index(path, "?"); idx != -1 {
		path = path[:idx]
	}
	dir := filepath.Dir(path)
	if dir != "." && dir != "/" {
		return os.MkdirAll(dir, 0755)
	}
	return nil
}