NixMsg/internal/store/db.go

package store

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

	_ "modernc.org/sqlite"
)

// DBFileName 是数据目录下的主库文件名。
const DBFileName = "nixmsg.db"

// DB 持有读写连接与写入队列。
type DB struct {
	Write *sql.DB
	Read  *sql.DB
	Queue *Queue
}

// Open 打开数据目录下的库：写连接、读连接池、迁移，并启动简单写入队列。
// 若 nixmsg.db 尚不存在则跳过迁移前备份。
func Open(dataDir, synchronous string) (*DB, error) {
	if err := os.MkdirAll(dataDir, 0o755); err != nil {
		return nil, fmt.Errorf("mkdir data_dir: %w", err)
	}
	dbPath := filepath.Join(dataDir, DBFileName)
	existed, err := fileExists(dbPath)
	if err != nil {
		return nil, err
	}

	write, err := OpenWriter(dataDir, synchronous)
	if err != nil {
		return nil, err
	}
	if migErr := Migrate(write, dataDir, existed); migErr != nil {
		_ = write.Close()
		return nil, migErr
	}
	read, err := OpenReader(dataDir, synchronous)
	if err != nil {
		_ = write.Close()
		return nil, err
	}
	q := NewQueue(write)
	return &DB{Write: write, Read: read, Queue: q}, nil
}

// Close 停止写入队列并关闭连接。
func (d *DB) Close() error {
	var first error
	if d.Queue != nil {
		if err := d.Queue.Close(); err != nil {
			first = err
		}
	}
	if d.Read != nil {
		if err := d.Read.Close(); err != nil && first == nil {
			first = err
		}
	}
	if d.Write != nil {
		if err := d.Write.Close(); err != nil && first == nil {
			first = err
		}
	}
	return first
}

// Checkpoint 在写 goroutine 上执行 wal_checkpoint(TRUNCATE)。
func (d *DB) Checkpoint(ctx context.Context) error {
	if d == nil || d.Queue == nil {
		return fmt.Errorf("store: not open")
	}
	return d.Queue.Checkpoint(ctx)
}

// OpenWriter 按 DEVELOPMENT 7.7 打开写连接（带 _txlock=immediate，MaxOpenConns=1）。
func OpenWriter(dataDir, synchronous string) (*sql.DB, error) {
	return openWriter(dataDir, synchronous, true)
}

// OpenExistingWriter 打开已有库的写连接，不创建数据目录；库文件不存在时返回错误。
func OpenExistingWriter(dataDir, synchronous string) (*sql.DB, error) {
	dbPath := filepath.Join(dataDir, DBFileName)
	ok, err := fileExists(dbPath)
	if err != nil {
		return nil, err
	}
	if !ok {
		abs, absErr := filepath.Abs(dbPath)
		if absErr != nil {
			abs = dbPath
		}
		return nil, fmt.Errorf("database not found: %s", abs)
	}
	return openWriter(dataDir, synchronous, false)
}

func openWriter(dataDir, synchronous string, mkdir bool) (*sql.DB, error) {
	if mkdir {
		if err := os.MkdirAll(dataDir, 0o755); err != nil {
			return nil, fmt.Errorf("mkdir data_dir: %w", err)
		}
	}
	dsn, err := buildDSN(dataDir, synchronous, true)
	if err != nil {
		return nil, err
	}
	db, err := sql.Open("sqlite", dsn)
	if err != nil {
		return nil, fmt.Errorf("open sqlite writer: %w", err)
	}
	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(1)
	if err := db.Ping(); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("ping sqlite writer: %w", err)
	}
	return db, nil
}

// OpenReader 打开读连接池（DSN 不含 _txlock=immediate）。
func OpenReader(dataDir, synchronous string) (*sql.DB, error) {
	dsn, err := buildDSN(dataDir, synchronous, false)
	if err != nil {
		return nil, err
	}
	db, err := sql.Open("sqlite", dsn)
	if err != nil {
		return nil, fmt.Errorf("open sqlite reader: %w", err)
	}
	db.SetMaxOpenConns(64)
	db.SetMaxIdleConns(16)
	db.SetConnMaxIdleTime(5 * time.Minute)
	if err := db.Ping(); err != nil {
		_ = db.Close()
		return nil, fmt.Errorf("ping sqlite reader: %w", err)
	}
	return db, nil
}

func buildDSN(dataDir, synchronous string, writer bool) (string, error) {
	sync := strings.ToUpper(strings.TrimSpace(synchronous))
	if sync == "" {
		sync = "FULL"
	}
	if sync != "FULL" && sync != "NORMAL" {
		return "", fmt.Errorf("invalid sqlite_synchronous: %s", synchronous)
	}
	dbPath := filepath.ToSlash(filepath.Join(dataDir, DBFileName))
	dsn := fmt.Sprintf(
		"file:%s?_pragma=journal_mode(WAL)&_pragma=busy_timeout(5000)&_pragma=synchronous(%s)&_pragma=foreign_keys(ON)&_pragma=secure_delete(ON)",
		dbPath,
		sync,
	)
	if writer {
		dsn += "&_txlock=immediate"
	} else {
		dsn += "&_pragma=query_only(1)"
	}
	return dsn, nil
}

func fileExists(path string) (bool, error) {
	st, err := os.Stat(path)
	if err == nil {
		if st.IsDir() {
			return false, fmt.Errorf("db path is a directory: %s", path)
		}
		return true, nil
	}
	if os.IsNotExist(err) {
		return false, nil
	}
	return false, err
}