Coverage for src/local_deep_research/database/thread_local_session.py: 93%
200 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Thread-local database session management.
3Each thread gets its own database session, reused until its owned cleanup boundary.
4"""
6import functools
7import threading
8from contextlib import ContextDecorator
9from typing import Optional, Dict, Tuple
10from sqlalchemy import text
11from sqlalchemy.exc import PendingRollbackError
12from sqlalchemy.orm import Session
13from loguru import logger
15from .encrypted_db import db_manager
18class ThreadLocalSessionManager:
19 """
20 Manages database sessions per thread.
21 Each thread gets its own session, reused until the request or task finishes.
22 """
24 def __init__(self):
25 # Thread-local storage for sessions
26 self._local = threading.local()
27 # Track credentials per thread ID (for cleanup)
28 self._thread_credentials: Dict[int, Tuple[str, str]] = {}
29 self._lock = threading.Lock()
31 def enter_scope(self) -> None:
32 """Mark entry into a ``get_user_db_session`` block on this thread.
34 ``get_session`` uses the depth to tell a left-over transaction
35 from a *completed* block (safe and necessary to roll back — it
36 holds a SQLite SHARED lock) apart from an *enclosing* block that
37 is still active (rolling back would destroy the caller's
38 uncommitted writes)."""
39 self._local.scope_depth = getattr(self._local, "scope_depth", 0) + 1
41 def exit_scope(self) -> None:
42 """Mark exit from a ``get_user_db_session`` block on this thread."""
43 self._local.scope_depth = max(
44 0, getattr(self._local, "scope_depth", 0) - 1
45 )
47 def in_scope(self) -> bool:
48 """Whether an enclosing ``get_user_db_session`` block is active
49 on this thread."""
50 return getattr(self._local, "scope_depth", 0) > 0
52 def get_session(self, username: str, password: str) -> Optional[Session]:
53 """
54 Get or create a database session for the current thread.
56 The session is created once per thread and reused for all subsequent calls.
57 This avoids the expensive SQLCipher decryption on every database access.
59 Returns None when the database cannot be opened with ``password``
60 (or its schema cannot be initialised). A ``RekeyOperationalError``
61 from opening the database or creating the session (the rekey
62 sidecar or an epoch file could not be read, or a pending rekey
63 needs recovery) is raised instead: it is not a credential verdict,
64 and ``get_user_db_session`` turns it into
65 ``DatabaseSessionUnavailableError`` so a request can answer a
66 retryable 503 rather than log the user out.
67 """
68 thread_id = threading.get_ident()
70 # Check if we already have a session for this thread
71 if hasattr(self._local, "session") and self._local.session:
72 # SECURITY: ensure cached session belongs to the requesting user
73 if getattr(self._local, "username", None) != username:
74 logger.warning(
75 f"Thread {thread_id}: Session username mismatch "
76 f"(cached={self._local.username!r}, requested={username!r}), "
77 "clearing stale cross-user session"
78 )
79 self._cleanup_thread_session()
80 else:
81 # Verify it's still valid
82 try:
83 self._local.session.execute(text("SELECT 1"))
84 if not self.in_scope():
85 # Under DEFERRED isolation the validation SELECT
86 # opens a transaction that holds a SHARED lock on
87 # SQLite until an explicit commit/rollback. A
88 # long-lived thread-local session reused across
89 # requests would keep that lock held and block
90 # the first writer. Roll it back so subsequent
91 # callers start fresh.
92 self._local.session.rollback()
93 # else: re-entrant call from inside an enclosing
94 # get_user_db_session block on this thread — the
95 # validation SELECT joined the caller's open
96 # transaction, and rolling back here would silently
97 # destroy the caller's uncommitted writes (concrete
98 # data loss in library_rag_service.index_document,
99 # whose nested helpers open their own sessions).
100 return self._local.session
101 except PendingRollbackError:
102 # Session has a pending rollback (e.g. from a previous database lock error).
103 # Attempt rollback to recover without destroying the session.
104 logger.debug(
105 f"Thread {thread_id}: PendingRollbackError, attempting rollback recovery"
106 )
107 try:
108 self._local.session.rollback()
109 self._local.session.execute(text("SELECT 1"))
110 self._local.session.rollback()
111 return self._local.session
112 except Exception:
113 logger.warning(
114 f"Thread {thread_id}: Rollback recovery failed, creating new session"
115 )
116 self._cleanup_thread_session()
117 except Exception:
118 # Session is invalid, will create a new one
119 logger.debug(
120 f"Thread {thread_id}: Existing session invalid, creating new one"
121 )
122 self._cleanup_thread_session()
124 # Create new session for this thread
125 logger.debug(
126 f"Thread {thread_id}: Creating new database session for user {username}"
127 )
129 # Ensure database is open. open_user_database returns None for
130 # credential failures and raises DatabaseInitializationError when
131 # the schema can't be initialised; from a worker-thread caller
132 # both mean "no usable session right now" so collapse them.
133 #
134 # A RekeyOperationalError is deliberately NOT collapsed into None,
135 # here or from create_thread_safe_session_for_metrics below. It is
136 # an operational fault (an fd-exhaustion spike on the rekey
137 # sidecar, an unreadable epoch, pending recovery), and a None would
138 # reach get_user_db_session as the same DatabaseSessionError a
139 # wrong password produces, which the home page answered by logging
140 # the user out. get_user_db_session catches it instead.
141 from .encrypted_db import DatabaseInitializationError
143 try:
144 engine = db_manager.open_user_database(username, password)
145 except DatabaseInitializationError:
146 # ``logger.warning`` (no traceback) rather than
147 # ``logger.exception``: ``password`` is a live local in this
148 # frame, so rendering a traceback under ``diagnose=True``
149 # would dump the plaintext SQLCipher master password
150 # (unrecoverable — TRUST.md §5). The redacted failure detail
151 # is already logged at the raise site in
152 # ``open_user_database`` (#4182).
153 logger.warning(
154 f"Thread {thread_id}: database init failed for user {username}"
155 )
156 return None
157 if not engine:
158 logger.error(
159 f"Thread {thread_id}: Failed to open database for user {username}"
160 )
161 return None
163 # Create session for this thread
164 session = db_manager.create_thread_safe_session_for_metrics(
165 username, password
166 )
167 if not session: 167 ↛ 168line 167 didn't jump to line 168 because the condition on line 167 was never true
168 logger.error(
169 f"Thread {thread_id}: Failed to create session for user {username}"
170 )
171 return None
173 # Store in thread-local storage
174 self._local.session = session
175 self._local.username = username
177 # Track credentials for cleanup
178 with self._lock:
179 self._thread_credentials[thread_id] = (username, password)
181 return session
183 def get_current_session(self) -> Optional[Session]:
184 """Get the current thread's session if it exists."""
185 if hasattr(self._local, "session"):
186 return self._local.session
187 return None
189 def get_scoped_session(self, username: str) -> Optional[Session]:
190 """Borrow the active scope's session only when it belongs to username.
192 Metrics callbacks must not call ``get_session`` here: its recovery
193 path may roll back or replace an enclosing caller-owned transaction.
194 A callback for a different research owner simply omits best-effort
195 telemetry instead of touching that other user's session.
196 """
197 if not self.in_scope():
198 return None
199 if getattr(self._local, "username", None) != username:
200 return None
201 return self.get_current_session()
203 def _cleanup_thread_session(self):
204 """Clean up the current thread's session.
206 Sessions are bound to the shared per-user QueuePool engine, so
207 closing the session returns its connection to the pool — there
208 are no per-thread engines to dispose.
209 """
210 thread_id = threading.get_ident()
212 if hasattr(self._local, "session") and self._local.session:
213 try:
214 self._local.session.rollback()
215 except Exception:
216 logger.warning(
217 f"Thread {thread_id}: Error rolling back session during cleanup"
218 )
219 try:
220 self._local.session.close()
221 logger.debug(f"Thread {thread_id}: Closed database session")
222 except Exception:
223 logger.warning(f"Thread {thread_id}: Error closing session")
224 finally:
225 self._local.session = None
227 if hasattr(self._local, "username"):
228 self._local.username = None
230 # Remove from tracking
231 with self._lock:
232 self._thread_credentials.pop(thread_id, None)
234 def reset_session_if_matches(self, session: "Session | None") -> bool:
235 """If ``session`` is the current thread's cached session, clear it.
237 Used by ``safe_rollback`` in ``database/session_context.py`` to heal
238 the thread-local session cache when an error tells us the session
239 is structurally unusable (a provisioning race or an invalidated
240 cursor). The next caller on this thread then gets a fresh session
241 from the shared ``QueuePool`` instead of inheriting the broken
242 one.
244 Identity-checked: a caller that hands in a session from a
245 different thread (or a session that isn't the cached one at all
246 — e.g. one borrowed from ``g.db_session``) won't accidentally
247 clear someone else's cache. Returns ``True`` if a reset
248 actually happened, ``False`` otherwise.
249 """
250 if session is None:
251 return False
252 if not hasattr(self._local, "session") or self._local.session is None:
253 return False
254 if self._local.session is not session:
255 return False
256 self._cleanup_thread_session()
257 return True
259 def cleanup_thread(self, thread_id: Optional[int] = None):
260 """
261 Clean up session for a specific thread or current thread.
262 Called when a thread is finishing.
263 """
264 if thread_id is None:
265 thread_id = threading.get_ident()
267 # If it's the current thread, we can clean up directly
268 if thread_id == threading.get_ident():
269 self._cleanup_thread_session()
270 else:
271 # For other threads, just remove from tracking
272 # The thread-local storage will be cleaned up when the thread ends
273 with self._lock:
274 self._thread_credentials.pop(thread_id, None)
276 def cleanup_dead_threads(self):
277 """Remove credential entries for threads that are no longer alive.
279 Handles the abnormal case of threads that died without triggering
280 their cleanup handler. Uses threading.enumerate() to identify
281 alive threads and removes credential entries for dead ones.
282 Sessions on dead threads are garbage-collected normally and
283 their connections return to the shared per-user QueuePool.
285 Called from:
286 - processor_v2.py: every ~60s in the queue loop
287 - fastapi_app.py: in the request middleware's post-request cleanup
288 - connection_cleanup.py: in cleanup_idle_connections (every ~300s)
289 """
290 alive_ids = {t.ident for t in threading.enumerate()}
291 with self._lock:
292 dead_ids = [
293 tid for tid in self._thread_credentials if tid not in alive_ids
294 ]
295 for tid in dead_ids:
296 del self._thread_credentials[tid]
297 if dead_ids:
298 logger.debug(f"Swept {len(dead_ids)} dead thread credential(s)")
300 def clear_credentials_for_user(self, username: str) -> int:
301 """Drop cached plaintext credentials for one user, on any thread.
303 Entries here hold the user's plaintext SQLCipher password so a
304 pooled worker can rebuild its session without a round-trip to the
305 password store. Synchronous routes normally drop their owner-thread
306 entry through ``WorkerCleanupAPIRoute``. This username-wide cleanup
307 is still required on logout because another worker may be executing
308 background work for the same user, and one thread cannot clear a
309 different thread's ``threading.local`` state.
311 Only the tracking dict is cleared. The owning thread's
312 ``threading.local`` session is deliberately left alone: one thread
313 cannot reach into another's local storage, and the session is
314 harmless once the credential is gone — the username re-validation at
315 the top of ``get_session()`` evicts it before any reuse. In-flight
316 work is unaffected too, since ``get_session()`` is passed the
317 password explicitly and never reads it back out of this dict.
319 Returns the number of entries dropped.
320 """
321 with self._lock:
322 thread_ids = [
323 tid
324 for tid, (user, _password) in self._thread_credentials.items()
325 if user == username
326 ]
327 for tid in thread_ids:
328 del self._thread_credentials[tid]
329 if thread_ids:
330 logger.debug(
331 f"Cleared cached credentials on {len(thread_ids)} thread(s)"
332 )
333 return len(thread_ids)
335 def cleanup_all(self):
336 """Clean up all tracked sessions (for shutdown)."""
337 with self._lock:
338 thread_ids = list(self._thread_credentials.keys())
340 for thread_id in thread_ids:
341 self.cleanup_thread(thread_id)
344# Global instance
345thread_session_manager = ThreadLocalSessionManager()
348def get_metrics_session(username: str, password: str) -> Optional[Session]:
349 """
350 Get a database session for metrics operations in the current thread.
351 The session is created once and reused until the thread's owned cleanup
352 boundary (request, task, or worker exit).
354 Note: This specifically uses create_thread_safe_session_for_metrics internally
355 and should only be used for metrics-related database operations.
357 Raises ``RekeyOperationalError`` when the database is unavailable for an
358 operational reason; see ``ThreadLocalSessionManager.get_session``.
359 """
360 return thread_session_manager.get_session(username, password)
363def get_current_thread_session() -> Optional[Session]:
364 """Get the current thread's session if it exists."""
365 return thread_session_manager.get_current_session()
368def cleanup_current_thread():
369 """Clean up every database session owned by the current thread.
371 Also clears any passwords cached by ``metrics_writer`` on this thread —
372 pooled worker threads must not retain plaintext credentials across tasks.
373 """
374 try:
375 thread_session_manager.cleanup_thread()
376 except Exception:
377 logger.debug(
378 "cleanup_current_thread: error closing thread-local DB session",
379 exc_info=True,
380 )
381 try:
382 # Lazy import avoids the encrypted_db -> thread_local_session import
383 # cycle during module bootstrap. This is the second session path:
384 # ambient SettingsManager callers use db_utils' per-thread LRU rather
385 # than ThreadLocalSessionManager.
386 from ..utilities.db_utils import (
387 cleanup_cached_user_sessions_current_thread,
388 )
390 cleanup_cached_user_sessions_current_thread()
391 except Exception:
392 # ERROR, not debug: this is the #6095 owner-thread registry
393 # cleanup that turns a GC-reclaimable session leak into a
394 # deterministic release. If it ever breaks (e.g. a cachetools
395 # upgrade renames ``cache_lock``), the fix silently reverts to the
396 # pre-#6095 behavior (an evicted session waits on the cyclic
397 # collector to reclaim its pooled connection) with nothing above
398 # debug level to notice.
399 #
400 # Uses logger.exception(), NOT logger.warning(..., exc_info=True):
401 # this module imports raw loguru (see the module imports above),
402 # which has no ``exc_info`` parameter. Passing it to warning()
403 # makes it an ordinary str.format kwarg that is silently discarded,
404 # so no traceback is emitted at all -- verified empirically.
405 # ``logger.exception()`` logs at ERROR, louder than WARNING, which
406 # suits a silent revert to pre-#6095 behaviour.
407 #
408 # It DOES attach the traceback: this is plain loguru, not
409 # SecureLogger, so nothing gates that on diagnose mode. The
410 # traceback still does not leave the process. The encrypted-DB sink
411 # persists ``_exception_context(record) + record["message"]`` -- an
412 # exception type/value prefix, never the frames -- and the frontend
413 # sink forwards only ``record["message"]``; see ``database_sink``,
414 # ``_exception_context`` and the frontend sink in
415 # ``utilities/log_utils.py``. It reaches the stderr/file sinks,
416 # which is exactly where it is wanted.
417 logger.exception(
418 "cleanup_current_thread: error closing cached DB sessions "
419 "(owner-thread registry cleanup failed — falling back to "
420 "GC-reclaimed cleanup for this thread's sessions)"
421 )
422 try:
423 from .thread_metrics import metrics_writer
425 metrics_writer.clear_passwords()
426 except Exception:
427 logger.debug(
428 "cleanup_current_thread: error clearing metrics_writer passwords",
429 exc_info=True,
430 )
433def clear_user_credentials(username: str) -> int:
434 """Drop ``username`` from the session manager's credential registry.
436 Called on logout so the plaintext SQLCipher password does not outlive
437 the session that authorised it. See
438 ``ThreadLocalSessionManager.clear_credentials_for_user`` for why the
439 registry cleanup differs from owner-thread cleanup.
441 This function clears only ``ThreadLocalSessionManager``'s process-wide
442 username-keyed registry. ``thread_metrics.metrics_writer`` storage is
443 owner-thread-local; normal owned worker boundaries call
444 ``cleanup_current_thread``, and direct integrations must do likewise.
445 """
446 try:
447 return thread_session_manager.clear_credentials_for_user(username)
448 except Exception:
449 logger.exception("Failed to clear cached credentials on logout")
450 return 0
453def cleanup_dead_threads():
454 """Sweep dead-thread credential entries from the session manager."""
455 try:
456 thread_session_manager.cleanup_dead_threads()
457 except Exception:
458 logger.warning("Dead-thread session sweep failed")
461class _ThreadCleanup(ContextDecorator):
462 """Context manager / decorator for thread-local resource cleanup."""
464 def __enter__(self):
465 return self
467 def __exit__(self, exc_type, exc_val, exc_tb):
468 try:
469 cleanup_current_thread()
470 except Exception:
471 logger.debug(
472 "thread_cleanup: error during DB session cleanup",
473 exc_info=True,
474 )
475 try:
476 from ..config.thread_settings import clear_settings_context
478 clear_settings_context()
479 except Exception:
480 logger.debug(
481 "thread_cleanup: error clearing settings context",
482 exc_info=True,
483 )
484 try:
485 from ..utilities.thread_context import clear_search_context
487 clear_search_context()
488 except Exception:
489 logger.debug(
490 "thread_cleanup: error clearing search context",
491 exc_info=True,
492 )
493 try:
494 # Defense-in-depth audit hook (PEP 578). Leaving an active
495 # EgressContext on a pooled worker thread would let the
496 # next, unrelated task inherit the previous run's scope.
497 from ..security.egress.audit_hook import clear_active_context
499 clear_active_context()
500 except Exception:
501 logger.debug(
502 "thread_cleanup: error clearing egress audit context",
503 exc_info=True,
504 )
505 return False
508def thread_cleanup(func=None):
509 """Ensure all thread-local resources are cleaned up when a function or block exits.
511 Works as a bare decorator, a decorator factory, or a context manager::
513 @thread_cleanup
514 def worker(): ...
516 @thread_cleanup()
517 def worker(): ...
519 with thread_cleanup():
520 ...
522 executor.submit(thread_cleanup(func), arg)
523 """
524 if func is not None:
526 @functools.wraps(func)
527 def wrapper(*args, **kwargs):
528 with _ThreadCleanup():
529 return func(*args, **kwargs)
531 return wrapper
532 return _ThreadCleanup()
535# Context manager for automatic cleanup
536class ThreadSessionContext:
537 """
538 Context manager that ensures thread session is cleaned up.
539 Usage:
540 with ThreadSessionContext(username, password) as session:
541 # Use session
542 """
544 def __init__(self, username: str, password: str):
545 self.username = username
546 self.password = password
547 self.session = None
549 def __enter__(self) -> Optional[Session]:
550 self.session = get_metrics_session(self.username, self.password)
551 return self.session
553 def __exit__(self, exc_type, exc_val, exc_tb):
554 # Don't cleanup here - let the thread keep its session
555 # Only cleanup when thread ends
556 pass