Coverage for src/local_deep_research/database/library_init.py: 100%
94 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +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 logger.info(f"Created default Library collection: {library_id}")
183 return library_id
185 except Exception:
186 logger.warning("Error ensuring default Library collection")
187 raise
190def ensure_research_history_collection(
191 username: str, password: str = None
192) -> str:
193 """
194 Ensure the "Research History" collection exists for a user.
195 This collection is used for semantic search over research reports and sources.
196 Creates it if it doesn't exist.
198 Args:
199 username: User to check/create collection for
200 password: User's password (optional, uses session context)
202 Returns:
203 UUID of the Research History collection
204 """
205 try:
206 with (
207 _get_user_init_lock(username),
208 get_user_db_session(username, password) as session,
209 ):
210 # Check if research history collection exists
211 collection = (
212 session.query(Collection)
213 .filter_by(collection_type="research_history")
214 .first()
215 )
217 if collection:
218 logger.debug(
219 f"Research History collection exists: {collection.id}"
220 )
221 return collection.id
223 # Create Research History collection
224 collection_id = str(uuid.uuid4())
225 collection = Collection(
226 id=collection_id,
227 name=RESEARCH_HISTORY_COLLECTION_NAME,
228 description=RESEARCH_HISTORY_COLLECTION_DESCRIPTION,
229 collection_type="research_history",
230 is_default=False,
231 )
232 session.add(collection)
233 session.commit()
235 logger.info(f"Created Research History collection: {collection_id}")
236 return collection_id
238 except Exception:
239 logger.warning("Error ensuring Research History collection")
240 raise
243def initialize_library_for_user(username: str, password: str = None) -> dict:
244 """
245 Complete initialization of library system for a user.
246 Seeds source types and ensures default Library and Research History collections exist.
248 Args:
249 username: User to initialize for
250 password: User's password (optional, uses session context)
252 Returns:
253 Dict with initialization results
254 """
255 results = {
256 "source_types_seeded": False,
257 "library_collection_id": None,
258 "research_history_collection_id": None,
259 "success": False,
260 }
262 try:
263 # Seed source types
264 seed_source_types(username, password)
265 results["source_types_seeded"] = True
267 # Ensure Library collection
268 library_id = ensure_default_library_collection(username, password)
269 results["library_collection_id"] = library_id
271 # Ensure Research History collection
272 research_history_id = ensure_research_history_collection(
273 username, password
274 )
275 results["research_history_collection_id"] = research_history_id
277 results["success"] = True
278 logger.info(f"Library initialization complete for user: {username}")
280 except Exception as e:
281 logger.warning(f"Library initialization failed for {username}")
282 results["error"] = str(e)
284 return results
287def get_default_library_id(username: str, password: str = None) -> str:
288 """
289 Get the ID of the default Library collection for a user.
290 Creates it if it doesn't exist.
292 Args:
293 username: User to get library for
294 password: User's password (optional, uses session context)
296 Returns:
297 UUID of the Library collection
298 """
299 return ensure_default_library_collection(username, password)
302def get_source_type_id(
303 username: str, type_name: str, password: str = None
304) -> str:
305 """
306 Get the ID of a source type by name.
308 Args:
309 username: User to query for
310 type_name: Name of source type (e.g., 'research_download', 'user_upload')
311 password: User's password (optional, uses session context)
313 Returns:
314 UUID of the source type
316 Raises:
317 ValueError: If source type not found
318 """
319 try:
320 with get_user_db_session(username, password) as session:
321 source_type = (
322 session.query(SourceType).filter_by(name=type_name).first()
323 )
325 if not source_type:
326 raise ValueError(f"Source type not found: {type_name}") # noqa: TRY301 — inside db session context, except logs and re-raises
328 return source_type.id
330 except Exception:
331 logger.warning("Error getting source type ID")
332 raise