GopherGraph/sqlite_checkpoint.go

package GopherGraph

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"regexp"
	"sync"
	"time"

	_ "modernc.org/sqlite"
)

// CheckpointVersion 保存多版本检查点的元数据信息。
type CheckpointVersion struct {
	ID         int64     `json:"id"`
	ThreadID   string    `json:"thread_id"`
	Version    int64     `json:"version"`
	NextNode   string    `json:"next_node"`
	IsPaused   bool      `json:"is_paused"`
	IsFinished bool      `json:"is_finished"`
	CreatedAt  time.Time `json:"created_at"`
}

// SQLiteCheckpointer 实现了基于 SQLite 数据库的多版本追加写入（Append-Only Multi-Version）状态持久化。
type SQLiteCheckpointer[S any] struct {
	db        *sql.DB
	tableName string
	pretty    bool
	ownsDB    bool
	mu        sync.RWMutex
}

// SQLiteCheckpointerOption 是 SQLiteCheckpointer 的可选配置项。
type SQLiteCheckpointerOption func(*sqliteCheckpointerConfig)

type sqliteCheckpointerConfig struct {
	tableName string
	pretty    bool
}

// WithSQLiteTableName 指定状态持久化存储的表名（默认 "checkpoints"）。
func WithSQLiteTableName(tableName string) SQLiteCheckpointerOption {
	return func(c *sqliteCheckpointerConfig) {
		if tableName != "" {
			c.tableName = tableName
		}
	}
}

// WithSQLitePrettyJSON 开启 JSON 缩进格式存储（适合开发调试阶段）。
func WithSQLitePrettyJSON() SQLiteCheckpointerOption {
	return func(c *sqliteCheckpointerConfig) {
		c.pretty = true
	}
}

var validTableNameRegex = regexp.MustCompile(`^[a-zA-Z0-9_]+$`)

func validateTableName(name string) error {
	if name == "" {
		return fmt.Errorf("table name must not be empty")
	}
	if !validTableNameRegex.MatchString(name) {
		return fmt.Errorf("invalid table name %q: must contain only alphanumeric characters and underscores", name)
	}
	return nil
}

// NewSQLiteCheckpointer 通过 SQLite 数据库文件路径创建一个多版本状态检查点管理器。
// 会自动打开 SQLite 连接并开启 WAL 模式以保障高并发读写性能。
func NewSQLiteCheckpointer[S any](dbPath string, opts ...SQLiteCheckpointerOption) (*SQLiteCheckpointer[S], error) {
	if dbPath == "" {
		return nil, fmt.Errorf("dbPath must not be empty")
	}

	// 连接 SQLite，配置 WAL 模式与 busy_timeout
	dsn := fmt.Sprintf("%s?_pragma=journal_mode(WAL)&_pragma=busy_timeout(5000)", dbPath)
	db, err := sql.Open("sqlite", dsn)
	if err != nil {
		return nil, fmt.Errorf("failed to open sqlite db: %w", err)
	}
	// 限制为单一共享连接：
	// 1. 确保 PRAGMA WAL / busy_timeout 对所有请求生效（不会因多连接而失效）；
	// 2. SQLite 单写者模型下，单连接消除 SQLITE_BUSY 并发写冲突。
	// 读操作在 WAL 模式下仍可与写操作并行，性能不受影响。
	db.SetMaxOpenConns(1)
	db.SetMaxIdleConns(1)

	sc, err := NewSQLiteCheckpointerWithDB[S](db, opts...)
	if err != nil {
		_ = db.Close()
		return nil, err
	}
	sc.ownsDB = true
	return sc, nil
}

// NewSQLiteCheckpointerWithDB 使用已有的 *sql.DB 实例创建一个 SQLite 检查点管理器。
func NewSQLiteCheckpointerWithDB[S any](db *sql.DB, opts ...SQLiteCheckpointerOption) (*SQLiteCheckpointer[S], error) {
	if db == nil {
		return nil, fmt.Errorf("db must not be nil")
	}

	cfg := &sqliteCheckpointerConfig{
		tableName: "checkpoints",
	}
	for _, opt := range opts {
		opt(cfg)
	}

	if err := validateTableName(cfg.tableName); err != nil {
		return nil, err
	}

	sc := &SQLiteCheckpointer[S]{
		db:        db,
		tableName: cfg.tableName,
		pretty:    cfg.pretty,
		ownsDB:    false,
	}

	if err := sc.initTable(); err != nil {
		return nil, err
	}

	return sc, nil
}

