Coverage for src/local_deep_research/database/initialize.py: 93%

1""" 

2Centralized database initialization module. 

3 

4This module provides a single entry point for database initialization, 

5using Alembic for schema migrations and version control. 

6""" 

7 

8import time 

9from typing import Any, Optional 

10from loguru import logger 

11from sqlalchemy import Engine, inspect 

12from sqlalchemy.orm import Session 

13 

14from ..database.models import Base 

15from .alembic_runner import run_migrations 

16 

17 

18def initialize_database( 

19 engine: Engine, 

20 db_session: Optional[Session] = None, 

21) -> None: 

22 """ 

23 Initialize database tables if they don't exist. 

24 

25 Uses Alembic migrations to create and update the database schema. 

26 

27 Args: 

28 engine: SQLAlchemy engine for the database 

29 db_session: Optional database session for settings initialization 

30 """ 

31 inspector = inspect(engine) 

32 existing_tables = inspector.get_table_names() 

33 

34 logger.info( 

35 f"Initializing database with {len(existing_tables)} existing tables" 

36 ) 

37 logger.debug( 

38 f"Base.metadata has {len(Base.metadata.tables)} tables defined" 

39 ) 

40 

41 # Run Alembic migrations (creates tables and applies schema changes) 

42 run_migrations(engine) 

43 

44 # Check what was created (need new inspector to avoid caching) 

45 new_inspector = inspect(engine) 

46 new_tables = new_inspector.get_table_names() 

47 logger.info(f"After initialization: {len(new_tables)} tables exist") 

48 

49 # Initialize default settings if session provided 

50 if db_session: 

51 try: 

52 _initialize_default_settings(db_session) 

53 except Exception: 

54 logger.warning("Could not initialize default settings") 

55 

56 logger.info("Database initialization complete") 

57 

58 

59def _initialize_default_settings(db_session: Session) -> None: 

60 """ 

61 Initialize default settings from the defaults file. 

62 

63 Args: 

64 db_session: Database session to use for settings initialization 

65 """ 

66 from ..settings.manager import SettingsManager 

67 

68 try: 

69 settings_mgr = SettingsManager(db_session) 

70 

71 # Check if we need to update settings 

72 if settings_mgr.db_version_matches_package(): 

73 logger.debug("Settings version matches package, skipping update") 

74 return 

75 

76 logger.info("Loading default settings into database") 

77 

78 # Load settings from defaults file 

79 # This will not overwrite existing settings but will add new ones. 

80 # Direct import_settings call: the load_from_defaults_file wrapper 

81 # no longer accepts override_locked (#5841), and this trusted 

82 # bootstrap must bypass the lock so a locked account still 

83 # receives settings a later release ships. Time it directly too: 

84 # routing through the wrapper used to log this for free. 

85 start = time.perf_counter() 

86 row_count = len(settings_mgr.default_settings) 

87 settings_mgr.import_settings( 

88 settings_mgr.default_settings, 

89 overwrite=False, 

90 delete_extra=True, 

91 override_locked=True, 

92 ) 

93 elapsed_ms = (time.perf_counter() - start) * 1000 

94 logger.info( 

95 f"Loaded {row_count} default settings in {elapsed_ms:.0f}ms" 

96 ) 

97 

98 # Update the saved version 

99 settings_mgr.update_db_version() 

100 

101 logger.info("Default settings initialized successfully") 

102 

103 except Exception: 

104 logger.exception("Error initializing default settings") 

105 

106 

107def check_database_schema(engine: Engine) -> dict: 

108 """ 

109 Check the current database schema and return information about tables. 

110 

111 Args: 

112 engine: SQLAlchemy engine for the database 

113 

114 Returns: 

115 Dictionary with schema information including tables and their columns 

116 """ 

117 inspector = inspect(engine) 

118 schema_info: dict[str, Any] = { 

119 "tables": {}, 

120 "missing_tables": [], 

121 "has_news_tables": False, 

122 } 

123 

124 # Tables that only exist in the auth database, not per-user databases 

125 auth_only_tables = {"users"} 

126 

127 # Check core tables 

128 for table_name in Base.metadata.tables.keys(): 

129 if table_name in auth_only_tables: 

130 continue 

131 if inspector.has_table(table_name): 

132 columns = [col["name"] for col in inspector.get_columns(table_name)] 

133 schema_info["tables"][table_name] = columns 

134 else: 

135 schema_info["missing_tables"].append(table_name) 

136 

137 # Check if news tables exist 

138 news_tables = ["news_subscription", "news_card", "news_interest"] 

139 for table_name in news_tables: 

140 if table_name in schema_info["tables"]: 140 ↛ 141line 140 didn't jump to line 141 because the condition on line 140 was never true

141 schema_info["has_news_tables"] = True 

142 break 

143 

144 return schema_info