diff --git a/README.md b/README.md index 02f0735..80b0084 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,7 @@ Go 言語（ジェネリクス）で書かれた、軽量・高性能・型安 - **死循环硬性熔断：** 通过 `Engine.WithMaxSteps` 设置步数上限，图中的任何意外环路都会被明确报错而非无限阻塞。 - **轻量级生命周期 Hooks（支持多重链式累加）：** `Engine.WithPreNodeHook` / `WithPostNodeHook` 支持多次调用链式累加（不覆盖已有钩子），无需修改 `NodeFn`，极易解耦日志、链路追踪（OpenTelemetry）或流式输出等逻辑。 - **编译期结构校验：** 图在编译 (`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) ``` -#### 5. 文件持久化 (Checkpointer) 示例 +#### 5. 状态持久化 (Checkpointer) 示例 + +##### A. 基于 JSON 本地文件的 Checkpointer ```go -// 创建一个文件存储器，指定存放状态文件的目录 +// 创建文件存储器，指定存放状态文件的目录 fc, _ := GopherGraph.NewFileCheckpointer[MyState]("./checkpoints") // 在遇到中断挂起时，将 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 # 完整的"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. - **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. -- **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. ### Quick Start @@ -270,6 +295,8 @@ close(streamCh) ``` #### 5. State Persistence (Checkpointer) Example + +##### A. JSON File Checkpointer ```go // Initialize a local directory storage 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 ```bash # Interactive "Translation -> Review -> Human Approval -> Publish" demo @@ -319,7 +369,7 @@ go test -v -race ./... - **無限ループのサーキットブレーカー:** `Engine.WithMaxSteps` でステップ数の上限を設定。グラフ内に意図しないサイクルが発生しても、永久ブロックせず明確なエラーを返します。 - **軽量なライフサイクルフック（複数チェーン登録対応）:** `Engine.WithPreNodeHook` / `WithPostNodeHook` により、`NodeFn` のシグネチャを変更することなくロギング、トレーシング（OpenTelemetry）、ストリーミング出力を実装できます。**複数回呼び出すことで複数のフックを順番に実行（チェーン蓄積）**可能です。 - **コンパイル時の構造検証:** `Compile()` 実行時に、すべてのノードや辺の存在をチェックするだけでなく、同一ノードに対する出辺タイプ（静的辺、条件付き辺、並行辺）の重複競合も厳密に検出し、意図しない曖昧な挙動を防ぎます。 -- **状態の永続化とパス安全性:** 抽象的な `Checkpointer[S]` インタフェースと、書き込み時の破損を防ぐアトミック書き込み（一時ファイルからリネーム）に対応した `FileCheckpointer[S]` を内蔵。`threadID` のパス・トラバーサル検証によりセキュアに状態の復旧が可能です。 +- **状態の永続化とパス安全性:** 抽象的な `Checkpointer[S]` インタフェースと、アトミック書き込み対応の `FileCheckpointer[S]`、およびマルチバージョン追記（Append-Only）・履歴復元に対応した `SQLiteCheckpointer[S]` を内蔵。`threadID` のパス・トラバーサル検証によりセキュアに状態の復旧が可能です。 - **ピュア Go & 高性能:** 外部依存関係ゼロ。メモリ内チャネルを用いた高速なコンテキスト切り替え。 ### クイックスタート @@ -416,6 +466,8 @@ close(streamCh) ``` #### 5. 状態の永続化（Checkpointer）のコード例 + +##### A. JSON ファイル Checkpointer ```go // 保存先ディレクトリを指定してローカルファイルチェッカーを初期化 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 # 「AI翻訳 -> 監査 -> 人間による確認 -> 公開」のインタラクティブデモ diff --git a/doc.go b/doc.go index b60e43f..81144b1 100644 --- a/doc.go +++ b/doc.go @@ -11,9 +11,5 @@ // - [Engine] 增强包装器：通过 [NewEngine] 创建，在 CompiledGraph 基础上叠加 // StateCloner（深拷贝）、MaxSteps（死循环熔断）、PreNodeHook/PostNodeHook（生命周期钩子）。 // - [Thread] 执行快照：承载当前状态、下一节点、暂停/结束标记，可序列化用于断点恢复。 -// - [Checkpointer] 持久化接口：提供 Save/Load/Delete 用于工作流进度存储和清理，内置 [FileCheckpointer] 实现基于 JSON 文件的状态存储。 -// -// # 零外部依赖 -// -// GopherGraph 仅使用 Go 标准库，无任何第三方依赖。 +// - [Checkpointer] 持久化接口：提供 Save/Load/Delete 用于工作流进度存储和清理，内置 [FileCheckpointer] 基于 JSON 文件存储，以及 [SQLiteCheckpointer] 支持多版本追加写入（Append-Only）与版本回溯。 package GopherGraph diff --git a/go.mod b/go.mod index 886323a..309f712 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,16 @@ module github.com/unclesam-ly/GopherGraph 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 +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..9869c9c --- /dev/null +++ b/go.sum @@ -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= diff --git a/sqlite_checkpoint.go b/sqlite_checkpoint.go new file mode 100644 index 0000000..15b7d32 --- /dev/null +++ b/sqlite_checkpoint.go @@ -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 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 +} diff --git a/sqlite_checkpoint_test.go b/sqlite_checkpoint_test.go new file mode 100644 index 0000000..a5faa46 --- /dev/null +++ b/sqlite_checkpoint_test.go @@ -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) + } + } +}