75 lines
2.8 KiB
Python
75 lines
2.8 KiB
Python
"""Database resources owned by the standalone calculation worker."""
|
|
from __future__ import annotations
|
|
|
|
from contextlib import asynccontextmanager
|
|
from uuid import UUID
|
|
|
|
from sqlalchemy import text
|
|
|
|
|
|
@asynccontextmanager
|
|
async def queue_repository(project_id):
|
|
"""Resolve the current project route and open its calculation queue."""
|
|
from app.core.config import settings
|
|
from app.infra.db.dynamic_manager import project_connection_manager
|
|
from app.infra.db.metadb.database import SessionLocal
|
|
from app.infra.db.metadb.repositories.metadata_repository import MetadataRepository
|
|
from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
|
|
|
|
project_uuid = UUID(str(project_id))
|
|
async with SessionLocal() as session:
|
|
metadata = MetadataRepository(session)
|
|
project = await metadata.get_project_by_id(project_uuid)
|
|
if not project or project.status != "active":
|
|
raise ValueError("Calculation project is not active")
|
|
routing = await metadata.get_project_db_routing(project_uuid, "biz_data")
|
|
if not routing or routing.db_type != "postgresql":
|
|
raise ValueError("Project business PostgreSQL routing is unavailable")
|
|
await session.commit()
|
|
pool_min_size = (
|
|
routing.pool_min_size
|
|
if routing.pool_min_size is not None
|
|
else settings.PROJECT_PG_POOL_MIN_SIZE
|
|
)
|
|
pool_max_size = (
|
|
routing.pool_max_size
|
|
if routing.pool_max_size is not None
|
|
else settings.PROJECT_PG_POOL_SIZE
|
|
)
|
|
async with project_connection_manager.pg_connection(
|
|
project_uuid,
|
|
"biz_data",
|
|
routing.dsn,
|
|
pool_min_size,
|
|
pool_max_size,
|
|
) as connection:
|
|
yield CalculationJobRepository(connection, project_id)
|
|
|
|
|
|
class CalculationSupervisorLease:
|
|
"""Keep the singleton advisory lock alive on one metadata connection."""
|
|
|
|
def __init__(self, connection):
|
|
self.connection = connection
|
|
|
|
async def ping(self) -> None:
|
|
await self.connection.execute(text("SELECT 1"))
|
|
await self.connection.commit()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def calculation_supervisor_lease():
|
|
"""Acquire and release the session-level calculation supervisor lock."""
|
|
from app.infra.db.metadb.database import engine
|
|
|
|
async with engine.connect() as connection:
|
|
acquired = await connection.scalar(text("SELECT pg_try_advisory_lock(1748531781, 2)"))
|
|
await connection.commit()
|
|
if not acquired:
|
|
raise RuntimeError("Another calculation supervisor is already active")
|
|
try:
|
|
yield CalculationSupervisorLease(connection)
|
|
finally:
|
|
await connection.execute(text("SELECT pg_advisory_unlock(1748531781, 2)"))
|
|
await connection.commit()
|