fix: enhance get_ui_roles method to handle exceptions and return error messages

This commit is contained in:
Théophile Diot 2024-11-15 13:47:21 +01:00
parent c3c5826f64
commit 23b7ef6ee1
No known key found for this signature in database
GPG key ID: FA995104A0BA376A
2 changed files with 22 additions and 18 deletions

View file

@ -254,28 +254,31 @@ class UIDatabase(Database):
return ""
def get_ui_roles(self, *, as_dict: bool = False) -> List[Union[Roles, dict]]:
def get_ui_roles(self, *, as_dict: bool = False) -> Union[str, List[Union[Roles, dict]]]:
"""Get ui roles."""
with self._db_session() as session:
roles = session.query(Roles).with_entities(Roles.name, Roles.description, Roles.update_datetime).all()
if not as_dict:
return roles
try:
roles = session.query(Roles).with_entities(Roles.name, Roles.description, Roles.update_datetime).all()
if not as_dict:
return roles
roles_data = []
for role in roles:
role_data = {
"name": role.name,
"description": role.description,
"update_datetime": role.update_datetime,
"permissions": [],
}
roles_data = []
for role in roles:
role_data = {
"name": role.name,
"description": role.description,
"update_datetime": role.update_datetime,
"permissions": [],
}
for permission in session.query(RolesPermissions).with_entities(RolesPermissions.permission_name).filter_by(role_name=role.name):
role_data["permissions"].append(permission.permission_name)
for permission in session.query(RolesPermissions).with_entities(RolesPermissions.permission_name).filter_by(role_name=role.name):
role_data["permissions"].append(permission.permission_name)
roles_data.append(role_data)
roles_data.append(role_data)
return roles_data
return roles_data
except BaseException as e:
return str(e)
def refresh_ui_user_recovery_codes(self, username: str, codes: List[str]) -> str:
"""Refresh ui user recovery codes."""

View file

@ -101,14 +101,15 @@ def on_starting(server):
ready = False
while not ready:
db_metadata = DB.get_metadata()
if isinstance(db_metadata, str) or not db_metadata["is_initialized"]:
ui_roles = DB.get_ui_roles(as_dict=True)
if isinstance(db_metadata, str) or not db_metadata["is_initialized"] or (isinstance(ui_roles, str) and "doesn't exist" in ui_roles):
LOGGER.warning("Database is not initialized, retrying in 5s ...")
else:
ready = True
continue
sleep(5)
if not DB.get_ui_roles(as_dict=True):
if not ui_roles:
ret = DB.create_ui_role("admin", "Admins can create new users, edit and read the data.", ["manage", "write", "read"])
if ret:
LOGGER.error(f"Couldn't create the admin role in the database: {ret}")