You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
80 lines
2.3 KiB
Python
80 lines
2.3 KiB
Python
"""DB helper utilities: upsert, bulk operations, health check."""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from sqlalchemy import select, text
|
|
from sqlalchemy.dialects.postgresql import insert
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from libs.db.models import Base
|
|
|
|
|
|
async def upsert(
|
|
session: AsyncSession,
|
|
model: type[Base],
|
|
values: dict[str, Any],
|
|
index_elements: list[str],
|
|
update_columns: list[str] | None = None,
|
|
) -> None:
|
|
"""Insert or update a single row using PostgreSQL ON CONFLICT."""
|
|
stmt = insert(model).values(**values)
|
|
if update_columns:
|
|
update_dict = {col: getattr(stmt.excluded, col) for col in update_columns}
|
|
stmt = stmt.on_conflict_do_update(
|
|
index_elements=index_elements, set_=update_dict
|
|
)
|
|
else:
|
|
stmt = stmt.on_conflict_do_nothing(index_elements=index_elements)
|
|
await session.execute(stmt)
|
|
|
|
|
|
async def bulk_upsert(
|
|
session: AsyncSession,
|
|
model: type[Base],
|
|
rows: list[dict[str, Any]],
|
|
index_elements: list[str],
|
|
update_columns: list[str] | None = None,
|
|
) -> int:
|
|
"""Bulk upsert; returns number of rows processed."""
|
|
if not rows:
|
|
return 0
|
|
stmt = insert(model).values(rows)
|
|
if update_columns:
|
|
update_dict = {col: getattr(stmt.excluded, col) for col in update_columns}
|
|
stmt = stmt.on_conflict_do_update(
|
|
index_elements=index_elements, set_=update_dict
|
|
)
|
|
else:
|
|
stmt = stmt.on_conflict_do_nothing(index_elements=index_elements)
|
|
await session.execute(stmt)
|
|
return len(rows)
|
|
|
|
|
|
async def get_or_create(
|
|
session: AsyncSession,
|
|
model: type[Base],
|
|
pk_value: Any,
|
|
pk_column: str,
|
|
defaults: dict[str, Any],
|
|
) -> tuple[Any, bool]:
|
|
"""Return (instance, created). Upserts if not found."""
|
|
col = getattr(model, pk_column)
|
|
result = await session.execute(select(model).where(col == pk_value))
|
|
row = result.scalar_one_or_none()
|
|
if row is not None:
|
|
return row, False
|
|
instance = model(**{pk_column: pk_value, **defaults})
|
|
session.add(instance)
|
|
await session.flush()
|
|
return instance, True
|
|
|
|
|
|
async def health_check(session: AsyncSession) -> bool:
|
|
"""Return True if DB is reachable."""
|
|
try:
|
|
await session.execute(text("SELECT 1"))
|
|
return True
|
|
except Exception:
|
|
return False
|