Coverage for src/local_deep_research/database/session_context.py: 98%
80 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1"""
2Database session context manager and decorator for encrypted databases.
3Ensures all database access has proper encryption context.
4"""
6import functools
7from contextlib import contextmanager
8from typing import Callable, Optional
10from loguru import logger
11from sqlalchemy.exc import SQLAlchemyError
12from sqlalchemy.orm import Session
14from ..utilities.thread_context import get_search_context
15from .encrypted_db import db_manager
16from .thread_local_session import thread_session_manager
18# Placeholder password used when accessing unencrypted databases.
19# This should only be used when LDR_ALLOW_UNENCRYPTED=true is set.
20UNENCRYPTED_DB_PLACEHOLDER = "unencrypted-mode"
23class DatabaseSessionError(Exception):
24 """Raised when database session cannot be established."""
26 pass
29def safe_rollback(session: Session, context: str = "") -> None:
30 """Roll back the session, swallowing and logging any rollback failure.
32 SQLAlchemy requires explicit rollback after a failed flush/commit before
33 the session is usable again. Skipping it leaves the session in
34 PendingRollbackError state and every subsequent ORM operation cascades.
36 This helper exists so call sites can recover the session in one line
37 without repeating the try/except/log boilerplate at every except handler.
38 ``context`` is included in the error log so failed rollbacks can be
39 traced back to the call site.
41 Two SQLAlchemy error shapes are treated as "the session is structurally
42 unusable — give up on it, drop the thread-local cache, and let the next
43 caller get a fresh one" rather than as loud failures:
45 * ``InvalidRequestError("...provisioning a new connection; concurrent
46 operations are not permitted...")`` — the second thread racing on the
47 per-user QueuePool.
48 * ``InterfaceError("Cursor needed to be reset because of commit/rollback
49 and can no longer be fetched from")`` — a cursor was invalidated by a
50 commit/rollback that fired between the original ``execute()`` and the
51 lazy-attribute fetch that followed.
53 In both cases SQLAlchemy has already marked the session unrecoverable, so
54 a no-op rollback is correct, the thread-local cache is cleared, and the
55 failure is logged at DEBUG so the production stderr stream stays clean.
56 Other ``SQLAlchemyError`` failures still hit the loud ``logger.exception``
57 path (unless their message matches a known broken-session signature) —
58 the session is normally recoverable via rollback and the operator needs
59 to see them.
60 """
61 if session is None:
62 return
63 log_msg = (
64 f"Failed to rollback session: {context}"
65 if context
66 else "Failed to rollback session"
67 )
68 try:
69 session.rollback()
70 except SQLAlchemyError as exc:
71 msg = str(exc)
72 # Message-substring matching pinned against SQLAlchemy 2.0+ QueuePool error messages
73 # ('provisioning a new connection', 'concurrent operations are not permitted').
74 is_provisioning_race = (
75 "provisioning a new connection" in msg
76 or "concurrent operations are not permitted" in msg
77 )
78 # Requires 'Cursor needed to be reset' message match so other InterfaceErrors
79 # (e.g., driver/connectivity failures) are not quietly swallowed into the DEBUG path.
80 is_cursor_invalidated = "Cursor needed to be reset" in msg
81 if is_provisioning_race or is_cursor_invalidated:
82 label = f": {context}" if context else ""
83 logger.debug(f"safe_rollback — resetting broken session{label}")
84 # Drop the thread-local cache so the next caller on this
85 # thread gets a fresh session. Identity-checked inside the
86 # helper, so a caller that hands in a session that ISN'T
87 # the cached one (e.g. borrowed from ``g.db_session`` or
88 # owned by a different thread) won't accidentally clear
89 # someone else's cache. The reset itself is best-effort —
90 # never let it raise past ``safe_rollback`` (call sites are
91 # themselves in except handlers).
92 try:
93 thread_session_manager.reset_session_if_matches(session)
94 except Exception:
95 logger.debug(
96 f"safe_rollback: reset_session_if_matches raised for{label}"
97 )
98 return
99 logger.exception(log_msg)
100 except Exception:
101 logger.exception(log_msg)
104@contextmanager
105def get_user_db_session(
106 username: Optional[str] = None,
107 password: Optional[str] = None,
108 session_id: Optional[str] = None,
109):
110 """
111 Context manager that ensures proper database session with encryption.
112 Now uses thread-local sessions for better performance.
114 Args:
115 username: Username (required; must be passed explicitly under FastAPI).
116 password: Password for encrypted database (required for first access).
117 session_id: Optional session ID for exact per-session password lookup.
118 Request handlers should pass the current request's session_id so
119 two concurrent sessions for the same user can't cross-pollinate.
120 If omitted, the resolver falls back to scanning any active session
121 for the user — fine for background threads, unsafe for request
122 handlers.
124 Yields:
125 Database session for the user
127 Raises:
128 DatabaseSessionError: If session cannot be established
129 """
130 # Import here to avoid circular imports
131 from .thread_local_session import get_metrics_session
132 from .session_passwords import session_password_store
134 if not username:
135 raise DatabaseSessionError("No authenticated user")
137 # Resolve password from the provided session_id, then fall back to
138 # any active session, then the thread context for background workers.
139 if not password and session_id:
140 password = session_password_store.get_session_password(
141 username, session_id
142 )
143 if password: 143 ↛ 146line 143 didn't jump to line 146 because the condition on line 143 was always true
144 logger.debug(f"Got password from session store for {username}")
146 if not password:
147 # Scan active sessions for this user. Safe for background threads;
148 # request handlers should pass session_id explicitly.
149 password = session_password_store.get_any_session_password(username)
151 if not password:
152 thread_context = get_search_context()
153 if thread_context and thread_context.get("user_password"):
154 password = thread_context["user_password"]
155 logger.debug(f"Got password from thread context for {username}")
157 if not password and db_manager.has_encryption:
158 raise DatabaseSessionError(
159 f"Encrypted database for {username} requires password"
160 )
161 if not password:
162 logger.warning(
163 f"Accessing unencrypted database for {username} - "
164 "ensure this is intentional (LDR_ALLOW_UNENCRYPTED=true)"
165 )
166 password = UNENCRYPTED_DB_PLACEHOLDER
168 session = get_metrics_session(username, password)
169 if not session:
170 raise DatabaseSessionError(
171 f"Could not establish session for {username}"
172 )
174 # Thread-local sessions are managed by the thread — do not close
175 # here. But we MUST wrap the yield in try/except so an exception
176 # inside the `with` block doesn't leave a half-committed
177 # transaction attached to this thread's session. The next caller
178 # on the same thread would otherwise inherit that dirty state.
179 # Actual connection close happens in `cleanup_current_thread()`,
180 # called by middleware / worker-loop finally blocks.
181 #
182 # The scope depth tells get_session's re-validation whether an
183 # enclosing block is still active on this thread: a nested
184 # get_user_db_session call (e.g. a helper opening its own session
185 # inside a caller's with-block) must NOT trigger the stale-lock
186 # rollback, or the caller's uncommitted writes are destroyed.
187 from .thread_local_session import thread_session_manager
189 thread_session_manager.enter_scope()
190 try:
191 yield session
192 except Exception:
193 # The yielded session is a *reused* thread-local session, not a
194 # fresh one closed on exit. If the caller's ``with`` block raised
195 # (most importantly a failed ``session.commit()``/``flush()``),
196 # the session is left in ``PendingRollbackError`` state and the
197 # next operation on this thread cascades. Roll it back here so an
198 # unguarded ``with`` block can't poison the thread, then re-raise
199 # so the original error still surfaces to the caller.
200 safe_rollback(session, "get_user_db_session")
201 raise
202 finally:
203 thread_session_manager.exit_scope()
206def with_user_database(func: Callable) -> Callable:
207 """
208 Decorator that ensures function has access to user's database.
209 Injects 'db_session' as first argument to the decorated function.
211 Usage:
212 @with_user_database
213 def get_user_settings(db_session, setting_key):
214 return db_session.query(Setting).filter_by(key=setting_key).first()
215 """
217 @functools.wraps(func)
218 def wrapper(*args, **kwargs):
219 # Check if username/password provided in kwargs
220 username = kwargs.pop("_username", None)
221 password = kwargs.pop("_password", None)
223 with get_user_db_session(username, password) as db_session:
224 return func(db_session, *args, **kwargs)
226 return wrapper
229class DatabaseAccessMixin:
230 """
231 Mixin class for services that need database access.
232 Provides convenient methods for database operations.
233 """
235 def get_db_session(
236 self, username: Optional[str] = None
237 ) -> Optional[Session]:
238 """
239 DEPRECATED: This method returns a closed session due to context manager exit.
241 Use `with get_user_db_session(username) as session:` instead.
243 Raises:
244 DeprecationWarning: Always raised to prevent usage of broken method.
245 """
246 raise DeprecationWarning(
247 "get_db_session() is deprecated and returns a closed session. "
248 "Use `with get_user_db_session(username) as session:` instead."
249 )
251 @with_user_database
252 def execute_with_db(
253 self, db_session: Session, query_func: Callable, *args, **kwargs
254 ):
255 """Execute a function with database session."""
256 return query_func(db_session, *args, **kwargs)