Coverage for core/init_db.py: 100.00%

1import logging 

2 

3from sqlalchemy import delete, select, text 

4 

5from core.config import settings 

6from core.database import AsyncSessionLocal 

7from core.permissions import get_attributes 

8from core.security import hash_password 

9from models.role_attributes import RoleAttributes 

10from models.role_mapper import RoleMapper 

11from models.roles import Roles 

12from models.users import Users 

13 

14logger = logging.getLogger("init_db") 

15 

16 

17async def is_already_initialized() -> bool: 

18 """Return True if super admin role and user already exist (seed done).""" 

19 async with AsyncSessionLocal() as db: 

20 # Super admin role must exist 

21 super_role_result = await db.execute( 

22 select(Roles).where(Roles.name == settings.DEFAULT_SUPER_ADMIN_ROLE) 

23 ) 

24 super_role = super_role_result.scalar_one_or_none() 

25 if not super_role: 

26 return False 

27 

28 # At least one attribute must exist 

29 attrs_result = await db.execute(select(RoleAttributes)) 

30 if not attrs_result.scalars().first(): 

31 return False 

32 

33 # If any user already mapped to super admin role, consider seeded 

34 super_users_result = await db.execute( 

35 select(Users) 

36 .join(RoleMapper, Users.id == RoleMapper.user_id) 

37 .where(RoleMapper.role_id == super_role.id) 

38 ) 

39 if super_users_result.scalars().first(): 

40 return True 

41 

42 return False 

43 

44 

45async def init_database(): 

46 """Initialize database with default data: role attributes, roles and admin account""" 

47 lock_name = "init_db_lock" 

48 lock_timeout_sec = 15 

49 

50 # Use DB advisory lock to avoid multiple workers seeding simultaneously 

51 async with AsyncSessionLocal() as lock_session: 

52 lock_result = await lock_session.execute( 

53 text("SELECT GET_LOCK(:name, :timeout)"), 

54 {"name": lock_name, "timeout": lock_timeout_sec}, 

55 ) 

56 lock_acquired = lock_result.scalar() 

57 

58 if lock_acquired != 1: 

59 logger.info("Skipping database initialization: another worker is seeding.") 

60 return 

61 

62 try: 

63 # Double-check after getting the lock, if already initialized, skip 

64 if await is_already_initialized(): 

65 logger.info("Database initialization already completed, skipping.") 

66 return 

67 

68 logger.info("Starting database initialization...") 

69 

70 await create_role_attributes() 

71 await create_default_roles() 

72 await create_default_admin() 

73 

74 logger.info("Database initialization completed") 

75 

76 except Exception as e: 

77 logger.error(f"Database initialization failed: {str(e)}") 

78 raise 

79 finally: 

80 # Always release the advisory lock 

81 await lock_session.execute(text("SELECT RELEASE_LOCK(:name)"), {"name": lock_name}) 

82 await lock_session.commit() 

83 

84 

85async def create_role_attributes(): 

86 """Create role attributes""" 

87 async with AsyncSessionLocal() as db: 

88 try: 

89 attributes = get_attributes() 

90 created_count = 0 

91 updated_count = 0 

92 

93 for attr_config in attributes: 

94 existing_attr = await db.execute( 

95 select(RoleAttributes).where(RoleAttributes.name == attr_config["name"]) 

96 ) 

97 existing = existing_attr.scalar_one_or_none() 

98 if existing: 

99 if ( 

100 getattr(existing, "group", None) is None 

101 and attr_config.get("group") is not None 

102 ): 

103 existing.group = attr_config.get("group") 

104 updated_count += 1 

105 if ( 

106 getattr(existing, "category", None) is None 

107 and attr_config.get("category") is not None 

108 ): 

109 existing.category = attr_config.get("category") 

110 updated_count += 1 

111 continue 

112 

113 attribute = RoleAttributes( 

114 name=attr_config["name"], 

115 group=attr_config.get("group"), 

116 category=attr_config.get("category"), 

117 ) 

118 db.add(attribute) 

119 created_count += 1 

120 

121 await db.commit() 

122 

123 except Exception as e: 

124 logger.error(f"Failed to create role attributes: {str(e)}") 

125 await db.rollback() 

126 raise 

127 

128 

129async def create_default_roles(): 