func (sc *SQLiteCheckpointer[S]) initTable() error {
	query := fmt.Sprintf(`
		CREATE TABLE IF NOT EXISTS %s (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			thread_id TEXT NOT NULL,
			version INTEGER NOT NULL,
			next_node TEXT NOT NULL,
			is_paused BOOLEAN NOT NULL DEFAULT 0,
			is_finished BOOLEAN NOT NULL DEFAULT 0,
			data TEXT NOT NULL,
			created_at DATETIME DEFAULT CURRENT_TIMESTAMP
		);
		CREATE INDEX IF NOT EXISTS idx_%s_thread_ver ON %s(thread_id, version DESC);
	`, sc.tableName, sc.tableName, sc.tableName)

	_, err := sc.db.Exec(query)
	if err != nil {
		return fmt.Errorf("failed to initialize checkpoints table %q: %w", sc.tableName, err)
	}
	return nil
}

// Save 将当前 Thread 追加保存为一个新的版本号记录（Append-Only）。
func (sc *SQLiteCheckpointer[S]) Save(ctx context.Context, threadID string, thread *Thread[S]) error {
	if err := validateThreadID(threadID); err != nil {
		return err
	}
	if thread == nil {
		return fmt.Errorf("thread must not be nil")
	}

	// Save 使用写锁（Lock）而非读锁：
	// SELECT MAX(version) + INSERT 是两步操作，必须在应用层互斥，
	// 防止并发 goroutine 读到相同的 maxVersion 并写入重复的版本号。
	sc.mu.Lock()
	defer sc.mu.Unlock()
	if sc.db == nil {
		return fmt.Errorf("SQLiteCheckpointer is already closed")
	}

	var data []byte
	var err error
	if sc.pretty {
		data, err = json.MarshalIndent(thread, "", "  ")
	} else {
		data, err = json.Marshal(thread)
	}
	if err != nil {
		return fmt.Errorf("failed to marshal thread: %w", err)
	}

	tx, err := sc.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("failed to begin tx: %w", err)
	}
	defer func() {
		_ = tx.Rollback()
	}()

	// 获取当前 threadID 的最大版本号
	var maxVersion sql.NullInt64
	queryMax := fmt.Sprintf("SELECT MAX(version) FROM %s WHERE thread_id = ?", sc.tableName)
	if err := tx.QueryRowContext(ctx, queryMax, threadID).Scan(&maxVersion); err != nil {
		return fmt.Errorf("failed to query max version: %w", err)
	}

	newVersion := int64(1)
	if maxVersion.Valid {
		newVersion = maxVersion.Int64 + 1
	}

	// 追加插入新版本记录
	insertQuery := fmt.Sprintf(`
		INSERT INTO %s (thread_id, version, next_node, is_paused, is_finished, data, created_at)
		VALUES (?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
	`, sc.tableName)

	_, err = tx.ExecContext(ctx, insertQuery, threadID, newVersion, thread.NextNode, thread.IsPaused, thread.IsFinished, string(data))
	if err != nil {
		return fmt.Errorf("failed to insert checkpoint version %d: %w", newVersion, err)
	}

	if err := tx.Commit(); err != nil {
		return fmt.Errorf("failed to commit checkpoint transaction: %w", err)
	}

	return nil
}

// Load 读取指定 threadID 的最新版本快照。
// 若最新线程状态为已完结 (IsFinished == true)，会包装 ErrAlreadyFinished 作为额外信息返回。
func (sc *SQLiteCheckpointer[S]) Load(ctx context.Context, threadID string) (*Thread[S], error) {
	if err := validateThreadID(threadID); err != nil {
		return nil, err
	}

	sc.mu.RLock()
	defer sc.mu.RUnlock()
	if sc.db == nil {
		return nil, fmt.Errorf("SQLiteCheckpointer is already closed")
	}

	query := fmt.Sprintf("SELECT data FROM %s WHERE thread_id = ? ORDER BY version DESC LIMIT 1", sc.tableName)
	var dataStr string
	err := sc.db.QueryRowContext(ctx, query, threadID).Scan(&dataStr)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, fmt.Errorf("checkpoint for thread %q not found", threadID)
		}
		return nil, fmt.Errorf("failed to load latest checkpoint: %w", err)
	}

	var thread Thread[S]
	if err := json.Unmarshal([]byte(dataStr), &thread); err != nil {
		return nil, fmt.Errorf("failed to unmarshal thread data: %w", err)
	}

	if thread.IsFinished {
		return &thread, fmt.Errorf("thread %q has already finished: %w", threadID, ErrAlreadyFinished)
	}

	return &thread, nil
}

