Coverage for src/local_deep_research/database/session_context.py: 94%
122 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-19 23:35 +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 flask import (
11 g,
12 has_app_context,
13 has_request_context,
14 jsonify,
15 session as flask_session,
16)
17from loguru import logger
18from sqlalchemy.orm import Session
20from ..utilities.thread_context import get_search_context
21from .encrypted_db import db_manager
23# Placeholder password used when accessing unencrypted databases.
24# This should only be used when LDR_ALLOW_UNENCRYPTED=true is set.
25UNENCRYPTED_DB_PLACEHOLDER = "unencrypted-mode"
28class DatabaseSessionError(Exception):
29 """Raised when database session cannot be established."""
31 pass
34def safe_rollback(session: Session, context: str = "") -> None:
35 """Roll back the session, swallowing and logging any rollback failure.
37 SQLAlchemy requires explicit rollback after a failed flush/commit before
38 the session is usable again. Skipping it leaves the session in
39 PendingRollbackError state and every subsequent ORM operation cascades.
41 This helper exists so call sites can recover the session in one line
42 without repeating the try/except/log boilerplate at every except handler.
43 ``context`` is included in the error log so failed rollbacks can be
44 traced back to the call site.
45 """
46 try:
47 session.rollback()
48 except Exception:
49 if context: 49 ↛ 52line 49 didn't jump to line 52 because the condition on line 49 was always true
50 logger.exception(f"Failed to rollback session: {context}")
51 else:
52 logger.exception("Failed to rollback session")
55def get_g_db_session() -> Optional[Session]:
56 """Lazily create and cache a DB session on Flask g for the current request.
58 Follows Flask's recommended pattern for lazy resource creation:
59 only check out a pool connection when a route actually needs one.
60 Returns None if no user is authenticated or DB is not connected.
61 """
62 if (
63 has_app_context()
64 and hasattr(g, "db_session")
65 and g.db_session is not None
66 ):
67 return g.db_session
68 username = getattr(g, "current_user", None) if has_app_context() else None
69 if not username or not db_manager.is_user_connected(username):
70 return None
71 try:
72 session = db_manager.get_session(username)
73 g.db_session = session
74 return session
75 except Exception:
76 logger.exception(f"Error lazily creating session for {username}")
77 return None
80@contextmanager
81def get_user_db_session(
82 username: Optional[str] = None, password: Optional[str] = None
83):
84 """
85 Context manager that ensures proper database session with encryption.
86 Now uses thread-local sessions for better performance.
88 Args:
89 username: Username (if not provided, gets from Flask session)
90 password: Password for encrypted database (required for first access)
92 Yields:
93 Database session for the user
95 Raises:
96 DatabaseSessionError: If session cannot be established
97 """
98 # Import here to avoid circular imports
99 from .thread_local_session import get_metrics_session
100 from .session_passwords import session_password_store
102 session = None
103 needs_close = False
105 try:
106 # Get username from Flask session if not provided (only in Flask context)
107 if not username and has_request_context():
108 username = flask_session.get("username")
110 if not username:
111 raise DatabaseSessionError("No authenticated user")
113 # First, try to get/create a session via Flask g (best performance)
114 if has_app_context():
115 cached = get_g_db_session()
116 if cached:
117 session = cached
118 needs_close = False
120 if not session:
121 # Get password if not provided
122 if not password and has_app_context():
123 # Try to get from g (works with app context)
124 if hasattr(g, "user_password"):
125 password = g.user_password
126 logger.debug(
127 f"Got password from g.user_password for {username}"
128 )
130 # Try session password store (requires request context for flask_session)
131 if not password and has_request_context():
132 session_id = flask_session.get("session_id")
133 if session_id:
134 logger.debug(
135 f"Trying session password store for {username}"
136 )
137 password = session_password_store.get_session_password(
138 username, session_id
139 )
140 if password:
141 logger.debug(
142 f"Got password from session store for {username}"
143 )
144 else:
145 logger.debug(
146 f"No password in session store for {username}"
147 )
149 # Try thread context (for background threads)
150 if not password:
151 thread_context = get_search_context()
152 if thread_context and thread_context.get("user_password"):
153 password = thread_context["user_password"]
154 logger.debug(
155 f"Got password from thread context for {username}"
156 )
158 if not password and db_manager.has_encryption:
159 raise DatabaseSessionError(
160 f"Encrypted database for {username} requires password"
161 )
162 elif not password:
163 logger.warning(
164 f"Accessing unencrypted database for {username} - "
165 "ensure this is intentional (LDR_ALLOW_UNENCRYPTED=true)"
166 )
167 password = UNENCRYPTED_DB_PLACEHOLDER
169 # Use thread-local session (will reuse existing or create new)
170 session = get_metrics_session(username, password)
171 if not session:
172 raise DatabaseSessionError(
173 f"Could not establish session for {username}"
174 )
175 # Thread-local sessions are managed by the thread, don't close them
176 needs_close = False
178 # Store the password we successfully used
179 if password and has_app_context():
180 g.user_password = password
182 # Wrap only the yield: the session is established and non-None here,
183 # and we want to recover *the caller's* block — not setup failures
184 # above, where there may be no usable session to roll back.
185 try:
186 yield session
187 except Exception:
188 # The yielded session is a *reused* thread-local session, not a
189 # fresh one closed on exit. If the caller's ``with`` block raised
190 # (most importantly a failed ``session.commit()``/``flush()``),
191 # the session is left in ``PendingRollbackError`` state and the
192 # next operation on this thread cascades. Roll it back here so an
193 # unguarded ``with`` block can't poison the thread, then re-raise
194 # so the original error still surfaces to the caller.
195 safe_rollback(session, "get_user_db_session")
196 raise
198 finally:
199 # Only close if we created a new session (which we don't anymore)
200 if session and needs_close: 200 ↛ 201line 200 didn't jump to line 201 because the condition on line 200 was never true
201 try:
202 session.close()
203 except Exception:
204 logger.debug("Failed to close session during cleanup")
207def with_user_database(func: Callable) -> Callable:
208 """
209 Decorator that ensures function has access to user's database.
210 Injects 'db_session' as first argument to the decorated function.
212 Usage:
213 @with_user_database
214 def get_user_settings(db_session, setting_key):
215 return db_session.query(Setting).filter_by(key=setting_key).first()
216 """
218 @functools.wraps(func)
219 def wrapper(*args, **kwargs):
220 # Check if username/password provided in kwargs
221 username = kwargs.pop("_username", None)
222 password = kwargs.pop("_password", None)
224 with get_user_db_session(username, password) as db_session:
225 return func(db_session, *args, **kwargs)
227 return wrapper
230def ensure_db_session(view_func: Callable) -> Callable:
231 """
232 Flask view decorator that ensures database session is available.
233 Sets g.db_session for use in the request.
235 Usage:
236 @app.route('/my-route')
237 @ensure_db_session
238 def my_view():
239 # g.db_session is available here
240 settings = g.db_session.query(Setting).all()
241 """
243 @functools.wraps(view_func)
244 def wrapper(*args, **kwargs):
245 username = flask_session.get("username")
247 if not username:
248 # Let the view handle unauthenticated users
249 return view_func(*args, **kwargs)
251 try:
252 # Try to get or create session
253 if db_manager.is_user_connected(username):
254 g.db_session = db_manager.get_session(username)
255 else:
256 # Database not open - for encrypted DBs this means user needs to re-login
257 if db_manager.has_encryption:
258 # Clear session to force re-login
259 flask_session.clear()
260 from flask import redirect, url_for
262 return redirect(url_for("auth.login"))
263 # Try to reopen unencrypted database
264 logger.warning(
265 f"Reopening unencrypted database for {username} - "
266 "ensure this is intentional"
267 )
268 engine = db_manager.open_user_database(
269 username, UNENCRYPTED_DB_PLACEHOLDER
270 )
271 if engine:
272 g.db_session = db_manager.get_session(username)
274 except Exception:
275 logger.exception(
276 f"Failed to ensure database session for {username}"
277 )
278 return jsonify({"error": "Database session unavailable"}), 500
280 return view_func(*args, **kwargs)
282 return wrapper
285class DatabaseAccessMixin:
286 """
287 Mixin class for services that need database access.
288 Provides convenient methods for database operations.
289 """
291 def get_db_session(
292 self, username: Optional[str] = None
293 ) -> Optional[Session]:
294 """
295 DEPRECATED: This method returns a closed session due to context manager exit.
297 Use `with get_user_db_session(username) as session:` instead.
299 Raises:
300 DeprecationWarning: Always raised to prevent usage of broken method.
301 """
302 raise DeprecationWarning(
303 "get_db_session() is deprecated and returns a closed session. "
304 "Use `with get_user_db_session(username) as session:` instead."
305 )
307 @with_user_database
308 def execute_with_db(
309 self, db_session: Session, query_func: Callable, *args, **kwargs
310 ):
311 """Execute a function with database session."""
312 return query_func(db_session, *args, **kwargs)