PMFlow/backend/app/utils/sequence.py

38 lines
1.2 KiB
Python

from __future__ import annotations
from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession
async def _ensure_table_id_sequence(db: AsyncSession, *, table_name: str) -> None:
bind = db.get_bind()
dialect_name = str(getattr(getattr(bind, "dialect", None), "name", "")).lower()
if dialect_name != "postgresql":
return
await db.execute(
text(
f"""
SELECT setval(
pg_get_serial_sequence('{table_name}', 'id'),
GREATEST(
COALESCE((SELECT MAX(id) FROM {table_name}), 0),
COALESCE((SELECT last_value FROM {table_name}_id_seq), 0),
1
),
true
)
"""
)
)
async def ensure_project_id_sequence(db: AsyncSession) -> None:
"""Keep the PostgreSQL `projects.id` sequence ahead of existing rows."""
await _ensure_table_id_sequence(db, table_name="projects")
async def ensure_signal_id_sequence(db: AsyncSession) -> None:
"""Keep the PostgreSQL `signals.id` sequence ahead of existing rows."""
await _ensure_table_id_sequence(db, table_name="signals")