Coverage for src/local_deep_research/database/library_init.py: 100%
95 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-19 23:35 +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_user_init_locks: dict[str, threading.Lock] = {}
30_user_init_locks_lock = threading.Lock()
33def _get_user_init_lock(username: str) -> threading.Lock:
34 """Get (or lazily create) the per-user lock used to serialise the
35 check-then-insert idempotent collection initialisers.
36 """
37 with _user_init_locks_lock:
38 lock = _user_init_locks.get(username)
39 if lock is None:
40 lock = threading.Lock()
41 _user_init_locks[username] = lock
42 return lock
45def pop_user_init_lock(username: str) -> None:
46 """Remove the per-user init lock for ``username`` from the registry.
48 Called from the user-close path (``db_manager.close_user_database``
49 callers in ``web/auth/connection_cleanup.py`` and ``web/auth/routes.py``)
50 so the module-level dict doesn't accumulate one entry per username
51 across the process lifetime. The next login lazily re-creates the
52 lock, which is fine — the lock has no state that needs to persist
53 across login/logout.
54 """
55 with _user_init_locks_lock:
56 _user_init_locks.pop(username, None)
59def seed_source_types(username: str, password: str = None) -> None:
60 """
61 Seed the source_types table with predefined document source types.
63 Args:
64 username: User to seed types for
65 password: User's password (optional, uses session context)
66 """
67 predefined_types = [
68 {
69 "name": "research_download",
70 "display_name": "Research Download",
71 "description": "Documents downloaded from research sessions (arXiv, PubMed, etc.)",
72 "icon": "download",
73 },
74 {
75 "name": "user_upload",
76 "display_name": "User Upload",
77 "description": "Documents manually uploaded by the user",
78 "icon": "upload",
79 },
80 {
81 "name": "manual_entry",
82 "display_name": "Manual Entry",
83 "description": "Documents manually created or entered",
84 "icon": "edit",
85 },
86 {
87 "name": "research_report",
88 "display_name": "Research Report",
89 "description": "Generated research reports (markdown) for semantic search",
90 "icon": "file-alt",
91 },
92 {
93 "name": "research_source",
94 "display_name": "Research Source",
95 "description": "Sources discovered during research with content for semantic search",
96 "icon": "link",
97 },
98 {
99 "name": "note",
100 "display_name": "Note",
101 "description": "User-created notes with AI-enhanced features",
102 "icon": "sticky-note",
103 },
104 {
105 "name": "zotero",
106 "display_name": "Zotero",
107 "description": "Documents imported from a Zotero library or collection",
108 "icon": "book",
109 },
110 ]
112 try:
113 with get_user_db_session(username, password) as session:
114 for type_data in predefined_types:
115 # Check if type already exists
116 existing = (
117 session.query(SourceType)
118 .filter_by(name=type_data["name"])
119 .first()
120 )
122 if not existing:
123 source_type = SourceType(id=str(uuid.uuid4()), **type_data)
124 session.add(source_type)
125 logger.info(f"Created source type: {type_data['name']}")
127 session.commit()
128 logger.info("Source types seeded successfully")
130 except IntegrityError:
131 logger.warning("Source types may already exist")
132 except Exception:
133 logger.warning("Error seeding source types")
134 raise
137def ensure_default_library_collection(
138 username: str, password: str = None
139) -> str:
140 """
141 Ensure the default "Library" collection exists for a user.
142 Creates it if it doesn't exist.
144 Args:
145 username: User to check/create library for
146 password: User's password (optional, uses session context)
148 Returns:
149 UUID of the Library collection
150 """
151 try:
152 with (
153 _get_user_init_lock(username),
154 get_user_db_session(username, password) as session,
155 ):
156 # Check if default library exists
157 library = (
158 session.query(Collection).filter_by(is_default=True).first()
159 )
161 if library:
162 logger.debug(f"Default Library collection exists: {library.id}")
163 return library.id
165 # Create default Library collection
166 library_id = str(uuid.uuid4())
167 library = Collection(
168 id=library_id,
169 name="Library",
170 description="Default collection for research downloads and documents",
171 collection_type="default_library",
172 is_default=True,
173 )
174 session.add(library)
175 session.commit()
177 logger.info(f"Created default Library collection: {library_id}")
178 return library_id
180 except Exception:
181 logger.warning("Error ensuring default Library collection")
182 raise
185def ensure_research_history_collection(
186 username: str, password: str = None
187) -> str:
188 """
189 Ensure the "Research History" collection exists for a user.
190 This collection is used for semantic search over research reports and sources.
191 Creates it if it doesn't exist.
193 Args:
194 username: User to check/create collection for
195 password: User's password (optional, uses session context)
197 Returns:
198 UUID of the Research History collection
199 """
200 try:
201 with (
202 _get_user_init_lock(username),
203 get_user_db_session(username, password) as session,
204 ):
205 # Check if research history collection exists
206 collection = (
207 session.query(Collection)
208 .filter_by(collection_type="research_history")
209 .first()
210 )
212 if collection:
213 logger.debug(
214 f"Research History collection exists: {collection.id}"
215 )
216 return collection.id
218 # Create Research History collection
219 collection_id = str(uuid.uuid4())
220 collection = Collection(
221 id=collection_id,
222 name=RESEARCH_HISTORY_COLLECTION_NAME,
223 description=RESEARCH_HISTORY_COLLECTION_DESCRIPTION,
224 collection_type="research_history",
225 is_default=False,
226 )
227 session.add(collection)
228 session.commit()
230 logger.info(f"Created Research History collection: {collection_id}")
231 return collection_id
233 except Exception:
234 logger.warning("Error ensuring Research History collection")
235 raise
238def initialize_library_for_user(username: str, password: str = None) -> dict:
239 """
240 Complete initialization of library system for a user.
241 Seeds source types and ensures default Library and Research History collections exist.
243 Args:
244 username: User to initialize for
245 password: User's password (optional, uses session context)
247 Returns:
248 Dict with initialization results
249 """
250 results = {
251 "source_types_seeded": False,
252 "library_collection_id": None,
253 "research_history_collection_id": None,
254 "success": False,
255 }
257 try:
258 # Seed source types
259 seed_source_types(username, password)
260 results["source_types_seeded"] = True
262 # Ensure Library collection
263 library_id = ensure_default_library_collection(username, password)
264 results["library_collection_id"] = library_id
266 # Ensure Research History collection
267 research_history_id = ensure_research_history_collection(
268 username, password
269 )
270 results["research_history_collection_id"] = research_history_id
272 results["success"] = True
273 logger.info(f"Library initialization complete for user: {username}")
275 except Exception as e:
276 logger.warning(f"Library initialization failed for {username}")
277 results["error"] = str(e)
279 return results
282def get_default_library_id(username: str, password: str = None) -> str:
283 """
284 Get the ID of the default Library collection for a user.
285 Creates it if it doesn't exist.
287 Args:
288 username: User to get library for
289 password: User's password (optional, uses session context)
291 Returns:
292 UUID of the Library collection
293 """
294 return ensure_default_library_collection(username, password)
297def get_source_type_id(
298 username: str, type_name: str, password: str = None
299) -> str:
300 """
301 Get the ID of a source type by name.
303 Args:
304 username: User to query for
305 type_name: Name of source type (e.g., 'research_download', 'user_upload')
306 password: User's password (optional, uses session context)
308 Returns:
309 UUID of the source type
311 Raises:
312 ValueError: If source type not found
313 """
314 try:
315 with get_user_db_session(username, password) as session:
316 source_type = (
317 session.query(SourceType).filter_by(name=type_name).first()
318 )
320 if not source_type:
321 raise ValueError(f"Source type not found: {type_name}") # noqa: TRY301 — inside db session context, except logs and re-raises
323 return source_type.id
325 except Exception:
326 logger.warning("Error getting source type ID")
327 raise