388 lines
15 KiB
Python
388 lines
15 KiB
Python
import contextvars
|
|
import logging
|
|
from typing import Any
|
|
|
|
import sentry_sdk
|
|
from sqlalchemy import MetaData, event, text
|
|
from sqlalchemy.exc import OperationalError
|
|
from sqlalchemy.exc import TimeoutError as SQLAlchemyTimeoutError
|
|
from sqlalchemy.ext.asyncio import (
|
|
AsyncSession,
|
|
async_sessionmaker,
|
|
create_async_engine,
|
|
)
|
|
from sqlalchemy.orm import declarative_base
|
|
from sqlalchemy.pool import NullPool, QueuePool
|
|
from tenacity import (
|
|
AsyncRetrying,
|
|
retry_if_exception_type,
|
|
stop_after_delay,
|
|
wait_exponential_jitter,
|
|
)
|
|
|
|
from src.config import settings
|
|
from src.telemetry.prometheus.metrics import (
|
|
db_queries_in_flight_gauge,
|
|
prometheus_metrics,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
connect_args = {"prepare_threshold": None}
|
|
|
|
# Context variable to store request context
|
|
request_context: contextvars.ContextVar[str | None] = contextvars.ContextVar(
|
|
"request_context", default=None
|
|
)
|
|
|
|
engine_kwargs = {}
|
|
|
|
if settings.DB.POOL_CLASS == "null":
|
|
engine_kwargs["poolclass"] = NullPool
|
|
else:
|
|
# Only add pool-related kwargs for pooled connections
|
|
engine_kwargs.update( # pyright: ignore
|
|
{
|
|
"pool_pre_ping": settings.DB.POOL_PRE_PING,
|
|
"pool_size": settings.DB.POOL_SIZE,
|
|
"max_overflow": settings.DB.MAX_OVERFLOW,
|
|
"pool_timeout": settings.DB.POOL_TIMEOUT,
|
|
"pool_recycle": settings.DB.POOL_RECYCLE,
|
|
"pool_use_lifo": settings.DB.POOL_USE_LIFO,
|
|
}
|
|
)
|
|
|
|
engine = create_async_engine(
|
|
settings.DB.CONNECTION_URI,
|
|
connect_args=connect_args,
|
|
echo=settings.DB.SQL_DEBUG,
|
|
**engine_kwargs,
|
|
)
|
|
|
|
# NOTE: SessionLocal is defined further down, after HonchoAsyncSession (its
|
|
# session class) and acquire_connection_with_retry (which that class calls).
|
|
|
|
# Errors worth retrying when acquiring a pooled connection: SQLAlchemy's local
|
|
# pool-checkout timeout, and OperationalError (how a saturated transaction
|
|
# pooler surfaces "too many clients" / connection refusals).
|
|
RETRYABLE_DB_CONNECTION_ERRORS = (SQLAlchemyTimeoutError, OperationalError)
|
|
|
|
# Identifies this process ("api" | "deriver") on DB metrics. Set once at startup
|
|
# by register_db_query_instrumentation; stays "unknown" if metrics are disabled.
|
|
_db_instance_type: str = "unknown"
|
|
|
|
|
|
def _record_acquisition_outcome(outcome: str) -> None:
|
|
"""Record a connection-acquisition outcome (no-op when metrics disabled)."""
|
|
if settings.METRICS.ENABLED:
|
|
prometheus_metrics.record_db_connection_acquisition(
|
|
instance_type=_db_instance_type, outcome=outcome
|
|
)
|
|
|
|
|
|
def get_pool_stats() -> dict[str, int]:
|
|
"""Return live connection-pool stats for this process.
|
|
|
|
``engine.pool`` is the AsyncEngine's pool (the same object as
|
|
``engine.sync_engine.pool``); its stat methods are synchronous counter
|
|
reads with no I/O, so they are safe to call without ``await``. Returns
|
|
zeros for pools that do not track connections (e.g. ``NullPool``).
|
|
"""
|
|
zeros = {"checked_out": 0, "checked_in": 0, "size": 0, "overflow": 0}
|
|
pool = engine.pool
|
|
# Only QueuePool (and its AsyncAdaptedQueuePool subclass) tracks connection
|
|
# counts; NullPool and others have no meaningful stats.
|
|
if not isinstance(pool, QueuePool):
|
|
return zeros
|
|
try:
|
|
# overflow() is negative until the base pool fills (it starts at
|
|
# -pool_size); clamp to the count of overflow connections actually open.
|
|
return {
|
|
"checked_out": pool.checkedout(),
|
|
"checked_in": pool.checkedin(),
|
|
"size": pool.size(),
|
|
"overflow": max(0, pool.overflow()),
|
|
}
|
|
except Exception:
|
|
return zeros
|
|
|
|
|
|
async def acquire_connection_with_retry(db: AsyncSession, context: str) -> None:
|
|
"""Force pool checkout (which ``SessionLocal()`` defers) with bounded backoff.
|
|
|
|
``SessionLocal()`` is lazy: the pool checkout — and any pooler rejection —
|
|
happens on the first query. We force it here inside a retry block so that
|
|
transient saturation of the transaction pooler is retried with exponential
|
|
backoff rather than surfacing as an immediate error. The checkout is wrapped
|
|
in a Sentry span so wait time is visible in traces; on budget exhaustion the
|
|
original error is reraised after capturing live pool stats to Sentry.
|
|
|
|
Each attempt rolls the session back on a retryable failure before retrying:
|
|
a failed checkout can leave the autobegun transaction in a pending-rollback
|
|
state, which would make the next ``db.connection()`` raise instead of
|
|
re-checking-out cleanly. The rollback is pure Python-side state cleanup when
|
|
no connection was bound, so it is cheap and safe.
|
|
"""
|
|
with sentry_sdk.start_span(op="db.pool.acquire", name=context):
|
|
if not settings.DB.CONNECTION_RETRY_ENABLED:
|
|
await db.connection()
|
|
return
|
|
attempts = 0
|
|
try:
|
|
async for attempt in AsyncRetrying(
|
|
wait=wait_exponential_jitter(
|
|
initial=settings.DB.CONNECTION_RETRY_BACKOFF_INITIAL_SECONDS,
|
|
max=settings.DB.CONNECTION_RETRY_BACKOFF_MAX_SECONDS,
|
|
),
|
|
stop=stop_after_delay(settings.DB.CONNECTION_RETRY_MAX_DELAY_SECONDS),
|
|
retry=retry_if_exception_type(RETRYABLE_DB_CONNECTION_ERRORS),
|
|
reraise=True,
|
|
):
|
|
with attempt:
|
|
attempts += 1
|
|
try:
|
|
await db.connection()
|
|
except RETRYABLE_DB_CONNECTION_ERRORS:
|
|
# Reset session state so the next attempt starts clean.
|
|
try:
|
|
await db.rollback()
|
|
except Exception:
|
|
logger.debug(
|
|
"rollback after failed checkout failed",
|
|
exc_info=True,
|
|
)
|
|
raise
|
|
except RETRYABLE_DB_CONNECTION_ERRORS as e:
|
|
_record_acquisition_outcome("exhausted")
|
|
if settings.SENTRY.ENABLED:
|
|
sentry_sdk.set_context("db_pool", get_pool_stats())
|
|
sentry_sdk.capture_exception(e)
|
|
raise
|
|
# "ok" on first try, "retried" if backoff was needed before success.
|
|
_record_acquisition_outcome("ok" if attempts <= 1 else "retried")
|
|
|
|
|
|
class HonchoAsyncSession(AsyncSession):
|
|
"""AsyncSession that lazily checks out its connection, with retry.
|
|
|
|
The pool checkout — and any pooler-rejection retry via
|
|
``acquire_connection_with_retry`` — happens on the FIRST DB-touching call,
|
|
not at construction. So a request handler that does non-DB work (embedding,
|
|
file processing, an LLM call) before its first query does NOT pin a
|
|
connection across that work, while still getting checkout retry on the
|
|
request path.
|
|
|
|
Only the checkout is retried; the SQL statement itself runs exactly once
|
|
(we never retry ``super().execute`` after a broad OperationalError), so
|
|
writes are never duplicated. The context for tracing is read from the
|
|
``request_context`` ContextVar, which the request/task scope has already set.
|
|
"""
|
|
|
|
# Class-level default; per-instance assignment shadows it (the subclass has
|
|
# a __dict__ even though AsyncSession declares __slots__).
|
|
_honcho_acquired: bool = False
|
|
|
|
async def _ensure_acquired(self) -> None:
|
|
if self._honcho_acquired:
|
|
return
|
|
context = request_context.get() or "unknown"
|
|
await acquire_connection_with_retry(self, context)
|
|
self._honcho_acquired = True
|
|
if settings.DB.TRACING:
|
|
# Forced checkout already happened above; this rides the same
|
|
# connection. super() to avoid re-entering _ensure_acquired.
|
|
await super().execute(
|
|
text("SELECT set_config('application_name', :name, false)"),
|
|
{"name": context},
|
|
)
|
|
|
|
# The overrides below are thin: ensure the connection is checked out (once,
|
|
# with retry) before delegating to AsyncSession. They cover every public
|
|
# DB-touching async method so the "lazy retry on first DB use" guarantee has
|
|
# no holes. Signatures are widened to *args/**kwargs because we only forward;
|
|
# call sites are typed against the AsyncSession base, so this does not weaken
|
|
# type-checking elsewhere. (connection() is intentionally NOT wrapped —
|
|
# acquire_connection_with_retry calls it, so wrapping would recurse.)
|
|
async def execute(self, *args: Any, **kwargs: Any) -> Any:
|
|
await self._ensure_acquired()
|
|
return await super().execute(*args, **kwargs)
|
|
|
|
async def scalar(self, *args: Any, **kwargs: Any) -> Any:
|
|
await self._ensure_acquired()
|
|
return await super().scalar(*args, **kwargs)
|
|
|
|
async def scalars(self, *args: Any, **kwargs: Any) -> Any:
|
|
await self._ensure_acquired()
|
|
return await super().scalars(*args, **kwargs)
|
|
|
|
async def get(self, *args: Any, **kwargs: Any) -> Any:
|
|
await self._ensure_acquired()
|
|
return await super().get(*args, **kwargs)
|
|
|
|
async def get_one(self, *args: Any, **kwargs: Any) -> Any:
|
|
await self._ensure_acquired()
|
|
return await super().get_one(*args, **kwargs)
|
|
|
|
async def stream(self, *args: Any, **kwargs: Any) -> Any:
|
|
await self._ensure_acquired()
|
|
return await super().stream(*args, **kwargs)
|
|
|
|
async def stream_scalars(self, *args: Any, **kwargs: Any) -> Any:
|
|
await self._ensure_acquired()
|
|
return await super().stream_scalars(*args, **kwargs)
|
|
|
|
async def flush(self, *args: Any, **kwargs: Any) -> None:
|
|
await self._ensure_acquired()
|
|
await super().flush(*args, **kwargs)
|
|
|
|
async def merge(self, *args: Any, **kwargs: Any) -> Any:
|
|
await self._ensure_acquired()
|
|
return await super().merge(*args, **kwargs)
|
|
|
|
async def refresh(self, *args: Any, **kwargs: Any) -> None:
|
|
await self._ensure_acquired()
|
|
await super().refresh(*args, **kwargs)
|
|
|
|
async def delete(self, *args: Any, **kwargs: Any) -> None:
|
|
await self._ensure_acquired()
|
|
await super().delete(*args, **kwargs)
|
|
|
|
async def commit(self) -> None:
|
|
# Ensures the add()->commit() path (autoflush on commit) also retries.
|
|
await self._ensure_acquired()
|
|
try:
|
|
await super().commit()
|
|
finally:
|
|
# Transaction ended; a later op must re-acquire (and re-wrap retry).
|
|
self._honcho_acquired = False
|
|
|
|
async def rollback(self) -> None:
|
|
try:
|
|
await super().rollback()
|
|
finally:
|
|
self._honcho_acquired = False
|
|
|
|
async def close(self) -> None:
|
|
try:
|
|
await super().close()
|
|
finally:
|
|
# The connection is released; a reused session must re-acquire.
|
|
self._honcho_acquired = False
|
|
|
|
async def reset(self) -> None:
|
|
try:
|
|
await super().reset()
|
|
finally:
|
|
self._honcho_acquired = False
|
|
|
|
|
|
SessionLocal = async_sessionmaker(
|
|
autocommit=False,
|
|
autoflush=False,
|
|
expire_on_commit=False,
|
|
bind=engine,
|
|
class_=HonchoAsyncSession,
|
|
)
|
|
|
|
|
|
class DBQueryInflightTracker:
|
|
"""Tracks statements executing on the wire via SQLAlchemy cursor events.
|
|
|
|
Drift-proof: marks ``Connection.info`` when a statement starts and clears it
|
|
on completion OR error, so the gauge can't leak upward (an errored statement
|
|
skips ``after_cursor_execute``) or go negative (a connect-time error has no
|
|
matching start). Bound to a pre-resolved labeled gauge child so the
|
|
per-statement hot path does no label resolution.
|
|
"""
|
|
|
|
# Marker on Connection.info recording that we incremented for the current
|
|
# statement, so we decrement exactly once on completion or error.
|
|
INFLIGHT_KEY: str = "_honcho_inflight"
|
|
|
|
def __init__(self, gauge_child: Any) -> None:
|
|
self._child: Any = gauge_child
|
|
|
|
def on_before(self, conn: Any, *_: Any) -> None:
|
|
try:
|
|
conn.info[self.INFLIGHT_KEY] = True
|
|
self._child.inc()
|
|
except Exception:
|
|
logger.debug("in-flight gauge inc failed", exc_info=True)
|
|
|
|
def on_after(self, conn: Any, *_: Any) -> None:
|
|
try:
|
|
if conn.info.pop(self.INFLIGHT_KEY, False):
|
|
self._child.dec()
|
|
except Exception:
|
|
logger.debug("in-flight gauge dec failed", exc_info=True)
|
|
|
|
def on_error(self, exception_context: Any) -> None:
|
|
try:
|
|
conn = exception_context.connection
|
|
if conn is not None and conn.info.pop(self.INFLIGHT_KEY, False):
|
|
self._child.dec()
|
|
except Exception:
|
|
logger.debug("in-flight gauge error-path dec failed", exc_info=True)
|
|
|
|
|
|
# Process-wide tracker, created at registration (None until then / if metrics off).
|
|
_inflight_tracker: DBQueryInflightTracker | None = None
|
|
|
|
|
|
_db_query_instrumentation_registered = False
|
|
|
|
|
|
def register_db_query_instrumentation(instance_type: str) -> None:
|
|
"""Attach per-statement in-flight tracking to the engine (no-op if off).
|
|
|
|
Gated on METRICS.ENABLED so there is zero overhead — not even attached event
|
|
listeners — when metrics are disabled. Idempotent: repeated calls (e.g. a
|
|
re-run lifespan or test startup) won't attach duplicate listeners, which
|
|
would double-count in-flight statements.
|
|
"""
|
|
global _db_instance_type, _inflight_tracker, _db_query_instrumentation_registered
|
|
_db_instance_type = instance_type
|
|
if not settings.METRICS.ENABLED or _db_query_instrumentation_registered:
|
|
return
|
|
child = db_queries_in_flight_gauge.labels(instance_type=instance_type)
|
|
_inflight_tracker = DBQueryInflightTracker(child)
|
|
sync_engine = engine.sync_engine
|
|
event.listen(sync_engine, "before_cursor_execute", _inflight_tracker.on_before)
|
|
event.listen(sync_engine, "after_cursor_execute", _inflight_tracker.on_after)
|
|
event.listen(sync_engine, "handle_error", _inflight_tracker.on_error)
|
|
_db_query_instrumentation_registered = True
|
|
|
|
|
|
# Define your naming convention
|
|
convention = {
|
|
"ix": "ix_%(table_name)s_%(column_0_N_name)s", # Index - supports multi-column
|
|
"uq": "uq_%(table_name)s_%(column_0_N_name)s", # Unique constraint - supports multi-column
|
|
"ck": "ck_%(table_name)s_%(constraint_name)s", # Check constraint
|
|
"fk": "fk_%(table_name)s_%(column_0_N_name)s_%(referred_table_name)s", # Foreign key - supports composite keys
|
|
"pk": "pk_%(table_name)s", # Primary key
|
|
}
|
|
|
|
table_schema = settings.DB.SCHEMA
|
|
# Note: column_0_N_name expands to include all columns in multi-column constraints
|
|
# e.g., "workspace_id_tenant_id" for a composite constraint on both columns
|
|
meta = MetaData(naming_convention=convention)
|
|
meta.schema = table_schema
|
|
Base = declarative_base(metadata=meta)
|
|
|
|
|
|
async def init_db():
|
|
"""Initialize the database using Alembic migrations"""
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
|
|
async with engine.connect() as connection:
|
|
# Create schema if it doesn't exist
|
|
await connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{table_schema}"'))
|
|
# Install pgvector extension if it doesn't exist
|
|
await connection.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
|
|
await connection.commit()
|
|
|
|
# Run Alembic migrations
|
|
alembic_cfg = Config("alembic.ini")
|
|
command.upgrade(alembic_cfg, "head")
|