MailUI4Agents/gateway/internal/db/db.go

// Package db 提供数据库连接与方言适配。
//
// AgentMail 默认用 SQLite（零依赖、单文件，配合 go:embed 的前端就是「一个二进制 + 一个 .db」），
// 用户显式给出 DATABASE_URL 时切换到外部 PostgreSQL。
//
// 两种方言的差异集中在本包处理，repo 层只写一份 SQL：
//   - 占位符：SQLite 也支持 $1/$2，无需改写
//   - NOW() / gen_random_uuid()：SQLite 侧注册同名函数补齐
//   - JSONB 包含判断：走 CCHas/CCArg 辅助函数（唯一必须分支的查询）
//   - 唯一约束冲突：IsUniqueViolation 统一识别
package db

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgconn"
	_ "github.com/jackc/pgx/v5/stdlib" // database/sql 驱动：pgx
	sqlite "modernc.org/sqlite"
)

type Dialect string

const (
	Postgres Dialect = "postgres"
	SQLite   Dialect = "sqlite"
)

var (
	DB *sql.DB
	// D 是当前生效的方言，repo 层据此选择 SQL 片段
	D Dialect
)

func init() {
	// SQLite 没有 NOW() 与 gen_random_uuid()，注册同名函数使 repo 层 SQL 与 PG 保持一致。
	// 函数名在 SQLite 中大小写不敏感，注册小写即可匹配 SQL 里的 NOW()。
	sqlite.MustRegisterDeterministicScalarFunction("gen_random_uuid", 0,
		func(*sqlite.FunctionContext, []driver.Value) (driver.Value, error) {
			return uuid.NewString(), nil
		})

	// NOW() 必须非确定性：同一语句内多次调用要各自取当前时刻，
	// 且格式与 SQLite 的 CURRENT_TIMESTAMP 一致，才能统一扫进 time.Time。
	sqlite.MustRegisterScalarFunction("now", 0,
		func(*sqlite.FunctionContext, []driver.Value) (driver.Value, error) {
			return time.Now().UTC().Format("2006-01-02 15:04:05"), nil
		})
}

// Connect 依据 DATABASE_URL 建立连接。空值时落到 SQLite。
func Connect(ctx context.Context, dsn string) error {
	driverName, connStr, dialect, err := resolve(dsn)
	if err != nil {
		return err
	}

	pool, err := sql.Open(driverName, connStr)
	if err != nil {
		return fmt.Errorf("open %s: %w", dialect, err)
	}

	switch dialect {
	case Postgres:
		pool.SetMaxOpenConns(20)
		pool.SetMaxIdleConns(4)
		pool.SetConnMaxLifetime(30 * time.Minute)
		pool.SetConnMaxIdleTime(5 * time.Minute)
	case SQLite:
		// SQLite 单写者：并发写靠 WAL + busy_timeout 排队，连接数放大只会加剧锁竞争。
		pool.SetMaxOpenConns(1)
		pool.SetMaxIdleConns(1)
		pool.SetConnMaxLifetime(0)
	}

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

	DB = pool
	D = dialect
	return nil
}

// resolve 把 DATABASE_URL 解析为 (驱动名, 连接串, 方言)。
func resolve(dsn string) (string, string, Dialect, error) {
	dsn = strings.TrimSpace(dsn)

	if dsn == "" {
		return "sqlite", sqliteDSN(defaultDBPath()), SQLite, nil
	}

	switch {
	case strings.HasPrefix(dsn, "postgres://"), strings.HasPrefix(dsn, "postgresql://"):
		return "pgx", dsn, Postgres, nil

	case strings.HasPrefix(dsn, "sqlite://"):
		return "sqlite", sqliteDSN(strings.TrimPrefix(dsn, "sqlite://")), SQLite, nil

	case strings.HasPrefix(dsn, "sqlite:"):
		return "sqlite", sqliteDSN(strings.TrimPrefix(dsn, "sqlite:")), SQLite, nil

	case strings.HasPrefix(dsn, "file:"):
		// 已是 SQLite URI，原样透传（调用方自带 pragma）
		return "sqlite", dsn, SQLite, nil

	case strings.HasSuffix(dsn, ".db"), strings.HasSuffix(dsn, ".sqlite"), strings.HasSuffix(dsn, ".sqlite3"):
		return "sqlite", sqliteDSN(dsn), SQLite, nil
	}

	return "", "", "", fmt.Errorf("无法识别的 DATABASE_URL %q：期望 postgres://…、sqlite:///path/x.db 或 /path/x.db", dsn)
}

// defaultDBPath 返回默认 SQLite 文件位置（AGENTMAIL_DATA_DIR 可覆盖）。
func defaultDBPath() string {
	dir := os.Getenv("AGENTMAIL_DATA_DIR")
	if dir == "" {
		dir = "data"
	}
	return filepath.Join(dir, "agentmail.db")
}

// sqliteDSN 把文件路径包装为带 pragma 的 SQLite URI，并确保父目录存在。
//
//   - journal_mode=WAL：读写不互斥，SSE 长连接查询不会被写入阻塞
//   - busy_timeout=5000：并发写时排队 5s 而不是立刻 SQLITE_BUSY
//   - foreign_keys=1：SQLite 默认不校验外键，必须显式打开
func sqliteDSN(path string) string {
	if dir := filepath.Dir(path); dir != "" && dir != "." {
		os.MkdirAll(dir, 0o755)
	}
	return "file:" + path +
		"?_pragma=journal_mode(WAL)" +
		"&_pragma=busy_timeout(5000)" +
		"&_pragma=foreign_keys(1)"
}

func Close() {
	if DB != nil {
		DB.Close()
	}
}

// ---------- 方言差异 ----------

// CCHas 返回「cc_list 是否抄送了某人」的 SQL 片段，argN 是该人名对应的占位符序号。
//
// 两个方言的实参都是【纯人名字符串】，不是 JSON 探针——因为多处查询把同一个
// 占位符同时用于 from_name/to_name 比较和抄送判断，两种实参约定必然出错。
// PG 侧在 SQL 里用 jsonb_build_* 现场构造探针；SQLite 侧用 json_each 展开逐项比对。
func CCHas(col string, argN int) string {
	if D == Postgres {
		return fmt.Sprintf("%s @> jsonb_build_array(jsonb_build_object('name', $%d::text))", col, argN)
	}
	return fmt.Sprintf(
		"EXISTS (SELECT 1 FROM json_each(%s) WHERE json_extract(value, '$.name') = $%d)",
		col, argN)
}

// JSONCast 返回把占位符转成 JSONB 的后缀（PG 需要 ::jsonb，SQLite 存 TEXT 无需转换）。
func JSONCast() string {
	if D == Postgres {
		return "::jsonb"
	}
	return ""
}

// IsUniqueViolation 判断错误是否为唯一约束冲突（用于别名撞名重试）。
func IsUniqueViolation(err error) bool {
	if err == nil {
		return false
	}
	var pgErr *pgconn.PgError
	if errors.As(err, &pgErr) {
		return pgErr.Code == "23505"
	}
	var liteErr *sqlite.Error
	if errors.As(err, &liteErr) {
		// SQLITE_CONSTRAINT_UNIQUE = 2067、SQLITE_CONSTRAINT_PRIMARYKEY = 1555
		code := liteErr.Code()
		return code == 2067 || code == 1555
	}
	return false
}