Coverage for src/local_deep_research/database/thread_local_session.py: 92%
194 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +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.
58 """
59 thread_id = threading.get_ident()
61 # Check if we already have a session for this thread
62 if hasattr(self._local, "session") and self._local.session:
63 # SECURITY: ensure cached session belongs to the requesting user
64 if getattr(self._local, "username", None) != username:
65 logger.warning(
66 f"Thread {thread_id}: Session username mismatch "
67 f"(cached={self._local.username!r}, requested={username!r}), "
68 "clearing stale cross-user session"
69 )
70 self._cleanup_thread_session()
71 else:
72 # Verify it's still valid
73 try:
74 self._local.session.execute(text("SELECT 1"))
75 if not self.in_scope():
76 # Under DEFERRED isolation the validation SELECT
77 # opens a transaction that holds a SHARED lock on
78 # SQLite until an explicit commit/rollback. A
79 # long-lived thread-local session reused across
80 # requests would keep that lock held and block
81 # the first writer. Roll it back so subsequent
82 # callers start fresh.
83 self._local.session.rollback()
84 # else: re-entrant call from inside an enclosing
85 # get_user_db_session block on this thread — the
86 # validation SELECT joined the caller's open
87 # transaction, and rolling back here would silently
88 # destroy the caller's uncommitted writes (concrete
89 # data loss in library_rag_service.index_document,
90 # whose nested helpers open their own sessions).
91 return self._local.session
92 except PendingRollbackError:
93 # Session has a pending rollback (e.g. from a previous database lock error).
94 # Attempt rollback to recover without destroying the session.
95 logger.debug(
96 f"Thread {thread_id}: PendingRollbackError, attempting rollback recovery"
97 )
98 try:
99 self._local.session.rollback()
100 self._local.session.execute(text("SELECT 1"))
101 self._local.session.rollback()
102 return self._local.session
103 except Exception:
104 logger.warning(
105 f"Thread {thread_id}: Rollback recovery failed, creating new session"
106 )
107 self._cleanup_thread_session()
108 except Exception:
109 # Session is invalid, will create a new one
110 logger.debug(
111 f"Thread {thread_id}: Existing session invalid, creating new one"
112 )
113 self._cleanup_thread_session()
115 # Create new session for this thread
116 logger.debug(
117 f"Thread {thread_id}: Creating new database session for user {username}"
118 )
120 # Ensure database is open. open_user_database returns None for
121 # credential failures and raises DatabaseInitializationError when
122 # the schema can't be initialised; from a worker-thread caller
123 # both mean "no usable session right now" so collapse them.
124 from .encrypted_db import DatabaseInitializationError
126 try:
127 engine = db_manager.open_user_database(username, password)
128 except DatabaseInitializationError:
129 # ``logger.warning`` (no traceback) rather than
130 # ``logger.exception``: ``password`` is a live local in this
131 # frame, so rendering a traceback under ``diagnose=True``
132 # would dump the plaintext SQLCipher master password
133 # (unrecoverable — TRUST.md §5). The redacted failure detail
134 # is already logged at the raise site in
135 # ``open_user_database`` (#4182).
136 logger.warning(
137 f"Thread {thread_id}: database init failed for user {username}"
138 )
139 return None
140 if not engine:
141 logger.error(
142 f"Thread {thread_id}: Failed to open database for user {username}"
143 )
144 return None
146 # Create session for this thread
147 session = db_manager.create_thread_safe_session_for_metrics(
148 username, password
149 )
150 if not session: 150 ↛ 151line 150 didn't jump to line 151 because the condition on line 150 was never true
151 logger.error(
152 f"Thread {thread_id}: Failed to create session for user {username}"
153 )
154 return None
156 # Store in thread-local storage
157 self._local.session = session
158 self._local.username = username
160 # Track credentials for cleanup
161 with self._lock:
162 self._thread_credentials[thread_id] = (username, password)
164 return session
166 def get_current_session(self) -> Optional[Session]:
167 """Get the current thread's session if it exists."""
168 if hasattr(self._local, "session"):
169 return self._local.session
170 return None
172 def _cleanup_thread_session(self):
173 """Clean up the current thread's session.
175 Sessions are bound to the shared per-user QueuePool engine, so
176 closing the session returns its connection to the pool — there
177 are no per-thread engines to dispose.
178 """
179 thread_id = threading.get_ident()
181 if hasattr(self._local, "session") and self._local.session:
182 try:
183 self._local.session.rollback()
184 except Exception:
185 logger.warning(
186 f"Thread {thread_id}: Error rolling back session during cleanup"
187 )
188 try:
189 self._local.session.close()
190 logger.debug(f"Thread {thread_id}: Closed database session")
191 except Exception:
192 logger.warning(f"Thread {thread_id}: Error closing session")
193 finally:
194 self._local.session = None
196 if hasattr(self._local, "username"):
197 self._local.username = None
199 # Remove from tracking
200 with self._lock:
201 self._thread_credentials.pop(thread_id, None)
203 def reset_session_if_matches(self, session: "Session | None") -> bool:
204 """If ``session`` is the current thread's cached session, clear it.
206 Used by ``safe_rollback`` in ``database/session_context.py`` to heal
207 the thread-local session cache when an error tells us the session
208 is structurally unusable (a provisioning race or an invalidated
209 cursor). The next caller on this thread then gets a fresh session
210 from the shared ``QueuePool`` instead of inheriting the broken
211 one.
213 Identity-checked: a caller that hands in a session from a
214 different thread (or a session that isn't the cached one at all
215 — e.g. one borrowed from ``g.db_session``) won't accidentally
216 clear someone else's cache. Returns ``True`` if a reset
217 actually happened, ``False`` otherwise.
218 """
219 if session is None:
220 return False
221 if not hasattr(self._local, "session") or self._local.session is None: 221 ↛ 222line 221 didn't jump to line 222 because the condition on line 221 was never true
222 return False
223 if self._local.session is not session:
224 return False
225 self._cleanup_thread_session()
226 return True
228 def cleanup_thread(self, thread_id: Optional[int] = None):
229 """
230 Clean up session for a specific thread or current thread.
231 Called when a thread is finishing.
232 """
233 if thread_id is None:
234 thread_id = threading.get_ident()
236 # If it's the current thread, we can clean up directly
237 if thread_id == threading.get_ident():
238 self._cleanup_thread_session()
239 else:
240 # For other threads, just remove from tracking
241 # The thread-local storage will be cleaned up when the thread ends
242 with self._lock:
243 self._thread_credentials.pop(thread_id, None)
245 def cleanup_dead_threads(self):
246 """Remove credential entries for threads that are no longer alive.
248 Handles the abnormal case of threads that died without triggering
249 their cleanup handler. Uses threading.enumerate() to identify
250 alive threads and removes credential entries for dead ones.
251 Sessions on dead threads are garbage-collected normally and
252 their connections return to the shared per-user QueuePool.
254 Called from:
255 - processor_v2.py: every ~60s in the queue loop
256 - fastapi_app.py: in the request middleware's post-request cleanup
257 - connection_cleanup.py: in cleanup_idle_connections (every ~300s)
258 """
259 alive_ids = {t.ident for t in threading.enumerate()}
260 with self._lock:
261 dead_ids = [
262 tid for tid in self._thread_credentials if tid not in alive_ids
263 ]
264 for tid in dead_ids:
265 del self._thread_credentials[tid]
266 if dead_ids:
267 logger.debug(f"Swept {len(dead_ids)} dead thread credential(s)")
269 def clear_credentials_for_user(self, username: str) -> int:
270 """Drop cached plaintext credentials for one user, on any thread.
272 Entries here hold the user's plaintext SQLCipher password so a
273 pooled worker can rebuild its session without a round-trip to the
274 password store. Synchronous routes normally drop their owner-thread
275 entry through ``WorkerCleanupAPIRoute``. This username-wide cleanup
276 is still required on logout because another worker may be executing
277 background work for the same user, and one thread cannot clear a
278 different thread's ``threading.local`` state.
280 Only the tracking dict is cleared. The owning thread's
281 ``threading.local`` session is deliberately left alone: one thread
282 cannot reach into another's local storage, and the session is
283 harmless once the credential is gone — the username re-validation at
284 the top of ``get_session()`` evicts it before any reuse. In-flight
285 work is unaffected too, since ``get_session()`` is passed the
286 password explicitly and never reads it back out of this dict.
288 Returns the number of entries dropped.
289 """
290 with self._lock:
291 thread_ids = [
292 tid
293 for tid, (user, _password) in self._thread_credentials.items()
294 if user == username
295 ]
296 for tid in thread_ids:
297 del self._thread_credentials[tid]
298 if thread_ids:
299 logger.debug(
300 f"Cleared cached credentials on {len(thread_ids)} thread(s)"
301 )
302 return len(thread_ids)
304 def cleanup_all(self):
305 """Clean up all tracked sessions (for shutdown)."""
306 with self._lock:
307 thread_ids = list(self._thread_credentials.keys())
309 for thread_id in thread_ids:
310 self.cleanup_thread(thread_id)
313# Global instance
314thread_session_manager = ThreadLocalSessionManager()
317def get_metrics_session(username: str, password: str) -> Optional[Session]:
318 """
319 Get a database session for metrics operations in the current thread.
320 The session is created once and reused until the thread's owned cleanup
321 boundary (request, task, or worker exit).
323 Note: This specifically uses create_thread_safe_session_for_metrics internally
324 and should only be used for metrics-related database operations.
325 """
326 return thread_session_manager.get_session(username, password)
329def get_current_thread_session() -> Optional[Session]:
330 """Get the current thread's session if it exists."""
331 return thread_session_manager.get_current_session()
334def cleanup_current_thread():
335 """Clean up every database session owned by the current thread.
337 Also clears any passwords cached by ``metrics_writer`` on this thread —
338 pooled worker threads must not retain plaintext credentials across tasks.
339 """
340 try:
341 thread_session_manager.cleanup_thread()
342 except Exception:
343 logger.debug(
344 "cleanup_current_thread: error closing thread-local DB session",
345 exc_info=True,
346 )
347 try:
348 # Lazy import avoids the encrypted_db -> thread_local_session import
349 # cycle during module bootstrap. This is the second session path:
350 # ambient SettingsManager callers use db_utils' per-thread LRU rather
351 # than ThreadLocalSessionManager.
352 from ..utilities.db_utils import (
353 cleanup_cached_user_sessions_current_thread,
354 )
356 cleanup_cached_user_sessions_current_thread()
357 except Exception:
358 # ERROR, not debug: this is the #6095 owner-thread registry
359 # cleanup that turns a GC-reclaimable session leak into a
360 # deterministic release. If it ever breaks (e.g. a cachetools
361 # upgrade renames ``cache_lock``), the fix silently reverts to the
362 # pre-#6095 behavior (an evicted session waits on the cyclic
363 # collector to reclaim its pooled connection) with nothing above
364 # debug level to notice.
365 #
366 # Uses logger.exception(), NOT logger.warning(..., exc_info=True):
367 # this module imports raw loguru (see the module imports above),
368 # which has no ``exc_info`` parameter. Passing it to warning()
369 # makes it an ordinary str.format kwarg that is silently discarded,
370 # so no traceback is emitted at all -- verified empirically.
371 # ``logger.exception()`` logs at ERROR, louder than WARNING, which
372 # suits a silent revert to pre-#6095 behaviour.
373 #
374 # It DOES attach the traceback: this is plain loguru, not
375 # SecureLogger, so nothing gates that on diagnose mode. The
376 # traceback still does not leave the process. The encrypted-DB sink
377 # persists ``_exception_context(record) + record["message"]`` -- an
378 # exception type/value prefix, never the frames -- and the frontend
379 # sink forwards only ``record["message"]``; see ``database_sink``,
380 # ``_exception_context`` and the frontend sink in
381 # ``utilities/log_utils.py``. It reaches the stderr/file sinks,
382 # which is exactly where it is wanted.
383 logger.exception(
384 "cleanup_current_thread: error closing cached DB sessions "
385 "(owner-thread registry cleanup failed — falling back to "
386 "GC-reclaimed cleanup for this thread's sessions)"
387 )
388 try:
389 from .thread_metrics import metrics_writer
391 metrics_writer.clear_passwords()
392 except Exception:
393 logger.debug(
394 "cleanup_current_thread: error clearing metrics_writer passwords",
395 exc_info=True,
396 )
399def clear_user_credentials(username: str) -> int:
400 """Drop ``username`` from the session manager's credential registry.
402 Called on logout so the plaintext SQLCipher password does not outlive
403 the session that authorised it. See
404 ``ThreadLocalSessionManager.clear_credentials_for_user`` for why the
405 registry cleanup differs from owner-thread cleanup.
407 This function clears only ``ThreadLocalSessionManager``'s process-wide
408 username-keyed registry. ``thread_metrics.metrics_writer`` storage is
409 owner-thread-local; normal owned worker boundaries call
410 ``cleanup_current_thread``, and direct integrations must do likewise.
411 """
412 try:
413 return thread_session_manager.clear_credentials_for_user(username)
414 except Exception:
415 logger.exception("Failed to clear cached credentials on logout")
416 return 0
419def cleanup_dead_threads():
420 """Sweep dead-thread credential entries from the session manager."""
421 try:
422 thread_session_manager.cleanup_dead_threads()
423 except Exception:
424 logger.warning("Dead-thread session sweep failed")
427class _ThreadCleanup(ContextDecorator):
428 """Context manager / decorator for thread-local resource cleanup."""
430 def __enter__(self):
431 return self
433 def __exit__(self, exc_type, exc_val, exc_tb):
434 try:
435 cleanup_current_thread()
436 except Exception:
437 logger.debug(
438 "thread_cleanup: error during DB session cleanup",
439 exc_info=True,
440 )
441 try:
442 from ..config.thread_settings import clear_settings_context
444 clear_settings_context()
445 except Exception:
446 logger.debug(
447 "thread_cleanup: error clearing settings context",
448 exc_info=True,
449 )
450 try:
451 from ..utilities.thread_context import clear_search_context
453 clear_search_context()
454 except Exception:
455 logger.debug(
456 "thread_cleanup: error clearing search context",
457 exc_info=True,
458 )
459 try:
460 # Defense-in-depth audit hook (PEP 578). Leaving an active
461 # EgressContext on a pooled worker thread would let the
462 # next, unrelated task inherit the previous run's scope.
463 from ..security.egress.audit_hook import clear_active_context
465 clear_active_context()
466 except Exception:
467 logger.debug(
468 "thread_cleanup: error clearing egress audit context",
469 exc_info=True,
470 )
471 return False
474def thread_cleanup(func=None):
475 """Ensure all thread-local resources are cleaned up when a function or block exits.
477 Works as a bare decorator, a decorator factory, or a context manager::
479 @thread_cleanup
480 def worker(): ...
482 @thread_cleanup()
483 def worker(): ...
485 with thread_cleanup():
486 ...
488 executor.submit(thread_cleanup(func), arg)
489 """
490 if func is not None:
492 @functools.wraps(func)
493 def wrapper(*args, **kwargs):
494 with _ThreadCleanup():
495 return func(*args, **kwargs)
497 return wrapper
498 return _ThreadCleanup()
501# Context manager for automatic cleanup
502class ThreadSessionContext:
503 """
504 Context manager that ensures thread session is cleaned up.
505 Usage:
506 with ThreadSessionContext(username, password) as session:
507 # Use session
508 """
510 def __init__(self, username: str, password: str):
511 self.username = username
512 self.password = password
513 self.session = None
515 def __enter__(self) -> Optional[Session]:
516 self.session = get_metrics_session(self.username, self.password)
517 return self.session
519 def __exit__(self, exc_type, exc_val, exc_tb):
520 # Don't cleanup here - let the thread keep its session
521 # Only cleanup when thread ends
522 pass