"""Async SQLAlchemy engine, session factory, and Base model.""" from __future__ import annotations from collections.abc import AsyncGenerator from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import DeclarativeBase from app.config import settings engine = create_async_engine(settings.database_url, echo=False, pool_pre_ping=True) SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) class Base(DeclarativeBase): pass async def get_db() -> AsyncGenerator[AsyncSession, None]: async with SessionLocal() as session: yield session async def init_db() -> None: """Create tables and apply lightweight dev migrations.""" from app import models # noqa: F401 from sqlalchemy import text async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) # Dev migration: locations → vessels (rename legacy table/columns) await conn.execute( text( """ DO $$ BEGIN IF EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'locations') AND NOT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'vessels') THEN ALTER TABLE locations RENAME TO vessels; END IF; IF EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'devices' AND column_name = 'location_id') THEN ALTER TABLE devices RENAME COLUMN location_id TO vessel_id; END IF; IF EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'tasks' AND column_name = 'location_id') THEN ALTER TABLE tasks RENAME COLUMN location_id TO vessel_id; END IF; IF NOT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'devices' AND column_name = 'catalog_key') THEN ALTER TABLE devices ADD COLUMN catalog_key VARCHAR(64); END IF; END $$; """ ) ) # Indexes for hot query paths (idempotent; covers pre-existing DBs). for stmt in ( "CREATE INDEX IF NOT EXISTS ix_tasks_vessel_id ON tasks (vessel_id)", "CREATE INDEX IF NOT EXISTS ix_task_events_task_id_id ON task_events (task_id, id)", "CREATE INDEX IF NOT EXISTS ix_approval_requests_task_status " "ON approval_requests (task_id, status)", ): await conn.execute(text(stmt))