Coverage for src / local_deep_research / database / library_init.py: 68%

1""" 

2Database initialization for Library - Unified Document Architecture. 

3 

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""" 

9 

10import uuid 

11from loguru import logger 

12from sqlalchemy.exc import IntegrityError 

13 

14from .models import SourceType, Collection 

15from .session_context import get_user_db_session 

16 

17 

18def seed_source_types(username: str, password: str = None) -> None: 

19 """ 

20 Seed the source_types table with predefined document source types. 

21 

22 Args: 

23 username: User to seed types for 

24 password: User's password (optional, uses session context) 

25 """ 

26 predefined_types = [ 

27 { 

28 "name": "research_download", 

29 "display_name": "Research Download", 

30 "description": "Documents downloaded from research sessions (arXiv, PubMed, etc.)", 

31 "icon": "download", 

32 }, 

33 { 

34 "name": "user_upload", 

35 "display_name": "User Upload", 

36 "description": "Documents manually uploaded by the user", 

37 "icon": "upload", 

38 }, 

39 { 

40 "name": "manual_entry", 

41 "display_name": "Manual Entry", 

42 "description": "Documents manually created or entered", 

43 "icon": "edit", 

44 }, 

45 ] 

46 

47 try: 

48 with get_user_db_session(username, password) as session: 

49 for type_data in predefined_types: 

50 # Check if type already exists 

51 existing = ( 

52 session.query(SourceType) 

53 .filter_by(name=type_data["name"]) 

54 .first() 

55 ) 

56 

57 if not existing: 

58 source_type = SourceType(id=str(uuid.uuid4()), **type_data) 

59 session.add(source_type) 

60 logger.info(f"Created source type: {type_data['name']}") 

61 

62 session.commit() 

63 logger.info("Source types seeded successfully") 

64 

65 except IntegrityError as e: 

66 logger.warning(f"Source types may already exist: {e}") 

67 except Exception: 

68 logger.exception("Error seeding source types") 

69 raise 

70 

71 

72def ensure_default_library_collection( 

73 username: str, password: str = None 

74) -> str: 

75 """ 

76 Ensure the default "Library" collection exists for a user. 

77 Creates it if it doesn't exist. 

78 

79 Args: 

80 username: User to check/create library for 

81 password: User's password (optional, uses session context) 

82 

83 Returns: 

84 UUID of the Library collection 

85 """ 

86 try: 

87 with get_user_db_session(username, password) as session: 

88 # Check if default library exists 

89 library = ( 

90 session.query(Collection).filter_by(is_default=True).first() 

91 ) 

92 

93 if library: 

94 logger.debug(f"Default Library collection exists: {library.id}") 

95 return library.id 

96 

97 # Create default Library collection 

98 library_id = str(uuid.uuid4()) 

99 library = Collection( 

100 id=library_id, 

101 name="Library", 

102 description="Default collection for research downloads and documents", 

103 collection_type="default_library", 

104 is_default=True, 

105 ) 

106 session.add(library) 

107 session.commit() 

108 

109 logger.info(f"Created default Library collection: {library_id}") 

110 return library_id 

111 

112 except Exception: 

113 logger.exception("Error ensuring default Library collection") 

114 raise 

115 

116 

117def initialize_library_for_user(username: str, password: str = None) -> dict: 

118 """ 

119 Complete initialization of library system for a user. 

120 Seeds source types and ensures default Library collection exists. 

121 

122 Args: 

123 username: User to initialize for 

124 password: User's password (optional, uses session context) 

125 

126 Returns: 

127 Dict with initialization results 

128 """ 

129 results = { 

130 "source_types_seeded": False, 

131 "library_collection_id": None, 

132 "success": False, 

133 } 

134 

135 try: 

136 # Seed source types 

137 seed_source_types(username, password) 

138 results["source_types_seeded"] = True 

139 

140 # Ensure Library collection 

141 library_id = ensure_default_library_collection(username, password) 

142 results["library_collection_id"] = library_id 

143 

144 results["success"] = True 

145 logger.info(f"Library initialization complete for user: {username}") 

146 

147 except Exception as e: 

148 logger.exception(f"Library initialization failed for {username}") 

149 results["error"] = str(e) 

150 

151 return results 

152 

153 

154def get_default_library_id(username: str, password: str = None) -> str: 

155 """ 

156 Get the ID of the default Library collection for a user. 

157 Creates it if it doesn't exist. 

158 

159 Args: 

160 username: User to get library for 

161 password: User's password (optional, uses session context) 

162 

163 Returns: 

164 UUID of the Library collection 

165 """ 

166 return ensure_default_library_collection(username, password) 

167 

168 

169def get_source_type_id( 

170 username: str, type_name: str, password: str = None 

171) -> str: 

172 """ 

173 Get the ID of a source type by name. 

174 

175 Args: 

176 username: User to query for 

177 type_name: Name of source type (e.g., 'research_download', 'user_upload') 

178 password: User's password (optional, uses session context) 

179 

180 Returns: 

181 UUID of the source type 

182 

183 Raises: 

184 ValueError: If source type not found 

185 """ 

186 try: 

187 with get_user_db_session(username, password) as session: 

188 source_type = ( 

189 session.query(SourceType).filter_by(name=type_name).first() 

190 ) 

191 

192 if not source_type: 

193 raise ValueError(f"Source type not found: {type_name}") 

194 

195 return source_type.id 

196 

197 except Exception: 

198 logger.exception("Error getting source type ID") 

199 raise