cornflow 1.2.3a3__py3-none-any.whl → 1.2.3a5__py3-none-any.whl

cornflow/app.py

@@ -164,7 +164,7 @@ def create_base_user(username, email, password, verbose):
 @click.option("-v", "--verbose", is_flag=True, default=False)
 @with_appcontext
 def register_roles(verbose):
-    register_roles_command(verbose)
+    register_roles_command(verbose=verbose)
 
 
 @click.command("register_actions")

cornflow/cli/roles.py

@@ -18,4 +18,4 @@ def roles():
 def init(verbose):
     app = get_app()
     with app.app_context():
-        register_roles_command(verbose)
+        register_roles_command(verbose=verbose)

cornflow/cli/service.py

@@ -2,8 +2,20 @@ import os
 import subprocess
 import sys
 import time
+import logging
 from logging import error
 
+# Configure a logger for the service module
+logger = logging.getLogger("cornflow.service")
+logger.setLevel(logging.INFO)
+
+# Add console handler if not already present
+if not logger.handlers:
+    handler = logging.StreamHandler(sys.stdout)
+    formatter = logging.Formatter("%(asctime)s [%(name)s] [%(levelname)s] %(message)s")
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    logger.propagate = False
 
 import click
 from .utils import get_db_conn
@@ -89,9 +101,10 @@ def init_cornflow_service():
 
         external_app_lib = import_module(external_app_module)
         app = external_app_lib.create_wsgi_app(environment, cornflow_db_conn)
-
         with app.app_context():
+
             _initialize_database(app, external_app_module)
+
             _create_initial_users(
                 config["auth"],
                 config["cornflow_admin_user"],
@@ -101,6 +114,7 @@ def init_cornflow_service():
                 config["cornflow_service_email"],
                 config["cornflow_service_pwd"],
             )
