Coverage for src / local_deep_research / database / session_context.py: 94%
94 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-02-25 01:07 +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 session as flask_session,
15)
16from loguru import logger
17from sqlalchemy.orm import Session
19from ..utilities.thread_context import get_search_context
20from .encrypted_db import db_manager
22# Placeholder password used when accessing unencrypted databases.
23# This should only be used when LDR_ALLOW_UNENCRYPTED=true is set.
24UNENCRYPTED_DB_PLACEHOLDER = "unencrypted-mode"
27class DatabaseSessionError(Exception):
28 """Raised when database session cannot be established."""
30 pass
33@contextmanager
34def get_user_db_session(
35 username: Optional[str] = None, password: Optional[str] = None
36):
37 """
38 Context manager that ensures proper database session with encryption.
39 Now uses thread-local sessions for better performance.
41 Args:
42 username: Username (if not provided, gets from Flask session)
43 password: Password for encrypted database (required for first access)
45 Yields:
46 Database session for the user
48 Raises:
49 DatabaseSessionError: If session cannot be established
50 """
51 # Import here to avoid circular imports
52 from .thread_local_session import get_metrics_session
53 from .session_passwords import session_password_store
55 session = None
56 needs_close = False
58 try:
59 # Get username from Flask session if not provided (only in Flask context)
60 if not username and has_request_context():
61 username = flask_session.get("username")
63 if not username:
64 raise DatabaseSessionError("No authenticated user")
66 # First, check if we have a session in Flask's g object (best performance)
67 if has_app_context() and hasattr(g, "db_session") and g.db_session:
68 # Use existing session from g - no need to close
69 session = g.db_session
70 needs_close = False
71 else:
72 # Get password if not provided
73 if not password and has_app_context():
74 # Try to get from g (works with app context)
75 if hasattr(g, "user_password"):
76 password = g.user_password
77 logger.debug(
78 f"Got password from g.user_password for {username}"
79 )
81 # Try session password store (requires request context for flask_session)
82 if not password and has_request_context():
83 session_id = flask_session.get("session_id")
84 if session_id:
85 logger.debug(
86 f"Trying session password store for {username}"
87 )
88 password = session_password_store.get_session_password(
89 username, session_id
90 )
91 if password: 91 ↛ 96line 91 didn't jump to line 96 because the condition on line 91 was always true
92 logger.debug(
93 f"Got password from session store for {username}"
94 )
95 else:
96 logger.debug(
97 f"No password in session store for {username}"
98 )
100 # Try thread context (for background threads)
101 if not password:
102 thread_context = get_search_context()
103 if thread_context and thread_context.get("user_password"):
104 password = thread_context["user_password"]
105 logger.debug(
106 f"Got password from thread context for {username}"
107 )
109 if not password and db_manager.has_encryption:
110 raise DatabaseSessionError(
111 f"Encrypted database for {username} requires password"
112 )
113 elif not password:
114 logger.warning(
115 f"Accessing unencrypted database for {username} - "
116 "ensure this is intentional (LDR_ALLOW_UNENCRYPTED=true)"
117 )
118 password = UNENCRYPTED_DB_PLACEHOLDER
120 # Use thread-local session (will reuse existing or create new)
121 session = get_metrics_session(username, password)
122 if not session:
123 raise DatabaseSessionError(
124 f"Could not establish session for {username}"
125 )
126 # Thread-local sessions are managed by the thread, don't close them
127 needs_close = False
129 # Store the password we successfully used
130 if password and has_app_context():
131 g.user_password = password
133 yield session
135 finally:
136 # Only close if we created a new session (which we don't anymore)
137 if session and needs_close: 137 ↛ 138line 137 didn't jump to line 138 because the condition on line 137 was never true
138 try:
139 session.close()
140 except Exception:
141 logger.debug("Failed to close session during cleanup")
144def with_user_database(func: Callable) -> Callable:
145 """
146 Decorator that ensures function has access to user's database.
147 Injects 'db_session' as first argument to the decorated function.
149 Usage:
150 @with_user_database
151 def get_user_settings(db_session, setting_key):
152 return db_session.query(Setting).filter_by(key=setting_key).first()
153 """
155 @functools.wraps(func)
156 def wrapper(*args, **kwargs):
157 # Check if username/password provided in kwargs
158 username = kwargs.pop("_username", None)
159 password = kwargs.pop("_password", None)
161 with get_user_db_session(username, password) as db_session:
162 return func(db_session, *args, **kwargs)
164 return wrapper
167def ensure_db_session(view_func: Callable) -> Callable:
168 """
169 Flask view decorator that ensures database session is available.
170 Sets g.db_session for use in the request.
172 Usage:
173 @app.route('/my-route')
174 @ensure_db_session
175 def my_view():
176 # g.db_session is available here
177 settings = g.db_session.query(Setting).all()
178 """
180 @functools.wraps(view_func)
181 def wrapper(*args, **kwargs):
182 username = flask_session.get("username")
184 if not username:
185 # Let the view handle unauthenticated users
186 return view_func(*args, **kwargs)
188 try:
189 # Try to get or create session
190 if db_manager.is_user_connected(username):
191 g.db_session = db_manager.get_session(username)
192 else:
193 # Database not open - for encrypted DBs this means user needs to re-login
194 if db_manager.has_encryption:
195 # Clear session to force re-login
196 flask_session.clear()
197 from flask import redirect, url_for
199 return redirect(url_for("auth.login"))
200 else:
201 # Try to reopen unencrypted database
202 logger.warning(
203 f"Reopening unencrypted database for {username} - "
204 "ensure this is intentional"
205 )
206 engine = db_manager.open_user_database(
207 username, UNENCRYPTED_DB_PLACEHOLDER
208 )
209 if engine:
210 g.db_session = db_manager.get_session(username)
212 except Exception:
213 logger.exception(
214 f"Failed to ensure database session for {username}"
215 )
217 return view_func(*args, **kwargs)
219 return wrapper
222class DatabaseAccessMixin:
223 """
224 Mixin class for services that need database access.
225 Provides convenient methods for database operations.
226 """
228 def get_db_session(
229 self, username: Optional[str] = None
230 ) -> Optional[Session]:
231 """
232 DEPRECATED: This method returns a closed session due to context manager exit.
234 Use `with get_user_db_session(username) as session:` instead.
236 Raises:
237 DeprecationWarning: Always raised to prevent usage of broken method.
238 """
239 raise DeprecationWarning(
240 "get_db_session() is deprecated and returns a closed session. "
241 "Use `with get_user_db_session(username) as session:` instead."
242 )
244 @with_user_database
245 def execute_with_db(
246 self, db_session: Session, query_func: Callable, *args, **kwargs
247 ):
248 """Execute a function with database session."""
249 return query_func(db_session, *args, **kwargs)