Coverage for src / local_deep_research / database / session_context.py: 94%
111 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-14 23:55 +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 get_g_db_session() -> Optional[Session]:
35 """Lazily create and cache a DB session on Flask g for the current request.
37 Follows Flask's recommended pattern for lazy resource creation:
38 only check out a pool connection when a route actually needs one.
39 Returns None if no user is authenticated or DB is not connected.
40 """
41 if (
42 has_app_context()
43 and hasattr(g, "db_session")
44 and g.db_session is not None
45 ):
46 return g.db_session
47 username = getattr(g, "current_user", None) if has_app_context() else None
48 if not username or not db_manager.is_user_connected(username):
49 return None
50 try:
51 session = db_manager.get_session(username)
52 g.db_session = session
53 return session
54 except Exception:
55 logger.exception(f"Error lazily creating session for {username}")
56 return None
59@contextmanager
60def get_user_db_session(
61 username: Optional[str] = None, password: Optional[str] = None
62):
63 """
64 Context manager that ensures proper database session with encryption.
65 Now uses thread-local sessions for better performance.
67 Args:
68 username: Username (if not provided, gets from Flask session)
69 password: Password for encrypted database (required for first access)
71 Yields:
72 Database session for the user
74 Raises:
75 DatabaseSessionError: If session cannot be established
76 """
77 # Import here to avoid circular imports
78 from .thread_local_session import get_metrics_session
79 from .session_passwords import session_password_store
81 session = None
82 needs_close = False
84 try:
85 # Get username from Flask session if not provided (only in Flask context)
86 if not username and has_request_context():
87 username = flask_session.get("username")
89 if not username:
90 raise DatabaseSessionError("No authenticated user")
92 # First, try to get/create a session via Flask g (best performance)
93 if has_app_context():
94 cached = get_g_db_session()
95 if cached:
96 session = cached
97 needs_close = False
99 if not session:
100 # Get password if not provided
101 if not password and has_app_context():
102 # Try to get from g (works with app context)
103 if hasattr(g, "user_password"):
104 password = g.user_password
105 logger.debug(
106 f"Got password from g.user_password for {username}"
107 )
109 # Try session password store (requires request context for flask_session)
110 if not password and has_request_context():
111 session_id = flask_session.get("session_id")
112 if session_id:
113 logger.debug(
114 f"Trying session password store for {username}"
115 )
116 password = session_password_store.get_session_password(
117 username, session_id
118 )
119 if password:
120 logger.debug(
121 f"Got password from session store for {username}"
122 )
123 else:
124 logger.debug(
125 f"No password in session store for {username}"
126 )
128 # Try thread context (for background threads)
129 if not password:
130 thread_context = get_search_context()
131 if thread_context and thread_context.get("user_password"):
132 password = thread_context["user_password"]
133 logger.debug(
134 f"Got password from thread context for {username}"
135 )
137 if not password and db_manager.has_encryption:
138 raise DatabaseSessionError(
139 f"Encrypted database for {username} requires password"
140 )
141 elif not password:
142 logger.warning(
143 f"Accessing unencrypted database for {username} - "
144 "ensure this is intentional (LDR_ALLOW_UNENCRYPTED=true)"
145 )
146 password = UNENCRYPTED_DB_PLACEHOLDER
148 # Use thread-local session (will reuse existing or create new)
149 session = get_metrics_session(username, password)
150 if not session:
151 raise DatabaseSessionError(
152 f"Could not establish session for {username}"
153 )
154 # Thread-local sessions are managed by the thread, don't close them
155 needs_close = False
157 # Store the password we successfully used
158 if password and has_app_context():
159 g.user_password = password
161 yield session
163 finally:
164 # Only close if we created a new session (which we don't anymore)
165 if session and needs_close: 165 ↛ 166line 165 didn't jump to line 166 because the condition on line 165 was never true
166 try:
167 session.close()
168 except Exception:
169 logger.debug("Failed to close session during cleanup")
172def with_user_database(func: Callable) -> Callable:
173 """
174 Decorator that ensures function has access to user's database.
175 Injects 'db_session' as first argument to the decorated function.
177 Usage:
178 @with_user_database
179 def get_user_settings(db_session, setting_key):
180 return db_session.query(Setting).filter_by(key=setting_key).first()
181 """
183 @functools.wraps(func)
184 def wrapper(*args, **kwargs):
185 # Check if username/password provided in kwargs
186 username = kwargs.pop("_username", None)
187 password = kwargs.pop("_password", None)
189 with get_user_db_session(username, password) as db_session:
190 return func(db_session, *args, **kwargs)
192 return wrapper
195def ensure_db_session(view_func: Callable) -> Callable:
196 """
197 Flask view decorator that ensures database session is available.
198 Sets g.db_session for use in the request.
200 Usage:
201 @app.route('/my-route')
202 @ensure_db_session
203 def my_view():
204 # g.db_session is available here
205 settings = g.db_session.query(Setting).all()
206 """
208 @functools.wraps(view_func)
209 def wrapper(*args, **kwargs):
210 username = flask_session.get("username")
212 if not username:
213 # Let the view handle unauthenticated users
214 return view_func(*args, **kwargs)
216 try:
217 # Try to get or create session
218 if db_manager.is_user_connected(username):
219 g.db_session = db_manager.get_session(username)
220 else:
221 # Database not open - for encrypted DBs this means user needs to re-login
222 if db_manager.has_encryption:
223 # Clear session to force re-login
224 flask_session.clear()
225 from flask import redirect, url_for
227 return redirect(url_for("auth.login"))
228 # Try to reopen unencrypted database
229 logger.warning(
230 f"Reopening unencrypted database for {username} - "
231 "ensure this is intentional"
232 )
233 engine = db_manager.open_user_database(
234 username, UNENCRYPTED_DB_PLACEHOLDER
235 )
236 if engine:
237 g.db_session = db_manager.get_session(username)
239 except Exception:
240 logger.exception(
241 f"Failed to ensure database session for {username}"
242 )
243 return jsonify({"error": "Database session unavailable"}), 500
245 return view_func(*args, **kwargs)
247 return wrapper
250class DatabaseAccessMixin:
251 """
252 Mixin class for services that need database access.
253 Provides convenient methods for database operations.
254 """
256 def get_db_session(
257 self, username: Optional[str] = None
258 ) -> Optional[Session]:
259 """
260 DEPRECATED: This method returns a closed session due to context manager exit.
262 Use `with get_user_db_session(username) as session:` instead.
264 Raises:
265 DeprecationWarning: Always raised to prevent usage of broken method.
266 """
267 raise DeprecationWarning(
268 "get_db_session() is deprecated and returns a closed session. "
269 "Use `with get_user_db_session(username) as session:` instead."
270 )
272 @with_user_database
273 def execute_with_db(
274 self, db_session: Session, query_func: Callable, *args, **kwargs
275 ):
276 """Execute a function with database session."""
277 return query_func(db_session, *args, **kwargs)