36 lines
1.1 KiB
Python
36 lines
1.1 KiB
Python
"""同一 task 内消息序号的统一事务分配入口。"""
|
||
from __future__ import annotations
|
||
|
||
from typing import Optional
|
||
from uuid import UUID
|
||
|
||
from sqlalchemy import func, select
|
||
from sqlalchemy.orm import Session as OrmSession
|
||
|
||
from .models import Message, Task
|
||
|
||
|
||
def allocate_message_idx(
|
||
session: OrmSession,
|
||
task_id: UUID,
|
||
*,
|
||
locked_task: Optional[Task] = None,
|
||
) -> int:
|
||
"""锁定 task 行并返回唯一递增 idx；max 校准兼容蓝绿旧实例。"""
|
||
task = locked_task
|
||
if task is None:
|
||
task = session.execute(
|
||
select(Task).where(Task.task_id == task_id).with_for_update()
|
||
).scalar_one()
|
||
elif task.task_id != task_id:
|
||
raise ValueError("locked_task 与 task_id 不一致")
|
||
|
||
max_idx = session.execute(
|
||
select(func.max(Message.idx)).where(Message.task_id == task_id)
|
||
).scalar_one()
|
||
persisted_next = int(getattr(task, "next_message_idx", 0) or 0)
|
||
observed_next = -1 if max_idx is None else int(max_idx)
|
||
next_idx = max(persisted_next, observed_next + 1)
|
||
task.next_message_idx = next_idx + 1
|
||
return next_idx
|