// LoadVersion 读取指定 threadID 的特定历史版本号 (version) 的快照。
func (sc *SQLiteCheckpointer[S]) LoadVersion(ctx context.Context, threadID string, version int64) (*Thread[S], error) {
	if err := validateThreadID(threadID); err != nil {
		return nil, err
	}
	if version <= 0 {
		return nil, fmt.Errorf("version must be positive")
	}

	sc.mu.RLock()
	defer sc.mu.RUnlock()
	if sc.db == nil {
		return nil, fmt.Errorf("SQLiteCheckpointer is already closed")
	}

	query := fmt.Sprintf("SELECT data FROM %s WHERE thread_id = ? AND version = ?", sc.tableName)
	var dataStr string
	err := sc.db.QueryRowContext(ctx, query, threadID, version).Scan(&dataStr)
	if err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, fmt.Errorf("checkpoint for thread %q version %d not found", threadID, version)
		}
		return nil, fmt.Errorf("failed to load checkpoint version %d: %w", version, err)
	}

	var thread Thread[S]
	if err := json.Unmarshal([]byte(dataStr), &thread); err != nil {
		return nil, fmt.Errorf("failed to unmarshal thread data: %w", err)
	}

	if thread.IsFinished {
		return &thread, fmt.Errorf("thread %q (version %d) has already finished: %w", threadID, version, ErrAlreadyFinished)
	}

	return &thread, nil
}

// ListVersions 获取指定 threadID 的所有历史版本元数据列表（按 version 升序排列）。
func (sc *SQLiteCheckpointer[S]) ListVersions(ctx context.Context, threadID string) ([]CheckpointVersion, error) {
	if err := validateThreadID(threadID); err != nil {
		return nil, err
	}

	sc.mu.RLock()
	defer sc.mu.RUnlock()
	if sc.db == nil {
		return nil, fmt.Errorf("SQLiteCheckpointer is already closed")
	}

	query := fmt.Sprintf("SELECT id, thread_id, version, next_node, is_paused, is_finished, created_at FROM %s WHERE thread_id = ? ORDER BY version ASC", sc.tableName)
	rows, err := sc.db.QueryContext(ctx, query, threadID)
	if err != nil {
		return nil, fmt.Errorf("failed to list versions: %w", err)
	}
	defer rows.Close()

	// 初始化为空 slice 而非 nil，使调用方无需区分 nil 与空列表两种情况
	versions := make([]CheckpointVersion, 0)
	for rows.Next() {
		var v CheckpointVersion
		if err := rows.Scan(&v.ID, &v.ThreadID, &v.Version, &v.NextNode, &v.IsPaused, &v.IsFinished, &v.CreatedAt); err != nil {
			return nil, fmt.Errorf("failed to scan checkpoint version: %w", err)
		}
		versions = append(versions, v)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}

	return versions, nil
}

// PruneVersions 对指定 threadID 的历史检查点进行剪枝清理，仅保留最新的 keepN 个版本。
func (sc *SQLiteCheckpointer[S]) PruneVersions(ctx context.Context, threadID string, keepN int) error {
	if err := validateThreadID(threadID); err != nil {
		return err
	}
	if keepN <= 0 {
		return fmt.Errorf("keepN must be greater than 0")
	}

	sc.mu.Lock()
	defer sc.mu.Unlock()
	if sc.db == nil {
		return fmt.Errorf("SQLiteCheckpointer is already closed")
	}

	query := fmt.Sprintf(`
		DELETE FROM %s 
		WHERE thread_id = ? 
		  AND id NOT IN (
			SELECT id FROM %s WHERE thread_id = ? ORDER BY version DESC LIMIT ?
		  )
	`, sc.tableName, sc.tableName)

	// threadID 传两次：第一个参数对应外层 WHERE，第二个对应子查询 WHERE
	_, err := sc.db.ExecContext(ctx, query, threadID, threadID, keepN)
	if err != nil {
		return fmt.Errorf("failed to prune versions: %w", err)
	}
	return nil
}

// Delete 删除指定 threadID 对应的所有版本历史快照记录（幂等操作）。
func (sc *SQLiteCheckpointer[S]) Delete(ctx context.Context, threadID string) error {
	if err := validateThreadID(threadID); err != nil {
		return err
	}

	sc.mu.Lock()
	defer sc.mu.Unlock()
	if sc.db == nil {
		return fmt.Errorf("SQLiteCheckpointer is already closed")
	}

	query := fmt.Sprintf("DELETE FROM %s WHERE thread_id = ?", sc.tableName)
	_, err := sc.db.ExecContext(ctx, query, threadID)
	if err != nil {
		return fmt.Errorf("failed to delete checkpoints for thread %q: %w", threadID, err)
	}
	return nil
}

// Close 关闭由 NewSQLiteCheckpointer 创建的底层 SQLite 数据库连接。
// 若数据库连接是由外部创建并传给 NewSQLiteCheckpointerWithDB 的，则调用 Close 为无操作（由外部管理）。
//
// Close 使用写锁（Lock），会等待所有进行中的 Save/Load 等操作（持有 RLock）全部完成后
// 才将 sc.db 置 nil 并关闭连接，保证不会出现 use-after-close 的竞态问题。
func (sc *SQLiteCheckpointer[S]) Close() error {
	sc.mu.Lock()
	defer sc.mu.Unlock()

	if sc.ownsDB && sc.db != nil {
		err := sc.db.Close()
		sc.db = nil
		return err
	}
	return nil
}