T01: core infrastructure + auth + multi-tenant + RLS
- 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)
This commit is contained in:
@@ -0,0 +1,134 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user