Coverage for src/local_deep_research/utilities/db_utils.py: 98%
77 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +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.
91 """
92 user_session = db_manager.get_session(username)
93 if user_session:
94 return user_session
95 raise NoUserDatabaseError(f"No database found for user {username}")
98def cleanup_cached_user_sessions_current_thread() -> int:
99 """Close sessions cached or acquired by the current worker thread.
101 ``_get_cached_user_session`` uses a process-wide LRU with the worker's
102 UUID in each key. FastAPI reuses those workers across requests, so the
103 entries must be removed and closed on the worker itself when a request or
104 owned background task finishes. Sessions evicted from the LRU are also
105 closed via the owner-thread registry populated by ``get_db_session``.
107 Returns the number of distinct sessions closed.
108 """
109 thread_key = getattr(g_thread_local_store, "thread_id", None)
110 tracked = getattr(g_thread_local_store, "db_sessions", {})
111 sessions = dict(tracked)
113 cache = _get_cached_user_session.cache
114 lock = _get_cached_user_session.cache_lock
115 if thread_key is not None:
116 # Detach under cachetools' own lock, but close outside it. Session
117 # close can perform driver work and must not serialize unrelated
118 # workers' cache access.
119 with lock:
120 current_keys = [
121 key for key in list(cache) if key and key[0] == thread_key
122 ]
123 for key in current_keys:
124 session = cache.pop(key)
125 sessions[id(session)] = session
127 if hasattr(g_thread_local_store, "db_sessions"):
128 del g_thread_local_store.db_sessions
130 for session in sessions.values():
131 try:
132 session.close()
133 except Exception:
134 # This session has already been dropped from both the LRU
135 # cache and the owner-thread registry above, so this is the
136 # only place the actual close failure is ever observable — a
137 # bare warning with no traceback makes a real driver-level
138 # failure indistinguishable from routine noise.
139 #
140 # logger.exception(), not logger.warning(..., exc_info=True):
141 # this module imports raw loguru (see the module imports
142 # above), which has no ``exc_info`` parameter, so that kwarg is
143 # silently swallowed as a format argument and yields no
144 # traceback whatsoever.
145 #
146 # ``logger.exception()`` DOES attach the traceback — this is
147 # plain loguru, not SecureLogger, so nothing gates it on
148 # diagnose mode. The traceback does not leave the process: the
149 # encrypted-DB sink persists ``_exception_context(record) +
150 # record["message"]`` (exception type/value prefix only, never
151 # the frames) and the frontend sink forwards only
152 # ``record["message"]``; see ``database_sink``,
153 # ``_exception_context`` and the frontend sink in
154 # ``utilities/log_utils.py``. It reaches the stderr/file sinks,
155 # which is the point.
156 logger.exception(
157 "Failed to close a cached DB session on worker cleanup"
158 )
160 return len(sessions)
163def get_settings_manager(
164 db_session: Session | None = None, username: str | None = None
165):
166 """
167 Get the settings manager for the current context.
169 Args:
170 db_session: Optional database session
171 username: Optional username for caching (required for SettingsManager)
173 Returns:
174 The appropriate settings manager instance.
175 """
176 # Track whether we are borrowing a caller-provided session we don't own.
177 # Borrowed sessions must NOT be closed by SettingsManager — their owner is
178 # responsible for cleanup.
179 borrowed_session = db_session is not None
181 # Resolve the current user from the request-context contextvar when the
182 # caller passed neither a session nor a username. Pre-migration this block
183 # reused the Flask request's g.db_session and read the username from
184 # flask.session; both are gone. get_db_session() below resolves the user
185 # via the same contextvar, so there is no Flask request session to borrow.
186 if db_session is None and username is None:
187 username = get_current_username()
189 if db_session is None:
190 try:
191 db_session = get_db_session(username=username)
192 except NoUserDatabaseError:
193 # This user has no database, so defaults are the only answer we
194 # have. Deliberately NOT `except RuntimeError`: that also caught
195 # the background-thread guard above and any failure from inside
196 # db_manager.get_session, and answering either of those with
197 # anonymous defaults is silently wrong rather than degraded --
198 # a worker would read another configuration than the user's, and
199 # an unopenable database would read as "every setting is at its
200 # default" with nothing logged. Both now propagate.
201 db_session = None
202 username = "anonymous"
204 # Import here to avoid circular imports
205 from ..settings import SettingsManager
207 logger.debug(
208 "get_settings_manager: session_source={}, owned={}",
209 "borrowed" if borrowed_session else ("new" if db_session else "None"),
210 not borrowed_session,
211 )
213 # Always use regular SettingsManager (now with built-in simple caching)
214 return SettingsManager(db_session, owns_session=not borrowed_session)
217def no_db_settings(func: Callable[..., Any]) -> Callable[..., Any]:
218 """
219 Decorator that runs the wrapped function with the settings database
220 completely disabled. This will prevent the function from accidentally
221 reading settings from the DB. Settings can only be read from environment
222 variables or the defaults file.
224 Args:
225 func: The function to wrap.
227 Returns:
228 The wrapped function.
230 """
232 @functools.wraps(func)
233 def wrapper(*args, **kwargs):
234 # Temporarily disable DB access in the settings manager.
235 manager = get_settings_manager()
236 db_session = manager.db_session
237 manager.db_session = None
239 try:
240 return func(*args, **kwargs)
241 finally:
242 # Restore the original database session.
243 manager.db_session = db_session
245 return wrapper