Agent Zero
efba5ceb9c
The CircuitBreakerMiddleware was blocking all requests (503 circuit_open) because every exception in get_db() — including 401 Unauthorized, 403 Forbidden, 404 Not Found — was calling record_failure() on the DB circuit breaker. This caused the circuit to trip after 5 non-DB errors (e.g. failed login attempts during security testing). Fix: Only call record_failure() when _is_transient_db_error(exc) returns True, filtering out HTTP exceptions that are not DB-related.
378 lines
12 KiB
Python
378 lines
12 KiB
Python
"""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 get_circuit, retry_db, _is_transient_db_error
|
|
|
|
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
|