feat: implement SQLite-based append-only multi-version state persistence for thread checkpoints

This commit is contained in:
unclesam committed 2026-08-12 13:14:18 +08:00
1 parent 7028571a4f
commit 3a4c1e9c84
6 files changed
+803 -10

No files matched your search

+80 -5
View File
@@ -26,7 +26,7 @@ Go 言語(ジェネリクス)で書かれた、軽量・高性能・型安
- **死循环硬性熔断:** 通过 `Engine.WithMaxSteps` 设置步数上限,图中的任何意外环路都会被明确报错而非无限阻塞。 - **死循环硬性熔断:** 通过 `Engine.WithMaxSteps` 设置步数上限,图中的任何意外环路都会被明确报错而非无限阻塞。
- **轻量级生命周期 Hooks(支持多重链式累加):** `Engine.WithPreNodeHook` / `WithPostNodeHook` 支持多次调用链式累加(不覆盖已有钩子),无需修改 `NodeFn`,极易解耦日志、链路追踪(OpenTelemetry)或流式输出等逻辑。 - **轻量级生命周期 Hooks(支持多重链式累加):** `Engine.WithPreNodeHook` / `WithPostNodeHook` 支持多次调用链式累加(不覆盖已有钩子),无需修改 `NodeFn`,极易解耦日志、链路追踪(OpenTelemetry)或流式输出等逻辑。
- **编译期结构校验:** 图在编译 (`Compile`) 时,不仅会验证所有节点和边是否存在,还会严格检测同一节点是否冲突地定义了多种出边类型(如同时定义静态边、条件边与并发边),规避多义性逻辑隐患。 - **编译期结构校验:** 图在编译 (`Compile`) 时,不仅会验证所有节点和边是否存在,还会严格检测同一节点是否冲突地定义了多种出边类型(如同时定义静态边、条件边与并发边),规避多义性逻辑隐患。
- **开箱即用持久化与安全校验:** 提供 `Checkpointer[S]` 接口与内置的 `FileCheckpointer[S]`,采用“先写临时文件再原子重命名”防损坏机制,并对 `threadID` 进行路径穿越防御性校验,实现进程意外退出后进度安全无损地恢复。 - **开箱即用持久化与安全校验:** 提供 `Checkpointer[S]` 接口,内置基于 JSON 文件的 `FileCheckpointer[S]`(原子重命名防损坏)以及基于 SQLite 的 `SQLiteCheckpointer[S]`(支持多版本追加写入 Append-Only、历史版本回溯与剪枝),并对 `threadID` 进行路径穿越防御性校验,实现进程意外退出后进度安全无损地恢复。
- **纯粹的高性能:** 基于内存流转,零外部依赖,极低的上下文切换开销。 - **纯粹的高性能:** 基于内存流转,零外部依赖,极低的上下文切换开销。
### 快速开始 ### 快速开始
@@ -123,9 +123,11 @@ engine.Start(ctx, "start", MyState{})
close(streamCh) close(streamCh)
``` ```
#### 5. 文件持久化 (Checkpointer) 示例 #### 5. 状态持久化 (Checkpointer) 示例
##### A. 基于 JSON 本地文件的 Checkpointer
```go ```go
// 创建一个文件存储器,指定存放状态文件的目录 // 创建文件存储器,指定存放状态文件的目录
fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints") fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints")
// 在遇到中断挂起时,将 thread 保存到磁盘 // 在遇到中断挂起时,将 thread 保存到磁盘
@@ -142,6 +144,29 @@ if thread.IsFinished {
} }
``` ```
##### B. 基于 SQLite 的多版本追加写入 (Append-Only) Checkpointer
```go
// 创建 SQLite 存储器,自动开启 WAL 高并发模式
sc, _ := GopherGraph.NewSQLiteCheckpointer[MyState]("./checkpoints.db")
defer sc.Close()
// 每次 Save 都会追加写入一个新的版本号 (v1, v2, v3...)
sessionID := "user-session-123"
sc.Save(ctx, sessionID, thread)
// Load 默认获取最新版本的线程快照
latestThread, _ := sc.Load(ctx, sessionID)
// LoadVersion 支持精确回溯并恢复指定历史版本的快照(版本重放/回滚)
v1Thread, _ := sc.LoadVersion(ctx, sessionID, 1)
// ListVersions 查询该 session 的所有历史版本元数据
versions, _ := sc.ListVersions(ctx, sessionID)
// PruneVersions 保留最近 N 个版本,自动清理过期历史
sc.PruneVersions(ctx, sessionID, 5)
```
#### 运行内置示例 #### 运行内置示例
```bash ```bash
# 完整的"AI翻译 -> 质量检测 -> 人工审核 -> 发布"交互式 Demo # 完整的"AI翻译 -> 质量检测 -> 人工审核 -> 发布"交互式 Demo
@@ -173,7 +198,7 @@ go test -v -race ./...
- **Infinite Loop Circuit Breaker:** Set a hard step limit with `Engine.WithMaxSteps`. Any accidental cycle in the graph returns a clear error instead of blocking forever. - **Infinite Loop Circuit Breaker:** Set a hard step limit with `Engine.WithMaxSteps`. Any accidental cycle in the graph returns a clear error instead of blocking forever.
- **Lightweight Lifecycle Hooks (Additive Chaining):** `Engine.WithPreNodeHook` / `WithPostNodeHook` fire before and after each node without touching any `NodeFn` signature. Supports multiple chained registrations (additive) instead of overwriting, making it easy to decouple concerns like logging, tracing (OpenTelemetry), and streaming. - **Lightweight Lifecycle Hooks (Additive Chaining):** `Engine.WithPreNodeHook` / `WithPostNodeHook` fire before and after each node without touching any `NodeFn` signature. Supports multiple chained registrations (additive) instead of overwriting, making it easy to decouple concerns like logging, tracing (OpenTelemetry), and streaming.
- **Compile-Time Structural Validation:** During `Compile()`, the engine validates node/edge existence and strictly checks for outgoing edge conflicts on each node (e.g., preventing a node from simultaneously defining static, conditional, and parallel edges) to ensure unambiguous flow execution. - **Compile-Time Structural Validation:** During `Compile()`, the engine validates node/edge existence and strictly checks for outgoing edge conflicts on each node (e.g., preventing a node from simultaneously defining static, conditional, and parallel edges) to ensure unambiguous flow execution.
- **State Persistence & Path Safety:** Generic `Checkpointer[S]` interface with built-in `FileCheckpointer[S]`. Uses an atomic write (temp-file write followed by rename) to prevent corruption, and includes path traversal validation on `threadID`. - **State Persistence & Path Safety:** Generic `Checkpointer[S]` interface with built-in `FileCheckpointer[S]` (atomic file rename) and `SQLiteCheckpointer[S]` (append-only multi-versioning, time travel history & pruning), featuring path traversal validation on `threadID`.
- **Zero-Dependency & High-Performance:** Written in pure Go with zero external dependencies, leveraging in-memory queues for lightning-fast orchestration. - **Zero-Dependency & High-Performance:** Written in pure Go with zero external dependencies, leveraging in-memory queues for lightning-fast orchestration.
### Quick Start ### Quick Start
@@ -270,6 +295,8 @@ close(streamCh)
``` ```
#### 5. State Persistence (Checkpointer) Example #### 5. State Persistence (Checkpointer) Example
##### A. JSON File Checkpointer
```go ```go
// Initialize a local directory storage // Initialize a local directory storage
fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints") fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints")
@@ -288,6 +315,29 @@ if thread.IsFinished {
} }
``` ```
##### B. SQLite Append-Only Multi-Version Checkpointer
```go
// Initialize SQLite checkpointer (WAL mode enabled automatically)
sc, _ := GopherGraph.NewSQLiteCheckpointer[MyState]("./checkpoints.db")
defer sc.Close()
// Append a new version snapshot on every Save (v1, v2, v3...)
sessionID := "user-session-123"
sc.Save(ctx, sessionID, thread)
// Load gets the latest version by default
latestThread, _ := sc.Load(ctx, sessionID)
// LoadVersion fetches a specific historical version for replay/rollback
v1Thread, _ := sc.LoadVersion(ctx, sessionID, 1)
// ListVersions retrieves version metadata for the thread
versions, _ := sc.ListVersions(ctx, sessionID)
// PruneVersions retains the N most recent versions and purges older records
sc.PruneVersions(ctx, sessionID, 5)
```
#### Run the Interactive Demo #### Run the Interactive Demo
```bash ```bash
# Interactive "Translation -> Review -> Human Approval -> Publish" demo # Interactive "Translation -> Review -> Human Approval -> Publish" demo
@@ -319,7 +369,7 @@ go test -v -race ./...
- **無限ループのサーキットブレーカー:** `Engine.WithMaxSteps` でステップ数の上限を設定。グラフ内に意図しないサイクルが発生しても、永久ブロックせず明確なエラーを返します。 - **無限ループのサーキットブレーカー:** `Engine.WithMaxSteps` でステップ数の上限を設定。グラフ内に意図しないサイクルが発生しても、永久ブロックせず明確なエラーを返します。
- **軽量なライフサイクルフック(複数チェーン登録対応):** `Engine.WithPreNodeHook` / `WithPostNodeHook` により、`NodeFn` のシグネチャを変更することなくロギング、トレーシング(OpenTelemetry)、ストリーミング出力を実装できます。**複数回呼び出すことで複数のフックを順番に実行(チェーン蓄積)**可能です。 - **軽量なライフサイクルフック(複数チェーン登録対応):** `Engine.WithPreNodeHook` / `WithPostNodeHook` により、`NodeFn` のシグネチャを変更することなくロギング、トレーシング(OpenTelemetry)、ストリーミング出力を実装できます。**複数回呼び出すことで複数のフックを順番に実行(チェーン蓄積)**可能です。
- **コンパイル時の構造検証:** `Compile()` 実行時に、すべてのノードや辺の存在をチェックするだけでなく、同一ノードに対する出辺タイプ(静的辺、条件付き辺、並行辺)の重複競合も厳密に検出し、意図しない曖昧な挙動を防ぎます。 - **コンパイル時の構造検証:** `Compile()` 実行時に、すべてのノードや辺の存在をチェックするだけでなく、同一ノードに対する出辺タイプ(静的辺、条件付き辺、並行辺)の重複競合も厳密に検出し、意図しない曖昧な挙動を防ぎます。
- **状態の永続化とパス安全性:** 抽象的な `Checkpointer[S]` インタフェースと、書き込み時の破損を防ぐアトミック書き込み(一時ファイルからリネーム)に対応した `FileCheckpointer[S]` を内蔵。`threadID` のパス・トラバーサル検証によりセキュアに状態の復旧が可能です。 - **状態の永続化とパス安全性:** 抽象的な `Checkpointer[S]` インタフェースと、アトミック書き込み対応の `FileCheckpointer[S]`、およびマルチバージョン追記(Append-Only)・履歴復元に対応した `SQLiteCheckpointer[S]` を内蔵。`threadID` のパス・トラバーサル検証によりセキュアに状態の復旧が可能です。
- **ピュア Go & 高性能:** 外部依存関係ゼロ。メモリ内チャネルを用いた高速なコンテキスト切り替え。 - **ピュア Go & 高性能:** 外部依存関係ゼロ。メモリ内チャネルを用いた高速なコンテキスト切り替え。
### クイックスタート ### クイックスタート
@@ -416,6 +466,8 @@ close(streamCh)
``` ```
#### 5. 状態の永続化(Checkpointer)のコード例 #### 5. 状態の永続化(Checkpointer)のコード例
##### A. JSON ファイル Checkpointer
```go ```go
// 保存先ディレクトリを指定してローカルファイルチェッカーを初期化 // 保存先ディレクトリを指定してローカルファイルチェッカーを初期化
fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints") fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints")
@@ -434,6 +486,29 @@ if thread.IsFinished {
} }
``` ```
##### B. SQLite マルチバージョン追記 (Append-Only) Checkpointer
```go
// SQLite チェッカーを初期化(WAL 高並行モードが自動有効化)
sc, _ := GopherGraph.NewSQLiteCheckpointer[MyState]("./checkpoints.db")
defer sc.Close()
// Save のたびに新しいバージョン(v1, v2, v3...)を追記保存
sessionID := "user-session-123"
sc.Save(ctx, sessionID, thread)
// Load はデフォルトで最新バージョンのスナップショットを取得
latestThread, _ := sc.Load(ctx, sessionID)
// LoadVersion は特定バージョンを正確に復元(タイムトラベル / 状態ロールバック)
v1Thread, _ := sc.LoadVersion(ctx, sessionID, 1)
// ListVersions で全履歴のメタデータ一覧を取得
versions, _ := sc.ListVersions(ctx, sessionID)
// PruneVersions で最新 N 件のみを保持し、古い履歴を削除
sc.PruneVersions(ctx, sessionID, 5)
```
#### インタラクティブデモの実行 #### インタラクティブデモの実行
```bash ```bash
# 「AI翻訳 -> 監査 -> 人間による確認 -> 公開」のインタラクティブデモ # 「AI翻訳 -> 監査 -> 人間による確認 -> 公開」のインタラクティブデモ
+1 -5
View File
@@ -11,9 +11,5 @@
// - [Engine] 增强包装器:通过 [NewEngine] 创建,在 CompiledGraph 基础上叠加 // - [Engine] 增强包装器:通过 [NewEngine] 创建,在 CompiledGraph 基础上叠加
// StateCloner(深拷贝)、MaxSteps(死循环熔断)、PreNodeHook/PostNodeHook(生命周期钩子)。 // StateCloner(深拷贝)、MaxSteps(死循环熔断)、PreNodeHook/PostNodeHook(生命周期钩子)。
// - [Thread] 执行快照:承载当前状态、下一节点、暂停/结束标记,可序列化用于断点恢复。 // - [Thread] 执行快照:承载当前状态、下一节点、暂停/结束标记,可序列化用于断点恢复。
// - [Checkpointer] 持久化接口:提供 Save/Load/Delete 用于工作流进度存储和清理,内置 [FileCheckpointer] 实现基于 JSON 文件的状态存储。 // - [Checkpointer] 持久化接口:提供 Save/Load/Delete 用于工作流进度存储和清理,内置 [FileCheckpointer] 基于 JSON 文件存储,以及 [SQLiteCheckpointer] 支持多版本追加写入(Append-Only)与版本回溯。
//
// # 零外部依赖
//
// GopherGraph 仅使用 Go 标准库,无任何第三方依赖。
package GopherGraph package GopherGraph
+13
View File
@@ -1,3 +1,16 @@
module github.com/unclesam-ly/GopherGraph module github.com/unclesam-ly/GopherGraph
go 1.26.2 go 1.26.2
require (
github.com/dustin/go-humanize v1.0.1 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/mattn/go-isatty v0.0.24 // indirect
github.com/ncruces/go-strftime v1.0.0 // indirect
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
golang.org/x/sys v0.47.0 // indirect
modernc.org/libc v1.74.4 // indirect
modernc.org/mathutil v1.7.1 // indirect
modernc.org/memory v1.11.0 // indirect
modernc.org/sqlite v1.56.0 // indirect
)
+20
View File
@@ -0,0 +1,20 @@
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/mattn/go-isatty v0.0.24 h1:tGZZoVgT/KiqK1c8ocVLeDS8BSWMRd47J3Lbz7vsReI=
github.com/mattn/go-isatty v0.0.24/go.mod h1:nMCL3Zebbrt45jsMDgnfIwz6ydEQApk5oEI3HqDio6A=
github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
modernc.org/libc v1.74.4 h1:fX1Omw4o2/1C2iRkkIsrQTasJQldLhRmuPreXLoWs9k=
modernc.org/libc v1.74.4/go.mod h1:eeQAS9W3sZeKYMFubydxJpII9ybHWshk+7or7bLG9co=
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
modernc.org/sqlite v1.56.0 h1:/D8e2RfFqoy/Zc6PuC76U28zFwmI/sYx1Kjm4yEn9e0=
modernc.org/sqlite v1.56.0/go.mod h1:yCJ2cmAaIkHQ25oXWrF8H4O1lIfPYPR26yCEDj2P3pQ=
+359
View File
@@ -0,0 +1,359 @@
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 !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)
}
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")
}
sc.mu.Lock()
defer sc.mu.Unlock()
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()
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()
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()
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()
var versions []CheckpointVersion
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()
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)
_, 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()
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 为无操作(由外部管理)。
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
}
+330
View File
@@ -0,0 +1,330 @@
package GopherGraph
import (
"context"
"database/sql"
"errors"
"path/filepath"
"sync"
"testing"
"time"
)
type SQLiteTestState struct {
Query string
Count int
}
func TestSQLiteCheckpointerBasic(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "test_checkpoints.db")
sc, err := NewSQLiteCheckpointer[SQLiteTestState](dbPath)
if err != nil {
t.Fatalf("Failed to create SQLiteCheckpointer: %v", err)
}
defer sc.Close()
ctx := context.Background()
threadID := "thread-1"
// 验证未存入时 Load 返回错误
_, err = sc.Load(ctx, threadID)
if err == nil {
t.Fatalf("Expected error when loading non-existent thread, got nil")
}
thread := &Thread[SQLiteTestState]{
State: SQLiteTestState{Query: "hello", Count: 1},
NextNode: "node1",
IsPaused: true,
IsFinished: false,
}
// 存入快照
if err := sc.Save(ctx, threadID, thread); err != nil {
t.Fatalf("Failed to save thread: %v", err)
}
// 读取快照
loaded, err := sc.Load(ctx, threadID)
if err != nil {
t.Fatalf("Failed to load thread: %v", err)
}
if loaded.State.Query != "hello" || loaded.State.Count != 1 {
t.Errorf("Loaded state mismatch, got %+v", loaded.State)
}
if loaded.NextNode != "node1" || !loaded.IsPaused || loaded.IsFinished {
t.Errorf("Loaded thread metadata mismatch, got %+v", loaded)
}
// 删除线程记录
if err := sc.Delete(ctx, threadID); err != nil {
t.Fatalf("Failed to delete thread: %v", err)
}
// 校验删除后 Load 失败
_, err = sc.Load(ctx, threadID)
if err == nil {
t.Fatalf("Expected error after delete, got nil")
}
}
func TestSQLiteCheckpointerMultiVersion(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "multiver.db")
sc, err := NewSQLiteCheckpointer[SQLiteTestState](dbPath, WithSQLitePrettyJSON())
if err != nil {
t.Fatalf("Failed to create SQLiteCheckpointer: %v", err)
}
defer sc.Close()
ctx := context.Background()
threadID := "thread-multi"
// 连续保存 3 个不同版本的线程状态
t1 := &Thread[SQLiteTestState]{State: SQLiteTestState{Query: "step1", Count: 10}, NextNode: "node1"}
t2 := &Thread[SQLiteTestState]{State: SQLiteTestState{Query: "step2", Count: 20}, NextNode: "node2"}
t3 := &Thread[SQLiteTestState]{State: SQLiteTestState{Query: "step3", Count: 30}, NextNode: "node3", IsFinished: true}
if err := sc.Save(ctx, threadID, t1); err != nil {
t.Fatalf("Save v1 failed: %v", err)
}
if err := sc.Save(ctx, threadID, t2); err != nil {
t.Fatalf("Save v2 failed: %v", err)
}
if err := sc.Save(ctx, threadID, t3); err != nil {
t.Fatalf("Save v3 failed: %v", err)
}
// 1. 验证 ListVersions
versions, err := sc.ListVersions(ctx, threadID)
if err != nil {
t.Fatalf("ListVersions failed: %v", err)
}
if len(versions) != 3 {
t.Fatalf("Expected 3 versions, got %d", len(versions))
}
for i, v := range versions {
expectedVer := int64(i + 1)
if v.Version != expectedVer {
t.Errorf("Version index %d expected version %d, got %d", i, expectedVer, v.Version)
}
}
// 2. 验证 Load 默认取最新的 v3 (并且因为 IsFinished == true,包装了 ErrAlreadyFinished)
latest, err := sc.Load(ctx, threadID)
if !errors.Is(err, ErrAlreadyFinished) {
t.Fatalf("Expected ErrAlreadyFinished when loading finished latest state, got: %v", err)
}
if latest.State.Query != "step3" || latest.State.Count != 30 {
t.Errorf("Latest state mismatch, got %+v", latest.State)
}
// 3. 验证 LoadVersion 取历史版本 v1 和 v2
v1State, err := sc.LoadVersion(ctx, threadID, 1)
if err != nil {
t.Fatalf("LoadVersion(1) failed: %v", err)
}
if v1State.State.Query != "step1" || v1State.State.Count != 10 || v1State.NextNode != "node1" {
t.Errorf("v1 state mismatch, got %+v", v1State)
}
v2State, err := sc.LoadVersion(ctx, threadID, 2)
if err != nil {
t.Fatalf("LoadVersion(2) failed: %v", err)
}
if v2State.State.Query != "step2" || v2State.State.Count != 20 || v2State.NextNode != "node2" {
t.Errorf("v2 state mismatch, got %+v", v2State)
}
// 4. 验证 LoadVersion(3) 也因为 IsFinished 报 ErrAlreadyFinished
v3State, err := sc.LoadVersion(ctx, threadID, 3)
if !errors.Is(err, ErrAlreadyFinished) {
t.Fatalf("Expected ErrAlreadyFinished for LoadVersion(3), got: %v", err)
}
if v3State.State.Query != "step3" {
t.Errorf("v3 state mismatch, got %+v", v3State)
}
}
func TestSQLiteCheckpointerPruneVersions(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "prune.db")
sc, err := NewSQLiteCheckpointer[SQLiteTestState](dbPath)
if err != nil {
t.Fatalf("Failed to create SQLiteCheckpointer: %v", err)
}
defer sc.Close()
ctx := context.Background()
threadID := "thread-prune"
// 存入 5 个版本
for i := 1; i <= 5; i++ {
tState := &Thread[SQLiteTestState]{
State: SQLiteTestState{Count: i},
NextNode: "node",
}
if err := sc.Save(ctx, threadID, tState); err != nil {
t.Fatalf("Save v%d failed: %v", i, err)
}
}
// 保留最新的 2 个版本 (v4, v5)
if err := sc.PruneVersions(ctx, threadID, 2); err != nil {
t.Fatalf("PruneVersions failed: %v", err)
}
versions, err := sc.ListVersions(ctx, threadID)
if err != nil {
t.Fatalf("ListVersions failed: %v", err)
}
if len(versions) != 2 {
t.Fatalf("Expected 2 versions after prune, got %d", len(versions))
}
if versions[0].Version != 4 || versions[1].Version != 5 {
t.Errorf("Expected remaining versions 4 and 5, got %d and %d", versions[0].Version, versions[1].Version)
}
// 验证旧版本 v1 已不可读取
_, err = sc.LoadVersion(ctx, threadID, 1)
if err == nil {
t.Fatalf("Expected error loading pruned version 1, got nil")
}
// 验证 v4 和 v5 依然能正常读取
v4, err := sc.LoadVersion(ctx, threadID, 4)
if err != nil || v4.State.Count != 4 {
t.Errorf("LoadVersion(4) failed or mismatch: %v, state: %+v", err, v4)
}
}
func TestSQLiteCheckpointerCustomTableAndExternalDB(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "custom.db")
rawDB, err := sql.Open("sqlite", dbPath)
if err != nil {
t.Fatalf("Failed to open sqlite raw DB: %v", err)
}
sc, err := NewSQLiteCheckpointerWithDB[SQLiteTestState](rawDB, WithSQLiteTableName("my_custom_checkpoints"))
if err != nil {
rawDB.Close()
t.Fatalf("Failed to create SQLiteCheckpointerWithDB: %v", err)
}
ctx := context.Background()
threadID := "thread-custom"
thread := &Thread[SQLiteTestState]{State: SQLiteTestState{Query: "custom table"}}
if err := sc.Save(ctx, threadID, thread); err != nil {
t.Fatalf("Save failed: %v", err)
}
loaded, err := sc.Load(ctx, threadID)
if err != nil {
t.Fatalf("Load failed: %v", err)
}
if loaded.State.Query != "custom table" {
t.Errorf("Loaded state mismatch, got %s", loaded.State.Query)
}
// 测试 Close 不会关闭外部传入的 rawDB (ownsDB == false)
if err := sc.Close(); err != nil {
t.Fatalf("Close failed: %v", err)
}
if err := rawDB.Ping(); err != nil {
t.Fatalf("External DB should remain open after sc.Close(): %v", err)
}
rawDB.Close()
}
func TestSQLiteCheckpointerThreadIDValidation(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "val.db")
sc, err := NewSQLiteCheckpointer[SQLiteTestState](dbPath)
if err != nil {
t.Fatalf("Failed to create SQLiteCheckpointer: %v", err)
}
defer sc.Close()
ctx := context.Background()
invalidIDs := []string{"", "../malicious", "foo/bar", "foo\\bar"}
dummyThread := &Thread[SQLiteTestState]{State: SQLiteTestState{Query: "test"}}
for _, id := range invalidIDs {
if err := sc.Save(ctx, id, dummyThread); err == nil {
t.Errorf("Expected Save error for invalid threadID %q, got nil", id)
}
if _, err := sc.Load(ctx, id); err == nil {
t.Errorf("Expected Load error for invalid threadID %q, got nil", id)
}
if err := sc.Delete(ctx, id); err == nil {
t.Errorf("Expected Delete error for invalid threadID %q, got nil", id)
}
if _, err := sc.ListVersions(ctx, id); err == nil {
t.Errorf("Expected ListVersions error for invalid threadID %q, got nil", id)
}
}
}
func TestSQLiteCheckpointerConcurrency(t *testing.T) {
tmpDir := t.TempDir()
dbPath := filepath.Join(tmpDir, "concurrent.db")
sc, err := NewSQLiteCheckpointer[SQLiteTestState](dbPath)
if err != nil {
t.Fatalf("Failed to create SQLiteCheckpointer: %v", err)
}
defer sc.Close()
ctx := context.Background()
threadID := "thread-concurrent"
var wg sync.WaitGroup
numRoutines := 10
savesPerRoutine := 5
for i := 0; i < numRoutines; i++ {
wg.Add(1)
go func(routineID int) {
defer wg.Done()
for j := 0; j < savesPerRoutine; j++ {
tState := &Thread[SQLiteTestState]{
State: SQLiteTestState{Count: routineID*100 + j},
NextNode: "step",
}
if err := sc.Save(ctx, threadID, tState); err != nil {
t.Errorf("Concurrent Save failed for routine %d step %d: %v", routineID, j, err)
}
time.Sleep(1 * time.Millisecond)
}
}(i)
}
wg.Wait()
versions, err := sc.ListVersions(ctx, threadID)
if err != nil {
t.Fatalf("ListVersions failed after concurrency test: %v", err)
}
expectedTotal := numRoutines * savesPerRoutine
if len(versions) != expectedTotal {
t.Fatalf("Expected %d total versions, got %d", expectedTotal, len(versions))
}
// 验证版本号单调递增且不重复 (1 ~ expectedTotal)
for i, v := range versions {
expectedVer := int64(i + 1)
if v.Version != expectedVer {
t.Errorf("Version mismatch at index %d, expected %d, got %d", i, expectedVer, v.Version)
}
}
}