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 }