leocrm/app/core/db/__init__.py

"""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,
        )
    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 for DDL operations. This engine connects as the table owner.
    Falls back to the main engine if MIGRATION_DATABASE_URL is not set.
    """
    global _migration_engine
    if _migration_engine is None:
        settings = get_settings()
        url = settings.migration_database_url or settings.database_url
        _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.
    """
    factory = get_session_factory()
    async with factory() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise


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