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

cornflow/endpoints/dag.py

@@ -2,6 +2,7 @@
 Internal endpoint for getting and posting execution data
 These are the endpoints used by airflow in its communication with cornflow
 """
+
 # Import from libraries
 from cornflow_client.constants import SOLUTION_SCHEMA
 from flask import current_app
@@ -17,7 +18,6 @@ from cornflow.schemas.execution import (
     ExecutionDagPostRequest,
     ExecutionDagRequest,
     ExecutionDetailsEndpointResponse,
-    ExecutionSchema,
 )
 
 from cornflow.shared.authentication import Auth, authenticate
@@ -85,7 +85,7 @@ class DAGDetailEndpoint(BaseMetaResource):
     @doc(description="Edit an execution", tags=["DAGs"])
     @authenticate(auth_class=Auth())
     @use_kwargs(ExecutionDagRequest, location="json")
-    def put(self, idx, **req_data):
+    def put(self, idx, **kwargs):
         """
         API method to write the results of the execution
         It requires authentication to be passed in the form of a token that has to be linked to
@@ -95,13 +95,20 @@ class DAGDetailEndpoint(BaseMetaResource):
         :return: A dictionary with a message (body) and an integer with the HTTP status code
         :rtype: Tuple(dict, integer)
         """
-        solution_schema = req_data.pop("solution_schema", "pulp")
+        execution = ExecutionModel.get_one_object(user=self.get_user(), idx=idx)
+        if execution is None:
+            err = "The execution does not exist."
+            raise ObjectDoesNotExist(
+                error=err,
+                log_txt=f"Error while user {self.get_user()} tries to edit execution {idx}."
+                + err,
+            )
+
+        solution_schema = execution.schema
 
-        # TODO: the solution_schema maybe we should get it from the created execution_id?
-        #  at least, check they have the same schema-name
         # Check data format
-        data = req_data.get("data")
-        checks = req_data.get("checks")
+        data = kwargs.get("data")
+        checks = kwargs.get("checks")
         if data is None:
             # only check format if executions_results exist
             solution_schema = None
@@ -111,24 +118,19 @@ class DAGDetailEndpoint(BaseMetaResource):
         if solution_schema is not None:
             config = current_app.config
 
-            solution_schema = DeployedDAG.get_one_schema(config, solution_schema, SOLUTION_SCHEMA)
+            solution_schema = DeployedDAG.get_one_schema(
+                config, solution_schema, SOLUTION_SCHEMA
+            )
             solution_errors = json_schema_validate_as_string(solution_schema, data)
 
             if solution_errors:
                 raise InvalidData(
                     payload=dict(jsonschema_errors=solution_errors),
                     log_txt=f"Error while user {self.get_user()} tries to edit execution {idx}. "
-                            f"Solution data do not match the jsonschema.",
+                    f"Solution data do not match the jsonschema.",
                 )
-        execution = ExecutionModel.get_one_object(user=self.get_user(), idx=idx)
-        if execution is None:
-            err = "The execution does not exist."
-            raise ObjectDoesNotExist(
-                error=err,
-                log_txt=f"Error while user {self.get_user()} tries to edit execution {idx}."
-                + err,
-            )
-        state = req_data.get("state", EXEC_STATE_CORRECT)
+
+        state = kwargs.get("state", EXEC_STATE_CORRECT)
         new_data = dict(
             state=state,
             state_message=EXECUTION_STATE_MESSAGE_DICT[state],
@@ -141,10 +143,9 @@ class DAGDetailEndpoint(BaseMetaResource):
             new_data["data"] = data
         if checks is not None:
             new_data["checks"] = checks
-        req_data.update(new_data)
-        execution.update(req_data)
-        # TODO: is this save necessary?
-        execution.save()
+        kwargs.update(new_data)
+        execution.update(kwargs)
+
         current_app.logger.info(f"User {self.get_user()} edits execution {idx}")
         return {"message": "results successfully saved"}, 200
 
@@ -207,7 +208,6 @@ class DAGEndpointManual(BaseMetaResource):
 
         # Check data format
         data = kwargs.get("data")
-        # TODO: create a function to validate and replace data/ execution_results
         if data is None:
             # only check format if executions_results exist
             solution_schema = None
@@ -215,14 +215,16 @@ class DAGEndpointManual(BaseMetaResource):
             solution_schema = "solve_model_dag"
         if solution_schema is not None:
             config = current_app.config
-            solution_schema = DeployedDAG.get_one_schema(config, solution_schema, SOLUTION_SCHEMA)
+            solution_schema = DeployedDAG.get_one_schema(
+                config, solution_schema, SOLUTION_SCHEMA
+            )
             solution_errors = json_schema_validate_as_string(solution_schema, data)
 
             if solution_errors:
                 raise InvalidData(
                     payload=dict(jsonschema_errors=solution_errors),
                     log_txt=f"Error while user {self.get_user()} tries to manually create an execution. "
-                            f"Solution data do not match the jsonschema.",
+                    f"Solution data do not match the jsonschema.",
                 )
 
         kwargs_copy = dict(kwargs)

cornflow/endpoints/data_check.py

@@ -14,25 +14,105 @@ from cornflow.models import InstanceModel, ExecutionModel, CaseModel, DeployedDA
 from cornflow.schemas.execution import ExecutionDetailsEndpointResponse
 from cornflow.shared.authentication import Auth, authenticate
 from cornflow.shared.const import (
+    AIRFLOW_ERROR_MSG,
+    AIRFLOW_NOT_REACHABLE_MSG,
+    DAG_PAUSED_MSG,
     EXEC_STATE_QUEUED,
     EXEC_STATE_ERROR,
     EXEC_STATE_ERROR_START,
     EXEC_STATE_NOT_RUN,
-    EXECUTION_STATE_MESSAGE_DICT, VIEWER_ROLE, PLANNER_ROLE, ADMIN_ROLE,
+    EXECUTION_STATE_MESSAGE_DICT,
+    PLANNER_ROLE,
+    ADMIN_ROLE,
 )
 from cornflow.shared.exceptions import (
     AirflowError,
     ObjectDoesNotExist,
     InvalidUsage,
-    InvalidData
+    InvalidData,
 )
 from cornflow.shared.validators import json_schema_validate_as_string
 
 
+def _run_airflow_data_check(
+    af_client: Airflow,
+    execution: ExecutionModel,
+    schema: str,
+    user,
+    context_str: str,
+    **run_dag_kwargs,
+):
+    """
+    Helper function to check Airflow status and run a data check DAG.
+
+    :param af_client: Initialized Airflow client.
+    :param execution: The ExecutionModel object to update.
+    :param schema: The name of the schema/DAG to run.
+    :param user: The user object performing the action.
+    :param context_str: A string describing the context (e.g., "execution {id}") for logging.
+    :param run_dag_kwargs: Additional keyword arguments for af_client.run_dag.
+    :return: None. Updates execution object in place and raises AirflowError on failure.
+    """
+    # Check Airflow liveness
+    if not af_client.is_alive():
+        current_app.logger.error(AIRFLOW_NOT_REACHABLE_MSG)
+        execution.update_state(EXEC_STATE_ERROR_START)
+        raise AirflowError(
+            error=AIRFLOW_NOT_REACHABLE_MSG,
+            payload=dict(
+                message=EXECUTION_STATE_MESSAGE_DICT[EXEC_STATE_ERROR_START],
+                state=EXEC_STATE_ERROR_START,
+            ),
+            log_txt=f"Error while user {user} tries to run data checks on {context_str}. "
+            + AIRFLOW_NOT_REACHABLE_MSG,
+        )
+
+    # Check if DAG is paused
+    schema_info = af_client.get_dag_info(schema)
+    info = schema_info.json()
+    if info.get("is_paused", False):
+        current_app.logger.error(DAG_PAUSED_MSG)
+        execution.update_state(EXEC_STATE_ERROR_START)
+        raise AirflowError(
+            error=DAG_PAUSED_MSG,
+            payload=dict(
+                message=EXECUTION_STATE_MESSAGE_DICT[EXEC_STATE_ERROR_START],
+                state=EXEC_STATE_ERROR_START,
+            ),
+            log_txt=f"Error while user {user} tries to run data checks on {context_str}. "
+            + DAG_PAUSED_MSG,
+        )
+
+    # Run the DAG
+    try:
+        response = af_client.run_dag(
+            execution.id, dag_name=schema, checks_only=True, **run_dag_kwargs
+        )
+    except AirflowError as err:
+        error = f"{AIRFLOW_ERROR_MSG} {err}"
+        current_app.logger.error(error)
+        execution.update_state(EXEC_STATE_ERROR)
+        raise AirflowError(
+            error=error,
+            payload=dict(
+                message=EXECUTION_STATE_MESSAGE_DICT[EXEC_STATE_ERROR],
+                state=EXEC_STATE_ERROR,
+            ),
+            log_txt=f"Error while user {user} tries to run data checks on {context_str}. "
+            + error,
+        )
+
+    # Update execution on success
+    af_data = response.json()
+    execution.dag_run_id = af_data.get("dag_run_id")
+    execution.update_state(EXEC_STATE_QUEUED)
+
+
 class DataCheckExecutionEndpoint(BaseMetaResource):
     """
     Endpoint used to execute the instance and solution checks on an execution
     """
+
     ROLES_WITH_ACCESS = [PLANNER_ROLE, ADMIN_ROLE]
 
     def __init__(self):
@@ -56,13 +136,16 @@ class DataCheckExecutionEndpoint(BaseMetaResource):
         :rtype: Tuple(dict, integer)
         """
         config = current_app.config