+
             _sync_with_airflow(
                 config["airflow_url"],
                 config["airflow_user"],
@@ -117,6 +131,7 @@ def init_cornflow_service():
 
 def _setup_environment_variables():
     """Reads environment variables, sets defaults, and returns config values."""
+
     environment = os.getenv("FLASK_ENV", "development")
     os.environ["FLASK_ENV"] = environment
 
@@ -218,7 +233,7 @@ def _configure_logging(cornflow_logging):
             if logrotate.returncode != 0:
                 error(f"Error configuring logrotate: {logrotate.stderr}")
             else:
-                print(logrotate.stdout)
+                logger.info(logrotate.stdout)
         except Exception as e:
             error(f"Exception during logrotate configuration: {e}")
 
@@ -236,6 +251,7 @@ def _initialize_database(app, external_app_module=None):
 
         Migrate(app=app, db=db, directory=migrations_path)
         upgrade()
+        logger.info("----------------Migrations applied----------------")
         access_init_command(verbose=False)
 
 
@@ -285,14 +301,21 @@ def _sync_with_airflow(
 
 def _setup_external_app():
     """Performs setup steps specific to external applications."""
+
     os.chdir(MAIN_WD)
+
     if _register_key():
+
         prefix = "CUSTOM_SSH_"
         env_variables = {
             key: value for key, value in os.environ.items() if key.startswith(prefix)
         }
         for _, value in env_variables.items():
             _register_ssh_host(value)
+    else:
+        logger.info(
+            "************************ NO SSH KEY TO REGISTER ************************"
+        )
 
     # Install requirements for the external app
     pip_install_cmd = "$(command -v pip) install --user -r requirements.txt"
@@ -301,10 +324,15 @@ def _setup_external_app():
     if result.returncode != 0:
         error(f"Error installing requirements: {result.stderr}")
     else:
-        print(result.stdout)
+        logger.info(result.stdout)
     time.sleep(5)  # Consider if this sleep is truly necessary
     sys.path.append(MAIN_WD)
 
+    # Add .local path to sys.path so pip --user packages can be found
+    local_lib_path = os.path.expanduser("~/.local/lib/python3.12/site-packages")
+    if local_lib_path not in sys.path:
+        sys.path.insert(0, local_lib_path)
+
 
 def _start_application(external_application, environment, external_app_module=None):
     """Starts the Gunicorn server."""
@@ -349,4 +377,4 @@ def _register_key():
         os.system(add_key)
         return True
     else:
-        return False
+        return False

cornflow/commands/access.py

@@ -1,3 +1,4 @@
+import logging
 import os
 
 from .actions import register_actions_command
@@ -5,18 +6,28 @@ from .permissions import register_base_permissions_command
 from .roles import register_roles_command
 from .views import register_views_command
 
+# Configure logger for access
+logger = logging.getLogger("cornflow.access")
+
 
 def access_init_command(verbose: bool = False):
+    """
+    Initialize the access to the system.
+    """
+
     external = int(os.getenv("EXTERNAL_APP", 0))
     external_app = os.getenv("EXTERNAL_APP_MODULE", "external_app")
 
     register_actions_command(verbose)
-    register_roles_command(verbose)
 
-    register_views_command(verbose=verbose)
     if external != 0:
+        register_roles_command(external_app=external_app, verbose=verbose)
         register_views_command(external_app=external_app, verbose=verbose)
+    else:
+        register_roles_command(verbose=verbose)
+        register_views_command(verbose=verbose)
 
-    register_base_permissions_command(verbose=verbose)
     if external != 0:
         register_base_permissions_command(external_app=external_app, verbose=verbose)
+    else:
+        register_base_permissions_command(verbose=verbose)

cornflow/commands/auxiliar.py

@@ -0,0 +1,106 @@
+import sys
+from importlib import import_module
+
+from flask import current_app
+
+from cornflow.endpoints import resources, alarms_resources
+from cornflow.models import RoleModel
+from cornflow.shared.const import (
+    EXTRA_PERMISSION_ASSIGNATION,
+    ALL_DEFAULT_ROLES,
+)
+from cornflow.shared.const import ROLES_MAP
+
+
+def get_all_external(external_app):
+    """
+    Get all resources, extra permissions, and custom roles actions.
+    external_app: If provided, it will get the resources and extra permissions for the external app.
+    """
+    if external_app is None:
+        resources_to_register = resources
+        extra_permissions = EXTRA_PERMISSION_ASSIGNATION
+        custom_roles_actions = {}
+        if current_app.config["ALARMS_ENDPOINTS"]:
+            resources_to_register = resources + alarms_resources
+    else:
+        sys.path.append("./")
+        external_module = import_module(external_app)
+        try:
+            extra_permissions = (
+                EXTRA_PERMISSION_ASSIGNATION
+                + external_module.shared.const.EXTRA_PERMISSION_ASSIGNATION
+            )
+        except AttributeError:
+            extra_permissions = EXTRA_PERMISSION_ASSIGNATION
+
+        try:
+            custom_roles_actions = external_module.shared.const.CUSTOM_ROLES_ACTIONS
+        except AttributeError:
+            custom_roles_actions = {}
+
+        if current_app.config["ALARMS_ENDPOINTS"]:
+            resources_to_register = (
+                external_module.endpoints.resources + resources + alarms_resources
+            )
+        else:
+            resources_to_register = external_module.endpoints.resources + resources
+    return resources_to_register, extra_permissions, custom_roles_actions
+
+
+def get_all_resources(resources_to_register):
+    """
+    Get all resources and roles with access.
+    resources_to_register: List of resources to register.
+    """
+
+    resources_roles_with_access = {
+        resource["endpoint"]: resource["resource"].ROLES_WITH_ACCESS
+        for resource in resources_to_register
+    }
+
+    return resources_roles_with_access
+
+
+def get_new_roles_to_add(extra_permissions, resources_roles_with_access):
+    """
+    Get the new roles to add.
+    extra_permissions: List of extra permissions.
+    resources_roles_with_access: Dictionary of resources and roles with access.
+    """
+
+    roles_with_access = list(
+        set([role for roles in resources_roles_with_access.values() for role in roles])
+    )
+    roles_in_extra_permissions = [role for role, _, _ in extra_permissions]
+    roles_with_access = list(set(roles_with_access + roles_in_extra_permissions))
+
+    # Add all default roles that are referenced in BASE_PERMISSION_ASSIGNATION
+    roles_with_access = list(set(roles_with_access + ALL_DEFAULT_ROLES))
+
+    # We extract the existing roles in the database
+    existing_roles = [role.id for role in RoleModel.get_all_objects()]
+    new_roles_to_add = []
+
+    for role_id in roles_with_access:
+        if role_id not in existing_roles:
+            if role_id in ALL_DEFAULT_ROLES:
+                # Create standard role with predefined name
+                role_name = ROLES_MAP[role_id]
+                new_role = RoleModel(
+                    {
+                        "id": role_id,
+                        "name": role_name,
+                    }
+                )
+            else:
+                # Create custom role with custom_role_<id> name
+                new_role = RoleModel(
+                    {
+                        "id": role_id,
+                        "name": f"custom_role_{role_id}",
+                    }
+                )
+            new_roles_to_add.append(new_role)
+
+    return new_roles_to_add

cornflow/commands/permissions.py

@@ -1,112 +1,99 @@
-import sys
-from importlib import import_module
+from flask import current_app
+from sqlalchemy.exc import DBAPIError, IntegrityError
 
+from cornflow.commands.auxiliar import (
+    get_all_external,
+    get_all_resources,
+)
+from cornflow.models import ViewModel, PermissionViewRoleModel
+from cornflow.shared import db
+from cornflow.shared.const import ALL_DEFAULT_ROLES, GET_ACTION
 from cornflow.shared.const import (
     BASE_PERMISSION_ASSIGNATION,
-    EXTRA_PERMISSION_ASSIGNATION,
-    ROLES_MAP,
-    GET_ACTION,
-    PATCH_ACTION,
-    POST_ACTION,
-    PUT_ACTION,
-    DELETE_ACTION,
 )
-from cornflow.models import ViewModel, PermissionViewRoleModel, RoleModel
-from cornflow.shared import db
-from flask import current_app
-from sqlalchemy.exc import DBAPIError, IntegrityError
 
 
 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
+    """
+    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, 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)
+
+    if len(permissions_to_register) > 0:
+        current_app.logger.info(f"Permissions registered: {permissions_to_register}")
     else:
-        resources_to_register = []
-        exit()
+        current_app.logger.info("No new permissions to register")
 
-    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]
-    roles_in_db = [role.id for role in RoleModel.get_all_objects()]
-    # Check which roles are not in ROLES_MAP
-    roles_not_in_map = [
-        role_id for role_id in roles_in_db if role_id not in ROLES_MAP.keys()
-    ]
-    complete_base_assignation = BASE_PERMISSION_ASSIGNATION.copy()
-    if len(roles_not_in_map) > 0:
-        # We add to the complete_base_assignation the roles that are not in ROLES_MAP
-        for role_id in roles_not_in_map:
-            for action in [
-                GET_ACTION,
-                PATCH_ACTION,
-                POST_ACTION,
-                PUT_ACTION,
-                DELETE_ACTION,
-            ]:
-                complete_base_assignation.append((role_id, action))
-
-    # 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 complete_base_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_delete) > 0:
+        current_app.logger.info(f"Permissions deleted: {permissions_to_delete}")
+    else:
+        current_app.logger.info("No permissions to delete")
 
-    permissions_in_app_keys = [
-        (perm.role_id, perm.action_id, perm.api_view_id) for perm in permissions_in_app
-    ]
 
-    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
-    ]
+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}")
 
-    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_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)
-
     try:
         db.session.commit()
     except IntegrityError as e:
@@ -129,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")