Coverage for src/local_deep_research/database/session_context.py: 98%
118 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Database session context manager and decorator for encrypted databases.
3Ensures all database access has proper encryption context.
4"""
6import functools
7from contextlib import contextmanager
8from typing import Callable, Iterator, Optional
10from loguru import logger
11from sqlalchemy.exc import SQLAlchemyError
12from sqlalchemy.orm import Session
14from ..security.log_sanitizer import scrub_error
15from ..utilities.thread_context import get_search_context
16from .encrypted_db import db_manager
17from .rekey_lock import RekeyOperationalError
18from .thread_local_session import thread_session_manager
20# Placeholder password used when accessing unencrypted databases.
21# This should only be used when LDR_ALLOW_UNENCRYPTED=true is set.
22UNENCRYPTED_DB_PLACEHOLDER = "unencrypted-mode"
25class DatabaseSessionError(Exception):
26 """Raised when database session cannot be established."""
28 pass
31class DatabaseSessionUnavailableError(DatabaseSessionError):
32 """The session could not be established for an operational reason.
34 ``get_user_db_session`` raises this instead of the base class when
35 opening or claiming the user's database failed with a
36 ``RekeyOperationalError``: the rekey sidecar could not be opened, locked
37 or released, an epoch file could not be read or first published (an
38 fd-exhaustion spike, EACCES, EIO), or the durable rekey state needs
39 recovery first. None of these is a verdict on the credential or on the
40 session. A request handler that would end the session on a
41 ``DatabaseSessionError`` answers this one with a retryable 503 and keeps
42 the session instead (``database_unavailable_exception`` in
43 ``web/dependencies/auth.py``).
45 ``reason`` is the underlying fault's type name, for diagnostics. The
46 fault itself is neither chained nor kept: the frames it was raised
47 through hold the plaintext password.
48 """
50 def __init__(self, reason: str) -> None:
51 self.reason = reason
52 super().__init__("Database session unavailable")
55def safe_rollback(session: Session, context: str = "") -> None:
56 """Roll back the session, swallowing and logging any rollback failure.
58 SQLAlchemy requires explicit rollback after a failed flush/commit before
59 the session is usable again. Skipping it leaves the session in
60 PendingRollbackError state and every subsequent ORM operation cascades.
62 This helper exists so call sites can recover the session in one line
63 without repeating the try/except/log boilerplate at every except handler.
64 ``context`` is included in the error log so failed rollbacks can be
65 traced back to the call site.
67 Two SQLAlchemy error shapes are treated as "the session is structurally
68 unusable — give up on it, drop the thread-local cache, and let the next
69 caller get a fresh one" rather than as loud failures:
71 * ``InvalidRequestError("...provisioning a new connection; concurrent
72 operations are not permitted...")`` — the second thread racing on the
73 per-user QueuePool.
74 * ``InterfaceError("Cursor needed to be reset because of commit/rollback
75 and can no longer be fetched from")`` — a cursor was invalidated by a
76 commit/rollback that fired between the original ``execute()`` and the
77 lazy-attribute fetch that followed.
79 In both cases SQLAlchemy has already marked the session unrecoverable, so
80 a no-op rollback is correct, the thread-local cache is cleared, and the
81 failure is logged at DEBUG so the production stderr stream stays clean.
82 Other ``SQLAlchemyError`` failures still hit the loud ``logger.exception``
83 path (unless their message matches a known broken-session signature) —
84 the session is normally recoverable via rollback and the operator needs
85 to see them.
86 """
87 if session is None:
88 return
89 log_msg = (
90 f"Failed to rollback session: {context}"
91 if context
92 else "Failed to rollback session"
93 )
94 try:
95 session.rollback()
96 except SQLAlchemyError as exc:
97 msg = str(exc)
98 # Message-substring matching pinned against SQLAlchemy 2.0+ QueuePool error messages
99 # ('provisioning a new connection', 'concurrent operations are not permitted').
100 is_provisioning_race = (
101 "provisioning a new connection" in msg
102 or "concurrent operations are not permitted" in msg
103 )
104 # Requires 'Cursor needed to be reset' message match so other InterfaceErrors
105 # (e.g., driver/connectivity failures) are not quietly swallowed into the DEBUG path.
106 is_cursor_invalidated = "Cursor needed to be reset" in msg
107 if is_provisioning_race or is_cursor_invalidated:
108 label = f": {context}" if context else ""
109 logger.debug(f"safe_rollback — resetting broken session{label}")
110 # Drop the thread-local cache so the next caller on this
111 # thread gets a fresh session. Identity-checked inside the
112 # helper, so a caller that hands in a session that ISN'T
113 # the cached one (e.g. borrowed from ``g.db_session`` or
114 # owned by a different thread) won't accidentally clear
115 # someone else's cache. The reset itself is best-effort —
116 # never let it raise past ``safe_rollback`` (call sites are
117 # themselves in except handlers).
118 try:
119 thread_session_manager.reset_session_if_matches(session)
120 except Exception:
121 logger.debug(
122 f"safe_rollback: reset_session_if_matches raised for{label}"
123 )
124 return
125 logger.exception(log_msg)
126 except Exception:
127 logger.exception(log_msg)
130@contextmanager
131def telemetry_savepoint(session: Session) -> Iterator[None]:
132 """Write best-effort telemetry in a SAVEPOINT on a borrowed session.
134 The telemetry joins the caller's transaction only when the caller
135 already holds an open write transaction; otherwise it commits at once
136 in a short transaction of its own. The signal is the DBAPI connection's
137 ``in_transaction`` flag. Production sqlite3/SQLCipher connections use
138 legacy deferred mode (``isolation_level=""``): the driver opens a
139 transaction only before INSERT, UPDATE or DELETE and closes it on commit
140 or rollback, so a SELECT never opens one. ``Session.in_transaction()``
141 and the pending ``new``/``dirty``/``deleted`` collections cannot make
142 this distinction: any query autobegins SQLAlchemy's logical transaction,
143 and a flushed write leaves those collections empty.
145 * Open (the caller has written, or began a transaction explicitly as
146 ``research_cascade._ensure_outer_db_transaction`` does): the SAVEPOINT
147 nests inside it, so the telemetry commits or rolls back with the
148 caller, and no second connection waits on this thread's own writer
149 lock.
150 * Idle (the scope has only read, or has already committed): SQLite
151 begins a transaction for the outermost SAVEPOINT and RELEASE commits
152 it. The telemetry is durable at once, survives the scope's exit
153 rollback, and holds the writer lock only while it writes. If the
154 telemetry fails, that transaction is rolled back here, because
155 ROLLBACK TO alone would keep it (and any writer lock it took) open
156 until the caller's scope ends.
158 Other dialects begin their transaction with the first statement, so the
159 savepoint always nests there. A connection that does not report a
160 definite ``False`` counts as open and is never rolled back here.
161 Failures propagate to the caller, which logs and skips the telemetry.
162 """
163 connection = session.connection()
164 owned_connection = None
165 if connection.dialect.name == "sqlite":
166 driver_connection = connection.connection.driver_connection
167 if getattr(driver_connection, "in_transaction", None) is False:
168 owned_connection = driver_connection
169 try:
170 with session.begin_nested():
171 yield
172 except BaseException:
173 if owned_connection is not None and owned_connection.in_transaction:
174 try:
175 owned_connection.rollback()
176 except Exception:
177 logger.debug("Could not end a failed telemetry transaction")
178 raise
181@contextmanager
182def get_user_db_session(
183 username: Optional[str] = None,
184 password: Optional[str] = None,
185 session_id: Optional[str] = None,
186):
187 """
188 Context manager that ensures proper database session with encryption.
189 Now uses thread-local sessions for better performance.
191 Args:
192 username: Username (required; must be passed explicitly under FastAPI).
193 password: Password for encrypted database (required for first access).
194 session_id: Optional session ID for exact per-session password lookup.
195 Request handlers should pass the current request's session_id so
196 two concurrent sessions for the same user can't cross-pollinate.
197 If omitted, the resolver falls back to scanning any active session
198 for the user — fine for background threads, unsafe for request
199 handlers.
201 Yields:
202 Database session for the user
204 Raises:
205 DatabaseSessionUnavailableError: If the session cannot be
206 established because of a ``RekeyOperationalError`` (a rekey
207 sidecar or epoch fault, or pending rekey recovery). A subclass
208 of ``DatabaseSessionError``.
209 DatabaseSessionError: If session cannot be established
210 """
211 # Import here to avoid circular imports
212 from .thread_local_session import get_metrics_session
213 from .session_passwords import session_password_store
215 if not username:
216 raise DatabaseSessionError("No authenticated user")
218 # Resolve password from the provided session_id, then fall back to
219 # any active session, then the thread context for background workers.
220 if not password and session_id:
221 password = session_password_store.get_session_password(
222 username, session_id
223 )
224 if password:
225 logger.debug(f"Got password from session store for {username}")
227 if not password:
228 # Scan active sessions for this user. Safe for background threads;
229 # request handlers should pass session_id explicitly.
230 password = session_password_store.get_any_session_password(username)
232 if not password:
233 thread_context = get_search_context()
234 if thread_context and thread_context.get("user_password"):
235 password = thread_context["user_password"]
236 logger.debug(f"Got password from thread context for {username}")
238 if not password and db_manager.has_encryption:
239 raise DatabaseSessionError(
240 f"Encrypted database for {username} requires password"
241 )
242 if not password:
243 logger.warning(
244 f"Accessing unencrypted database for {username} - "
245 "ensure this is intentional (LDR_ALLOW_UNENCRYPTED=true)"
246 )
247 password = UNENCRYPTED_DB_PLACEHOLDER
249 # A RekeyOperationalError from opening or claiming the database (in
250 # open_user_database or in create_thread_safe_session_for_metrics, which
251 # both let it through) is not a credential verdict, so it gets its own
252 # subclass: request handlers answer it with a retryable 503 that keeps
253 # the session, where the base class may end it. Raised after the except
254 # block, so the new error does not even carry the fault as its
255 # suppressed __context__: the frames it came through hold the password.
256 session = None
257 unavailable_reason = None
258 try:
259 session = get_metrics_session(username, password)
260 except RekeyOperationalError as error:
261 unavailable_reason = type(error).__name__
262 if unavailable_reason is not None:
263 logger.warning(
264 f"Database session unavailable for {username}: {unavailable_reason}"
265 )
266 raise DatabaseSessionUnavailableError(unavailable_reason)
267 if not session:
268 raise DatabaseSessionError(
269 f"Could not establish session for {username}"
270 )
272 # The session is a reused thread-local session. On normal scope exit
273 # (outermost block) the checked-out connection is released in the
274 # ``finally`` below, so the pool isn't pinned for the thread's lifetime.
275 # We MUST wrap the yield in try/except so an exception inside the ``with``
276 # block doesn't leave a half-committed transaction attached to this
277 # thread's session: the next caller on the same thread would otherwise
278 # inherit that dirty state. Thread cleanup (``cleanup_current_thread()``)
279 # remains as a safety net for paths that never exit a scope.
280 #
281 # The scope depth tells get_session's re-validation whether an
282 # enclosing block is still active on this thread: a nested
283 # get_user_db_session call (e.g. a helper opening its own session
284 # inside a caller's with-block) must NOT trigger the stale-lock
285 # rollback, or the caller's uncommitted writes are destroyed.
286 from .thread_local_session import thread_session_manager
288 thread_session_manager.enter_scope()
289 try:
290 yield session
291 except Exception:
292 # The yielded session is a *reused* thread-local session, not a
293 # fresh one closed on exit. If the caller's ``with`` block raised
294 # (most importantly a failed ``session.commit()``/``flush()``),
295 # the session is left in ``PendingRollbackError`` state and the
296 # next operation on this thread cascades. Roll it back here so an
297 # unguarded ``with`` block can't poison the thread, then re-raise
298 # so the original error still surfaces to the caller.
299 safe_rollback(session, "get_user_db_session")
300 raise
301 finally:
302 thread_session_manager.exit_scope()
303 if not thread_session_manager.in_scope():
304 try:
305 # If a transaction was left open (e.g. from read queries autobeginning
306 # a transaction and checking out a connection from QueuePool, or uncommitted writes),
307 # roll back on outermost exit to return the checked-out connection to QueuePool
308 # and release SQLite SHARED read locks.
309 if (
310 hasattr(session, "in_transaction")
311 and session.in_transaction()
312 ):
313 safe_rollback(session, "get_user_db_session exit")
314 except Exception as exc:
315 logger.debug(
316 f"safe_rollback on get_user_db_session exit raised: {scrub_error(exc)}"
317 )
320def with_user_database(func: Callable) -> Callable:
321 """
322 Decorator that ensures function has access to user's database.
323 Injects 'db_session' as first argument to the decorated function.
325 Usage:
326 @with_user_database
327 def get_user_settings(db_session, setting_key):
328 return db_session.query(Setting).filter_by(key=setting_key).first()
329 """
331 @functools.wraps(func)
332 def wrapper(*args, **kwargs):
333 # Check if username/password provided in kwargs
334 username = kwargs.pop("_username", None)
335 password = kwargs.pop("_password", None)
337 with get_user_db_session(username, password) as db_session:
338 return func(db_session, *args, **kwargs)
340 return wrapper
343class DatabaseAccessMixin:
344 """
345 Mixin class for services that need database access.
346 Provides convenient methods for database operations.
347 """
349 def get_db_session(
350 self, username: Optional[str] = None
351 ) -> Optional[Session]:
352 """
353 DEPRECATED: This method returns a closed session due to context manager exit.
355 Use `with get_user_db_session(username) as session:` instead.
357 Raises:
358 DeprecationWarning: Always raised to prevent usage of broken method.
359 """
360 raise DeprecationWarning(
361 "get_db_session() is deprecated and returns a closed session. "
362 "Use `with get_user_db_session(username) as session:` instead."
363 )
365 @with_user_database
366 def execute_with_db(
367 self, db_session: Session, query_func: Callable, *args, **kwargs
368 ):
369 """Execute a function with database session."""
370 return query_func(db_session, *args, **kwargs)