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

1import os 

2from collections.abc import Callable 

3from pathlib import Path 

4from time import sleep 

5 

6from alembic import command, config, script 

7from alembic.config import Config 

8from alembic.runtime import migration 

9from sqlalchemy import engine, orm, text 

10 

11from mealie.core import root_logger 

12from mealie.core.config import get_app_settings 

13from mealie.db.db_setup import session_context 

14from mealie.db.fixes.fix_group_with_no_name import fix_group_with_no_name 

15from mealie.db.fixes.fix_migration_data import fix_migration_data 

16from mealie.db.fixes.fix_slug_foods import fix_slug_food_names 

17from mealie.repos.all_repositories import get_repositories 

18from mealie.repos.repository_factory import AllRepositories 

19from mealie.repos.seed.init_users import default_user_init 

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

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

22from mealie.services.group_services.group_service import GroupService 

23from mealie.services.household_services.household_service import HouseholdService 

24 

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

26 

27logger = root_logger.get_logger() 

28 

29 

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

31 settings = get_app_settings() 

32 

33 instance_repos = get_repositories(session) 

34 default_group = default_group_init(instance_repos, settings.DEFAULT_GROUP) 

35 

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

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

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

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) 

41 default_user_init(household_repos) 

42 

43 

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

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

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

47 

48 

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

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: 

56 settings = get_app_settings() 

57 url = settings.DB_URL 

58 

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

60 raise ValueError("No database url found") 

61 

62 connectable = engine.create_engine(url) 

63 directory = script.ScriptDirectory.from_config(alembic_cfg) 

64 with connectable.begin() as connection: 

65 context = migration.MigrationContext.configure(connection) 

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

67 

68 

69def safe_try(func: Callable): 

70 try: 

71 func() 

72 except Exception: 

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

74 

75 

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

77 try: 

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

79 return True 

80 except Exception: 

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

82 return False 

83 

84 

85def main(): 

86 # Wait for database to connect 

87 max_retry = 10 

88 wait_seconds = 1 

89 

90 with session_context() as session: 

91 while True: 

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

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

94 break 

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

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 true

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

108 

109 run_fixes = False 

110 alembic_cfg = Config(alembic_cfg_path) 

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 true

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

113 else: 

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

115 command.upgrade(alembic_cfg, "head") 

116 run_fixes = True 

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 true

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

120 

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

122 

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

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

132 init_db(session) 

133 

134 

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

136 main()