Coverage for src / local_deep_research / utilities / db_utils.py: 85%

1import functools 

2from typing import Any, Callable, Dict 

3 

4from cachetools import LRUCache 

5from flask import ( 

6 g, 

7 has_app_context, 

8 has_request_context, 

9 session as flask_session, 

10) 

11from loguru import logger 

12from sqlalchemy.orm import Session 

13 

14from ..settings.env_registry import use_fallback_llm 

15 

16from ..config.paths import get_data_directory 

17from ..database.encrypted_db import db_manager 

18from .threading_utils import thread_specific_cache 

19 

20# Database paths using new centralized configuration 

21DATA_DIR = get_data_directory() 

22# DB_PATH removed - use per-user encrypted databases instead 

23 

24 

25@thread_specific_cache(cache=LRUCache(maxsize=10)) 

26def get_db_session( 

27 _namespace: str = "", username: str | None = None 

28) -> Session: 

29 """ 

30 Get database session - uses encrypted per-user database if authenticated. 

31 

32 Args: 

33 _namespace: This can be specified to an arbitrary string in order to 

34 force the caching mechanism to create separate settings even in 

35 the same thread. Usually it does not need to be specified. 

36 username: Optional username for thread context (e.g., background research threads). 

37 If not provided, will try to get from Flask context. 

38 

39 Returns: 

40 The database session for the current user/context. 

41 """ 

42 # CRITICAL: Detect if we're in a background thread and raise an error 

43 # This helps identify code that's trying to access the database from threads 

44 import threading 

45 

46 # Check if we're in a background thread (not in Flask request context) 

47 # We check for request context specifically because app context might exist 

48 # during startup but we still shouldn't access the database from background threads 

49 thread_name = threading.current_thread().name 

50 

51 # Allow MainThread during startup, but not other threads 

52 if not has_app_context() and thread_name != "MainThread": 

53 thread_id = threading.get_ident() 

54 raise RuntimeError( 

55 f"Database access attempted from background thread '{thread_name}' (ID: {thread_id}). " 

56 f"Database access from threads is not allowed due to SQLite thread safety constraints. " 

57 f"Use settings_snapshot or pass all required data to the thread at creation time." 

58 ) 

59 

60 # If username is explicitly provided (e.g., from background thread) 

61 if username: 

62 user_session = db_manager.get_session(username) 

63 if user_session: 

64 return user_session 

65 raise RuntimeError(f"No database found for user {username}") 

66 

67 # Otherwise, check Flask request context 

68 try: 

69 # Check if we have a database session in Flask's g object 

70 if hasattr(g, "db_session") and g.db_session: 70 ↛ 74line 70 didn't jump to line 74 because the condition on line 70 was always true

71 return g.db_session 

72 

73 # Check if we have a username in the Flask session 

74 username = flask_session.get("username") 

75 if username: 

76 user_session = db_manager.get_session(username) 

77 if user_session: 

78 return user_session 

79 except Exception: 

80 # Error accessing Flask context 

81 pass 

82 

83 # No shared database - return None to allow SettingsManager to work without DB 

84 logger.warning( 

85 "get_db_session() is deprecated. Use get_user_db_session() from database.session_context" 

86 ) 

87 return None 

88 

89 

90def get_settings_manager( 

91 db_session: Session | None = None, username: str | None = None 

92): 

93 """ 

94 Get the settings manager for the current context. 

95 

96 Args: 

97 db_session: Optional database session 

98 username: Optional username for caching (required for SettingsManager) 

99 

100 Returns: 

101 The appropriate settings manager instance. 

102 """ 

103 # If db_session not provided, try to get one 

104 if db_session is None and username is None and has_request_context(): 

105 username = flask_session.get("username") 

106 

107 if db_session is None: 

108 try: 

109 db_session = get_db_session(username=username) 

110 except RuntimeError: 

111 # No authenticated user - settings manager will use defaults 

112 db_session = None 

113 username = "anonymous" 

114 

115 # Import here to avoid circular imports 

116 from ..settings import SettingsManager 

117 

118 # Always use regular SettingsManager (now with built-in simple caching) 

119 return SettingsManager(db_session) 

120 

121 

122def no_db_settings(func: Callable[..., Any]) -> Callable[..., Any]: 

123 """ 

124 Decorator that runs the wrapped function with the settings database 

125 completely disabled. This will prevent the function from accidentally 

126 reading settings from the DB. Settings can only be read from environment 

127 variables or the defaults file. 

128 

129 Args: 

130 func: The function to wrap. 

131 

132 Returns: 

133 The wrapped function. 

134 

135 """ 

136 

137 @functools.wraps(func) 

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

139 # Temporarily disable DB access in the settings manager. 

140 manager = get_settings_manager() 

141 db_session = manager.db_session 

142 manager.db_session = None 

143 

144 try: 

145 return func(*args, **kwargs) 

146 finally: 

147 # Restore the original database session. 

148 manager.db_session = db_session 

149 

150 return wrapper 

151 

152 

153def get_setting_from_db_main_thread( 

154 key: str, default_value: Any | None = None, username: str | None = None 

155) -> str | Dict[str, Any] | None: 

156 """ 

157 Get a setting from the database with fallback to default value 

158 

159 Args: 

160 key: The setting key. 

161 default_value: If the setting is not found, it will return this instead. 

162 username: Optional username for thread context (e.g., background research threads). 

163 

164 Returns: 

165 The setting value. 

166 

167 """ 

168 # In fallback LLM mode, always return default values without database access 

169 if use_fallback_llm(): 

170 logger.debug( 

171 f"Using default value for {key} in fallback LLM environment" 

172 ) 

173 return default_value 

174 

175 # CRITICAL: Detect if we're in a background thread and raise an error 

176 import threading 

177 

178 # Check if we're in a background thread 

179 thread_name = threading.current_thread().name 

180 

181 # Allow MainThread during startup, but not other threads 

182 if not has_app_context() and thread_name != "MainThread": 

183 thread_id = threading.get_ident() 

184 raise RuntimeError( 

185 f"get_db_setting('{key}') called from background thread '{thread_name}' (ID: {thread_id}). " 

186 f"Database access from threads is not allowed. Use settings_snapshot or thread-local settings context." 

187 ) 

188 

189 try: 

190 # Use the new session context to ensure proper database access 

191 from ..database.session_context import get_user_db_session 

192 

193 try: 

194 with get_user_db_session(username) as db_session: 

195 if db_session: 195 ↛ 197line 195 didn't jump to line 197 because the condition on line 195 was never true

196 # Use the unified settings manager 

197 settings_manager = get_settings_manager( 

198 db_session, username 

199 ) 

200 return settings_manager.get_setting( 

201 key, default=default_value 

202 ) 

203 except Exception: 

204 # If we can't get a session, fall back to default 

205 pass 

206 

207 except Exception: 

208 logger.exception(f"Error getting setting {key} from database") 

209 

210 logger.warning(f"Could not read setting '{key}' from the database.") 

211 return default_value