Coverage for src/local_deep_research/utilities/db_utils.py: 98%
77 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1import functools
2from typing import Any, Callable
4from cachetools import LRUCache
5from loguru import logger
6from sqlalchemy.orm import Session
8from ..config.paths import get_data_directory
9from ..exceptions import NoUserDatabaseError
10from ..database.encrypted_db import db_manager
11from .threading_utils import g_thread_local_store, thread_specific_cache
12from .request_context import get_current_username
14# Database paths using new centralized configuration
15DATA_DIR = get_data_directory()
16# DB_PATH removed - use per-user encrypted databases instead
19def get_db_session(
20 _namespace: str = "", username: str | None = None
21) -> Session:
22 """
23 Get database session - uses encrypted per-user database if authenticated.
25 Args:
26 _namespace: This can be specified to an arbitrary string in order to
27 force the caching mechanism to create separate settings even in
28 the same thread. Usually it does not need to be specified.
29 username: Optional username for thread context (e.g., background research threads).
30 If not provided, will be resolved from the request-context contextvar.
32 Returns:
33 The database session for the current user/context.
34 """
35 import threading
37 if not username:
38 # Resolve the authenticated user from the request contextvar.
39 # DatabaseMiddleware sets it per request, and Starlette copies the
40 # request context into the threadpool workers that run sync route
41 # handlers — so this resolves there too, not just on MainThread.
42 username = get_current_username()
44 # Only refuse callers with NO request context off the main thread:
45 # true background threads must use get_user_db_session() (or receive a
46 # settings_snapshot). Under Flask the equivalent guard was
47 # has_app_context(); the FastAPI port initially refused EVERY
48 # non-MainThread caller — but sync route handlers run in the anyio
49 # threadpool, so every no-arg settings read from a sync route silently
50 # fell back to anonymous defaults instead of the user's settings.
51 if not username and threading.current_thread().name != "MainThread":
52 thread_id = threading.get_ident()
53 raise RuntimeError(
54 f"Database access attempted from background thread "
55 f"'{threading.current_thread().name}' (ID: {thread_id}) with no "
56 f"request context. Use get_user_db_session() or pass all "
57 f"required data to the thread at creation time."
58 )
60 if not username:
61 # MainThread without an authenticated user (CLI, startup).
62 logger.warning(
63 "get_db_session() is deprecated. Use get_user_db_session() from database.session_context"
64 )
65 return None
67 session = _get_cached_user_session(username, _namespace)
69 # Keep an owner-thread registry in addition to the bounded shared cache.
70 # The cache can evict this thread's entry while the session is still in
71 # use, so scanning the cache alone at request teardown is insufficient:
72 # an evicted SQLAlchemy Session retains its QueuePool connection until a
73 # cyclic-GC pass happens to reclaim it. The registry lets the worker that
74 # acquired the session close it deterministically at its cleanup boundary.
75 tracked = getattr(g_thread_local_store, "db_sessions", None)
76 if tracked is None:
77 tracked = {}
78 g_thread_local_store.db_sessions = tracked
79 tracked[id(session)] = session
80 return session
83@thread_specific_cache(cache=LRUCache(maxsize=10))
84def _get_cached_user_session(username: str, _namespace: str = "") -> Session:
85 """Per-thread session cache, keyed by the RESOLVED username.
87 The cache must sit below username resolution: caching at the outer
88 call (where username is usually None) would key two different users'
89 requests served by the same threadpool worker to one entry, handing
90 user B a session opened for user A.
92 ``NoUserDatabaseError`` (which ``get_settings_manager`` answers with
93 the defaults file) is raised only for ``get_session``'s None: no usable
94 engine is cached. A rekey sidecar or epoch fault on a cached engine is
95 raised by ``get_session`` itself as a ``RekeyOperationalError`` and
96 passes through here untouched, so it never reads as "no database".
97 """
98 user_session = db_manager.get_session(username)
99 if user_session:
100 return user_session
101 raise NoUserDatabaseError(f"No database found for user {username}")
104def cleanup_cached_user_sessions_current_thread() -> int:
105 """Close sessions cached or acquired by the current worker thread.
107 ``_get_cached_user_session`` uses a process-wide LRU with the worker's
108 UUID in each key. FastAPI reuses those workers across requests, so the
109 entries must be removed and closed on the worker itself when a request or
110 owned background task finishes. Sessions evicted from the LRU are also
111 closed via the owner-thread registry populated by ``get_db_session``.
113 Returns the number of distinct sessions closed.
114 """
115 thread_key = getattr(g_thread_local_store, "thread_id", None)
116 tracked = getattr(g_thread_local_store, "db_sessions", {})
117 sessions = dict(tracked)
119 cache = _get_cached_user_session.cache
120 lock = _get_cached_user_session.cache_lock
121 if thread_key is not None:
122 # Detach under cachetools' own lock, but close outside it. Session
123 # close can perform driver work and must not serialize unrelated
124 # workers' cache access.
125 with lock:
126 current_keys = [
127 key for key in list(cache) if key and key[0] == thread_key
128 ]
129 for key in current_keys:
130 session = cache.pop(key)
131 sessions[id(session)] = session
133 if hasattr(g_thread_local_store, "db_sessions"):
134 del g_thread_local_store.db_sessions
136 for session in sessions.values():
137 try:
138 session.close()
139 except Exception:
140 # This session has already been dropped from both the LRU
141 # cache and the owner-thread registry above, so this is the
142 # only place the actual close failure is ever observable — a
143 # bare warning with no traceback makes a real driver-level
144 # failure indistinguishable from routine noise.
145 #
146 # logger.exception(), not logger.warning(..., exc_info=True):
147 # this module imports raw loguru (see the module imports
148 # above), which has no ``exc_info`` parameter, so that kwarg is
149 # silently swallowed as a format argument and yields no
150 # traceback whatsoever.
151 #
152 # ``logger.exception()`` DOES attach the traceback — this is
153 # plain loguru, not SecureLogger, so nothing gates it on
154 # diagnose mode. The traceback does not leave the process: the
155 # encrypted-DB sink persists ``_exception_context(record) +
156 # record["message"]`` (exception type/value prefix only, never
157 # the frames) and the frontend sink forwards only
158 # ``record["message"]``; see ``database_sink``,
159 # ``_exception_context`` and the frontend sink in
160 # ``utilities/log_utils.py``. It reaches the stderr/file sinks,
161 # which is the point.
162 logger.exception(
163 "Failed to close a cached DB session on worker cleanup"
164 )
166 return len(sessions)
169def get_settings_manager(
170 db_session: Session | None = None, username: str | None = None
171):
172 """
173 Get the settings manager for the current context.
175 Args:
176 db_session: Optional database session
177 username: Optional username for caching (required for SettingsManager)
179 Returns:
180 The appropriate settings manager instance.
181 """
182 # Track whether we are borrowing a caller-provided session we don't own.
183 # Borrowed sessions must NOT be closed by SettingsManager — their owner is
184 # responsible for cleanup.
185 borrowed_session = db_session is not None
187 # Resolve the current user from the request-context contextvar when the
188 # caller passed neither a session nor a username. Pre-migration this block
189 # reused the Flask request's g.db_session and read the username from
190 # flask.session; both are gone. get_db_session() below resolves the user
191 # via the same contextvar, so there is no Flask request session to borrow.
192 if db_session is None and username is None:
193 username = get_current_username()
195 if db_session is None:
196 try:
197 db_session = get_db_session(username=username)
198 except NoUserDatabaseError:
199 # This user has no database, so defaults are the only answer we
200 # have. Deliberately NOT `except RuntimeError`: that also caught
201 # the background-thread guard above and any failure from inside
202 # db_manager.get_session, and answering either of those with
203 # anonymous defaults is silently wrong rather than degraded --
204 # a worker would read another configuration than the user's, and
205 # an unopenable database would read as "every setting is at its
206 # default" with nothing logged. Both now propagate. So does a
207 # rekey sidecar or epoch fault on the user's cached engine (a
208 # RekeyOperationalError from db_manager.get_session, e.g. an
209 # fd-exhaustion spike): read as defaults, a stored "strict"
210 # egress scope would let a public download through. The app
211 # answers it with a retryable 503.
212 db_session = None
213 username = "anonymous"
215 # Import here to avoid circular imports
216 from ..settings import SettingsManager
218 logger.debug(
219 "get_settings_manager: session_source={}, owned={}",
220 "borrowed" if borrowed_session else ("new" if db_session else "None"),
221 not borrowed_session,
222 )
224 # Always use regular SettingsManager (now with built-in simple caching)
225 return SettingsManager(db_session, owns_session=not borrowed_session)
228def no_db_settings(func: Callable[..., Any]) -> Callable[..., Any]:
229 """
230 Decorator that runs the wrapped function with the settings database
231 completely disabled. This will prevent the function from accidentally
232 reading settings from the DB. Settings can only be read from environment
233 variables or the defaults file.
235 Args:
236 func: The function to wrap.
238 Returns:
239 The wrapped function.
241 """
243 @functools.wraps(func)
244 def wrapper(*args, **kwargs):
245 # Temporarily disable DB access in the settings manager.
246 manager = get_settings_manager()
247 db_session = manager.db_session
248 manager.db_session = None
250 try:
251 return func(*args, **kwargs)
252 finally:
253 # Restore the original database session.
254 manager.db_session = db_session
256 return wrapper