cornflow 2.0.0a8__tar.gz → 2.0.0a10__tar.gz

{cornflow-2.0.0a8/cornflow.egg-info → cornflow-2.0.0a10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: cornflow
-Version: 2.0.0a8
+Version: 2.0.0a10
 Summary: Cornflow is an open source multi-solver optimization server with a REST API built using flask.
 Home-page: https://github.com/baobabsoluciones/cornflow
 Author: baobab soluciones

{cornflow-2.0.0a8 → cornflow-2.0.0a10}/cornflow/cli/service.py

@@ -16,7 +16,13 @@ from cornflow.commands import (
     update_schemas_command,
     update_dag_registry_command,
 )
-from cornflow.shared.const import AUTH_DB, ADMIN_ROLE, DATABRICKS_BACKEND, SERVICE_ROLE, AIRFLOW_BACKEND
+from cornflow.shared.const import (
+    AUTH_DB,
+    ADMIN_ROLE,
+    DATABRICKS_BACKEND,
+    SERVICE_ROLE,
+    AIRFLOW_BACKEND,
+)
 from cornflow.shared import db
 from cryptography.fernet import Fernet
 from flask_migrate import Migrate, upgrade
@@ -106,7 +112,7 @@ def init_cornflow_service():
 
     # Check LDAP parameters for active directory and show message
     if os.getenv("AUTH_TYPE") == 2:
-        print(
+        click.echo(
             "WARNING: Cornflow will be deployed with LDAP Authorization. Please review your ldap auth configuration."
         )
 
@@ -129,10 +135,10 @@ def init_cornflow_service():
                 f"cat > /etc/logrotate.d/cornflow <<EOF\n {conf} \nEOF", shell=True
             )
             out_logrotate = logrotate.stdout
-            print(out_logrotate)
+            click.echo(out_logrotate)
 
         except error:
-            print(error)
+            click.echo(error)
 
     external_application = int(os.getenv("EXTERNAL_APP", 0))
     if external_application == 0:

{cornflow-2.0.0a8 → cornflow-2.0.0a10}/cornflow/commands/permissions.py

@@ -14,13 +14,22 @@ 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
+        extra_permissions = EXTRA_PERMISSION_ASSIGNATION
         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
+        try:
+            extra_permissions = (
+                EXTRA_PERMISSION_ASSIGNATION
+                + external_module.shared.const.EXTRA_PERMISSION_ASSIGNATION
+            )
+        except AttributeError:
+            extra_permissions = EXTRA_PERMISSION_ASSIGNATION
     else:
         resources_to_register = []
         exit()
@@ -52,7 +61,7 @@ def register_base_permissions_command(external_app: str = None, verbose: bool =
                 "api_view_id": views_in_db[endpoint],
             }
         )
