cornflow 1.2.1__py3-none-any.whl → 1.2.3__py3-none-any.whl

cornflow/commands/permissions.py

@@ -1,88 +1,99 @@
-import sys
-from importlib import import_module
+from flask import current_app
+from sqlalchemy.exc import DBAPIError, IntegrityError
 
-from cornflow.shared.const import (
-    BASE_PERMISSION_ASSIGNATION,
-    EXTRA_PERMISSION_ASSIGNATION,
+from cornflow.commands.auxiliar import (
+    get_all_external,
+    get_all_resources,
 )
 from cornflow.models import ViewModel, PermissionViewRoleModel
 from cornflow.shared import db
-from flask import current_app
-from sqlalchemy.exc import DBAPIError, IntegrityError
+from cornflow.shared.const import ALL_DEFAULT_ROLES, GET_ACTION
+from cornflow.shared.const import (
+    BASE_PERMISSION_ASSIGNATION,
+)
 
 
 def register_base_permissions_command(external_app: str = None, verbose: bool = False):
-    if external_app is None:
-        from cornflow.endpoints import resources, alarms_resources
-        resources_to_register = resources
-        if current_app.config["ALARMS_ENDPOINTS"]:
-            resources_to_register = resources + alarms_resources
-    elif external_app is not None:
-        sys.path.append("./")
-        external_module = import_module(external_app)
-        resources_to_register = external_module.endpoints.resources
-    else:
-        resources_to_register = []
-        exit()
-
+    """
+    Register base permissions for the application.
+    external_app: If provided, it will register the permissions for the external app.
+    verbose: If True, it will print the permissions that are being registered.
+    """
+    # Get all resources, extra permissions, and custom roles actions
+    resources_to_register, extra_permissions, custom_roles_actions = get_all_external(
+        external_app
+    )
+
+    # Get all views in the database
     views_in_db = {view.name: view.id for view in ViewModel.get_all_objects()}
-    permissions_in_db = [perm for perm in PermissionViewRoleModel.get_all_objects()]
-    permissions_in_db_keys = [
-        (perm.role_id, perm.action_id, perm.api_view_id) for perm in permissions_in_db
-    ]
-    resources_names = [resource["endpoint"] for resource in resources_to_register]
+    permissions_in_db, permissions_in_db_keys = get_db_permissions()
+
+    # Get all resources and roles with access
+    resources_roles_with_access = get_all_resources(resources_to_register)
+
+    # Get the new roles and base permissions assignation
+    base_permissions_assignation = get_base_permissions(
+        resources_roles_with_access, custom_roles_actions
+    )
+    # Get the permissions to register and delete
+    permissions_tuples = get_permissions_in_code_as_tuples(
+        resources_to_register,
+        views_in_db,
+        base_permissions_assignation,
+        extra_permissions,
+    )
+    permissions_to_register = get_permissions_to_register(
+        permissions_tuples, permissions_in_db_keys
+    )
+    permissions_to_delete = get_permissions_to_delete(
+        permissions_tuples, resources_roles_with_access.keys(), permissions_in_db
+    )
+
+    # Save the new permissions in the data
+    save_and_delete_permissions(permissions_to_register, permissions_to_delete)
 
-    # Create base permissions
-    permissions_in_app = [
-        PermissionViewRoleModel(
-            {
-                "role_id": role,
-                "action_id": action,
-                "api_view_id": views_in_db[view["endpoint"]],
-            }
-        )
-        for role, action in BASE_PERMISSION_ASSIGNATION
-        for view in resources_to_register
-        if role in view["resource"].ROLES_WITH_ACCESS
-    ] + [
-        PermissionViewRoleModel(
-            {
-                "role_id": role,
-                "action_id": action,
-                "api_view_id": views_in_db[endpoint],
-            }
-        )
-        for role, action, endpoint in EXTRA_PERMISSION_ASSIGNATION
-    ]
+    if len(permissions_to_register) > 0:
+        current_app.logger.info(f"Permissions registered: {permissions_to_register}")
+    else:
+        current_app.logger.info("No new permissions to register")
 
-    permissions_in_app_keys = [
-        (perm.role_id, perm.action_id, perm.api_view_id) for perm in permissions_in_app
-    ]
+    if len(permissions_to_delete) > 0:
+        current_app.logger.info(f"Permissions deleted: {permissions_to_delete}")
+    else:
+        current_app.logger.info("No permissions to delete")
 
-    permissions_to_register = [
-        permission
-        for permission in permissions_in_app
-        if (permission.role_id, permission.action_id, permission.api_view_id)
-        not in permissions_in_db_keys
-    ]
 
-    permissions_to_delete = [
-        permission
-        for permission in permissions_in_db
-        if (permission.role_id, permission.action_id, permission.api_view_id)
-        not in permissions_in_app_keys
-        and permission.api_view.name in resources_names
-    ]
+def save_new_roles(new_roles_to_add):
+    """
+    Save the new roles in the database.
+    new_roles_to_add: List of new roles to add.
+    """
+    if len(new_roles_to_add) > 0:
+        db.session.bulk_save_objects(new_roles_to_add)
+        try:
+            db.session.commit()
+        except IntegrityError as e:
+            db.session.rollback()
+            current_app.logger.error(
+                f"Integrity error on base permissions register: {e}"
+            )
+        except DBAPIError as e:
+            db.session.rollback()
+            current_app.logger.error(f"Unknown error on base permissions register: {e}")
+
 
+def save_and_delete_permissions(permissions_to_register, permissions_to_delete):
+    """
+    Save and delete permissions in the database.
+    permissions_to_register: List of permissions to register.
+    permissions_to_delete: List of permissions to delete.
+    """
     if len(permissions_to_register) > 0:
         db.session.bulk_save_objects(permissions_to_register)
 
-    # TODO: for now the permission are not going to get deleted just in case.
-    #  We are just going to register new permissions
-    # if len(permissions_to_delete) > 0:
-    #     for permission in permissions_to_delete:
-    #         db.session.delete(permission)
-
+    if len(permissions_to_delete) > 0:
+        for permission in permissions_to_delete:
+            db.session.delete(permission)
     try:
         db.session.commit()
     except IntegrityError as e:
@@ -105,25 +116,119 @@ def register_base_permissions_command(external_app: str = None, verbose: bool =
                 f"Unknown error on base permissions sequence updating: {e}"
             )
 
-    if verbose:
-        if len(permissions_to_register) > 0:
-            current_app.logger.info(
-                f"Permissions registered: {permissions_to_register}"
-            )
-        else:
-            current_app.logger.info("No new permissions to register")
 
-        if len(permissions_to_delete) > 0:
-            current_app.logger.info(f"Permissions deleted: {permissions_to_delete}")
-        else:
-            current_app.logger.info("No permissions to delete")
+def get_permissions_to_delete(permissions_tuples, resources_names, permissions_in_db):
+    """
+    Get the permissions to delete.
+    """
+    permissions_to_delete = [
+        permission
+        for permission in permissions_in_db
+        if (permission.role_id, permission.action_id, permission.api_view_id)
+        not in permissions_tuples
+    ]
+
+    return permissions_to_delete
+
+
+def get_permissions_to_register(permissions_tuples, permissions_in_db_keys):
+    """
+    Get the permissions to register.
+    """
+    # Convert set of tuples to list of PermissionViewRoleModel objects
+    return [
+        PermissionViewRoleModel(
+            {
+                "role_id": role_id,
+                "action_id": action_id,
+                "api_view_id": api_view_id,
+            }
+        )
+        for role_id, action_id, api_view_id in permissions_tuples
+        if (role_id, action_id, api_view_id) not in permissions_in_db_keys
+    ]
+
+
+def get_permissions_in_code_as_tuples(
+    resources_to_register, views_in_db, base_permissions_assignation, extra_permissions
+):
+    """
+    Get the permissions in code as tuples.
+    """
+    # Create base permissions using a set to avoid duplicates
+    permissions_tuples = set()
+
+    # Add permissions from ROLES_WITH_ACCESS
+    for role, action in base_permissions_assignation:
+        for view in resources_to_register:
+            if role in view["resource"].ROLES_WITH_ACCESS:
+                permissions_tuples.add((role, action, views_in_db[view["endpoint"]]))
+
+    # Add permissions from extra_permissions
+    for role, action, endpoint in extra_permissions:
+        if endpoint in views_in_db:
+            permissions_tuples.add((role, action, views_in_db[endpoint]))
+
+    return permissions_tuples
+
+
+def get_base_permissions(resources_roles_with_access, custom_roles_actions):
+    """
+    Get the new roles and base permissions assignation.
+    resources_roles_with_access: Dictionary of resources and roles with access.
+    custom_roles_actions: Dictionary mapping custom roles to their allowed actions.
+    """
+    # Get all custom roles (both new and existing) that appear in ROLES_WITH_ACCESS
+    all_custom_roles_in_access = set(
+        [
+            role
+            for roles in resources_roles_with_access.values()
+            for role in roles
+            if role not in ALL_DEFAULT_ROLES
+        ]
+    )
+
+    # Validate that all custom roles are defined in custom_roles_actions
+    undefined_roles = all_custom_roles_in_access - set(custom_roles_actions.keys())
+    if undefined_roles:
+        raise ValueError(
+            f"The following custom roles are used in code but not defined in CUSTOM_ROLES_ACTIONS: {undefined_roles}. "
+            f"Please define their allowed actions in the CUSTOM_ROLES_ACTIONS dictionary in shared/const.py."
+        )
+
+    # Create extended permission assignation including all custom roles
+    # For custom roles, use the actions defined in custom_roles_actions
+    custom_permissions = [
+        (custom_role, action)
+        for custom_role in all_custom_roles_in_access
+        for action in custom_roles_actions[custom_role]
+    ]
+
+    base_permissions_assignation = BASE_PERMISSION_ASSIGNATION + custom_permissions
+
+    return base_permissions_assignation
+
+
+def get_db_permissions():
+    """
+    Get all permissions in the database.
+    """
+    permissions_in_db = [perm for perm in PermissionViewRoleModel.get_all_objects()]
+    permissions_in_db_keys = [
+        (perm.role_id, perm.action_id, perm.api_view_id) for perm in permissions_in_db
+    ]
 
-    return True
+    return permissions_in_db, permissions_in_db_keys
 
 
 def register_dag_permissions_command(
     open_deployment: int = None, verbose: bool = False
 ):
+    """
+    Register DAG permissions.
+    open_deployment: If 1, it will register the permissions for the open deployment.
+    verbose: If True, it will print the permissions that are being registered.
+    """
 
     from flask import current_app
     from sqlalchemy.exc import DBAPIError, IntegrityError

cornflow/commands/roles.py

@@ -1,22 +1,23 @@
-def register_roles_command(verbose: bool = True):
+def register_roles_command(external_app: str = None, verbose: bool = True):
 
     from sqlalchemy.exc import DBAPIError, IntegrityError
     from flask import current_app
 
-    from cornflow.models import RoleModel
-    from cornflow.shared.const import ROLES_MAP
     from cornflow.shared import db
+    from cornflow.commands.auxiliar import (
+        get_all_external,
+        get_all_resources,
+        get_new_roles_to_add,
+    )
 
-    roles_registered = [role.name for role in RoleModel.get_all_objects()]
+    resources_to_register, extra_permissions, _ = get_all_external(external_app)
+    resources_roles_with_access = get_all_resources(resources_to_register)
+    new_roles_to_add = get_new_roles_to_add(
+        extra_permissions, resources_roles_with_access
+    )
 
-    roles_to_register = [
-        RoleModel({"id": key, "name": value})
-        for key, value in ROLES_MAP.items()
-        if value not in roles_registered
-    ]
-
-    if len(roles_to_register) > 0:
-        db.session.bulk_save_objects(roles_to_register)
+    if len(new_roles_to_add) > 0:
+        db.session.bulk_save_objects(new_roles_to_add)
 
     try:
         db.session.commit()
@@ -38,8 +39,8 @@ def register_roles_command(verbose: bool = True):
             current_app.logger.error(f"Unknown error on roles sequence updating: {e}")
 
     if verbose:
-        if len(roles_to_register) > 0:
-            current_app.logger.info(f"Roles registered: {roles_to_register}")
+        if len(new_roles_to_add) > 0:
+            current_app.logger.info(f"Roles registered: {new_roles_to_add}")
         else:
             current_app.logger.info("No new roles to be registered")
 

cornflow/commands/schemas.py

@@ -3,6 +3,7 @@ def update_schemas_command(url, user, pwd, verbose: bool = False):
     from flask import current_app
 
     from cornflow_client.airflow.api import Airflow
+    from cornflow.shared.const import AIRFLOW_NOT_REACHABLE_MSG
 
     af_client = Airflow(url, user, pwd)
     max_attempts = 20
@@ -10,12 +11,12 @@ def update_schemas_command(url, user, pwd, verbose: bool = False):
     while not af_client.is_alive() and attempts < max_attempts:
         attempts += 1
         if verbose == 1:
-            current_app.logger.info(f"Airflow is not reachable (attempt {attempts})")
+            current_app.logger.info(f"{AIRFLOW_NOT_REACHABLE_MSG} (attempt {attempts})")
         time.sleep(15)
 
     if not af_client.is_alive():
         if verbose == 1:
-            current_app.logger.info("Airflow is not reachable")
+            current_app.logger.info(f"{AIRFLOW_NOT_REACHABLE_MSG}")
         return False
 
     response = af_client.update_schemas()
@@ -34,6 +35,7 @@ def update_dag_registry_command(url, user, pwd, verbose: bool = False):
     from flask import current_app
 
     from cornflow_client.airflow.api import Airflow
+    from cornflow.shared.const import AIRFLOW_NOT_REACHABLE_MSG
 
     af_client = Airflow(url, user, pwd)
     max_attempts = 20
@@ -41,12 +43,12 @@ def update_dag_registry_command(url, user, pwd, verbose: bool = False):
     while not af_client.is_alive() and attempts < max_attempts:
         attempts += 1
         if verbose == 1:
-            current_app.logger.info(f"Airflow is not reachable (attempt {attempts})")
+            current_app.logger.info(f"{AIRFLOW_NOT_REACHABLE_MSG} (attempt {attempts})")
         time.sleep(15)
 
     if not af_client.is_alive():
         if verbose == 1:
-            current_app.logger.info("Airflow is not reachable")
+            current_app.logger.info(f"{AIRFLOW_NOT_REACHABLE_MSG}")
         return False
 
     response = af_client.update_dag_registry()

cornflow/commands/users.py

@@ -16,19 +16,16 @@ def create_user_with_role(
             current_app.logger.info(
                 f"User {username} is created and assigned {role_name} role"
             )
-        return True
+        return
 
     user_roles = UserRoleModel.get_all_objects(user_id=user.id)
     user_actual_roles = [ur.role for ur in user_roles]
-    if (
-        user_roles is not None
-        and RoleModel.get_one_object(role) in user_actual_roles
-    ):
+    if user_roles is not None and RoleModel.get_one_object(role) in user_actual_roles:
         if verbose:
             current_app.logger.info(
                 f"User {username} exists and already has {role_name} role assigned"
             )
-        return True
+        return
 
     user_role = UserRoleModel({"user_id": user.id, "role_id": role})
     user_role.save()
@@ -36,7 +33,6 @@ def create_user_with_role(
         current_app.logger.info(
             f"User {username} already exists and is assigned a {role_name} role"
         )
-    return True
 
 
 def create_service_user_command(username, email, password, verbose: bool = True):
@@ -45,8 +41,9 @@ def create_service_user_command(username, email, password, verbose: bool = True)
 
     if username is None or email is None or password is None:
         current_app.logger.info("Missing required arguments")
-        return False
-    return create_user_with_role(
+        return
+
+    create_user_with_role(
         username, email, password, "serviceuser", SERVICE_ROLE, verbose
     )
 
@@ -57,10 +54,9 @@ def create_admin_user_command(username, email, password, verbose: bool = True):
 
     if username is None or email is None or password is None:
         current_app.logger.info("Missing required arguments")
-        return False
-    return create_user_with_role(
-        username, email, password, "admin", ADMIN_ROLE, verbose
-    )
+        return
+
+    create_user_with_role(username, email, password, "admin", ADMIN_ROLE, verbose)
 
 
 def create_planner_user_command(username, email, password, verbose: bool = True):
@@ -69,7 +65,6 @@ def create_planner_user_command(username, email, password, verbose: bool = True)
 
     if username is None or email is None or password is None:
         current_app.logger.info("Missing required arguments")
-        return False
-    return create_user_with_role(
-        username, email, password, "planner", PLANNER_ROLE, verbose
-    )
+        return
+
+    create_user_with_role(username, email, password, "planner", PLANNER_ROLE, verbose)