Coverage for src / local_deep_research / database / session_context.py: 57%

1""" 

2Database session context manager and decorator for encrypted databases. 

3Ensures all database access has proper encryption context. 

4""" 

5 

6import functools 

7from contextlib import contextmanager 

8from typing import Callable, Optional 

9 

10from flask import g, has_app_context, session as flask_session 

11from loguru import logger 

12from sqlalchemy.orm import Session 

13 

14from ..utilities.thread_context import get_search_context 

15from .encrypted_db import db_manager 

16 

17# Placeholder password used when accessing unencrypted databases. 

18# This should only be used when LDR_ALLOW_UNENCRYPTED=true is set. 

19UNENCRYPTED_DB_PLACEHOLDER = "unencrypted-mode" 

20 

21 

22class DatabaseSessionError(Exception): 

23 """Raised when database session cannot be established.""" 

24 

25 pass 

26 

27 

28@contextmanager 

29def get_user_db_session( 

30 username: Optional[str] = None, password: Optional[str] = None 

31): 

32 """ 

33 Context manager that ensures proper database session with encryption. 

34 Now uses thread-local sessions for better performance. 

35 

36 Args: 

37 username: Username (if not provided, gets from Flask session) 

38 password: Password for encrypted database (required for first access) 

39 

40 Yields: 

41 Database session for the user 

42 

43 Raises: 

44 DatabaseSessionError: If session cannot be established 

45 """ 

46 # Import here to avoid circular imports 

47 from .thread_local_session import get_metrics_session 

48 from .session_passwords import session_password_store 

49 

50 session = None 

51 needs_close = False 

52 

53 try: 

54 # Get username from Flask session if not provided (only in Flask context) 

55 if not username and has_app_context(): 

56 username = flask_session.get("username") 

57 

58 if not username: 58 ↛ 59line 58 didn't jump to line 59 because the condition on line 58 was never true

59 raise DatabaseSessionError("No authenticated user") 

60 

61 # First, check if we have a session in Flask's g object (best performance) 

62 if has_app_context() and hasattr(g, "db_session") and g.db_session: 

63 # Use existing session from g - no need to close 

64 session = g.db_session 

65 needs_close = False 

66 else: 

67 # Get password if not provided 

68 if not password and has_app_context(): 

69 # Try to get from g 

70 if hasattr(g, "user_password"): 

71 password = g.user_password 

72 logger.debug( 

73 f"Got password from g.user_password for {username}" 

74 ) 

75 # Try session password store 

76 elif flask_session.get("session_id"): 76 ↛ 77line 76 didn't jump to line 77 because the condition on line 76 was never true

77 session_id = flask_session.get("session_id") 

78 logger.debug( 

79 f"Trying session password store for {username}" 

80 ) 

81 password = session_password_store.get_session_password( 

82 username, session_id 

83 ) 

84 if password: 

85 logger.debug( 

86 f"Got password from session store for {username}" 

87 ) 

88 else: 

89 logger.debug( 

90 f"No password in session store for {username}" 

91 ) 

92 

93 # Try thread context (for background threads) 

94 if not password: 

95 thread_context = get_search_context() 

96 if thread_context and thread_context.get("user_password"): 

97 password = thread_context["user_password"] 

98 logger.debug( 

99 f"Got password from thread context for {username}" 

100 ) 

101 

102 if not password and db_manager.has_encryption: 

103 raise DatabaseSessionError( 

104 f"Encrypted database for {username} requires password" 

105 ) 

106 elif not password: 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true

107 logger.warning( 

108 f"Accessing unencrypted database for {username} - " 

109 "ensure this is intentional (LDR_ALLOW_UNENCRYPTED=true)" 

110 ) 

111 password = UNENCRYPTED_DB_PLACEHOLDER 

112 

113 # Use thread-local session (will reuse existing or create new) 

114 session = get_metrics_session(username, password) 

115 if not session: 

116 raise DatabaseSessionError( 

117 f"Could not establish session for {username}" 

118 ) 

119 # Thread-local sessions are managed by the thread, don't close them 

120 needs_close = False 

121 

122 # Store the password we successfully used 

123 if password and has_app_context(): 

124 g.user_password = password 

125 

126 yield session 

127 

128 finally: 

129 # Only close if we created a new session (which we don't anymore) 

130 if session and needs_close: 130 ↛ 131line 130 didn't jump to line 131 because the condition on line 130 was never true

131 try: 

132 session.close() 

133 except Exception: 

134 pass 

135 

136 

137def with_user_database(func: Callable) -> Callable: 

138 """ 

139 Decorator that ensures function has access to user's database. 

140 Injects 'db_session' as first argument to the decorated function. 

141 

142 Usage: 

143 @with_user_database 

144 def get_user_settings(db_session, setting_key): 

145 return db_session.query(Setting).filter_by(key=setting_key).first() 

146 """ 

147 

148 @functools.wraps(func) 

149 def wrapper(*args, **kwargs): 

150 # Check if username/password provided in kwargs 

151 username = kwargs.pop("_username", None) 

152 password = kwargs.pop("_password", None) 

153 

154 with get_user_db_session(username, password) as db_session: 

155 return func(db_session, *args, **kwargs) 

156 

157 return wrapper 

158 

159 

160def ensure_db_session(view_func: Callable) -> Callable: 

161 """ 

162 Flask view decorator that ensures database session is available. 

163 Sets g.db_session for use in the request. 

164 

165 Usage: 

166 @app.route('/my-route') 

167 @ensure_db_session 

168 def my_view(): 

169 # g.db_session is available here 

170 settings = g.db_session.query(Setting).all() 

171 """ 

172 

173 @functools.wraps(view_func) 

174 def wrapper(*args, **kwargs): 

175 username = flask_session.get("username") 

176 

177 if not username: 

178 # Let the view handle unauthenticated users 

179 return view_func(*args, **kwargs) 

180 

181 try: 

182 # Try to get or create session 

183 if username in db_manager.connections: 

184 g.db_session = db_manager.get_session(username) 

185 else: 

186 # Database not open - for encrypted DBs this means user needs to re-login 

187 if db_manager.has_encryption: 

188 # Clear session to force re-login 

189 flask_session.clear() 

190 from flask import redirect, url_for 

191 

192 return redirect(url_for("auth.login")) 

193 else: 

194 # Try to reopen unencrypted database 

195 logger.warning( 

196 f"Reopening unencrypted database for {username} - " 

197 "ensure this is intentional" 

198 ) 

199 engine = db_manager.open_user_database( 

200 username, UNENCRYPTED_DB_PLACEHOLDER 

201 ) 

202 if engine: 

203 g.db_session = db_manager.get_session(username) 

204 

205 except Exception: 

206 logger.exception( 

207 f"Failed to ensure database session for {username}" 

208 ) 

209 

210 return view_func(*args, **kwargs) 

211 

212 return wrapper 

213 

214 

215class DatabaseAccessMixin: 

216 """ 

217 Mixin class for services that need database access. 

218 Provides convenient methods for database operations. 

219 """ 

220 

221 def get_db_session( 

222 self, username: Optional[str] = None 

223 ) -> Optional[Session]: 

224 """Get database session for user.""" 

225 try: 

226 with get_user_db_session(username) as session: 

227 return session 

228 except DatabaseSessionError: 

229 return None 

230 

231 @with_user_database 

232 def execute_with_db( 

233 self, db_session: Session, query_func: Callable, *args, **kwargs 

234 ): 

235 """Execute a function with database session.""" 

236 return query_func(db_session, *args, **kwargs)