227 lines
8.0 KiB
Go
227 lines
8.0 KiB
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 只有秒，
|
||
// 同一秒内插入的多封邮件排序就不确定 —— 「会话里最早那封」（决定联系人身份）
|
||
// 与「最后那封」（决定最新进展）都会取错行。实测同秒插 5 封，
|
||
// 按 created_at 排出来的顺序是乱的（由随机 UUID 决定）。
|
||
//
|
||
// 毫秒还不够：一次插入只要几十到几百微秒，循环里连插几封会落在同一毫秒。
|
||
// 微秒是实测确认驱动能原样扫回 time.Time 的精度（纳秒也行，但没必要）。
|
||
//
|
||
// 格式仍是 SQLite 认得的 "YYYY-MM-DD HH:MM:SS.ffffff"，因此：
|
||
// - 驱动能扫进 time.Time（列声明为 DATETIME 时）
|
||
// - 与老数据（秒精度）的文本比较依然正确：前缀相同时短的排前面，
|
||
// 而 ":31" 确实早于 ":31.767000"
|
||
sqlite.MustRegisterScalarFunction("now", 0,
|
||
func(*sqlite.FunctionContext, []driver.Value) (driver.Value, error) {
|
||
return time.Now().UTC().Format("2006-01-02 15:04:05.000000"), 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)" +
|
||
// _time_format / _timezone 决定 time.Time 参数怎么落成文本。
|
||
//
|
||
// 驱动的**默认行为是 Go 的 t.String()**，写出来是
|
||
// 2026-09-10 15:10:36.122781994 +0800 HKT m=+607182.882153215
|
||
// SQLite 的 datetime() 解析不了这种串（返回 NULL），而它跟我们注册的
|
||
// NOW()（"2006-01-02 15:04:05.000000" UTC）做的是**字符串**比较。
|
||
// 后果实测有两条，都不是「显示不好看」级别的：
|
||
// 1. 安全：`expires_at > NOW()` 比较两种格式且时区不同（+0800 vs UTC），
|
||
// 用户会话永不过期，`DELETE ... WHERE expires_at < NOW()` 删 0 行。
|
||
// 2. 功能：日历 `datetime(event_time,'-N minutes') <= NOW()` 恒为假，
|
||
// 一条提醒都发不出去。
|
||
//
|
||
// _time_format=sqlite 给 "2006-01-02 15:04:05.999999999-07:00"（驱动
|
||
// parseTimeFormats[0]，读回时原样认得）；_timezone=UTC 让偏移固定为
|
||
// +00:00，与 NOW() 同一时间轴。
|
||
//
|
||
// 不用 _time_format=datetime（"2006-01-02 15:04:05"）：它把亚秒截断，
|
||
// 会重新引入同秒插入多封邮件排序不确定的老问题 —— 「会话里最早那封」
|
||
// （决定联系人身份）与「最后那封」（决定最新进展）都会取错行。
|
||
"&_time_format=sqlite" +
|
||
"&_timezone=UTC"
|
||
}
|
||
|
||
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
|
||
}
|