62 lines
2.7 KiB
Python
62 lines
2.7 KiB
Python
"""统一编号服务 —— 原子取号，替代各处分散扫描式生成（§3.3 编号约定）。
|
||
|
||
取号基于 bre_sequence（name/current_seq）表 + PostgreSQL 事务级 advisory lock：
|
||
pg_advisory_xact_lock 在事务内串行化同名取号，锁随事务提交/回滚自动释放，
|
||
与 CRUD 同事务提交 ⇒ 并发请求不重复、失败自动回滚不留空号。
|
||
|
||
存量数据迁移：首次使用某序列名时由 seed_fn（服务端传入）从既有记录推算
|
||
起始号（原扫描逻辑），此后全部走原子递增，不再扫描。
|
||
"""
|
||
from typing import Any, Awaitable, Callable
|
||
|
||
from sqlalchemy import select, text
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.api.v1.module_bre.sequence.model import SequenceModel
|
||
|
||
|
||
class NumberGenService:
|
||
"""原子序号发生器。实例化于 HTTP/服务层，随当前事务取号。"""
|
||
|
||
def __init__(self, db: AsyncSession) -> None:
|
||
self.db = db
|
||
|
||
@staticmethod
|
||
async def _acquire_lock(db: AsyncSession, name: str) -> None:
|
||
from app.config.setting import settings
|
||
|
||
if settings.DATABASE_TYPE != "postgres":
|
||
return # sqlite 测试库无 advisory lock，退化为普通自增（单进程无并发竞争）
|
||
# hashtextextended 对文本做稳定 64 位哈希，同名序列恒映射同一把锁
|
||
await db.execute(text("SELECT pg_advisory_xact_lock(hashtextextended(:n, 0))"), {"n": name})
|
||
|
||
async def next_seq(
|
||
self,
|
||
name: str,
|
||
*,
|
||
start: int = 1,
|
||
seed_fn: Callable[[], Awaitable[int | None]] | None = None,
|
||
) -> int:
|
||
"""原子取号。返回本序列名下一个应使用的序号。
|
||
|
||
name: 序列名（如 "combination:2026" / "tree:3"），按编号维度天然分片；
|
||
start: 新序列的默认起始值；
|
||
seed_fn: 仅当序列行尚不存在时调用，从存量数据推算起始号（兼容既有数据迁移）。
|
||
"""
|
||
await self._acquire_lock(self.db, name)
|
||
row = await self.db.scalar(select(SequenceModel).where(SequenceModel.name == name))
|
||
if row is None:
|
||
init = await seed_fn() if seed_fn is not None else start
|
||
if init is None or init < 1:
|
||
init = start
|
||
self.db.add(SequenceModel(name=name, current_seq=init))
|
||
await self.db.flush()
|
||
return init
|
||
row.current_seq += 1
|
||
await self.db.flush()
|
||
return int(row.current_seq)
|
||
|
||
async def exists(self, name: str) -> bool:
|
||
"""序列行是否已初始化（供服务端判断是否需要 seed_fn 兜底）。"""
|
||
return (await self.db.scalar(select(SequenceModel.id).where(SequenceModel.name == name))) is not None
|