-        for role, action, endpoint in EXTRA_PERMISSION_ASSIGNATION
+        for role, action, endpoint in extra_permissions
     ]
 
     permissions_in_app_keys = [
@@ -124,7 +133,7 @@ def register_base_permissions_command(external_app: str = None, verbose: bool =
 def register_dag_permissions_command(
     open_deployment: int = None, verbose: bool = False
 ):
-    click.echo(f"Checkpoint 1")
+
     from flask import current_app
     from sqlalchemy.exc import DBAPIError, IntegrityError
 
@@ -138,7 +147,7 @@ def register_dag_permissions_command(
         (permission.dag_id, permission.user_id)
         for permission in PermissionsDAG.get_all_objects()
     ]
-    click.echo(f"Checkpoint 2")
+
     try:
         db.session.commit()
     except DBAPIError as e:
@@ -149,14 +158,13 @@ def register_dag_permissions_command(
     all_dags = DeployedOrch.get_all_objects().all()
 
     if open_deployment == 1:
-        click.echo(f"Checkpoint 3")
+
         permissions = [
             PermissionsDAG({"dag_id": dag.id, "user_id": user.id})
             for user in all_users
             for dag in all_dags
             if (dag.id, user.id) not in existing_permissions
         ]
-        click.echo(f"Checkpoint 4")
 
     else:
         permissions = [
@@ -165,10 +173,10 @@ def register_dag_permissions_command(
             for dag in all_dags
             if (dag.id, user.id) not in existing_permissions and user.is_service_user()
         ]
-        click.echo(f"Checkpoint 5")
+
     if len(permissions) > 1:
         db.session.bulk_save_objects(permissions)
-    click.echo(f"Checkpoint 6")
+
     try:
         db.session.commit()
     except IntegrityError as e:
@@ -177,7 +185,7 @@ def register_dag_permissions_command(
     except DBAPIError as e:
         db.session.rollback()
         current_app.logger.error(f"Unknown error on dag permissions register: {e}")
-    click.echo(f"Checkpoint 7")
+
     if "postgres" in str(db.session.get_bind()):
         db.engine.execute(
             "SELECT setval(pg_get_serial_sequence('permission_dag', 'id'), MAX(id)) FROM permission_dag;"
@@ -190,11 +198,12 @@ def register_dag_permissions_command(
             current_app.logger.error(
                 f"Unknown error on dag permissions sequence updating: {e}"
             )
-    click.echo(f"Checkpoint 7")
+
     if verbose:
+        click.echo(f"DAG permissions registered")
         if len(permissions) > 1:
             current_app.logger.info(f"DAG permissions registered: {permissions}")
         else:
             current_app.logger.info("No new DAG permissions")
-    click.echo(f"Checkpoint 8")
+    
     pass

{cornflow-2.0.0a8 → cornflow-2.0.0a10}/cornflow/config.py

@@ -41,7 +41,7 @@ class DefaultConfig(object):
     DATABRICKS_EP_CLUSTERS = os.getenv("DATABRICKS_EP_CLUSTERS")
     DATABRICKS_CLIENT_ID = os.getenv("DATABRICKS_CLIENT_ID")
 
-    # If service user is allow to log with username and password
+    # If service user is allowed to log with username and password
     SERVICE_USER_ALLOW_PASSWORD_LOGIN = int(
         os.getenv("SERVICE_USER_ALLOW_PASSWORD_LOGIN", 1)
     )

{cornflow-2.0.0a8 → cornflow-2.0.0a10}/cornflow/endpoints/execution_databricks.py

@@ -5,29 +5,17 @@ These endpoints hve different access url, but manage the same data entities
 """
 
 # Import from libraries
-import datetime
-import logging
-import time
-from databricks.sdk import WorkspaceClient
-import databricks.sdk.service.jobs as j
-from cornflow.constants import *
-from cornflow.shared.const import (
-    AIRFLOW_BACKEND,
-    DATABRICKS_BACKEND,
-    STATUS_HEALTHY,
-    STATUS_UNHEALTHY,
-)
-# TODO AGA: Modificar import para sacarlo de cornflow_client
-from cornflow.shared.databricks import Databricks
-from cornflow_client.constants import INSTANCE_SCHEMA, CONFIG_SCHEMA, SOLUTION_SCHEMA
 from cornflow_client.airflow.api import Airflow
-# TODO AGA: Porqué el import no funcina correctamente 
+from cornflow_client.constants import INSTANCE_SCHEMA, CONFIG_SCHEMA, SOLUTION_SCHEMA
+
+# TODO AGA: Porqué el import no funcina correctamente
 from flask import request, current_app
 from flask_apispec import marshal_with, use_kwargs, doc
 
 # Import from internal modules
 from cornflow.endpoints.meta_resource import BaseMetaResource
 from cornflow.models import InstanceModel, DeployedOrch, ExecutionModel
+from cornflow.orchestrator_constants import config_orchestrator
 from cornflow.schemas.execution import (
     ExecutionDetailsEndpointResponse,
     ExecutionDetailsEndpointWithIndicatorsResponse,
@@ -39,10 +27,14 @@ from cornflow.schemas.execution import (
     ExecutionEditRequest,
     QueryFiltersExecution,
     ReLaunchExecutionRequest,
-    ExecutionDetailsWithIndicatorsAndLogResponse
+    ExecutionDetailsWithIndicatorsAndLogResponse,
 )
 from cornflow.shared.authentication import Auth, authenticate
 from cornflow.shared.compress import compressed
+from cornflow.shared.const import (
+    AIRFLOW_BACKEND,
+    DATABRICKS_BACKEND,
+)
 from cornflow.shared.const import (
     EXEC_STATE_RUNNING,
     EXEC_STATE_ERROR,
@@ -52,16 +44,23 @@ from cornflow.shared.const import (
     EXECUTION_STATE_MESSAGE_DICT,
     AIRFLOW_TO_STATE_MAP,
     DATABRICKS_TO_STATE_MAP,
-    DATABRICKS_FINISH_TO_STATE_MAP,
     EXEC_STATE_STOPPED,
     EXEC_STATE_QUEUED,
 )
-from cornflow.shared.exceptions import AirflowError, DatabricksError, ObjectDoesNotExist, InvalidData
+
+# TODO AGA: Modificar import para sacarlo de cornflow_client
+from cornflow.shared.databricks import Databricks
+from cornflow.shared.exceptions import (
+    AirflowError,
+    DatabricksError,
+    ObjectDoesNotExist,
+    InvalidData,
+    EndpointNotImplemented,
+)
 from cornflow.shared.validators import (
     json_schema_validate_as_string,
     json_schema_extend_and_validate_as_string,
 )
-from cornflow.orchestrator_constants import config_orchestrator
 
 
 class ExecutionEndpoint(BaseMetaResource):
@@ -159,13 +158,13 @@ class ExecutionEndpoint(BaseMetaResource):
         # TODO: should the schema field be cross validated with the instance schema field?
 
         ORQ_TYPE = current_app.config["CORNFLOW_BACKEND"]
-        if ORQ_TYPE==AIRFLOW_BACKEND:
-            orq_const= config_orchestrator["airflow"]
-            ORQ_ERROR=AirflowError
-        elif ORQ_TYPE==DATABRICKS_BACKEND:
-            orq_const= config_orchestrator["databricks"]
+        if ORQ_TYPE == AIRFLOW_BACKEND:
+            orq_const = config_orchestrator["airflow"]
+            ORQ_ERROR = AirflowError
+        elif ORQ_TYPE == DATABRICKS_BACKEND:
+            orq_const = config_orchestrator["databricks"]
             # TODO AGA: Revisar si esto funcionaría correctamente
-            ORQ_ERROR=DatabricksError
+            ORQ_ERROR = DatabricksError
 
         if "schema" not in kwargs:
             kwargs["schema"] = orq_const["def_schema"]
@@ -189,13 +188,15 @@ class ExecutionEndpoint(BaseMetaResource):
         # We try to create an orch client
         # Note schema is a string with the name of the job/dag
         schema = execution.schema
-        # If we are dealing with DataBricks, the schema will 
+        # If we are dealing with DataBricks, the schema will
         #   be the job id
-        orch_client, schema_info, execution= get_orch_client(schema,ORQ_TYPE,execution)
+        orch_client, schema_info, execution = get_orch_client(
+            schema, ORQ_TYPE, execution
+        )
         # endregion
 
         # region VALIDACIONES
-        # We check if the job/dag exists 
+        # We check if the job/dag exists
         orch_client.get_orch_info(schema)
         # Validate config before running the dag
         config_schema = DeployedOrch.get_one_schema(config, schema, CONFIG_SCHEMA)
@@ -249,8 +250,8 @@ class ExecutionEndpoint(BaseMetaResource):
                 execution.update_log_txt(f"{solution_errors}")
                 raise InvalidData(payload=dict(jsonschema_errors=solution_errors))
         # endregion
-        
-        if ORQ_TYPE==AIRFLOW_BACKEND:
+
+        if ORQ_TYPE == AIRFLOW_BACKEND:
             info = schema_info.json()
             if info["is_paused"]:
                 err = "The dag exists but it is paused in airflow"
@@ -268,7 +269,7 @@ class ExecutionEndpoint(BaseMetaResource):
         # TODO AGA: revisar si hay que hacer alguna verificación a los JOBS
 
         try:
-            # TODO AGA: Hay que genestionar la posible eliminación de execution.id como 
+            # TODO AGA: Hay que genestionar la posible eliminación de execution.id como
             #   parámetro, ya que no se puede seleccionar el id en databricks
             #   revisar las consecuencias que puede tener
             response = orch_client.run_workflow(execution.id, orch_name=schema)
@@ -319,13 +320,13 @@ class ExecutionRelaunchEndpoint(BaseMetaResource):
         :rtype: Tuple(dict, integer)
         """
         ORQ_TYPE = current_app.config["CORNFLOW_BACKEND"]
-        if ORQ_TYPE==AIRFLOW_BACKEND:
-            orq_const= config_orchestrator["airflow"]
-            ORQ_ERROR=AirflowError
-        elif ORQ_TYPE==DATABRICKS_BACKEND:
-            orq_const= config_orchestrator["databricks"]
+        if ORQ_TYPE == AIRFLOW_BACKEND:
+            orq_const = config_orchestrator["airflow"]
+            ORQ_ERROR = AirflowError
+        elif ORQ_TYPE == DATABRICKS_BACKEND:
+            orq_const = config_orchestrator["databricks"]
             # TODO AGA: Revisar si esto funcionaría correctamente
-            ORQ_ERROR=DatabricksError
+            ORQ_ERROR = DatabricksError
 
         config = current_app.config
         if "schema" not in kwargs:
@@ -370,10 +371,12 @@ class ExecutionRelaunchEndpoint(BaseMetaResource):
                 log_txt=f"Error while user {self.get_user()} tries to relaunch execution {idx}. "
                 f"Configuration data does not match the jsonschema.",
             )
-        orch_client, schema_info, execution = get_orch_client(schema,ORQ_TYPE,execution)
-        
+        orch_client, schema_info, execution = get_orch_client(
+            kwargs["schema"], ORQ_TYPE, execution
+        )
+
         if not orch_client.is_alive():
-            err = orq_const["name"]+" is not accessible"
+            err = orq_const["name"] + " is not accessible"
             current_app.logger.error(err)
             execution.update_state(EXEC_STATE_ERROR_START)
             raise ORQ_ERROR(
@@ -385,13 +388,13 @@ class ExecutionRelaunchEndpoint(BaseMetaResource):
                 log_txt=f"Error while user {self.get_user()} tries to relaunch execution {idx}. "
                 + err,
             )
-        
+
         # ask airflow if dag_name exists
         schema = execution.schema
         schema_info = orch_client.get_orch_info(schema)
 
         info = schema_info.json()
-        if ORQ_TYPE==AIRFLOW_BACKEND:
+        if ORQ_TYPE == AIRFLOW_BACKEND:
             if info["is_paused"]:
                 err = "The dag exists but it is paused in airflow"
                 current_app.logger.error(err)
@@ -409,7 +412,7 @@ class ExecutionRelaunchEndpoint(BaseMetaResource):
         try:
             response = orch_client.run_workflow(execution.id, orch_name=schema)
         except ORQ_ERROR as err:
-            error = orq_const["name"]+" responded with an error: {}".format(err)
+            error = orq_const["name"] + " responded with an error: {}".format(err)
             current_app.logger.error(error)
             execution.update_state(EXEC_STATE_ERROR)
             raise ORQ_ERROR(
@@ -437,6 +440,7 @@ class ExecutionDetailsEndpointBase(BaseMetaResource):
     """
     Endpoint used to get the information of a certain execution. But not the data!
     """
+
     # TODO AGA DUDA: Se usa? Qué debería devolver?
     def __init__(self):
         super().__init__()
@@ -565,13 +569,13 @@ class ExecutionStatusEndpoint(BaseMetaResource):
         :rtype: Tuple(dict, integer)
         """
         ORQ_TYPE = current_app.config["CORNFLOW_BACKEND"]
-        if ORQ_TYPE==AIRFLOW_BACKEND:
-            orq_const= config_orchestrator["airflow"]
-            ORQ_ERROR=AirflowError
-        elif ORQ_TYPE==DATABRICKS_BACKEND:
-            orq_const= config_orchestrator["databricks"]
+        if ORQ_TYPE == AIRFLOW_BACKEND:
+            orq_const = config_orchestrator["airflow"]
+            ORQ_ERROR = AirflowError
+        elif ORQ_TYPE == DATABRICKS_BACKEND:
+            orq_const = config_orchestrator["databricks"]
             # TODO AGA: Revisar si esto funcionaría correctamente
-            ORQ_ERROR=DatabricksError
+            ORQ_ERROR = DatabricksError
         execution = self.data_model.get_one_object(user=self.get_user(), idx=idx)
         if execution is None:
             raise ObjectDoesNotExist(
@@ -594,6 +598,7 @@ class ExecutionStatusEndpoint(BaseMetaResource):
             raise ORQ_ERROR(
                 error=error, payload=dict(message=message, state=state), log_txt=log_txt
             )
+
         print("The execution is ", execution)
         print("The execution user is ", self.get_user())
         print("The execution id is ", idx)
@@ -611,10 +616,12 @@ class ExecutionStatusEndpoint(BaseMetaResource):
             )
         schema = execution.schema
         # TODO AGA: Revisar si merece la pena hacer una funcion que solo
-        orch_client, schema_info, execution= get_orch_client(schema ,ORQ_TYPE,execution)
+        orch_client, schema_info, execution = get_orch_client(
+            schema, ORQ_TYPE, execution
+        )
 
         if not orch_client.is_alive():
-            err = orq_const["name"] +" is not accessible"
+            err = orq_const["name"] + " is not accessible"
             _raise_af_error(
                 execution,
                 err,
@@ -624,11 +631,9 @@ class ExecutionStatusEndpoint(BaseMetaResource):
 
         try:
             # TODO: get the dag_name from somewhere!
-            state = orch_client.get_run_status(
-                dag_run_id
-            )
+            state = orch_client.get_run_status(schema, dag_run_id)
         except ORQ_ERROR as err:
-            error = orq_const["name"] +f" responded with an error: {err}"
+            error = orq_const["name"] + f" responded with an error: {err}"
             _raise_af_error(
                 execution,
                 error,
@@ -723,38 +728,9 @@ class ExecutionLogEndpoint(ExecutionDetailsEndpointBase):
         current_app.logger.info(f"User {self.get_user()} gets log of execution {idx}")
         return self.get_detail(user=self.get_user(), idx=idx)
 
+
 # region aux_functions
-def submit_one_job(cid):
-    # trigger one-time-run job and get waiter object
-    waiter = w.jobs.submit(run_name=f'cornflow-job-{time.time()}', tasks=[
-        j.SubmitTask(
-            task_key='nippon_production_scheduling',
-            existing_cluster_id=cid,
-            libraries=[],
-            spark_python_task=j.SparkPythonTask(
-                python_file='/Workspace/Repos/nippon/nippon_production_scheduling/main.py',
-            ),
-            timeout_seconds=0,
-        )
-    ])
-    logging.info(f'starting to poll: {waiter.run_id}')
-    # callback, that receives a polled entity between state updates
-    # If you want to perform polling in a separate thread, process, or service,
-    # you can use w.jobs.wait_get_run_job_terminated_or_skipped(
-    #   run_id=waiter.run_id,
-    #   timeout=datetime.timedelta(minutes=15),
-    #   callback=print_status) to achieve the same results.
-    #
-    # Waiter interface allows for `w.jobs.submit(..).result()` simplicity in
-    # the scenarios, where you need to block the calling thread for the job to finish.
-    run = waiter.result(timeout=datetime.timedelta(minutes=15),
-                        callback=print_status)
-    logging.info(f'job finished: {run.run_page_url}')
-    return waiter.run_id
-
-def print_status(run: j.Run):
-    statuses = [f'{t.task_key}: {t.state.life_cycle_state}' for t in run.tasks]
-    logging.info(f'workflow intermediate status: {", ".join(statuses)}')
+
 
 def get_orch_client(schema, orq_type, execution):
     """
@@ -763,10 +739,10 @@ def get_orch_client(schema, orq_type, execution):
     if orq_type == AIRFLOW_BACKEND:
         return get_airflow(schema, execution=execution)
     elif orq_type == DATABRICKS_BACKEND:
-        return get_databricks(schema,execution=execution)
+        return get_databricks(schema, execution=execution)
     else:
         raise EndpointNotImplemented()
-            
+
 
 def get_airflow(schema, execution):
     """
@@ -784,12 +760,12 @@ def get_airflow(schema, execution):
                 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 create an execution "
+            log_txt=f"Error while user {execution.user_id} tries to create an execution "
             + err,
         )
-    # TODO AGA: revisar si tiene sentido que se devuelva execution o si 
+    # TODO AGA: revisar si tiene sentido que se devuelva execution o si
     #  es un puntero
-    return af_client,schema_info, execution
+    return af_client, schema_info, execution
 
 
 def get_databricks(schema, execution):
@@ -808,21 +784,23 @@ def get_databricks(schema, execution):
                 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 create an execution "
+            log_txt=f"Error while user {execution.user_id} tries to create an execution "
             + err,
         )
     return db_client, schema_info, execution
     # endregion
 
-def map_run_state(state,ORQ_TYPE):
+
+def map_run_state(state, ORQ_TYPE):
     """
     Maps the state of the execution in the orchestrator to the state of the execution in cornflow
     """
-    if ORQ_TYPE==AIRFLOW_BACKEND:
+    if ORQ_TYPE == AIRFLOW_BACKEND:
+        state = state.json()["state"]
         return AIRFLOW_TO_STATE_MAP.get(state, EXEC_STATE_UNKNOWN)
-    elif ORQ_TYPE==DATABRICKS_BACKEND:
+    elif ORQ_TYPE == DATABRICKS_BACKEND:
         print("The state is ", state)
-        preliminar_state = DATABRICKS_TO_STATE_MAP.get(state,EXEC_STATE_UNKNOWN)
+        preliminar_state = DATABRICKS_TO_STATE_MAP.get(state, EXEC_STATE_UNKNOWN)
         # print("The preliminar state is ", preliminar_state)
         # if preliminar_state =="TERMINATED":
         #     # TODO AGA DUDA: Revisar si es correcto el error predeterminado

{cornflow-2.0.0a8 → cornflow-2.0.0a10}/cornflow/models/execution.py

@@ -117,7 +117,6 @@ class ExecutionModel(BaseDataModel):
         :param str message: Message for the error
         :return: nothing
         """
-        print("Updating state to ", code)
         self.state = code
         if message is None:
             self.state_message = EXECUTION_STATE_MESSAGE_DICT[code]

{cornflow-2.0.0a8 → cornflow-2.0.0a10}/cornflow/shared/const.py

@@ -1,5 +1,5 @@
 """
-In this files we import the values for different constants on cornflow server
+In this file we import the values for different constants on cornflow server
 """
 
 # CORNFLOW BACKEND
@@ -62,7 +62,7 @@ DATABRICKS_TO_STATE_MAP = dict(
 DATABRICKS_FINISH_TO_STATE_MAP = dict(
     SUCCESS=EXEC_STATE_CORRECT,
     USER_CANCELED=EXEC_STATE_STOPPED,
-    )
+)
 
 DATABRICKS_TERMINATE_STATE = "TERMINATED"
 # These codes and names are inherited from flask app builder in order to have the same names and values