leocrm-bot
7a7daf8100
- 10 models: tenants, users, user_tenants, roles, sessions, audit_log, deletion_log, notifications, password_reset_tokens, api_tokens - Session-based auth (Redis + PostgreSQL audit trail) - Multi-tenant with ORM-level filtering + PostgreSQL RLS (set_config) - RBAC with roles/permissions + field-level permissions - CSRF protection via Origin header validation - Auth rate limiting (Redis counters with TTL) - CORS with explicit origins (no wildcard) - Health endpoint (no auth required) - Notification service + audit log middleware - 29 tests, 26 ACs, all passing - Coverage: 62% (infrastructure modules pending coverage in later tasks)
135 lines
4.0 KiB
Python
135 lines
4.0 KiB
Python
"""Database engine, session management, and base model."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
from collections.abc import AsyncGenerator
|
|
from typing import Any
|
|
|
|
from sqlalchemy import event as sa_event
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncEngine,
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
)
|
|
from sqlalchemy.orm import DeclarativeBase, Mapped, declared_attr, mapped_column
|
|
from sqlalchemy import text, String, DateTime, func
|
|
from sqlalchemy.dialects.postgresql import UUID as PGUUID
|
|
import uuid
|
|
from datetime import datetime
|
|
|
|
from app.config import get_settings
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
"""Declarative base with shared columns and tenant-scoping support."""
|
|
|
|
@declared_attr.directive
|
|
def __tablename__(cls) -> str:
|
|
return cls.__name__.lower() + "s"
|
|
|
|
|
|
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):
|
|
"""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 engine and session factory
|
|
_engine: AsyncEngine | None = None
|
|
_session_factory: async_sessionmaker[AsyncSession] | None = None
|
|
|
|
|
|
def get_engine() -> AsyncEngine:
|
|
"""Get or create the global async engine."""
|
|
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."""
|
|
global _session_factory
|
|
if _session_factory is None:
|
|
_session_factory = async_sessionmaker(
|
|
bind=get_engine(),
|
|
expire_on_commit=False,
|
|
class_=AsyncSession,
|
|
)
|
|
return _session_factory
|
|
|
|
|
|
async def get_db() -> AsyncGenerator[AsyncSession, None]:
|
|
"""FastAPI dependency: yield an async database session."""
|
|
factory = get_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."""
|
|
await session.execute(
|
|
text("SELECT set_config('app.current_tenant_id', :tid, true)"),
|
|
{"tid": str(tenant_id)},
|
|
)
|
|
|
|
|
|
@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 the global engine (for shutdown)."""
|
|
global _engine, _session_factory
|
|
if _engine is not None:
|
|
await _engine.dispose()
|
|
_engine = None
|
|
_session_factory = None
|
|
|
|
|
|
def reset_engine_for_testing(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]:
|
|
"""Replace the global engine with a test engine. Returns a session factory."""
|
|
global _engine, _session_factory
|
|
_engine = engine
|
|
_session_factory = async_sessionmaker(
|
|
bind=engine, expire_on_commit=False, class_=AsyncSession
|
|
)
|
|
return _session_factory
|