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

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 

8from typing import Any, Optional 

9from loguru import logger 

10from sqlalchemy import Engine, inspect 

11from sqlalchemy.orm import Session 

12 

13from ..database.models import Base 

14from .alembic_runner import run_migrations 

15 

16 

17def initialize_database( 

18 engine: Engine, 

19 db_session: Optional[Session] = None, 

20) -> None: 

21 """ 

22 Initialize database tables if they don't exist. 

23 

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

25 

26 Args: 

27 engine: SQLAlchemy engine for the database 

28 db_session: Optional database session for settings initialization 

29 """ 

30 inspector = inspect(engine) 

31 existing_tables = inspector.get_table_names() 

32 

33 logger.info( 

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

35 ) 

36 logger.debug( 

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

38 ) 

39 

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

41 run_migrations(engine) 

42 

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

44 new_inspector = inspect(engine) 

45 new_tables = new_inspector.get_table_names() 

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

47 

48 # Initialize default settings if session provided 

49 if db_session: 

50 try: 

51 _initialize_default_settings(db_session) 

52 except Exception: 

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

54 

55 logger.info("Database initialization complete") 

56 

57 

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

59 """ 

60 Initialize default settings from the defaults file. 

61 

62 Args: 

63 db_session: Database session to use for settings initialization 

64 """ 

65 from ..settings.manager import SettingsManager 

66 

67 try: 

68 settings_mgr = SettingsManager(db_session) 

69 

70 # Check if we need to update settings 

71 if settings_mgr.db_version_matches_package(): 

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

73 return 

74 

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

76 

77 # Load settings from defaults file 

78 # This will not overwrite existing settings but will add new ones 

79 settings_mgr.load_from_defaults_file( 

80 overwrite=False, delete_extra=True, override_locked=True 

81 ) 

82 

83 # Update the saved version 

84 settings_mgr.update_db_version() 

85 

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

87 

88 except Exception: 

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

90 

91 

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

93 """ 

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

95 

96 Args: 

97 engine: SQLAlchemy engine for the database 

98 

99 Returns: 

100 Dictionary with schema information including tables and their columns 

101 """ 

102 inspector = inspect(engine) 

103 schema_info: dict[str, Any] = { 

104 "tables": {}, 

105 "missing_tables": [], 

106 "has_news_tables": False, 

107 } 

108 

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

110 auth_only_tables = {"users"} 

111 

112 # Check core tables 

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

114 if table_name in auth_only_tables: 

115 continue 

116 if inspector.has_table(table_name): 

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

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

119 else: 

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

121 

122 # Check if news tables exist 

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

124 for table_name in news_tables: 

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

126 schema_info["has_news_tables"] = True 

127 break 

128 

129 return schema_info