Coverage for src/local_deep_research/database/session_context.py: 93%
118 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-03 23:15 +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 yield session
184 finally:
185 # Only close if we created a new session (which we don't anymore)
186 if session and needs_close: 186 ↛ 187line 186 didn't jump to line 187 because the condition on line 186 was never true
187 try:
188 session.close()
189 except Exception:
190 logger.debug("Failed to close session during cleanup")
193def with_user_database(func: Callable) -> Callable:
194 """
195 Decorator that ensures function has access to user's database.
196 Injects 'db_session' as first argument to the decorated function.
198 Usage:
199 @with_user_database
200 def get_user_settings(db_session, setting_key):
201 return db_session.query(Setting).filter_by(key=setting_key).first()
202 """
204 @functools.wraps(func)
205 def wrapper(*args, **kwargs):
206 # Check if username/password provided in kwargs
207 username = kwargs.pop("_username", None)
208 password = kwargs.pop("_password", None)
210 with get_user_db_session(username, password) as db_session:
211 return func(db_session, *args, **kwargs)
213 return wrapper
216def ensure_db_session(view_func: Callable) -> Callable:
217 """
218 Flask view decorator that ensures database session is available.
219 Sets g.db_session for use in the request.
221 Usage:
222 @app.route('/my-route')
223 @ensure_db_session
224 def my_view():
225 # g.db_session is available here
226 settings = g.db_session.query(Setting).all()
227 """
229 @functools.wraps(view_func)
230 def wrapper(*args, **kwargs):
231 username = flask_session.get("username")
233 if not username:
234 # Let the view handle unauthenticated users
235 return view_func(*args, **kwargs)
237 try:
238 # Try to get or create session
239 if db_manager.is_user_connected(username):
240 g.db_session = db_manager.get_session(username)
241 else:
242 # Database not open - for encrypted DBs this means user needs to re-login
243 if db_manager.has_encryption:
244 # Clear session to force re-login
245 flask_session.clear()
246 from flask import redirect, url_for
248 return redirect(url_for("auth.login"))
249 # Try to reopen unencrypted database
250 logger.warning(
251 f"Reopening unencrypted database for {username} - "
252 "ensure this is intentional"
253 )
254 engine = db_manager.open_user_database(
255 username, UNENCRYPTED_DB_PLACEHOLDER
256 )
257 if engine:
258 g.db_session = db_manager.get_session(username)
260 except Exception:
261 logger.exception(
262 f"Failed to ensure database session for {username}"
263 )
264 return jsonify({"error": "Database session unavailable"}), 500
266 return view_func(*args, **kwargs)
268 return wrapper
271class DatabaseAccessMixin:
272 """
273 Mixin class for services that need database access.
274 Provides convenient methods for database operations.
275 """
277 def get_db_session(
278 self, username: Optional[str] = None
279 ) -> Optional[Session]:
280 """
281 DEPRECATED: This method returns a closed session due to context manager exit.
283 Use `with get_user_db_session(username) as session:` instead.
285 Raises:
286 DeprecationWarning: Always raised to prevent usage of broken method.
287 """
288 raise DeprecationWarning(
289 "get_db_session() is deprecated and returns a closed session. "
290 "Use `with get_user_db_session(username) as session:` instead."
291 )
293 @with_user_database
294 def execute_with_db(
295 self, db_session: Session, query_func: Callable, *args, **kwargs
296 ):
297 """Execute a function with database session."""
298 return query_func(db_session, *args, **kwargs)