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

1import functools 

2from typing import Any, Callable, Dict 

3 

4from cachetools import LRUCache 

5from flask import g, has_app_context, session as flask_session 

6from loguru import logger 

7from sqlalchemy.orm import Session 

8 

9from ..settings.env_registry import use_fallback_llm 

10 

11from ..config.paths import get_data_directory 

12from ..database.encrypted_db import db_manager 

13from .threading_utils import thread_specific_cache 

14 

15# Database paths using new centralized configuration 

16DATA_DIR = get_data_directory() 

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

18 

19 

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

21def get_db_session( 

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

23) -> Session: 

24 """ 

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

26 

27 Args: 

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

29 force the caching mechanism to create separate settings even in 

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

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

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

33 

34 Returns: 

35 The database session for the current user/context. 

36 """ 

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

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

39 import threading 

40 

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

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

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

44 thread_name = threading.current_thread().name 

45 

46 # Allow MainThread during startup, but not other threads 

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

48 thread_id = threading.get_ident() 

49 raise RuntimeError( 

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

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

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

53 ) 

54 

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

56 if username: 

57 user_session = db_manager.get_session(username) 

58 if user_session: 

59 return user_session 

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

61 

62 # Otherwise, check Flask request context 

63 try: 

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

65 if hasattr(g, "db_session") and g.db_session: 

66 return g.db_session 

67 

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

69 username = flask_session.get("username") 

70 if username: 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true

71 user_session = db_manager.get_session(username) 

72 if user_session: 

73 return user_session 

74 except Exception: 

75 # Error accessing Flask context 

76 pass 

77 

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

79 logger.warning( 

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

81 ) 

82 return None 

83 

84 

85def get_settings_manager( 

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

87): 

88 """ 

89 Get the settings manager for the current context. 

90 

91 Args: 

92 db_session: Optional database session 

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

94 

95 Returns: 

96 The appropriate settings manager instance. 

97 """ 

98 # If db_session not provided, try to get one 

99 if db_session is None and username is None and has_app_context(): 

100 username = flask_session.get("username") 

101 

102 if db_session is None: 

103 try: 

104 db_session = get_db_session(username=username) 

105 except RuntimeError: 

106 # No authenticated user - settings manager will use defaults 

107 db_session = None 

108 username = "anonymous" 

109 

110 # Import here to avoid circular imports 

111 from ..settings import SettingsManager 

112 

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

114 return SettingsManager(db_session) 

115 

116 

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

118 """ 

119 Decorator that runs the wrapped function with the settings database 

120 completely disabled. This will prevent the function from accidentally 

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

122 variables or the defaults file. 

123 

124 Args: 

125 func: The function to wrap. 

126 

127 Returns: 

128 The wrapped function. 

129 

130 """ 

131 

132 @functools.wraps(func) 

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

134 # Temporarily disable DB access in the settings manager. 

135 manager = get_settings_manager() 

136 db_session = manager.db_session 

137 manager.db_session = None 

138 

139 try: 

140 return func(*args, **kwargs) 

141 finally: 

142 # Restore the original database session. 

143 manager.db_session = db_session 

144 

145 return wrapper 

146 

147 

148def get_setting_from_db_main_thread( 

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

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

151 """ 

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

153 

154 Args: 

155 key: The setting key. 

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

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

158 

159 Returns: 

160 The setting value. 

161 

162 """ 

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

164 if use_fallback_llm(): 

165 logger.debug( 

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

167 ) 

168 return default_value 

169 

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

171 import threading 

172 

173 # Check if we're in a background thread 

174 thread_name = threading.current_thread().name 

175 

176 # Allow MainThread during startup, but not other threads 

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

178 thread_id = threading.get_ident() 

179 raise RuntimeError( 

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

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

182 ) 

183 

184 try: 

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

186 from ..database.session_context import get_user_db_session 

187 

188 try: 

189 with get_user_db_session(username) as db_session: 

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

191 # Use the unified settings manager 

192 settings_manager = get_settings_manager( 

193 db_session, username 

194 ) 

195 return settings_manager.get_setting( 

196 key, default=default_value 

197 ) 

198 except Exception: 

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

200 pass 

201 

202 except Exception: 

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

204 

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

206 return default_value