Coverage for src/local_deep_research/database/thread_local_session.py: 93%
156 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-19 23:35 +0000
1"""
2Thread-local database session management.
3Each thread gets its own database session that persists for the thread's lifetime.
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 that is reused throughout the thread's lifetime.
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 get_session(self, username: str, password: str) -> Optional[Session]:
32 """
33 Get or create a database session for the current thread.
35 The session is created once per thread and reused for all subsequent calls.
36 This avoids the expensive SQLCipher decryption on every database access.
37 """
38 thread_id = threading.get_ident()
40 # Check if we already have a session for this thread
41 if hasattr(self._local, "session") and self._local.session:
42 # SECURITY: ensure cached session belongs to the requesting user
43 if getattr(self._local, "username", None) != username:
44 logger.warning(
45 f"Thread {thread_id}: Session username mismatch "
46 f"(cached={self._local.username!r}, requested={username!r}), "
47 "clearing stale cross-user session"
48 )
49 self._cleanup_thread_session()
50 else:
51 # Verify it's still valid
52 try:
53 self._local.session.execute(text("SELECT 1"))
54 # Under DEFERRED isolation the validation SELECT opens
55 # a transaction that holds a SHARED lock on SQLite
56 # until an explicit commit/rollback. A long-lived
57 # thread-local session reused across requests would
58 # keep that lock held and block the first writer.
59 # Roll it back so subsequent callers start fresh.
60 self._local.session.rollback()
61 return self._local.session
62 except PendingRollbackError:
63 # Session has a pending rollback (e.g. from a previous database lock error).
64 # Attempt rollback to recover without destroying the session.
65 logger.debug(
66 f"Thread {thread_id}: PendingRollbackError, attempting rollback recovery"
67 )
68 try:
69 self._local.session.rollback()
70 self._local.session.execute(text("SELECT 1"))
71 self._local.session.rollback()
72 return self._local.session
73 except Exception:
74 logger.warning(
75 f"Thread {thread_id}: Rollback recovery failed, creating new session"
76 )
77 self._cleanup_thread_session()
78 except Exception:
79 # Session is invalid, will create a new one
80 logger.debug(
81 f"Thread {thread_id}: Existing session invalid, creating new one"
82 )
83 self._cleanup_thread_session()
85 # Create new session for this thread
86 logger.debug(
87 f"Thread {thread_id}: Creating new database session for user {username}"
88 )
90 # Ensure database is open. open_user_database returns None for
91 # credential failures and raises DatabaseInitializationError when
92 # the schema can't be initialised; from a worker-thread caller
93 # both mean "no usable session right now" so collapse them.
94 from .encrypted_db import DatabaseInitializationError
96 try:
97 engine = db_manager.open_user_database(username, password)
98 except DatabaseInitializationError:
99 # ``logger.warning`` (no traceback) rather than
100 # ``logger.exception``: ``password`` is a live local in this
101 # frame, so rendering a traceback under ``diagnose=True``
102 # would dump the plaintext SQLCipher master password
103 # (unrecoverable — TRUST.md §5). The redacted failure detail
104 # is already logged at the raise site in
105 # ``open_user_database`` (#4182).
106 logger.warning(
107 f"Thread {thread_id}: database init failed for user {username}"
108 )
109 return None
110 if not engine:
111 logger.error(
112 f"Thread {thread_id}: Failed to open database for user {username}"
113 )
114 return None
116 # Create session for this thread
117 session = db_manager.create_thread_safe_session_for_metrics(
118 username, password
119 )
120 if not session: 120 ↛ 121line 120 didn't jump to line 121 because the condition on line 120 was never true
121 logger.error(
122 f"Thread {thread_id}: Failed to create session for user {username}"
123 )
124 return None
126 # Store in thread-local storage
127 self._local.session = session
128 self._local.username = username
130 # Track credentials for cleanup
131 with self._lock:
132 self._thread_credentials[thread_id] = (username, password)
134 return session
136 def get_current_session(self) -> Optional[Session]:
137 """Get the current thread's session if it exists."""
138 if hasattr(self._local, "session"):
139 return self._local.session
140 return None
142 def _cleanup_thread_session(self):
143 """Clean up the current thread's session.
145 Sessions are bound to the shared per-user QueuePool engine, so
146 closing the session returns its connection to the pool — there
147 are no per-thread engines to dispose.
148 """
149 thread_id = threading.get_ident()
151 if hasattr(self._local, "session") and self._local.session:
152 try:
153 self._local.session.rollback()
154 except Exception:
155 logger.warning(
156 f"Thread {thread_id}: Error rolling back session during cleanup"
157 )
158 try:
159 self._local.session.close()
160 logger.debug(f"Thread {thread_id}: Closed database session")
161 except Exception:
162 logger.warning(f"Thread {thread_id}: Error closing session")
163 finally:
164 self._local.session = None
166 if hasattr(self._local, "username"):
167 self._local.username = None
169 # Remove from tracking
170 with self._lock:
171 self._thread_credentials.pop(thread_id, None)
173 def cleanup_thread(self, thread_id: Optional[int] = None):
174 """
175 Clean up session for a specific thread or current thread.
176 Called when a thread is finishing.
177 """
178 if thread_id is None:
179 thread_id = threading.get_ident()
181 # If it's the current thread, we can clean up directly
182 if thread_id == threading.get_ident():
183 self._cleanup_thread_session()
184 else:
185 # For other threads, just remove from tracking
186 # The thread-local storage will be cleaned up when the thread ends
187 with self._lock:
188 self._thread_credentials.pop(thread_id, None)
190 def cleanup_dead_threads(self):
191 """Remove credential entries for threads that are no longer alive.
193 Handles the abnormal case of threads that died without triggering
194 their cleanup handler. Uses threading.enumerate() to identify
195 alive threads and removes credential entries for dead ones.
196 Sessions on dead threads are garbage-collected normally and
197 their connections return to the shared per-user QueuePool.
199 Called from:
200 - processor_v2.py: every ~60s in the queue loop
201 - app_factory.py: in teardown_appcontext
202 - connection_cleanup.py: in cleanup_idle_connections (every ~300s)
203 """
204 alive_ids = {t.ident for t in threading.enumerate()}
205 with self._lock:
206 dead_ids = [
207 tid for tid in self._thread_credentials if tid not in alive_ids
208 ]
209 for tid in dead_ids:
210 del self._thread_credentials[tid]
211 if dead_ids:
212 logger.debug(f"Swept {len(dead_ids)} dead thread credential(s)")
214 def cleanup_all(self):
215 """Clean up all tracked sessions (for shutdown)."""
216 with self._lock:
217 thread_ids = list(self._thread_credentials.keys())
219 for thread_id in thread_ids:
220 self.cleanup_thread(thread_id)
223# Global instance
224thread_session_manager = ThreadLocalSessionManager()
227def get_metrics_session(username: str, password: str) -> Optional[Session]:
228 """
229 Get a database session for metrics operations in the current thread.
230 The session is created once and reused for the thread's lifetime.
232 Note: This specifically uses create_thread_safe_session_for_metrics internally
233 and should only be used for metrics-related database operations.
234 """
235 return thread_session_manager.get_session(username, password)
238def get_current_thread_session() -> Optional[Session]:
239 """Get the current thread's session if it exists."""
240 return thread_session_manager.get_current_session()
243def cleanup_current_thread():
244 """Clean up the current thread's database session and cached credentials.
246 Also clears any passwords cached by ``metrics_writer`` on this thread —
247 pooled worker threads must not retain plaintext credentials across tasks.
248 """
249 thread_session_manager.cleanup_thread()
250 try:
251 from .thread_metrics import metrics_writer
253 metrics_writer.clear_passwords()
254 except Exception:
255 logger.debug(
256 "cleanup_current_thread: error clearing metrics_writer passwords",
257 exc_info=True,
258 )
261def cleanup_dead_threads():
262 """Sweep dead-thread credential entries from the session manager."""
263 try:
264 thread_session_manager.cleanup_dead_threads()
265 except Exception:
266 logger.warning("Dead-thread session sweep failed")
269class _ThreadCleanup(ContextDecorator):
270 """Context manager / decorator for thread-local resource cleanup."""
272 def __enter__(self):
273 return self
275 def __exit__(self, exc_type, exc_val, exc_tb):
276 try:
277 cleanup_current_thread()
278 except Exception:
279 logger.debug(
280 "thread_cleanup: error during DB session cleanup",
281 exc_info=True,
282 )
283 try:
284 from ..config.thread_settings import clear_settings_context
286 clear_settings_context()
287 except Exception:
288 logger.debug(
289 "thread_cleanup: error clearing settings context",
290 exc_info=True,
291 )
292 try:
293 from ..utilities.thread_context import clear_search_context
295 clear_search_context()
296 except Exception:
297 logger.debug(
298 "thread_cleanup: error clearing search context",
299 exc_info=True,
300 )
301 try:
302 # Defense-in-depth audit hook (PEP 578). Leaving an active
303 # EgressContext on a pooled worker thread would let the
304 # next, unrelated task inherit the previous run's scope.
305 from ..security.egress.audit_hook import clear_active_context
307 clear_active_context()
308 except Exception:
309 logger.debug(
310 "thread_cleanup: error clearing egress audit context",
311 exc_info=True,
312 )
313 return False
316def thread_cleanup(func=None):
317 """Ensure all thread-local resources are cleaned up when a function or block exits.
319 Works as a bare decorator, a decorator factory, or a context manager::
321 @thread_cleanup
322 def worker(): ...
324 @thread_cleanup()
325 def worker(): ...
327 with thread_cleanup():
328 ...
330 executor.submit(thread_cleanup(func), arg)
331 """
332 if func is not None:
334 @functools.wraps(func)
335 def wrapper(*args, **kwargs):
336 with _ThreadCleanup():
337 return func(*args, **kwargs)
339 return wrapper
340 return _ThreadCleanup()
343# Context manager for automatic cleanup
344class ThreadSessionContext:
345 """
346 Context manager that ensures thread session is cleaned up.
347 Usage:
348 with ThreadSessionContext(username, password) as session:
349 # Use session
350 """
352 def __init__(self, username: str, password: str):
353 self.username = username
354 self.password = password
355 self.session = None
357 def __enter__(self) -> Optional[Session]:
358 self.session = get_metrics_session(self.username, self.password)
359 return self.session
361 def __exit__(self, exc_type, exc_val, exc_tb):
362 # Don't cleanup here - let the thread keep its session
363 # Only cleanup when thread ends
364 pass