add read db (#773)
* feat: implement read DB and fix queue stale cleanup
* fix: use read_db in internal methods
* fix: mention read db in the CLAUDE.md
* fix: make TRACING checkout hook autocommit-safe; sample cleanup-gate jitter once
The DB.TRACING checkout hook ran `SELECT set_config(...)` at pool checkout,
before the dialect applies the read engine's AUTOCOMMIT isolation level. That
statement autobegins a transaction, and psycopg then refuses to switch the
connection into AUTOCOMMIT ("can't change 'autocommit' now: connection in
transaction status INTRANS"), so every read_only session 500s under TRACING and
the INTRANS connection leaks back to poison later write checkouts. Run the hook
in autocommit and restore the prior mode so it never leaves an open transaction;
set_config(..., is_local=false) is session-scoped and survives the boundary.
Add a regression test (fails without the fix) covering read_only + TRACING.
Also sample the stale-cleanup gate's jittered interval once per attempt instead
of re-rolling it every poll, so the spacing is a fixed deadline per cycle rather
than a random walk (and is testable at non-zero jitter ratios).
* fix: reset request_context in TRACING checkout-hook test
---------
Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
This commit is contained in:
parent
f75b336a3c
commit
bf494257b8
|
|
@ -116,6 +116,7 @@ cd sdks/typescript && bun run tsc --noEmit
|
|||
- Explicit error handling with appropriate exception types
|
||||
- Docstrings: Use Google style docstrings
|
||||
- **Never hold a DB session during external calls** (LLM, embedding, HTTP). If a function needs both a DB session and an external call result, compute the external result first and pass it as a parameter. This avoids tying up DB connections during slow network I/O. Use `tracked_db` for short-lived, DB-only operations; pass a shared session when multiple DB-only calls can reuse one connection.
|
||||
- **Never write through a read-only session** (`tracked_db(..., read_only=True)`, `get_read_db`, `ReadSessionLocal`). These run in AUTOCOMMIT mode with no transaction: writes are NOT blocked by the database — they silently commit immediately, and `begin_nested()` savepoints break. There is no runtime guard; this is enforced by convention only. Use `read_only=True` strictly for SELECT-only windows; anything that mutates (including get-or-create paths) must use a regular write session.
|
||||
|
||||
### Runtime Architecture
|
||||
|
||||
|
|
|
|||
|
|
@ -98,6 +98,10 @@ POLLING_BACKOFF_MULTIPLIER = 2.0
|
|||
POLLING_STARTUP_JITTER_SECONDS = 30.0
|
||||
POLLING_JITTER_RATIO = 0.5
|
||||
STALE_SESSION_TIMEOUT_MINUTES = 5
|
||||
# Minimum (jittered) spacing between stale-work-unit cleanup runs per instance.
|
||||
# Staleness is a minutes-timescale condition, so cleanup doesn't need to run on
|
||||
# every seconds-scale poll (0.0 = run every poll, legacy behavior).
|
||||
STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS = 60.0
|
||||
# QUEUE_ERROR_RETENTION_SECONDS = 2592000 # 30 days
|
||||
DEDUPLICATE = true
|
||||
LOG_OBSERVATIONS = false
|
||||
|
|
|
|||
|
|
@ -769,6 +769,10 @@ class DeriverSettings(HonchoSettings):
|
|||
# to 0.0 to disable.
|
||||
POLLING_JITTER_RATIO: Annotated[float, Field(default=0.5, ge=0.0, le=1.0)] = 0.5
|
||||
STALE_SESSION_TIMEOUT_MINUTES: Annotated[int, Field(default=5, gt=0, le=1440)] = 5
|
||||
# Minimum (jittered) spacing between stale-work-unit cleanup runs
|
||||
STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS: Annotated[
|
||||
float, Field(default=60.0, ge=0.0, le=3600.0)
|
||||
] = 60.0
|
||||
|
||||
# Retention window (seconds) for keeping errored items in the queue
|
||||
QUEUE_ERROR_RETENTION_SECONDS: Annotated[
|
||||
|
|
|
|||
|
|
@ -369,7 +369,7 @@ async def query_documents(
|
|||
max_distance,
|
||||
top_k,
|
||||
)
|
||||
async with tracked_db("query_documents.pgvector") as managed_db:
|
||||
async with tracked_db("query_documents.pgvector", read_only=True) as managed_db:
|
||||
docs = await _query_documents_pgvector(
|
||||
managed_db,
|
||||
workspace_name,
|
||||
|
|
@ -407,7 +407,7 @@ async def query_documents(
|
|||
document_ids=document_ids,
|
||||
filters=filters,
|
||||
)
|
||||
async with tracked_db("query_documents.fetch") as managed_db:
|
||||
async with tracked_db("query_documents.fetch", read_only=True) as managed_db:
|
||||
docs = await fetch_documents_by_ids(
|
||||
db=managed_db,
|
||||
workspace_name=workspace_name,
|
||||
|
|
|
|||
|
|
@ -817,7 +817,7 @@ async def _semantic_search_messages(
|
|||
# Pre-fetch peer session scope if needed (short-lived DB session)
|
||||
allowed_session_names: list[str] | None = None
|
||||
if observer and not session_name:
|
||||
async with tracked_db(f"{operation_name}.peer_scope") as db:
|
||||
async with tracked_db(f"{operation_name}.peer_scope", read_only=True) as db:
|
||||
allowed_session_names = await get_peer_session_names(
|
||||
db, workspace_name, observer
|
||||
)
|
||||
|
|
@ -837,7 +837,7 @@ async def _semantic_search_messages(
|
|||
if not message_ids:
|
||||
return []
|
||||
|
||||
async with tracked_db(operation_name) as db:
|
||||
async with tracked_db(operation_name, read_only=True) as db:
|
||||
matched_messages = (
|
||||
await _fetch_messages_by_ids(
|
||||
db,
|
||||
|
|
@ -853,7 +853,7 @@ async def _semantic_search_messages(
|
|||
_expunge_snippets(db, snippets)
|
||||
return snippets
|
||||
|
||||
async with tracked_db(operation_name) as db:
|
||||
async with tracked_db(operation_name, read_only=True) as db:
|
||||
snippets = await _search_messages_pgvector(
|
||||
db,
|
||||
workspace_name,
|
||||
|
|
@ -985,7 +985,7 @@ async def grep_messages(
|
|||
List of tuples: (matched_messages, context_messages)
|
||||
Each snippet may contain multiple matches if they were close together.
|
||||
"""
|
||||
async with tracked_db("message.grep_messages") as db:
|
||||
async with tracked_db("message.grep_messages", read_only=True) as db:
|
||||
# Pre-fetch peer session scope if needed
|
||||
allowed_session_names = None
|
||||
if observer and not session_name:
|
||||
|
|
|
|||
|
|
@ -276,7 +276,7 @@ class RepresentationManager:
|
|||
)
|
||||
|
||||
async with tracked_db(
|
||||
"representation_manager.get_working_representation"
|
||||
"representation_manager.get_working_representation", read_only=True
|
||||
) as new_db:
|
||||
return await self._get_working_representation_internal(
|
||||
new_db,
|
||||
|
|
|
|||
42
src/db.py
42
src/db.py
|
|
@ -66,6 +66,24 @@ SessionLocal = async_sessionmaker(
|
|||
class_=AsyncSession,
|
||||
)
|
||||
|
||||
# Read-only engine: shares `engine`'s pool, but checks connections out in DBAPI
|
||||
# AUTOCOMMIT mode, so psycopg emits NO BEGIN — a SELECT never autobegins a
|
||||
# transaction. The backend therefore returns to state 'idle' (not 'idle in
|
||||
# transaction') the moment a statement completes.
|
||||
read_engine = engine.execution_options(isolation_level="AUTOCOMMIT")
|
||||
|
||||
# Sessions for SELECT-only work (same lazy-checkout semantics as SessionLocal).
|
||||
# MUST NOT be used for writes: with no enclosing transaction, begin_nested()
|
||||
# savepoints (see the crud get-or-create paths) break, and every flush would
|
||||
# commit immediately. Use SessionLocal for anything that mutates.
|
||||
ReadSessionLocal = async_sessionmaker(
|
||||
autocommit=False,
|
||||
autoflush=False,
|
||||
expire_on_commit=False,
|
||||
bind=read_engine,
|
||||
class_=AsyncSession,
|
||||
)
|
||||
|
||||
|
||||
def _set_application_name_on_checkout(
|
||||
dbapi_connection: Any, _connection_record: Any, _connection_proxy: Any
|
||||
|
|
@ -76,16 +94,30 @@ def _set_application_name_on_checkout(
|
|||
reused pooled connection is re-tagged for the new caller), reading the
|
||||
per-task ``request_context`` the request/task scope has already set.
|
||||
Best-effort: a failure here must never break the checkout.
|
||||
|
||||
Runs in autocommit so it never leaves the connection 'idle in transaction'
|
||||
at checkout: this hook fires BEFORE the dialect applies execution-option
|
||||
isolation levels, and psycopg refuses to switch a connection into AUTOCOMMIT
|
||||
(which the read engine does) while a transaction opened by this statement is
|
||||
still in progress. set_config(..., is_local=false) is session-scoped, so it
|
||||
persists past the autocommit boundary.
|
||||
"""
|
||||
context = request_context.get() or "unknown"
|
||||
try:
|
||||
cursor = dbapi_connection.cursor()
|
||||
previous_autocommit = dbapi_connection.autocommit
|
||||
if not previous_autocommit:
|
||||
dbapi_connection.autocommit = True
|
||||
try:
|
||||
cursor.execute(
|
||||
"SELECT set_config('application_name', %s, false)", (context,)
|
||||
)
|
||||
cursor = dbapi_connection.cursor()
|
||||
try:
|
||||
cursor.execute(
|
||||
"SELECT set_config('application_name', %s, false)", (context,)
|
||||
)
|
||||
finally:
|
||||
cursor.close()
|
||||
finally:
|
||||
cursor.close()
|
||||
if not previous_autocommit:
|
||||
dbapi_connection.autocommit = False
|
||||
except Exception:
|
||||
logger.debug("setting application_name on checkout failed", exc_info=True)
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from contextlib import asynccontextmanager
|
|||
from fastapi import Depends
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src.db import SessionLocal, request_context
|
||||
from src.db import ReadSessionLocal, SessionLocal, request_context
|
||||
|
||||
|
||||
async def get_db():
|
||||
|
|
@ -32,12 +32,39 @@ async def get_db():
|
|||
await db.close()
|
||||
|
||||
|
||||
async def get_read_db():
|
||||
"""FastAPI Dependency Generator for SELECT-only handlers.
|
||||
|
||||
Same lazy-checkout semantics as get_db, but the session is bound to the
|
||||
AUTOCOMMIT read engine: no BEGIN is ever emitted, so the connection can not
|
||||
sit 'idle in transaction' between the query and this teardown — a delayed
|
||||
finally here is harmless (the backend is plain 'idle'). close() is still
|
||||
required to release the connection itself back to the pool.
|
||||
|
||||
MUST only be used by handlers that never mutate; see ReadSessionLocal.
|
||||
"""
|
||||
db: AsyncSession = ReadSessionLocal()
|
||||
try:
|
||||
yield db
|
||||
finally:
|
||||
# rollback is a wire-level no-op under AUTOCOMMIT; kept to reset any
|
||||
# Python-side session state before close, mirroring get_db.
|
||||
await db.rollback()
|
||||
await db.close()
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def tracked_db(operation_name: str | None = None):
|
||||
async def tracked_db(operation_name: str | None = None, *, read_only: bool = False):
|
||||
"""Context manager for tracked database sessions.
|
||||
|
||||
Sets a task-scoped request_context so the lazy session picks it up for
|
||||
tracing/attribution, then yields a lazy session (see get_db).
|
||||
|
||||
Pass read_only=True for SELECT-only windows: the session is then bound to
|
||||
the AUTOCOMMIT read engine, so the work inside the block never holds an
|
||||
open transaction (no idle-in-transaction parking; the pooler can reclaim
|
||||
the backend between statements). Never use read_only=True on a path that
|
||||
mutates — see ReadSessionLocal.
|
||||
"""
|
||||
# Get request ID if available, or create operation-specific one
|
||||
context = request_context.get()
|
||||
|
|
@ -47,14 +74,15 @@ async def tracked_db(operation_name: str | None = None):
|
|||
context = f"task:{operation_name}:{str(uuid.uuid4())[:8]}"
|
||||
token = request_context.set(context)
|
||||
|
||||
db = SessionLocal()
|
||||
db = (ReadSessionLocal if read_only else SessionLocal)()
|
||||
try:
|
||||
yield db
|
||||
except Exception:
|
||||
await db.rollback()
|
||||
raise
|
||||
finally:
|
||||
# Always send ROLLBACK unconditionally — see get_db() comment.
|
||||
# Always send ROLLBACK unconditionally — see get_db() comment. (Under
|
||||
# read_only/AUTOCOMMIT it is a wire-level no-op.)
|
||||
await db.rollback()
|
||||
await db.close()
|
||||
if token: # Only reset if we set it
|
||||
|
|
@ -62,3 +90,4 @@ async def tracked_db(operation_name: str | None = None):
|
|||
|
||||
|
||||
db: AsyncSession = Depends(get_db)
|
||||
read_db: AsyncSession = Depends(get_read_db)
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import asyncio
|
|||
import contextlib
|
||||
import random
|
||||
import signal
|
||||
import time
|
||||
from asyncio import Task
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass, field
|
||||
|
|
@ -134,6 +135,16 @@ class QueueManager:
|
|||
settings.DERIVER.POLLING_SLEEP_INTERVAL_SECONDS
|
||||
)
|
||||
|
||||
# Monotonic timestamp of the last stale-work-unit cleanup ATTEMPT.
|
||||
# None -> the first poll always runs cleanup (recovers rows left stale
|
||||
# by a crashed predecessor immediately).
|
||||
self._last_stale_cleanup_attempt: float | None = None
|
||||
# Jittered gate width (seconds) sampled ONCE per attempt, so the deadline
|
||||
# for the next run is fixed when the timestamp is set rather than
|
||||
# re-rolled on every poll (which would make the effective spacing a
|
||||
# random walk and untestable at non-zero jitter ratios).
|
||||
self._stale_cleanup_gate_seconds: float = 0.0
|
||||
|
||||
# Initialize from settings
|
||||
self.workers: int = settings.DERIVER.WORKERS
|
||||
self.semaphore: asyncio.Semaphore = asyncio.Semaphore(self.workers)
|
||||
|
|
@ -258,6 +269,35 @@ class QueueManager:
|
|||
# Polling and Scheduling #
|
||||
##########################
|
||||
|
||||
async def _maybe_cleanup_stale_work_units(self) -> None:
|
||||
"""Run stale-work-unit cleanup at most once per (jittered) interval.
|
||||
|
||||
Staleness is a minutes-timescale condition (STALE_SESSION_TIMEOUT_MINUTES),
|
||||
but the polling loop fires on a seconds timescale on every deriver
|
||||
instance — running cleanup unconditionally per poll multiplies into
|
||||
unnecessary write transactions. Gate it locally:
|
||||
concurrent cleaners on other instances remain safe via FOR UPDATE SKIP
|
||||
LOCKED, so no cross-instance coordination is required, and the jittered
|
||||
gate (sampled once per attempt) keeps instances from re-synchronizing
|
||||
their cleanup runs. The gate tracks the last ATTEMPT (set before
|
||||
running), so a failing cleanup waits a full interval instead of retrying
|
||||
every poll against a DB that is already struggling. An interval of 0
|
||||
preserves run-every-poll behavior.
|
||||
"""
|
||||
interval = settings.DERIVER.STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS
|
||||
if (
|
||||
interval > 0.0
|
||||
and self._last_stale_cleanup_attempt is not None
|
||||
and time.monotonic() - self._last_stale_cleanup_attempt
|
||||
< self._stale_cleanup_gate_seconds
|
||||
):
|
||||
return
|
||||
# Record the attempt and fix the next deadline before running, so the
|
||||
# gate width is stable for this cycle and a failing cleanup still waits.
|
||||
self._last_stale_cleanup_attempt = time.monotonic()
|
||||
self._stale_cleanup_gate_seconds = self._jitter(interval)
|
||||
await self.cleanup_stale_work_units()
|
||||
|
||||
async def cleanup_stale_work_units(self) -> None:
|
||||
"""Clean up stale work units"""
|
||||
async with tracked_db("cleanup_stale_work_units") as db:
|
||||
|
|
@ -457,7 +497,7 @@ class QueueManager:
|
|||
continue
|
||||
|
||||
try:
|
||||
await self.cleanup_stale_work_units()
|
||||
await self._maybe_cleanup_stale_work_units()
|
||||
claimed_work_units = await self.get_and_claim_work_units()
|
||||
if claimed_work_units:
|
||||
self._reset_poll_interval()
|
||||
|
|
|
|||
|
|
@ -40,7 +40,7 @@ async def agentic_chat(
|
|||
The synthesized answer string
|
||||
"""
|
||||
# Short-lived DB session for validation + config
|
||||
async with tracked_db("dialectic.preflight") as db:
|
||||
async with tracked_db("dialectic.preflight", read_only=True) as db:
|
||||
await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observer))
|
||||
if observer != observed:
|
||||
await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observed))
|
||||
|
|
@ -101,7 +101,7 @@ async def agentic_chat_stream(
|
|||
Chunks of the response text as they are generated
|
||||
"""
|
||||
# Short-lived DB session for validation + config
|
||||
async with tracked_db("dialectic.preflight") as db:
|
||||
async with tracked_db("dialectic.preflight", read_only=True) as db:
|
||||
await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observer))
|
||||
if observer != observed:
|
||||
await crud.get_peer(db, workspace_name, schemas.PeerCreate(name=observed))
|
||||
|
|
|
|||
|
|
@ -121,7 +121,7 @@ class DialecticAgent:
|
|||
token_limit=max_tokens,
|
||||
reverse=False, # chronological order
|
||||
)
|
||||
async with tracked_db("dialectic.session_history") as db:
|
||||
async with tracked_db("dialectic.session_history", read_only=True) as db:
|
||||
result = await db.execute(stmt)
|
||||
messages = result.scalars().all()
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from fastapi_pagination.ext.sqlalchemy import apaginate
|
|||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from src import crud, schemas
|
||||
from src.dependencies import db
|
||||
from src.dependencies import db, read_db
|
||||
from src.exceptions import ResourceNotFoundException, ValidationException
|
||||
from src.security import require_auth
|
||||
from src.telemetry.events import EmbeddingCallPurpose
|
||||
|
|
@ -67,7 +67,7 @@ async def list_conclusions(
|
|||
False,
|
||||
description="Whether to reverse the order of results",
|
||||
),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
):
|
||||
"""
|
||||
List Conclusions using optional filters, ordered by recency unless `reverse` is true. Results are paginated.
|
||||
|
|
@ -97,7 +97,7 @@ async def query_conclusions(
|
|||
...,
|
||||
description="Semantic search parameters for Conclusions",
|
||||
),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
) -> list[schemas.Conclusion]:
|
||||
"""
|
||||
Query Conclusions using semantic search. Use `top_k` to control the number of results returned.
|
||||
|
|
|
|||
|
|
@ -18,7 +18,7 @@ from sqlalchemy.orm.attributes import flag_modified
|
|||
|
||||
from src import crud, schemas
|
||||
from src.config import settings
|
||||
from src.dependencies import db
|
||||
from src.dependencies import db, read_db
|
||||
from src.deriver import enqueue
|
||||
from src.exceptions import FileTooLargeError, ResourceNotFoundException
|
||||
from src.security import require_auth
|
||||
|
|
@ -260,7 +260,7 @@ async def get_messages(
|
|||
reverse: bool | None = Query(
|
||||
False, description="Whether to reverse the order of results"
|
||||
),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
):
|
||||
"""Get all messages for a Session with optional filters. Results are paginated."""
|
||||
try:
|
||||
|
|
@ -288,7 +288,7 @@ async def get_message(
|
|||
workspace_id: str = Path(...),
|
||||
session_id: str = Path(...),
|
||||
message_id: str = Path(...),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
):
|
||||
"""Get a single message by ID from a Session."""
|
||||
honcho_message = await crud.get_message(
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
|||
|
||||
from src import crud, schemas
|
||||
from src.config import settings
|
||||
from src.dependencies import db, tracked_db
|
||||
from src.dependencies import db, read_db, tracked_db
|
||||
from src.dialectic.chat import agentic_chat, agentic_chat_stream
|
||||
from src.embedding_client import embedding_client
|
||||
from src.exceptions import AuthenticationException, ResourceNotFoundException
|
||||
|
|
@ -43,7 +43,7 @@ async def get_peers(
|
|||
None, description="Filtering options for the peers list"
|
||||
),
|
||||
reverse: bool = Query(False, description="Whether to reverse the order of results"),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
):
|
||||
"""Get all Peers for a Workspace, paginated with optional filters."""
|
||||
filter_param = None
|
||||
|
|
@ -134,7 +134,7 @@ async def get_sessions_for_peer(
|
|||
None, description="Filtering options for the sessions list"
|
||||
),
|
||||
reverse: bool = Query(False, description="Whether to reverse the order of results"),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
):
|
||||
"""Get all Sessions for a Peer, paginated with optional filters."""
|
||||
filter_param = None
|
||||
|
|
@ -318,7 +318,7 @@ async def get_peer_card(
|
|||
None,
|
||||
description="Optional target peer to retrieve a card for, from the observer's perspective. If not provided, returns the observer's own card",
|
||||
),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
):
|
||||
"""Get a peer card for a specific peer relationship.
|
||||
|
||||
|
|
@ -412,7 +412,6 @@ async def get_peer_context(
|
|||
le=100,
|
||||
description="Maximum number of conclusions to include in the representation",
|
||||
),
|
||||
db: AsyncSession = db,
|
||||
):
|
||||
"""
|
||||
Get context for a peer, including their representation and peer card.
|
||||
|
|
@ -459,10 +458,12 @@ async def get_peer_context(
|
|||
parent_category="api",
|
||||
)
|
||||
|
||||
# Get the peer card
|
||||
peer_card = await crud.get_peer_card(
|
||||
db, workspace_id, observer=peer_id, observed=observed
|
||||
)
|
||||
async with tracked_db(
|
||||
"peers.get_peer_context.peer_card", read_only=True
|
||||
) as card_db:
|
||||
peer_card = await crud.get_peer_card(
|
||||
card_db, workspace_id, observer=peer_id, observed=observed
|
||||
)
|
||||
|
||||
response = schemas.PeerContext(
|
||||
peer_id=peer_id,
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
|||
from src import config, crud, schemas
|
||||
from src.cache.client import safe_cache_delete
|
||||
from src.crud.session import session_cache_key
|
||||
from src.dependencies import db
|
||||
from src.dependencies import db, read_db
|
||||
from src.deriver.enqueue import enqueue_deletion
|
||||
from src.embedding_client import embedding_client
|
||||
from src.exceptions import (
|
||||
|
|
@ -251,7 +251,7 @@ async def get_sessions(
|
|||
None, description="Filtering and pagination options for the sessions list"
|
||||
),
|
||||
reverse: bool = Query(False, description="Whether to reverse the order of results"),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
):
|
||||
"""Get all Sessions for a Workspace, paginated with optional filters."""
|
||||
filter_param = None
|
||||
|
|
@ -544,7 +544,7 @@ async def get_peer_config(
|
|||
workspace_id: str = Path(...),
|
||||
session_id: str = Path(...),
|
||||
peer_id: str = Path(...),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
):
|
||||
"""Get the configuration for a Peer in a Session."""
|
||||
return await crud.get_peer_config(
|
||||
|
|
@ -599,7 +599,7 @@ async def set_peer_config(
|
|||
async def get_session_peers(
|
||||
workspace_id: str = Path(...),
|
||||
session_id: str = Path(...),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
):
|
||||
"""Get all Peers in a Session. Results are paginated."""
|
||||
try:
|
||||
|
|
@ -622,7 +622,7 @@ async def get_session_peers(
|
|||
async def get_session_context(
|
||||
workspace_id: str = Path(...),
|
||||
session_id: str = Path(...),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
tokens: int | None = Query(
|
||||
None,
|
||||
le=config.settings.GET_CONTEXT_MAX_TOKENS,
|
||||
|
|
@ -814,7 +814,7 @@ async def get_session_context(
|
|||
async def get_session_summaries(
|
||||
workspace_id: str = Path(...),
|
||||
session_id: str = Path(...),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
) -> schemas.SessionSummaries:
|
||||
"""
|
||||
Get available summaries for a Session.
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
|||
|
||||
from src import crud, schemas
|
||||
from src.config import settings
|
||||
from src.dependencies import db
|
||||
from src.dependencies import db, read_db
|
||||
from src.deriver.enqueue import enqueue_deletion, enqueue_dream
|
||||
from src.exceptions import AuthenticationException
|
||||
from src.security import JWTParams, require_auth
|
||||
|
|
@ -67,7 +67,7 @@ async def get_all_workspaces(
|
|||
None, description="Filtering and pagination options for the workspaces list"
|
||||
),
|
||||
reverse: bool = Query(False, description="Whether to reverse the order of results"),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
):
|
||||
"""Get all Workspaces, paginated with optional filters."""
|
||||
filter_param = None
|
||||
|
|
@ -169,7 +169,7 @@ async def get_queue_status(
|
|||
session_id: str | None = Query(
|
||||
None, description="Optional session ID to filter by"
|
||||
),
|
||||
db: AsyncSession = db,
|
||||
db: AsyncSession = read_db,
|
||||
):
|
||||
"""
|
||||
Get the processing queue status for a Workspace, optionally scoped to an observer, sender,
|
||||
|
|
|
|||
|
|
@ -1607,7 +1607,7 @@ async def _handle_get_recent_history(
|
|||
) -> "str | ToolResult":
|
||||
"""Handle get_recent_history tool."""
|
||||
_ = tool_input
|
||||
async with tracked_db("tool.get_recent_history") as db:
|
||||
async with tracked_db("tool.get_recent_history", read_only=True) as db:
|
||||
history: list[models.Message] = await get_recent_history(
|
||||
db,
|
||||
workspace_name=ctx.workspace_name,
|
||||
|
|
@ -1723,7 +1723,7 @@ async def _handle_get_observation_context(
|
|||
ctx: ToolContext, tool_input: dict[str, Any]
|
||||
) -> "str | ToolResult":
|
||||
"""Handle get_observation_context tool."""
|
||||
async with tracked_db("tool.get_observation_context") as db:
|
||||
async with tracked_db("tool.get_observation_context", read_only=True) as db:
|
||||
messages = await get_observation_context(
|
||||
db,
|
||||
workspace_name=ctx.workspace_name,
|
||||
|
|
@ -1862,7 +1862,7 @@ async def _handle_get_messages_by_date_range(
|
|||
if isinstance(before_date, str):
|
||||
return before_date # Error message
|
||||
|
||||
async with tracked_db("tool.get_messages_by_date_range") as db:
|
||||
async with tracked_db("tool.get_messages_by_date_range", read_only=True) as db:
|
||||
messages = await crud.get_messages_by_date_range(
|
||||
db,
|
||||
workspace_name=ctx.workspace_name,
|
||||
|
|
@ -1980,7 +1980,7 @@ async def _handle_get_recent_observations(
|
|||
) -> str:
|
||||
"""Handle get_recent_observations tool."""
|
||||
session_only = tool_input.get("session_only", False)
|
||||
async with tracked_db("tool.get_recent_observations") as db:
|
||||
async with tracked_db("tool.get_recent_observations", read_only=True) as db:
|
||||
documents = await crud.query_documents_recent(
|
||||
db=db,
|
||||
workspace_name=ctx.workspace_name,
|
||||
|
|
@ -2006,7 +2006,7 @@ async def _handle_get_most_derived_observations(
|
|||
ctx: ToolContext, tool_input: dict[str, Any]
|
||||
) -> str:
|
||||
"""Handle get_most_derived_observations tool."""
|
||||
async with tracked_db("tool.get_most_derived_observations") as db:
|
||||
async with tracked_db("tool.get_most_derived_observations", read_only=True) as db:
|
||||
documents = await crud.query_documents_most_derived(
|
||||
db=db,
|
||||
workspace_name=ctx.workspace_name,
|
||||
|
|
@ -2038,7 +2038,7 @@ async def _handle_get_session_summary(
|
|||
if summary_type == "long"
|
||||
else summarizer.SummaryType.SHORT
|
||||
)
|
||||
async with tracked_db("tool.get_session_summary") as db:
|
||||
async with tracked_db("tool.get_session_summary", read_only=True) as db:
|
||||
summary = await summarizer.get_summary(
|
||||
db, ctx.workspace_name, ctx.session_name, st
|
||||
)
|
||||
|
|
@ -2050,7 +2050,7 @@ async def _handle_get_session_summary(
|
|||
async def _handle_get_peer_card(ctx: ToolContext, tool_input: dict[str, Any]) -> str:
|
||||
"""Handle get_peer_card tool."""
|
||||
_ = tool_input
|
||||
async with tracked_db("tool.get_peer_card") as db:
|
||||
async with tracked_db("tool.get_peer_card", read_only=True) as db:
|
||||
peer_card = await crud.get_peer_card(
|
||||
db,
|
||||
workspace_name=ctx.workspace_name,
|
||||
|
|
@ -2209,7 +2209,7 @@ async def _handle_get_reasoning_chain(
|
|||
return f"ERROR: Invalid direction '{direction}'. Must be 'premises', 'conclusions', or 'both'"
|
||||
|
||||
# Get the observation itself
|
||||
async with tracked_db("tool.get_reasoning_chain") as db:
|
||||
async with tracked_db("tool.get_reasoning_chain", read_only=True) as db:
|
||||
docs = await crud.get_documents_by_ids(db, ctx.workspace_name, [observation_id])
|
||||
if not docs or not docs[0]:
|
||||
return f"ERROR: Observation '{observation_id}' not found"
|
||||
|
|
|
|||
|
|
@ -448,7 +448,7 @@ async def search(
|
|||
return search_results[0][:limit]
|
||||
return []
|
||||
|
||||
async with tracked_db("search.messages") as managed_db:
|
||||
async with tracked_db("search.messages", read_only=True) as managed_db:
|
||||
combined_results = await _run_search(managed_db)
|
||||
for message in combined_results:
|
||||
managed_db.expunge(message)
|
||||
|
|
|
|||
|
|
@ -32,7 +32,7 @@ from src import models
|
|||
from src.cache.client import cache
|
||||
from src.config import settings
|
||||
from src.db import Base
|
||||
from src.dependencies import get_db
|
||||
from src.dependencies import get_db, get_read_db
|
||||
from src.exceptions import HonchoException
|
||||
from src.main import app
|
||||
from src.models import Peer, Workspace
|
||||
|
|
@ -340,6 +340,10 @@ async def client(
|
|||
yield db_session
|
||||
|
||||
app.dependency_overrides[get_db] = override_get_db
|
||||
# Read-only routes use get_read_db (AUTOCOMMIT engine) in production; in
|
||||
# tests they must see the same per-test database/session as writes, both
|
||||
# for isolation and so data written by a test is visible to its reads.
|
||||
app.dependency_overrides[get_read_db] = override_get_db
|
||||
|
||||
# No-op the startup embedding-schema validator inside the lifespan. The
|
||||
# global `engine` it would inspect points to a DB that isn't migrated in
|
||||
|
|
@ -792,7 +796,10 @@ def mock_tracked_db(request: pytest.FixtureRequest):
|
|||
session_factory = async_sessionmaker(bind=db_engine, expire_on_commit=False)
|
||||
|
||||
@asynccontextmanager
|
||||
async def mock_tracked_db_context(_: str | None = None):
|
||||
async def mock_tracked_db_context(_: str | None = None, *, read_only: bool = False):
|
||||
# read_only is accepted (and ignored): in tests both engines resolve to
|
||||
# the same per-test database session.
|
||||
del read_only
|
||||
async with session_factory() as session:
|
||||
yield session
|
||||
|
||||
|
|
|
|||
|
|
@ -474,7 +474,10 @@ async def test_search_messages_external_lookup_happens_before_tracked_db(
|
|||
return [([message], [message])]
|
||||
|
||||
@asynccontextmanager
|
||||
async def fake_tracked_db(_operation_name: str | None = None):
|
||||
async def fake_tracked_db(
|
||||
_operation_name: str | None = None, *, read_only: bool = False
|
||||
):
|
||||
del read_only
|
||||
call_order.append("enter")
|
||||
yield fake_db
|
||||
call_order.append("exit")
|
||||
|
|
@ -578,7 +581,10 @@ async def test_search_messages_temporal_external_lookup_happens_before_tracked_d
|
|||
return [([message], [message])]
|
||||
|
||||
@asynccontextmanager
|
||||
async def fake_tracked_db(_operation_name: str | None = None):
|
||||
async def fake_tracked_db(
|
||||
_operation_name: str | None = None, *, read_only: bool = False
|
||||
):
|
||||
del read_only
|
||||
call_order.append("enter")
|
||||
yield fake_db
|
||||
call_order.append("exit")
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
|
|||
from uvicorn.config import Config
|
||||
from uvicorn.server import Server
|
||||
|
||||
from src.dependencies import get_db
|
||||
from src.dependencies import get_db, get_read_db
|
||||
from src.main import app
|
||||
|
||||
|
||||
|
|
@ -95,6 +95,9 @@ def ts_test_server(
|
|||
yield session
|
||||
|
||||
app.dependency_overrides[get_db] = override_get_db
|
||||
# Read-only routes use get_read_db (AUTOCOMMIT engine) in production; in
|
||||
# tests they must resolve to the same per-test database.
|
||||
app.dependency_overrides[get_read_db] = override_get_db
|
||||
|
||||
# No-op the lifespan's startup embedding-schema validator — same
|
||||
# reasoning as the `client` fixture in tests/conftest.py: the module-
|
||||
|
|
@ -133,7 +136,9 @@ def mock_tracked_db(ts_db_session: async_sessionmaker[AsyncSession]):
|
|||
|
||||
# Create a tracked_db that uses fresh sessions (not shared)
|
||||
@asynccontextmanager
|
||||
async def ts_tracked_db(_: str | None = None):
|
||||
async def ts_tracked_db(_: str | None = None, *, read_only: bool = False):
|
||||
# read_only accepted (and ignored): tests use one per-test database.
|
||||
del read_only
|
||||
async with ts_db_session() as session:
|
||||
yield session
|
||||
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from sqlalchemy import text
|
|||
from sqlalchemy.exc import OperationalError
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
|
||||
from src.config import settings
|
||||
from src.startup.embedding_validator import (
|
||||
StartupValidationError,
|
||||
_assert_pgvector_dims_match, # pyright: ignore[reportPrivateUsage]
|
||||
|
|
@ -120,18 +121,24 @@ async def test_validator_fails_closed_when_introspection_keeps_failing(
|
|||
@pytest.mark.asyncio
|
||||
async def test_validator_passes_against_test_database(
|
||||
db_engine: AsyncEngine,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""The test DB is provisioned at the default dim (1536); the validator
|
||||
should accept it without raising."""
|
||||
# conftest provisions the test tables in `public`; pin the validator to it
|
||||
# so a developer's local .env DB_SCHEMA can't point it at another schema.
|
||||
monkeypatch.setattr(settings.DB, "SCHEMA", "public")
|
||||
await validate_embedding_schema(db_engine)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validator_raises_when_schema_dim_diverges_from_settings(
|
||||
db_engine: AsyncEngine,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""ALTER one of the embedding columns to a non-1536 dim and confirm the
|
||||
validator raises with an actionable message."""
|
||||
monkeypatch.setattr(settings.DB, "SCHEMA", "public") # see test above
|
||||
async with db_engine.begin() as conn:
|
||||
await conn.execute(
|
||||
text(
|
||||
|
|
@ -181,8 +188,13 @@ def test_non_1536_pgvector_without_migrated_no_longer_raises_at_config_time() ->
|
|||
"""The dim-vs-MIGRATED guard has been removed. Constructing AppSettings
|
||||
with non-1536 + default pgvector + MIGRATED=false should now succeed
|
||||
(the runtime schema validator at startup is the safety net)."""
|
||||
# Minimal env, NOT a copy of os.environ: load_dotenv() in the app mutates
|
||||
# the parent pytest process's environ, so inheriting it would leak a
|
||||
# developer's local .env (DB_SCHEMA, VECTOR_STORE_MIGRATED, ...) into the
|
||||
# child despite PYTHON_DOTENV_DISABLED. The child must see pure defaults
|
||||
# plus exactly the overrides below.
|
||||
env = {
|
||||
**os.environ,
|
||||
"PATH": os.environ.get("PATH", ""),
|
||||
"PYTHON_DOTENV_DISABLED": "1",
|
||||
"EMBEDDING_VECTOR_DIMENSIONS": "768",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -50,6 +50,11 @@ class _FakeCursor:
|
|||
class _FakeDBAPIConn:
|
||||
def __init__(self, recorder: list[Any], raise_exc: Exception | None = None) -> None:
|
||||
self._cursor: _FakeCursor = _FakeCursor(recorder, raise_exc)
|
||||
# Real pooled connections are checked out in non-autocommit mode; the
|
||||
# hook flips this to True for its statement then restores it so it never
|
||||
# leaves an open transaction that would block the read engine's
|
||||
# AUTOCOMMIT switch.
|
||||
self.autocommit: bool = False
|
||||
|
||||
def cursor(self) -> _FakeCursor:
|
||||
return self._cursor
|
||||
|
|
@ -69,6 +74,8 @@ def test_checkout_hook_sets_application_name_from_request_context() -> None:
|
|||
assert "set_config" in sql and "application_name" in sql
|
||||
assert params == ("request:trace-ctx",)
|
||||
assert conn._cursor.closed is True # pyright: ignore[reportPrivateUsage]
|
||||
# The hook restored the original (non-autocommit) mode after its statement.
|
||||
assert conn.autocommit is False
|
||||
|
||||
|
||||
def test_checkout_hook_defaults_to_unknown_without_context() -> None:
|
||||
|
|
@ -198,3 +205,106 @@ def test_inflight_gauge_no_drift(monkeypatch: pytest.MonkeyPatch) -> None:
|
|||
# gauge negative.
|
||||
tracker.on_error(SimpleNamespace(connection=SimpleNamespace(info={})))
|
||||
assert value() == start
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_cleanup_time_gate(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""cleanup_stale_work_units runs at most once per gate interval per
|
||||
instance (staleness is a minutes-timescale condition; per-poll cleanup
|
||||
multiplies into needless fleet-wide write transactions). First poll always
|
||||
runs it so a crashed predecessor's stale rows are recovered immediately."""
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0)
|
||||
monkeypatch.setattr(
|
||||
settings.DERIVER, "STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS", 60.0
|
||||
)
|
||||
|
||||
from src.deriver import queue_manager as qm_mod
|
||||
|
||||
qm = qm_mod.QueueManager()
|
||||
runs = {"n": 0}
|
||||
|
||||
async def fake_cleanup() -> None:
|
||||
runs["n"] += 1
|
||||
|
||||
monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup)
|
||||
|
||||
clock = {"now": 1_000.0}
|
||||
monkeypatch.setattr(
|
||||
"src.deriver.queue_manager.time.monotonic", lambda: clock["now"]
|
||||
)
|
||||
|
||||
# First call runs (no prior attempt recorded).
|
||||
await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage]
|
||||
assert runs["n"] == 1
|
||||
|
||||
# Inside the gate window: skipped.
|
||||
clock["now"] += 10.0
|
||||
await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage]
|
||||
assert runs["n"] == 1
|
||||
|
||||
# Past the gate window: runs again.
|
||||
clock["now"] += 60.0
|
||||
await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage]
|
||||
assert runs["n"] == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_cleanup_gate_failed_attempt_waits_full_interval(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""The gate records the ATTEMPT before running, so a failing cleanup is not
|
||||
retried on every poll against a DB that is already struggling."""
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0)
|
||||
monkeypatch.setattr(
|
||||
settings.DERIVER, "STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS", 60.0
|
||||
)
|
||||
|
||||
from src.deriver import queue_manager as qm_mod
|
||||
|
||||
qm = qm_mod.QueueManager()
|
||||
attempts = {"n": 0}
|
||||
|
||||
async def failing_cleanup() -> None:
|
||||
attempts["n"] += 1
|
||||
raise RuntimeError("db unavailable")
|
||||
|
||||
monkeypatch.setattr(qm, "cleanup_stale_work_units", failing_cleanup)
|
||||
|
||||
clock = {"now": 1_000.0}
|
||||
monkeypatch.setattr(
|
||||
"src.deriver.queue_manager.time.monotonic", lambda: clock["now"]
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage]
|
||||
assert attempts["n"] == 1
|
||||
|
||||
# Immediately after the failure: still gated, no hammering.
|
||||
clock["now"] += 1.0
|
||||
await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage]
|
||||
assert attempts["n"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stale_cleanup_gate_zero_interval_runs_every_poll(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Interval 0.0 preserves legacy run-on-every-poll behavior."""
|
||||
monkeypatch.setattr(settings.DERIVER, "POLLING_JITTER_RATIO", 0.0)
|
||||
monkeypatch.setattr(
|
||||
settings.DERIVER, "STALE_WORK_UNIT_CLEANUP_INTERVAL_SECONDS", 0.0
|
||||
)
|
||||
|
||||
from src.deriver import queue_manager as qm_mod
|
||||
|
||||
qm = qm_mod.QueueManager()
|
||||
runs = {"n": 0}
|
||||
|
||||
async def fake_cleanup() -> None:
|
||||
runs["n"] += 1
|
||||
|
||||
monkeypatch.setattr(qm, "cleanup_stale_work_units", fake_cleanup)
|
||||
|
||||
await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage]
|
||||
await qm._maybe_cleanup_stale_work_units() # pyright: ignore[reportPrivateUsage]
|
||||
assert runs["n"] == 2
|
||||
|
|
|
|||
|
|
@ -150,3 +150,161 @@ async def test_tracked_db_rolls_back_open_transaction_on_exit(
|
|||
|
||||
assert fake_db.rollback_calls == 1
|
||||
assert fake_db.close_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tracked_db_read_only_uses_read_sessionmaker(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# read_only=True must construct the session from ReadSessionLocal (the
|
||||
# AUTOCOMMIT engine) — and never touch SessionLocal — while keeping the
|
||||
# same rollback/close teardown.
|
||||
read_fake = FakeSession()
|
||||
monkeypatch.setattr(dependencies_module, "ReadSessionLocal", lambda: read_fake)
|
||||
monkeypatch.setattr(
|
||||
dependencies_module,
|
||||
"SessionLocal",
|
||||
lambda: pytest.fail("read_only window constructed a write session"),
|
||||
)
|
||||
|
||||
async with real_tracked_db("read_op", read_only=True) as db:
|
||||
assert db is read_fake
|
||||
|
||||
assert read_fake.rollback_calls == 1
|
||||
assert read_fake.close_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_read_db_rolls_back_and_closes(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
read_fake = FakeSession()
|
||||
monkeypatch.setattr(dependencies_module, "ReadSessionLocal", lambda: read_fake)
|
||||
|
||||
dep_gen = dependencies_module.get_read_db()
|
||||
try:
|
||||
db = await anext(dep_gen)
|
||||
assert db is read_fake
|
||||
assert read_fake.connection_calls == 0 # still lazy, no eager checkout
|
||||
finally:
|
||||
await dep_gen.aclose()
|
||||
|
||||
assert read_fake.rollback_calls == 1
|
||||
assert read_fake.close_calls == 1
|
||||
|
||||
|
||||
def test_read_engine_is_autocommit_and_shares_pool() -> None:
|
||||
# The read engine must differ from the write engine ONLY by isolation
|
||||
# level: AUTOCOMMIT (so reads never autobegin a transaction) on the same
|
||||
# underlying pool (no second connection budget).
|
||||
from src.db import engine, read_engine
|
||||
|
||||
assert (
|
||||
read_engine.sync_engine._execution_options.get( # pyright: ignore[reportPrivateUsage]
|
||||
"isolation_level"
|
||||
)
|
||||
== "AUTOCOMMIT"
|
||||
)
|
||||
assert read_engine.sync_engine.pool is engine.sync_engine.pool
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_only_session_runs_in_autocommit_on_the_wire() -> None:
|
||||
# Wire-level guarantee behind the whole read-path fix: a read_only session's
|
||||
# connection has the DBAPI autocommit flag set, so psycopg emits no BEGIN
|
||||
# and the backend sits in state 'idle' (not 'idle in transaction') after a
|
||||
# statement returns. NOTE: get_isolation_level() can NOT verify this — it
|
||||
# reports the server's transaction_isolation GUC (READ COMMITTED), because
|
||||
# autocommit is a driver behavior, not a server isolation level.
|
||||
from sqlalchemy import text
|
||||
|
||||
from src.db import read_engine
|
||||
|
||||
async with real_tracked_db("read_op", read_only=True) as db:
|
||||
pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar()
|
||||
conn = await db.connection()
|
||||
raw = (await conn.get_raw_connection()).driver_connection
|
||||
assert raw is not None
|
||||
assert raw.autocommit is True
|
||||
|
||||
# Definitive check, from a second connection: after the SELECT above,
|
||||
# the session's backend must be plain 'idle' — an open transaction
|
||||
# would report 'idle in transaction' and be reapable in production.
|
||||
async with read_engine.connect() as observer:
|
||||
state = (
|
||||
await observer.execute(
|
||||
text("SELECT state FROM pg_stat_activity WHERE pid = :p"),
|
||||
{"p": pid},
|
||||
)
|
||||
).scalar()
|
||||
assert state == "idle"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_write_session_holds_idle_in_transaction_after_select() -> None:
|
||||
# Contrast guard documenting WHY the read engine exists: the default
|
||||
# (transactional) session autobegins on the first statement and leaves the
|
||||
# backend 'idle in transaction' until rollback/close — the state that
|
||||
# Postgres's idle_in_transaction_session_timeout reaps and that pins a
|
||||
# transaction-mode pooler backend.
|
||||
from sqlalchemy import text
|
||||
|
||||
from src.db import read_engine
|
||||
|
||||
async with real_tracked_db("write_op") as db:
|
||||
pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar()
|
||||
async with read_engine.connect() as observer:
|
||||
state = (
|
||||
await observer.execute(
|
||||
text("SELECT state FROM pg_stat_activity WHERE pid = :p"),
|
||||
{"p": pid},
|
||||
)
|
||||
).scalar()
|
||||
assert state == "idle in transaction"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_only_session_works_with_tracing_checkout_hook() -> None:
|
||||
# Regression: the DB.TRACING checkout hook runs set_config() at pool
|
||||
# checkout, BEFORE the dialect applies the read engine's AUTOCOMMIT
|
||||
# isolation level. If that statement is allowed to autobegin a transaction,
|
||||
# psycopg then refuses to switch the connection into AUTOCOMMIT
|
||||
# ("can't change 'autocommit' now: connection in transaction") and every
|
||||
# read_only session 500s under TRACING. The hook must run in autocommit so
|
||||
# it leaves the connection idle. This combination is otherwise untested
|
||||
# because DB.TRACING defaults to false.
|
||||
from sqlalchemy import event, text
|
||||
|
||||
from src.db import (
|
||||
_set_application_name_on_checkout, # pyright: ignore[reportPrivateUsage]
|
||||
engine,
|
||||
read_engine,
|
||||
)
|
||||
|
||||
context_token = request_context.set("tracing-regression")
|
||||
event.listen(engine.sync_engine, "checkout", _set_application_name_on_checkout)
|
||||
try:
|
||||
async with real_tracked_db("read_op", read_only=True) as db:
|
||||
pid = (await db.execute(text("SELECT pg_backend_pid()"))).scalar()
|
||||
app_name = (await db.execute(text("SHOW application_name"))).scalar()
|
||||
conn = await db.connection()
|
||||
raw = (await conn.get_raw_connection()).driver_connection
|
||||
assert raw is not None
|
||||
# AUTOCOMMIT was applied despite the checkout hook running first.
|
||||
assert raw.autocommit is True
|
||||
# The hook still tagged the connection (set_config is session-scoped,
|
||||
# so it survives the autocommit boundary).
|
||||
assert app_name == "tracing-regression"
|
||||
# Backend is idle, not idle-in-transaction: the no-BEGIN guarantee
|
||||
# holds even with the hook firing.
|
||||
async with read_engine.connect() as observer:
|
||||
state = (
|
||||
await observer.execute(
|
||||
text("SELECT state FROM pg_stat_activity WHERE pid = :p"),
|
||||
{"p": pid},
|
||||
)
|
||||
).scalar()
|
||||
assert state == "idle"
|
||||
finally:
|
||||
event.remove(engine.sync_engine, "checkout", _set_application_name_on_checkout)
|
||||
request_context.reset(context_token)
|
||||
|
|
|
|||
Loading…
Reference in New Issue