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(overwrite=False, delete_extra=True) 

80 

81 # Update the saved version 

82 settings_mgr.update_db_version() 

83 

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

85 

86 except Exception: 

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

88 

89 

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

91 """ 

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

93 

94 Args: 

95 engine: SQLAlchemy engine for the database 

96 

97 Returns: 

98 Dictionary with schema information including tables and their columns 

99 """ 

100 inspector = inspect(engine) 

101 schema_info: dict[str, Any] = { 

102 "tables": {}, 

103 "missing_tables": [], 

104 "has_news_tables": False, 

105 } 

106 

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

108 auth_only_tables = {"users"} 

109 

110 # Check core tables 

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

112 if table_name in auth_only_tables: 

113 continue 

114 if inspector.has_table(table_name): 

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

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

117 else: 

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

119 

120 # Check if news tables exist 

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

122 for table_name in news_tables: 

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

124 schema_info["has_news_tables"] = True 

125 break 

126 

127 return schema_info