boc/landvex-admin-backend/app/db/session.py

"""
LandveX Admin Backend — Database Session Management
Async SQLAlchemy session handling with tenant context
"""
from contextlib import asynccontextmanager
from typing import AsyncGenerator, Optional

from sqlalchemy.ext.asyncio import (
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlalchemy import text

from app.core.config import settings
from app.core.exceptions import TenantSchemaException


# Create async engine
engine = create_async_engine(
    settings.DATABASE_URL,
    pool_size=settings.DATABASE_POOL_SIZE,
    max_overflow=settings.DATABASE_MAX_OVERFLOW,
    pool_pre_ping=True,
    echo=settings.DEBUG,
)

# Create async session factory
AsyncSessionLocal = async_sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autoflush=False,
    autocommit=False,
)


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """Dependency to get database session."""
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()


@asynccontextmanager
async def get_db_context() -> AsyncGenerator[AsyncSession, None]:
    """Context manager for database sessions."""
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()


class TenantContext:
    """Context manager for tenant schema isolation.
    
    Uses PostgreSQL row-level security and schema search path
    to ensure data isolation between tenants.
    """
    
    def __init__(self, session: AsyncSession, tenant_schema: Optional[str] = None):
        self.session = session
        self.tenant_schema = tenant_schema or settings.DEFAULT_TENANT_SCHEMA
        self._original_search_path = None
    
    async def __aenter__(self):
        """Set tenant context on database connection."""
        # Set search path to tenant schema first, then public
        await self.session.execute(
            text(f'SET search_path TO "{self.tenant_schema}", public')
        )
        # Set application context for RLS policies
        await self.session.execute(
            text("SELECT set_config('app.current_tenant', :tenant, true)"),
            {"tenant": self.tenant_schema}
        )
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Reset tenant context."""
        # Reset search path
        await self.session.execute(text('SET search_path TO public'))
        # Clear application context
        await self.session.execute(
            text("SELECT set_config('app.current_tenant', '', true)")
        )
        return False


async def set_tenant_schema(session: AsyncSession, tenant_schema: str) -> None:
    """Set the tenant schema for the current session."""
    try:
        await session.execute(
            text(f'SET search_path TO "{tenant_schema}", public')
        )
        await session.execute(
            text("SELECT set_config('app.current_tenant', :tenant, true)"),
            {"tenant": tenant_schema}
        )
    except Exception as e:
        raise TenantSchemaException(
            tenant_id=tenant_schema,
            operation="set_schema",
            original_error=str(e)
        )


async def reset_tenant_schema(session: AsyncSession) -> None:
    """Reset schema to default."""
    await session.execute(text('SET search_path TO public'))
    await session.execute(
        text("SELECT set_config('app.current_tenant', '', true)")
    )


async def create_tenant_schema(schema_name: str) -> None:
    """Create a new schema for a tenant."""
    async with engine.begin() as conn:
        # Create schema if not exists
        await conn.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{schema_name}"'))
        # Grant usage
        await conn.execute(text(f'GRANT USAGE ON SCHEMA "{schema_name}" TO PUBLIC'))


async def drop_tenant_schema(schema_name: str) -> None:
    """Drop a tenant schema."""
    async with engine.begin() as conn:
        await conn.execute(text(f'DROP SCHEMA IF EXISTS "{schema_name}" CASCADE'))


async def schema_exists(schema_name: str) -> bool:
    """Check if a schema exists."""
    async with engine.connect() as conn:
        result = await conn.execute(
            text(
                "SELECT EXISTS(SELECT 1 FROM information_schema.schemata WHERE schema_name = :schema)"
            ),
            {"schema": schema_name}
        )
        return result.scalar()