Coverage for opt/mealie/lib/python3.12/site-packages/mealie/db/init_db.py: 68%

97 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-11-25 15:32 +0000

1import os 1a

2from collections.abc import Callable 1a

3from pathlib import Path 1a

4from time import sleep 1a

5 

6from alembic import command, config, script 1a

7from alembic.config import Config 1a

8from alembic.runtime import migration 1a

9from sqlalchemy import engine, orm, text 1a

10 

11from mealie.core import root_logger 1a

12from mealie.core.config import get_app_settings 1a

13from mealie.db.db_setup import session_context 1a

14from mealie.db.fixes.fix_group_with_no_name import fix_group_with_no_name 1a

15from mealie.db.fixes.fix_migration_data import fix_migration_data 1a

16from mealie.db.fixes.fix_slug_foods import fix_slug_food_names 1a

17from mealie.repos.all_repositories import get_repositories 1a

18from mealie.repos.repository_factory import AllRepositories 1a

19from mealie.repos.seed.init_users import default_user_init 1a

20from mealie.schema.household.household import HouseholdCreate, HouseholdInDB 1a

21from mealie.schema.user.user import GroupBase, GroupInDB 1a

22from mealie.services.group_services.group_service import GroupService 1a

23from mealie.services.household_services.household_service import HouseholdService 1a

24 

25ALEMBIC_DIR = Path(__file__).parent.parent / "alembic" 1a

26 

27logger = root_logger.get_logger() 1a

28 

29 

30def init_db(session: orm.Session) -> None: 1a

31 settings = get_app_settings() 1a

32 

33 instance_repos = get_repositories(session) 1a

34 default_group = default_group_init(instance_repos, settings.DEFAULT_GROUP) 1a

35 

36 group_repos = get_repositories(session, group_id=default_group.id, household_id=None) 1a

37 default_household = group_repos.households.get_by_name(settings.DEFAULT_HOUSEHOLD) 1a

38 if default_household is None: 38 ↛ 39line 38 didn't jump to line 39 because the condition on line 38 was never true1a

39 default_household = default_household_init(group_repos, settings.DEFAULT_HOUSEHOLD) 

40 household_repos = get_repositories(session, group_id=default_group.id, household_id=default_household.id) 1a

41 default_user_init(household_repos) 1a

42 

43 

44def default_group_init(repos: AllRepositories, name: str) -> GroupInDB: 1a

45 logger.info("Generating Default Group and Household") 1a

46 return GroupService.create_group(repos, GroupBase(name=name)) 1a

47 

48 

49def default_household_init(repos: AllRepositories, name: str) -> HouseholdInDB: 1a

50 logger.info("Generating Default Household") 

51 return HouseholdService.create_household(repos, HouseholdCreate(name=name)) 

52 

53 

54# Adapted from https://alembic.sqlalchemy.org/en/latest/cookbook.html#test-current-database-revision-is-at-head-s 

55def db_is_at_head(alembic_cfg: config.Config) -> bool: 1a

56 settings = get_app_settings() 1a

57 url = settings.DB_URL 1a

58 

59 if not url: 59 ↛ 60line 59 didn't jump to line 60 because the condition on line 59 was never true1a

60 raise ValueError("No database url found") 

61 

62 connectable = engine.create_engine(url) 1a

63 directory = script.ScriptDirectory.from_config(alembic_cfg) 1a

64 with connectable.begin() as connection: 1a

65 context = migration.MigrationContext.configure(connection) 1a

66 return set(context.get_current_heads()) == set(directory.get_heads()) 1a

67 

68 

69def safe_try(func: Callable): 1a

70 try: 

71 func() 

72 except Exception: 

73 logger.exception(f"Error calling '{func.__name__}'") 

74 

75 

76def connect(session: orm.Session) -> bool: 1a

77 try: 1a

78 session.execute(text("SELECT 1")) 1a

79 return True 1a

80 except Exception: 

81 logger.exception("Error connecting to database") 

82 return False 

83 

84 

85def main(): 1a

86 # Wait for database to connect 

87 max_retry = 10 1a

88 wait_seconds = 1 1a

89 

90 with session_context() as session: 1a

91 while True: 1a

92 if connect(session): 92 ↛ 96line 92 didn't jump to line 96 because the condition on line 92 was always true1a

93 logger.info("Database connection established.") 1a

94 break 1a

95 

96 logger.error(f"Database connection failed. Retrying in {wait_seconds} seconds...") 

97 max_retry -= 1 

98 

99 sleep(wait_seconds) 

100 

101 if max_retry == 0: 

102 raise ConnectionError("Database connection failed - exiting application.") 

103 

104 alembic_cfg_path = os.getenv("ALEMBIC_CONFIG_FILE", default=str(ALEMBIC_DIR / "alembic.ini")) 1a

105 

106 if not os.path.isfile(alembic_cfg_path): 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true1a

107 raise Exception("Provided alembic config path doesn't exist") 

108 

109 run_fixes = False 1a

110 alembic_cfg = Config(alembic_cfg_path) 1a

111 if db_is_at_head(alembic_cfg): 111 ↛ 112line 111 didn't jump to line 112 because the condition on line 111 was never true1a

112 logger.debug("Migration not needed.") 

113 else: 

114 logger.info("Migration needed. Performing migration...") 1a

115 command.upgrade(alembic_cfg, "head") 1a

116 run_fixes = True 1a

117 

118 if session.get_bind().name == "postgresql": # needed for fuzzy search and fast GIN text indices 118 ↛ 119line 118 didn't jump to line 119 because the condition on line 118 was never true1a

119 session.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm;")) 

120 

121 db = get_repositories(session, group_id=None, household_id=None) 1a

122 

123 if db.users.get_all(): 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true1a

124 logger.debug("Database exists") 

125 if run_fixes: 

126 safe_try(lambda: fix_migration_data(session)) 

127 safe_try(lambda: fix_slug_food_names(db)) 

128 safe_try(lambda: fix_group_with_no_name(session)) 

129 

130 else: 

131 logger.info("Database contains no users, initializing...") 1a

132 init_db(session) 1a

133 

134 

135if __name__ == "__main__": 135 ↛ 136line 135 didn't jump to line 136 because the condition on line 135 was never true1a

136 main()