+        user = self.get_user()
+        context_str = f"execution {idx}"
 
         execution = ExecutionModel.get_one_object(user=self.get_user(), idx=idx)
         if execution is None:
             err = "The execution to check does not exist"
             raise ObjectDoesNotExist(
                 error=err,
-                log_txt=f"Error while user {self.get_user()} tries to run data checks on execution {idx}. " + err
+                log_txt=f"Error while user {user} tries to run data checks on {context_str}. "
+                + err,
             )
 
         schema = execution.schema
@@ -72,7 +155,8 @@ class DataCheckExecutionEndpoint(BaseMetaResource):
             err = "The execution is still running"
             raise InvalidUsage(
                 error=err,
-                log_txt=f"Error while user {self.get_user()} tries to run data checks on execution {idx}. " + err
+                log_txt=f"Error while user {user} tries to run data checks on {context_str}. "
+                + err,
             )
 
         # this allows testing without airflow interaction:
@@ -80,58 +164,12 @@ class DataCheckExecutionEndpoint(BaseMetaResource):
             execution.update_state(EXEC_STATE_NOT_RUN)
             return execution, 201
 
-        # We now try to launch the task in airflow
+        # Initialize Airflow client
         af_client = Airflow.from_config(config)
-        if not af_client.is_alive():
-            err = "Airflow is not accessible"
-            current_app.logger.error(err)
-            execution.update_state(EXEC_STATE_ERROR_START)
-            raise AirflowError(
-                error=err,
-                payload=dict(
-                    message=EXECUTION_STATE_MESSAGE_DICT[EXEC_STATE_ERROR_START],
-                    state=EXEC_STATE_ERROR_START,
-                ),
-                log_txt=f"Error while user {self.get_user()} tries to run data checks on execution {idx}. " + err
-            )
-        # ask airflow if dag_name exists
-        schema_info = af_client.get_dag_info(schema)
-
-        info = schema_info.json()
-        if info["is_paused"]:
-            err = "The dag exists but it is paused in airflow"
-            current_app.logger.error(err)
-            execution.update_state(EXEC_STATE_ERROR_START)
-            raise AirflowError(
-                error=err,
-                payload=dict(
-                    message=EXECUTION_STATE_MESSAGE_DICT[EXEC_STATE_ERROR_START],
-                    state=EXEC_STATE_ERROR_START,
-                ),
-                log_txt=f"Error while user {self.get_user()} tries to run data checks on execution {idx}. " + err
-            )
 