130 """Create system super-admin role and a basic user role.""" 

131 async with AsyncSessionLocal() as db: 

132 try: 

133 default_roles = [ 

134 { 

135 "name": settings.DEFAULT_SUPER_ADMIN_ROLE, 

136 "description": "System super administrator (full access bypass)", 

137 "level": settings.DEFAULT_SUPER_ADMIN_LEVEL, 

138 }, 

139 { 

140 "name": "user", 

141 "description": "Regular user role with basic permissions", 

142 "level": settings.DEFAULT_USER_ROLE_LEVEL, 

143 }, 

144 ] 

145 

146 created_count = 0 

147 

148 for role_config in default_roles: 

149 existing_role = await db.execute( 

150 select(Roles).where(Roles.name == role_config["name"]) 

151 ) 

152 if existing_role.scalar_one_or_none(): 

153 continue 

154 

155 role = Roles( 

156 name=role_config["name"], 

157 description=role_config["description"], 

158 level=role_config["level"], 

159 ) 

160 db.add(role) 

161 created_count += 1 

162 

163 await db.commit() 

164 

165 except Exception as e: 

166 logger.error(f"Failed to create roles: {str(e)}") 

167 await db.rollback() 

168 raise 

169 

170 

171async def create_default_admin(): 

172 """Create default super-admin account from ENV settings.""" 

173 async with AsyncSessionLocal() as db: 

174 try: 

175 super_role_result = await db.execute( 

176 select(Roles).where(Roles.name == settings.DEFAULT_SUPER_ADMIN_ROLE) 

177 ) 

178 super_role = super_role_result.scalar_one_or_none() 

179 

180 if not super_role: 

181 logger.error("Super admin role not found, please run role initialization first") 

182 return 

183 

184 super_users_result = await db.execute( 

185 select(Users) 

186 .join(RoleMapper, Users.id == RoleMapper.user_id) 

187 .join(Roles, RoleMapper.role_id == Roles.id) 

188 .where(Roles.name == settings.DEFAULT_SUPER_ADMIN_ROLE) 

189 ) 

190 existing_super_users = super_users_result.scalars().all() 

191 

192 if existing_super_users: 

193 logger.info( 

194 "Super admin users already exist: " 

195 f"{[user.email for user in existing_super_users]}" 

196 ) 

197 await db.commit() 

198 return 

199 

200 existing_user_result = await db.execute( 

201 select(Users).where(Users.email == settings.DEFAULT_ADMIN_EMAIL) 

202 ) 

203 existing_user = existing_user_result.scalar_one_or_none() 

204 

205 if existing_user: 

206 existing_role_mapping = await db.execute( 

207 select(RoleMapper).where( 

208 RoleMapper.user_id == existing_user.id, RoleMapper.role_id == super_role.id 

209 ) 

210 ) 

211 if not existing_role_mapping.scalar_one_or_none(): 

212 role_mapping = RoleMapper(user_id=existing_user.id, role_id=super_role.id) 

213 db.add(role_mapping) 

214 logger.info( 

215 f"Assigned super admin role to existing user: {existing_user.email}" 

216 ) 

217 await db.execute( 

218 delete(RoleMapper).where( 

219 RoleMapper.user_id == existing_user.id, 

220 RoleMapper.role_id != super_role.id, 

221 ) 

222 ) 

223 else: 

224 admin_user = Users( 

225 first_name=settings.DEFAULT_ADMIN_FIRST_NAME, 

226 last_name=settings.DEFAULT_ADMIN_LAST_NAME, 

227 email=settings.DEFAULT_ADMIN_EMAIL, 

228 phone=settings.DEFAULT_ADMIN_PHONE, 

229 hash_password=await hash_password(settings.DEFAULT_ADMIN_PASSWORD), 

230 status=True, 

231 password_reset_required=False, 

232 email_verified=True, 

233 ) 

234 

235 db.add(admin_user) 

236 await db.commit() 

237 await db.refresh(admin_user) 

238 

239 role_mapping = RoleMapper(user_id=admin_user.id, role_id=super_role.id) 

240 db.add(role_mapping) 

241 

242 await db.commit() 

243 logger.info("Admin account initialization completed") 

244 

245 except Exception as e: 

246 logger.error(f"Failed to create admin account: {str(e)}") 

247 await db.rollback() 

248 raise