Coverage for src/local_deep_research/database/library_init.py: 100%
98 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Database initialization for Library - Unified Document Architecture.
4This module handles:
5- Seeding source_types table with predefined types
6- Creating the default "Library" collection
7- Must be called on app startup for each user
8"""
10import threading
11import uuid
12from loguru import logger
13from sqlalchemy.exc import IntegrityError
15from .models import SourceType, Collection
16from .session_context import get_user_db_session
17from ..constants import (
18 RESEARCH_HISTORY_COLLECTION_NAME,
19 RESEARCH_HISTORY_COLLECTION_DESCRIPTION,
20)
23# Per-user locks serialise the check-then-insert critical sections below.
24# Under IMMEDIATE isolation this was unnecessary; under DEFERRED, two
25# concurrent invocations (e.g. two logins of the same user from two
26# browser tabs) could both see the absent row and both insert, creating
27# duplicate default collections. An application-level lock is simpler
28# than a migration adding a partial UNIQUE constraint, and cheap.
29#
30# Entries deliberately retain stable identity for the process lifetime. A
31# caller receives its lock before acquiring it, so deleting an apparently idle
32# entry has a lookup-to-acquire race: that caller can later acquire the old lock
33# while a second caller creates and acquires a replacement.
34_user_init_locks: dict[str, threading.Lock] = {}
35_user_init_locks_lock = threading.Lock()
38def _get_user_init_lock(username: str) -> threading.Lock:
39 """Get (or lazily create) the per-user lock used to serialise the
40 check-then-insert idempotent collection initialisers.
41 """
42 with _user_init_locks_lock:
43 lock = _user_init_locks.get(username)
44 if lock is None:
45 lock = threading.Lock()
46 _user_init_locks[username] = lock
47 return lock
50def pop_user_init_lock(username: str) -> None:
51 """Compatibility no-op: per-user init locks keep stable identity.
53 User-close cleanup can race both a held lock and a caller that has looked
54 up the lock but not acquired it yet. Removing the registry entry in either
55 case permits a later initializer to create a second lock and enter the same
56 user's check-then-insert critical section concurrently.
58 Keep one small lock per distinct username for the process lifetime. The
59 function remains as a no-op for existing user-close cleanup callers.
60 """
61 del username
64def seed_source_types(username: str, password: str = None) -> None:
65 """
66 Seed the source_types table with predefined document source types.
68 Args:
69 username: User to seed types for
70 password: User's password (optional, uses session context)
71 """
72 predefined_types = [
73 {
74 "name": "research_download",
75 "display_name": "Research Download",
76 "description": "Documents downloaded from research sessions (arXiv, PubMed, etc.)",
77 "icon": "download",
78 },
79 {
80 "name": "user_upload",
81 "display_name": "User Upload",
82 "description": "Documents manually uploaded by the user",
83 "icon": "upload",
84 },
85 {
86 "name": "manual_entry",
87 "display_name": "Manual Entry",
88 "description": "Documents manually created or entered",
89 "icon": "edit",
90 },
91 {
92 "name": "research_report",
93 "display_name": "Research Report",
94 "description": "Generated research reports (markdown) for semantic search",
95 "icon": "file-alt",
96 },
97 {
98 "name": "research_source",
99 "display_name": "Research Source",
100 "description": "Sources discovered during research with content for semantic search",
101 "icon": "link",
102 },
103 {
104 "name": "note",
105 "display_name": "Note",
106 "description": "User-created notes with AI-enhanced features",
107 "icon": "sticky-note",
108 },
109 {
110 "name": "zotero",
111 "display_name": "Zotero",
112 "description": "Documents imported from a Zotero library or collection",
113 "icon": "book",
114 },
115 ]
117 try:
118 with get_user_db_session(username, password) as session:
119 for type_data in predefined_types:
120 # Check if type already exists
121 existing = (
122 session.query(SourceType)
123 .filter_by(name=type_data["name"])
124 .first()
125 )
127 if not existing:
128 source_type = SourceType(id=str(uuid.uuid4()), **type_data)
129 session.add(source_type)
130 logger.info(f"Created source type: {type_data['name']}")
132 session.commit()
133 logger.info("Source types seeded successfully")
135 except IntegrityError:
136 logger.warning("Source types may already exist")
137 except Exception:
138 logger.warning("Error seeding source types")
139 raise
142def ensure_default_library_collection(
143 username: str, password: str = None
144) -> str:
145 """
146 Ensure the default "Library" collection exists for a user.
147 Creates it if it doesn't exist.
149 Args:
150 username: User to check/create library for
151 password: User's password (optional, uses session context)
153 Returns:
154 UUID of the Library collection
155 """
156 try:
157 with (
158 _get_user_init_lock(username),
159 get_user_db_session(username, password) as session,
160 ):
161 # Check if default library exists
162 library = (
163 session.query(Collection).filter_by(is_default=True).first()
164 )
166 if library:
167 logger.debug(f"Default Library collection exists: {library.id}")
168 return library.id
170 # Create default Library collection
171 library_id = str(uuid.uuid4())
172 library = Collection(
173 id=library_id,
174 name="Library",
175 description="Default collection for research downloads and documents",
176 collection_type="default_library",
177 is_default=True,
178 )
179 session.add(library)
180 session.commit()
182 # The collection-engine cache may already hold this user's
183 # (empty) engine list from a search within the TTL; drop it so
184 # the next research sees the new collection as a tool (#6609).
185 # Lazy import: search_engines_config pulls in the search stack,
186 # which imports this package — module level would be circular.
187 from ..web_search_engines.search_engines_config import (
188 invalidate_collection_engines_cache,
189 )
191 invalidate_collection_engines_cache(username)
193 logger.info(f"Created default Library collection: {library_id}")
194 return library_id
196 except Exception:
197 logger.warning("Error ensuring default Library collection")
198 raise
201def ensure_research_history_collection(
202 username: str, password: str = None
203) -> str:
204 """
205 Ensure the "Research History" collection exists for a user.
206 This collection is used for semantic search over research reports and sources.
207 Creates it if it doesn't exist.
209 Args:
210 username: User to check/create collection for
211 password: User's password (optional, uses session context)
213 Returns:
214 UUID of the Research History collection
215 """
216 try:
217 with (
218 _get_user_init_lock(username),
219 get_user_db_session(username, password) as session,
220 ):
221 # Check if research history collection exists
222 collection = (
223 session.query(Collection)
224 .filter_by(collection_type="research_history")
225 .first()
226 )
228 if collection:
229 logger.debug(
230 f"Research History collection exists: {collection.id}"
231 )
232 return collection.id
234 # Create Research History collection
235 collection_id = str(uuid.uuid4())
236 collection = Collection(
237 id=collection_id,
238 name=RESEARCH_HISTORY_COLLECTION_NAME,
239 description=RESEARCH_HISTORY_COLLECTION_DESCRIPTION,
240 collection_type="research_history",
241 is_default=False,
242 )
243 session.add(collection)
244 session.commit()
246 # Same cache-invalidation contract as the default Library above
247 # (#6609): a search within the TTL must not hide this collection.
248 from ..web_search_engines.search_engines_config import (
249 invalidate_collection_engines_cache,
250 )
252 invalidate_collection_engines_cache(username)
254 logger.info(f"Created Research History collection: {collection_id}")
255 return collection_id
257 except Exception:
258 logger.warning("Error ensuring Research History collection")
259 raise
262def initialize_library_for_user(username: str, password: str = None) -> dict:
263 """
264 Complete initialization of library system for a user.
265 Seeds source types and ensures default Library and Research History collections exist.
267 Args:
268 username: User to initialize for
269 password: User's password (optional, uses session context)
271 Returns:
272 Dict with initialization results
273 """
274 results = {
275 "source_types_seeded": False,
276 "library_collection_id": None,
277 "research_history_collection_id": None,
278 "success": False,
279 }
281 try:
282 # Seed source types
283 seed_source_types(username, password)
284 results["source_types_seeded"] = True
286 # Ensure Library collection
287 library_id = ensure_default_library_collection(username, password)
288 results["library_collection_id"] = library_id
290 # Ensure Research History collection
291 research_history_id = ensure_research_history_collection(
292 username, password
293 )
294 results["research_history_collection_id"] = research_history_id
296 results["success"] = True
297 logger.info(f"Library initialization complete for user: {username}")
299 except Exception as e:
300 logger.warning(f"Library initialization failed for {username}")
301 results["error"] = str(e)
303 return results
306def get_default_library_id(username: str, password: str = None) -> str:
307 """
308 Get the ID of the default Library collection for a user.
309 Creates it if it doesn't exist.
311 Args:
312 username: User to get library for
313 password: User's password (optional, uses session context)
315 Returns:
316 UUID of the Library collection
317 """
318 return ensure_default_library_collection(username, password)
321def get_source_type_id(
322 username: str, type_name: str, password: str = None
323) -> str:
324 """
325 Get the ID of a source type by name.
327 Args:
328 username: User to query for
329 type_name: Name of source type (e.g., 'research_download', 'user_upload')
330 password: User's password (optional, uses session context)
332 Returns:
333 UUID of the source type
335 Raises:
336 ValueError: If source type not found
337 """
338 try:
339 with get_user_db_session(username, password) as session:
340 source_type = (
341 session.query(SourceType).filter_by(name=type_name).first()
342 )
344 if not source_type:
345 raise ValueError(f"Source type not found: {type_name}") # noqa: TRY301 — inside db session context, except logs and re-raises
347 return source_type.id
349 except Exception:
350 logger.warning("Error getting source type ID")
351 raise