starmen1265
279ae684ea
- Move pipeline.py, schemas.py, config.py, state.py, errors.py - Move retry.py, idempotency.py, link_parser.py - Create pipeline/__init__.py with updated imports - Update all internal imports to use kb_workflow.pipeline prefix Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
153 lines
4.6 KiB
Python
153 lines
4.6 KiB
Python
"""
|
||
WorkflowStateManager - 步骤级状态持久化管理器
|
||
|
||
提供工作流状态的原子写入、备份轮转和崩溃恢复功能。
|
||
|
||
Purpose: 为工作流提供可靠的状态管理，支持断点续传和崩溃恢复
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import tempfile
|
||
from pathlib import Path
|
||
from typing import Optional
|
||
|
||
from kb_workflow.pipeline.schemas import WorkflowState
|
||
|
||
|
||
class WorkflowStateManager:
|
||
"""工作流状态管理器
|
||
|
||
管理工作流状态的持久化，包括：
|
||
- 原子写入（使用临时文件 + os.replace）
|
||
- 备份轮转（workflow_state.json -> workflow_state.json.bak）
|
||
- 主文件损坏时从备份恢复
|
||
|
||
Attributes:
|
||
STATE_FILENAME: 状态文件名
|
||
content_dir: 内容目录路径
|
||
state_path: 状态文件完整路径
|
||
backup_path: 备份文件完整路径
|
||
"""
|
||
|
||
STATE_FILENAME = "workflow_state.json"
|
||
|
||
def __init__(self, content_dir: Path):
|
||
"""初始化状态管理器
|
||
|
||
Args:
|
||
content_dir: 内容目录路径，状态文件将保存在此目录下
|
||
"""
|
||
self.content_dir = Path(content_dir)
|
||
self.state_path = self.content_dir / self.STATE_FILENAME
|
||
self.backup_path = self.content_dir / f"{self.STATE_FILENAME}.bak"
|
||
|
||
def load(self) -> Optional[WorkflowState]:
|
||
"""加载工作流状态
|
||
|
||
加载策略：
|
||
1. 尝试加载主文件 (workflow_state.json)
|
||
2. 如果主文件损坏或不存在，尝试加载备份文件
|
||
3. 如果都不存在，返回 None
|
||
|
||
Returns:
|
||
WorkflowState 实例，如果文件不存在则返回 None
|
||
|
||
Raises:
|
||
ValueError: 如果文件格式无效且备份也无法恢复
|
||
"""
|
||
def _try_load(path: Path) -> Optional[WorkflowState]:
|
||
"""尝试从指定路径加载状态"""
|
||
if not path.exists():
|
||
return None
|
||
|
||
try:
|
||
with open(path, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
return WorkflowState.model_validate(data)
|
||
except (json.JSONDecodeError, ValueError):
|
||
return None
|
||
|
||
# 先尝试主文件
|
||
state = _try_load(self.state_path)
|
||
if state is not None:
|
||
return state
|
||
|
||
# 主文件失败，尝试备份
|
||
state = _try_load(self.backup_path)
|
||
return state
|
||
|
||
def save(self, state: WorkflowState) -> None:
|
||
"""原子保存工作流状态
|
||
|
||
保存策略（原子写入 + 备份轮转）：
|
||
1. 创建父目录（如果不存在）
|
||
2. 如果状态文件存在，将其重命名为备份文件
|
||
3. 写入临时文件
|
||
4. 原子替换（os.replace）
|
||
|
||
原子写入确保：
|
||
- 读取者永远不会看到部分写入的文件
|
||
- 崩溃时要么是旧文件，要么是新文件
|
||
|
||
Args:
|
||
state: 要保存的工作流状态
|
||
|
||
Raises:
|
||
OSError: 文件操作失败（权限、磁盘空间等）
|
||
"""
|
||
# 1. 创建父目录
|
||
self.content_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 2. 备份轮转：如果状态文件存在，移动到备份
|
||
if self.state_path.exists():
|
||
# 删除旧备份
|
||
if self.backup_path.exists():
|
||
self.backup_path.unlink()
|
||
# 将当前状态移动到备份
|
||
self.state_path.rename(self.backup_path)
|
||
|
||
# 3. 写入临时文件（在相同目录下，确保原子替换有效）
|
||
fd, temp_path = tempfile.mkstemp(dir=self.content_dir, suffix='.tmp')
|
||
try:
|
||
with os.fdopen(fd, 'w', encoding='utf-8') as f:
|
||
f.write(state.model_dump_json(indent=2))
|
||
# 4. 原子替换
|
||
os.replace(temp_path, self.state_path)
|
||
except Exception:
|
||
# 清理临时文件
|
||
if os.path.exists(temp_path):
|
||
os.unlink(temp_path)
|
||
raise
|
||
|
||
def get_or_create(
|
||
self,
|
||
content_id: str,
|
||
platform: str,
|
||
source_url: str
|
||
) -> WorkflowState:
|
||
"""获取现有状态或创建新状态
|
||
|
||
幂等操作：如果状态已存在，返回现有状态；否则创建新状态。
|
||
|
||
Args:
|
||
content_id: 内容唯一标识
|
||
platform: 平台名称
|
||
source_url: 原始链接
|
||
|
||
Returns:
|
||
WorkflowState 实例
|
||
"""
|
||
existing = self.load()
|
||
if existing is not None:
|
||
return existing
|
||
|
||
# 创建新状态
|
||
return WorkflowState(
|
||
content_id=content_id,
|
||
platform=platform,
|
||
source_url=source_url
|
||
)
|