-        try:
-            response = af_client.run_dag(
-                execution.id, dag_name=schema, checks_only=True
-            )
-        except AirflowError as err:
-            error = "Airflow responded with an error: {}".format(err)
-            current_app.logger.error(error)
-            execution.update_state(EXEC_STATE_ERROR)
-            raise AirflowError(
-                error=error,
-                payload=dict(
-                    message=EXECUTION_STATE_MESSAGE_DICT[EXEC_STATE_ERROR],
-                    state=EXEC_STATE_ERROR,
-                ),
-                log_txt=f"Error while user {self.get_user()} tries to run data checks on execution {idx}. " + error
-            )
+        # Call the shared Airflow execution logic
+        _run_airflow_data_check(af_client, execution, schema, user, context_str)
 
-        # if we succeed, we register the dag_run_id in the execution table:
-        af_data = response.json()
-        execution.dag_run_id = af_data["dag_run_id"]
-        execution.update_state(EXEC_STATE_QUEUED)
         current_app.logger.info(
             "User {} launches checks of execution {}".format(
                 self.get_user_id(), execution.id
@@ -144,7 +182,9 @@ class DataCheckInstanceEndpoint(BaseMetaResource):
     """
     Endpoint used to execute the instance and solution checks on an execution
     """
+
     ROLES_WITH_ACCESS = [PLANNER_ROLE, ADMIN_ROLE]
+
     def __init__(self):
         super().__init__()
         self.model = ExecutionModel
@@ -169,13 +209,16 @@ class DataCheckInstanceEndpoint(BaseMetaResource):
         :rtype: Tuple(dict, integer)
         """
         config = current_app.config
+        user = self.get_user()
+        context_str = f"instance {idx}"
 
         instance = InstanceModel.get_one_object(user=self.get_user(), idx=idx)
         if instance is None:
             err = "The instance to check does not exist"
             raise ObjectDoesNotExist(
                 error=err,
-                log_txt=f"Error while user {self.get_user()} tries to run data checks on instance {idx}. " + err
+                log_txt=f"Error while user {user} tries to run data checks on {context_str}. "
+                + err,
             )
         payload = dict(
             config=dict(checks_only=True),
@@ -185,67 +228,19 @@ class DataCheckInstanceEndpoint(BaseMetaResource):
         )
         schema = instance.schema
 
-        execution, status_code = self.post_list(data=payload)
+        execution, _ = self.post_list(data=payload)
 
         # this allows testing without airflow interaction:
         if request.args.get("run", "1") == "0":
             execution.update_state(EXEC_STATE_NOT_RUN)
             return execution, 201
 
-        # We now try to launch the task in airflow
+        # Initialize Airflow client
         af_client = Airflow.from_config(config)
-        if not af_client.is_alive():
-            err = "Airflow is not accessible"
-            current_app.logger.error(err)
-            execution.update_state(EXEC_STATE_ERROR_START)
-            raise AirflowError(
-                error=err,
-                payload=dict(
-                    message=EXECUTION_STATE_MESSAGE_DICT[EXEC_STATE_ERROR_START],
-                    state=EXEC_STATE_ERROR_START,
-                ),
-                log_txt=f"Error while user {self.get_user()} tries to run data checks on instance {idx}. " + err
 
-            )
-        # ask airflow if dag_name exists
-        schema_info = af_client.get_dag_info(schema)
-
-        info = schema_info.json()
-        if info["is_paused"]:
-            err = "The dag exists but it is paused in airflow"
-            current_app.logger.error(err)
-            execution.update_state(EXEC_STATE_ERROR_START)
-            raise AirflowError(
-                error=err,
-                payload=dict(
-                    message=EXECUTION_STATE_MESSAGE_DICT[EXEC_STATE_ERROR_START],
-                    state=EXEC_STATE_ERROR_START,
-                ),
-                log_txt=f"Error while user {self.get_user()} tries to run data checks on instance {idx}. " + err
-
-            )
+        # Call the shared Airflow execution logic
+        _run_airflow_data_check(af_client, execution, schema, user, context_str)
 
-        try:
-            response = af_client.run_dag(
-                execution.id, dag_name=schema, checks_only=True
-            )
-        except AirflowError as err:
-            error = "Airflow responded with an error: {}".format(err)
-            current_app.logger.error(error)
-            execution.update_state(EXEC_STATE_ERROR)
-            raise AirflowError(
-                error=error,
-                payload=dict(
-                    message=EXECUTION_STATE_MESSAGE_DICT[EXEC_STATE_ERROR],
-                    state=EXEC_STATE_ERROR,
-                ),
-                log_txt=f"Error while user {self.get_user()} tries to run data checks on instance {idx}. " + error
-            )
-
-        # if we succeed, we register the dag_run_id in the execution table:
-        af_data = response.json()
-        execution.dag_run_id = af_data["dag_run_id"]
-        execution.update_state(EXEC_STATE_QUEUED)
         current_app.logger.info(
             "User {} creates instance check execution {}".format(
                 self.get_user_id(), execution.id
@@ -258,7 +253,9 @@ class DataCheckCaseEndpoint(BaseMetaResource):
     """
     Endpoint used to execute the instance and solution checks on an execution
     """
+
     ROLES_WITH_ACCESS = [PLANNER_ROLE, ADMIN_ROLE]
+
     def __init__(self):
         super().__init__()
         self.model = ExecutionModel
@@ -283,13 +280,16 @@ class DataCheckCaseEndpoint(BaseMetaResource):
         :rtype: Tuple(dict, integer)
         """
         config = current_app.config
+        user = self.get_user()
+        context_str = f"case {idx}"
 
         case = CaseModel.get_one_object(user=self.get_user(), idx=idx)
         if case is None:
             err = "The case to check does not exist"
             raise ObjectDoesNotExist(
                 error=err,
-                log_txt=f"Error while user {self.get_user()} tries to run data checks on case {idx}. " + err
+                log_txt=f"Error while user {user} tries to run data checks on {context_str}. "
+                + err,
             )
 
         schema = case.schema or "solve_model_dag"
@@ -309,14 +309,18 @@ class DataCheckCaseEndpoint(BaseMetaResource):
             if schema == "pulp":
                 validation_schema = "solve_model_dag"
 
-            data_jsonschema = DeployedDAG.get_one_schema(config, validation_schema, INSTANCE_SCHEMA)
-            validation_errors = json_schema_validate_as_string(data_jsonschema, instance_payload["data"])
+            data_jsonschema = DeployedDAG.get_one_schema(
+                config, validation_schema, INSTANCE_SCHEMA
+            )
+            validation_errors = json_schema_validate_as_string(
+                data_jsonschema, instance_payload["data"]
+            )
 
             if validation_errors:
                 raise InvalidData(
                     payload=dict(jsonschema_errors=validation_errors),
-                    log_txt=f"Error while user {self.get_user()} tries to run data checks on case {idx}.  "
-                            f"Instance data does not match the jsonschema.",
+                    log_txt=f"Error while user {user} tries to run data checks on {context_str}.  "
+                    f"Instance data does not match the jsonschema.",
                 )
 
             instance, _ = self.post_list(data=instance_payload)
@@ -334,14 +338,18 @@ class DataCheckCaseEndpoint(BaseMetaResource):
 
             payload["data"] = case.solution
 
-            data_jsonschema = DeployedDAG.get_one_schema(config, validation_schema, SOLUTION_SCHEMA)
-            validation_errors = json_schema_validate_as_string(data_jsonschema, payload["data"])
+            data_jsonschema = DeployedDAG.get_one_schema(
+                config, validation_schema, SOLUTION_SCHEMA
+            )
+            validation_errors = json_schema_validate_as_string(
+                data_jsonschema, payload["data"]
+            )
 
             if validation_errors:
                 raise InvalidData(
                     payload=dict(jsonschema_errors=validation_errors),
-                    log_txt=f"Error while user {self.get_user()} tries to run data checks on case {idx}.  "
-                            f"Solution data does not match the jsonschema.",
+                    log_txt=f"Error while user {user} tries to run data checks on {context_str}.  "
+                    f"Solution data does not match the jsonschema.",
                 )
 
         self.data_model = ExecutionModel
@@ -354,59 +362,14 @@ class DataCheckCaseEndpoint(BaseMetaResource):
             execution.update_state(EXEC_STATE_NOT_RUN)
             return execution, 201
 
-        # We now try to launch the task in airflow
+        # Initialize Airflow client
         af_client = Airflow.from_config(config)
-        if not af_client.is_alive():
-            err = "Airflow is not accessible"
-            current_app.logger.error(err)
-            execution.update_state(EXEC_STATE_ERROR_START)
-            raise AirflowError(
-                error=err,
-                payload=dict(
-                    message=EXECUTION_STATE_MESSAGE_DICT[EXEC_STATE_ERROR_START],
-                    state=EXEC_STATE_ERROR_START,
-                ),
-                log_txt=f"Error while user {self.get_user()} tries to run data checks on case {idx}. " + err
-            )
-        # ask airflow if dag_name exists
-        schema_info = af_client.get_dag_info(schema)
-
-        info = schema_info.json()
-        if info["is_paused"]:
-            err = "The dag exists but it is paused in airflow"
-            current_app.logger.error(err)
-            execution.update_state(EXEC_STATE_ERROR_START)
-            raise AirflowError(
-                error=err,
-                payload=dict(
-                    message=EXECUTION_STATE_MESSAGE_DICT[EXEC_STATE_ERROR_START],
-                    state=EXEC_STATE_ERROR_START,
-                ),
-                log_txt=f"Error while user {self.get_user()} tries to run data checks on case {idx}. " + err
-            )
-
-        try:
-            response = af_client.run_dag(
-                execution.id, dag_name=schema, checks_only=True, case_id=idx
-            )
 
-        except AirflowError as err:
-            error = "Airflow responded with an error: {}".format(err)
-            current_app.logger.error(error)
-            execution.update_state(EXEC_STATE_ERROR)
-            raise AirflowError(
-                error=error,
-                payload=dict(
-                    message=EXECUTION_STATE_MESSAGE_DICT[EXEC_STATE_ERROR],
-                    state=EXEC_STATE_ERROR,
-                ),
-                log_txt=f"Error while user {self.get_user()} tries to run data checks on case {idx}. " + error
-            )
+        # Call the shared Airflow execution logic, passing case_id
+        _run_airflow_data_check(
+            af_client, execution, schema, user, context_str, case_id=idx
+        )
 
-        # if we succeed, we register the dag_run_id in the execution table:
-        af_data = response.json()
-        execution.dag_run_id = af_data["dag_run_id"]
-        execution.update_state(EXEC_STATE_QUEUED)
         current_app.logger.info(
             "User {} creates case check execution {}".format(
                 self.get_user_id(), execution.id

cornflow/endpoints/example_data.py

@@ -1,6 +1,7 @@
 """
 Endpoints to get the example data from a DAG
 """
+
 import json
 
 from cornflow_client.airflow.api import Airflow
@@ -11,7 +12,12 @@ from cornflow.endpoints.meta_resource import BaseMetaResource
 from cornflow.models import PermissionsDAG
 from cornflow.schemas.example_data import ExampleListData, ExampleDetailData
 from cornflow.shared.authentication import Auth, authenticate
-from cornflow.shared.const import VIEWER_ROLE, PLANNER_ROLE, ADMIN_ROLE
+from cornflow.shared.const import (
+    AIRFLOW_NOT_REACHABLE_MSG,
+    VIEWER_ROLE,
+    PLANNER_ROLE,
+    ADMIN_ROLE,
+)
 from cornflow.shared.exceptions import AirflowError, NoPermission, ObjectDoesNotExist
 
 
@@ -44,7 +50,7 @@ class ExampleDataListEndpoint(BaseMetaResource):
                 current_app.logger.error(
                     "Airflow not accessible when getting data {}".format(dag_name)
                 )
-                raise AirflowError(error="Airflow is not accessible")
+                raise AirflowError(error=f"{AIRFLOW_NOT_REACHABLE_MSG}")
 
             # try airflow and see if dag_name exists
             af_client.get_dag_info(dag_name)
@@ -87,7 +93,7 @@ class ExampleDataDetailEndpoint(BaseMetaResource):
                 current_app.logger.error(
                     "Airflow not accessible when getting data {}".format(dag_name)
                 )
-                raise AirflowError(error="Airflow is not accessible")
+                raise AirflowError(error=f"{AIRFLOW_NOT_REACHABLE_MSG}")
 
             # try airflow and see if dag_name exists
             af_client.get_dag_info(dag_name)