"""Database engine, session management, and base model.""" from __future__ import annotations import contextlib import uuid from collections.abc import AsyncGenerator from datetime import datetime from typing import Any # noqa: F401 from sqlalchemy import DateTime, String, func, text # noqa: F401 from sqlalchemy import event as sa_event # noqa: F401 from sqlalchemy.dialects.postgresql import UUID as PGUUID from sqlalchemy.ext.asyncio import ( AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine, ) from sqlalchemy.orm import DeclarativeBase, Mapped, declared_attr, mapped_column from app.config import get_settings class Base(DeclarativeBase): """Declarative base with shared columns and tenant-scoping support.""" @declared_attr.directive def __tablename__(self) -> str: return self.__name__.lower() + "s" class SoftDeleteMixin: """Adds deleted_at column for soft-delete support.""" deleted_at: Mapped[datetime | None] = mapped_column( DateTime(timezone=True), nullable=True, default=None ) class TimestampMixin: """Adds created_at and updated_at timestamps.""" created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), nullable=False, server_default=func.now() ) updated_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), nullable=False, server_default=func.now(), onupdate=func.now() ) class TenantMixin(TimestampMixin, SoftDeleteMixin): """Adds tenant_id column and enables ORM-level auto-filtering.""" tenant_id: Mapped[uuid.UUID] = mapped_column(PGUUID(as_uuid=True), nullable=False, index=True) # Global engines and session factories — separate per role _engine: AsyncEngine | None = None _session_factory: async_sessionmaker[AsyncSession] | None = None _auth_engine: AsyncEngine | None = None _auth_session_factory: async_sessionmaker[AsyncSession] | None = None _worker_engine: AsyncEngine | None = None _worker_session_factory: async_sessionmaker[AsyncSession] | None = None _migration_engine: AsyncEngine | None = None _migration_session_factory: async_sessionmaker[AsyncSession] | None = None def get_engine() -> AsyncEngine: """Get or create the global async engine (crm_api role).""" global _engine if _engine is None: settings = get_settings() _engine = create_async_engine( settings.database_url, pool_size=settings.db_pool_size, max_overflow=settings.db_max_overflow, echo=settings.db_echo, connect_args={ "server_settings": { "statement_timeout": "300000", # 5min — only kills truly stuck queries (deadlocks, infinite loops) }, }, ) return _engine def get_session_factory() -> async_sessionmaker[AsyncSession]: """Get or create the global session factory (crm_api role).""" global _session_factory if _session_factory is None: _session_factory = async_sessionmaker( bind=get_engine(), expire_on_commit=False, class_=AsyncSession, ) return _session_factory def get_auth_engine() -> AsyncEngine: """Get or create the auth engine (crm_auth role). Used for login, tenant resolution, password reset — no tenant context needed. Falls back to the main engine if AUTH_DATABASE_URL is not set. """ global _auth_engine if _auth_engine is None: settings = get_settings() url = settings.auth_database_url or settings.database_url _auth_engine = create_async_engine( url, pool_size=settings.db_pool_size, max_overflow=settings.db_max_overflow, echo=settings.db_echo, ) return _auth_engine def get_auth_session_factory() -> async_sessionmaker[AsyncSession]: """Get or create the auth session factory (crm_auth role).""" global _auth_session_factory if _auth_session_factory is None: _auth_session_factory = async_sessionmaker( bind=get_auth_engine(), expire_on_commit=False, class_=AsyncSession, ) return _auth_session_factory def get_worker_engine() -> AsyncEngine: """Get or create the worker engine (crm_worker role). Used by ARQ worker for background job processing. Falls back to the main engine if WORKER_DATABASE_URL is not set. """ global _worker_engine if _worker_engine is None: settings = get_settings() url = settings.worker_database_url or settings.database_url _worker_engine = create_async_engine( url, pool_size=settings.db_pool_size, max_overflow=settings.db_max_overflow, echo=settings.db_echo, ) return _worker_engine def get_worker_session_factory() -> async_sessionmaker[AsyncSession]: """Get or create the worker session factory (crm_worker role).""" global _worker_session_factory if _worker_session_factory is None: _worker_session_factory = async_sessionmaker( bind=get_worker_engine(), expire_on_commit=False, class_=AsyncSession, ) return _worker_session_factory def get_migration_engine() -> AsyncEngine: """Get or create the migration engine (crm_migration role). Used by Alembic and plugin migrations for DDL operations. This engine connects as the table owner with BYPASSRLS. Raises: RuntimeError: If MIGRATION_DATABASE_URL is not set. """ global _migration_engine if _migration_engine is None: settings = get_settings() url = settings.migration_database_url if not url: raise RuntimeError( "MIGRATION_DATABASE_URL is not set. " "Plugin migrations and Alembic require a dedicated migration " "database connection (crm_migration role). " "The application cannot start without it." ) _migration_engine = create_async_engine( url, pool_size=2, max_overflow=0, echo=settings.db_echo, ) return _migration_engine def get_migration_session_factory() -> async_sessionmaker[AsyncSession]: """Get or create the migration session factory (crm_migration role).""" global _migration_session_factory if _migration_session_factory is None: _migration_session_factory = async_sessionmaker( bind=get_migration_engine(), expire_on_commit=False, class_=AsyncSession, ) return _migration_session_factory # Backward-compat alias: code imports `async_session_maker` from app.core.db. # Behaves like the session factory — calling it returns an AsyncSession. # We use a wrapper class so `async with async_session_maker() as db:` works. class _AsyncSessionMakerWrapper: """Lazy proxy for the global async_sessionmaker.""" def __call__(self) -> AsyncSession: return get_session_factory()() def __getattr__(self, name: str) -> Any: return getattr(get_session_factory(), name) async_session_maker = _AsyncSessionMakerWrapper() async def get_db() -> AsyncGenerator[AsyncSession, None]: """FastAPI dependency: yield an async database session (crm_api role). Used for normal API requests with tenant context set via RLS. Includes retry logic for transient connection errors. """ from app.core.resilience import _is_transient_db_error, get_circuit, retry_db async def _get_session(): factory = get_session_factory() return factory() session = await retry_db(_get_session) try: yield session await session.commit() await get_circuit("db").record_success() except Exception as exc: await session.rollback() # Only record DB circuit failure for transient DB errors, not HTTP exceptions if _is_transient_db_error(exc): await get_circuit("db").record_failure() raise finally: await session.close() async def get_auth_db() -> AsyncGenerator[AsyncSession, None]: """FastAPI dependency: yield an auth database session (crm_auth role). Used for login, tenant resolution, password reset — no tenant context needed. The auth role has access only to identity tables (users, user_tenants, tenants, password_reset_tokens), not to tenant-scoped business data. """ factory = get_auth_session_factory() async with factory() as session: try: yield session await session.commit() except Exception: await session.rollback() raise async def get_worker_db() -> AsyncGenerator[AsyncSession, None]: """FastAPI dependency: yield a worker database session (crm_worker role). Used by ARQ worker for background job processing with tenant context. """ factory = get_worker_session_factory() async with factory() as session: try: yield session await session.commit() except Exception: await session.rollback() raise async def set_tenant_context(session: AsyncSession, tenant_id: uuid.UUID | str) -> None: """Set PostgreSQL session variable for RLS tenant context. Sets app.current_tenant_id (the only standard tenant context variable). The legacy app.tenant_id has been removed — all RLS policies now use app.current_tenant_id exclusively. """ tid = str(tenant_id) await session.execute( text("SELECT set_config('app.current_tenant_id', :tid, true)"), {"tid": tid}, ) async def set_user_context( session: AsyncSession, user_id: uuid.UUID | str, group_ids: list[uuid.UUID] | None = None, is_system_admin: bool = False, ) -> None: """Set PostgreSQL session variables for RLS user context. Sets: - app.current_user_id: the user's UUID - app.current_user_groups: comma-separated group UUIDs - app.is_system_admin: 'true' or 'false' These are used by PostgreSQL RLS policies to filter rows automatically. """ await session.execute( text("SELECT set_config('app.current_user_id', :uid, true)"), {"uid": str(user_id)}, ) groups_str = ",".join(str(g) for g in group_ids) if group_ids else "" await session.execute( text("SELECT set_config('app.current_user_groups', :groups, true)"), {"groups": groups_str}, ) await session.execute( text("SELECT set_config('app.is_system_admin', :admin, true)"), {"admin": "true" if is_system_admin else "false"}, ) @contextlib.asynccontextmanager async def create_db_session( tenant_id: uuid.UUID | str | None = None, ) -> AsyncGenerator[AsyncSession, None]: """Create a standalone session outside FastAPI (e.g. for tests/workers).""" factory = get_session_factory() async with factory() as session: if tenant_id is not None: await set_tenant_context(session, tenant_id) try: yield session await session.commit() except Exception: await session.rollback() raise async def close_engine() -> None: """Dispose all global engines (for shutdown).""" global _engine, _session_factory, _auth_engine, _auth_session_factory global _worker_engine, _worker_session_factory, _migration_engine, _migration_session_factory for eng in (_engine, _auth_engine, _worker_engine, _migration_engine): if eng is not None: await eng.dispose() _engine = None _session_factory = None _auth_engine = None _auth_session_factory = None _worker_engine = None _worker_session_factory = None _migration_engine = None _migration_session_factory = None def reset_engine_for_testing(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]: """Replace all global engines with a test engine. Returns a session factory. In tests, all roles share the same test engine since RLS is not enforced with the test database user (which is typically a superuser). """ global _engine, _session_factory global _auth_engine, _auth_session_factory global _worker_engine, _worker_session_factory global _migration_engine, _migration_session_factory _engine = engine _session_factory = async_sessionmaker(bind=engine, expire_on_commit=False, class_=AsyncSession) # In tests, all engines point to the same test DB _auth_engine = engine _auth_session_factory = _session_factory _worker_engine = engine _worker_session_factory = _session_factory _migration_engine = engine _migration_session_factory = _session_factory return _session_factory