zcbot/core/storage/utils.py

56 lines
1.6 KiB
Python

"""Storage 辅助工具:idempotent task 行创建、本地形态简化封装。"""
from __future__ import annotations
from typing import Optional
from uuid import UUID
from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert
from .engine import session_scope
from .models import SENTINEL_USER_ID, Task
def ensure_local_task_row(
task_id: UUID,
task_dir: str = "",
mode: str = "",
description: str = "",
model: str = "",
model_profile: str = "",
reasoning_effort: str = "",
user_id: UUID = SENTINEL_USER_ID,
) -> None:
"""本地形态 idempotent INSERT tasks 行。
用于 Session.append 首次写消息前打底,Step 2 阶段字段都是占位值;
Step 3 引入 TaskState ORM 后,TaskState.save 会把字段更新成真实值。
PG `INSERT ... ON CONFLICT DO NOTHING` 保证幂等且单 SQL,无 SELECT-then-INSERT
竞态。
"""
stmt = (
insert(Task)
.values(
task_id=task_id,
user_id=user_id,
task_dir=task_dir,
mode=mode,
description=description,
model=model,
model_profile=model_profile,
reasoning_effort=reasoning_effort,
)
.on_conflict_do_nothing(index_elements=["task_id"])
)
with session_scope() as s:
s.execute(stmt)
def get_task(task_id: UUID) -> Optional[Task]:
"""读 tasks 行,不存在返回 None。"""
with session_scope() as s:
return s.execute(
select(Task).where(Task.task_id == task_id)
).scalar_one_or_none()