cornflow 2.0.0a9__py3-none-any.whl → 2.0.0a10__py3-none-any.whl

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/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/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/endpoints/execution_databricks.py

@@ -5,28 +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.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,
@@ -38,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,
@@ -51,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):
@@ -158,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"]
@@ -188,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)
@@ -248,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"
@@ -267,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)
@@ -318,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:
@@ -369,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(
@@ -384,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)
@@ -408,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(
@@ -436,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__()
@@ -564,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(
@@ -593,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)
@@ -610,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,
@@ -623,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,
@@ -722,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):
     """
@@ -762,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):
     """
@@ -783,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):
@@ -807,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/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/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

cornflow/shared/databricks.py

@@ -1,24 +1,31 @@
 """
 Python class to implement the Databricks client wrapper
 """
+
 import requests
 from databricks.sdk import WorkspaceClient
 from flask import current_app
 from cornflow.orchestrator_constants import config_orchestrator
+
 # TODO AGA: CODIGO REPETIDO
 # TODO AGA: revisar si el import está bien
 from cornflow_client.constants import DatabricksError
-from cornflow.shared.const import DATABRICKS_TO_STATE_MAP,DATABRICKS_TERMINATE_STATE, DATABRICKS_FINISH_TO_STATE_MAP
+from cornflow.shared.const import (
+    DATABRICKS_TO_STATE_MAP,
+    DATABRICKS_TERMINATE_STATE,
+    DATABRICKS_FINISH_TO_STATE_MAP,
+)
+
 
 class Databricks:
     def __init__(self, url, auth_secret, token_endpoint, ep_clusters, client_id):
         self.url = url
-        self.constants=config_orchestrator["databricks"]
-        self.auth_secret=auth_secret
-        self.token_endpoint=token_endpoint
-        self.ep_clusters=ep_clusters
+        self.constants = config_orchestrator["databricks"]
+        self.auth_secret = auth_secret
+        self.token_endpoint = token_endpoint
+        self.ep_clusters = ep_clusters
         self.client_id = client_id
-        
+
     @classmethod
     def from_config(cls, config):
         data = dict(
@@ -32,61 +39,68 @@ class Databricks:
 
     def get_token(self):
         import requests
-        url = f'{self.url}{self.token_endpoint}'
-        data = {
-            "grant_type": "client_credentials",
-            "scope": "all-apis"
-        }
-        auth = (self.client_id,self.auth_secret)
-        oauth_response = requests.post(url,data=data,auth=auth)
+
+        url = f"{self.url}{self.token_endpoint}"
+        data = {"grant_type": "client_credentials", "scope": "all-apis"}
+        auth = (self.client_id, self.auth_secret)
+        oauth_response = requests.post(url, data=data, auth=auth)
         oauth_response.json()
         oauth_token = oauth_response.json()["access_token"]
         return oauth_token
-    
+
     def is_alive(self):
         try:
-          # TODO: this url is project specific. Either it has to be a config option or some other way has to be found
-            path="/Workspace/Repos/nippon/nippon_production_scheduling/requirements.txt"
+            # TODO: this url is project specific. Either it has to be a config option or some other way has to be found
+            path = (
+                "/Workspace/Repos/nippon/nippon_production_scheduling/requirements.txt"
+            )
             url = f"{self.url}/api/2.0/workspace/get-status?path={path}"
-            response = self.request_headers_auth(method="GET", url=url) 
-            if "error_code"  in response.json().keys():
+            response = self.request_headers_auth(method="GET", url=url)
+            if "error_code" in response.json().keys():
                 return False
             return True
-    
+
         except Exception as err:
-          current_app.logger.error(f"Error: {err}")
-          return False
-    
+            current_app.logger.error(f"Error: {err}")
+            return False
+
     def get_orch_info(self, orch_name, method="GET"):
         """
         Get information about a job in Databricks
         https://docs.databricks.com/api/workspace/jobs/get
         """
         url = f"{self.url}/api/2.1/jobs/get/?job_id={orch_name}"
-        schema_info = self.request_headers_auth(method=method, url=url) 
+        schema_info = self.request_headers_auth(method=method, url=url)
         if "error_code" in schema_info.json().keys():
             raise DatabricksError("JOB not available")
         return schema_info
-    # TODO AGA: incluir un id de job por defecto o hacer obligatorio el uso el parámetro. 
+
+    # TODO AGA: incluir un id de job por defecto o hacer obligatorio el uso el parámetro.
     #   Revisar los efectos secundarios de eliminar execution_id y usar el predeterminado
     def run_workflow(
-            self, execution_id, orch_name=config_orchestrator["databricks"]["def_schema"], checks_only=False, case_id=None
-        ):
+        self,
+        execution_id,
+        orch_name=config_orchestrator["databricks"]["def_schema"],
+        checks_only=False,
+        case_id=None,
+    ):
         """
         Run a job in Databricks
         """
         # TODO AGA: revisar si la url esta bien/si acepta asi los parámetros
         url = f"{self.url}/api/2.1/jobs/run-now/"
-        # TODO AGA: revisar si deben ser notebook parameters o job parameters. 
+        # TODO AGA: revisar si deben ser notebook parameters o job parameters.
         #   Entender cómo se usa checks_only
-        payload = dict(job_id=orch_name, notebook_parameters=dict(checks_only=checks_only))
+        payload = dict(
+            job_id=orch_name, notebook_parameters=dict(checks_only=checks_only)
+        )
         return self.request_headers_auth(method="POST", url=url, json=payload)
-    
-    def get_run_status(self, run_id):
+
+    def get_run_status(self, schema, run_id):
         """
         Get the status of a run in Databricks
         """
-        print( "asking for run id ", run_id)
+        print("asking for run id ", run_id)
         url = f"{self.url}/api/2.1/jobs/runs/get"
         payload = dict(run_id=run_id)
         info = self.request_headers_auth(method="GET", url=url, json=payload)
@@ -94,18 +108,22 @@ class Databricks:
         print("info is ", info)
         state = info["status"]["state"]
         if state == DATABRICKS_TERMINATE_STATE:
-            if info["status"]["termination_details"]["code"] in DATABRICKS_FINISH_TO_STATE_MAP.keys():
+            if (
+                info["status"]["termination_details"]["code"]
+                in DATABRICKS_FINISH_TO_STATE_MAP.keys()
+            ):
                 return info["status"]["termination_details"]["code"]
-            else: 
+            else:
                 return "OTHER_FINISH_ERROR"
         return state
+
     def request_headers_auth(self, status=200, **kwargs):
-        token =self.get_token()
-        def_headers = {"Authorization": "Bearer "+ str(token)}
+        token = self.get_token()
+        def_headers = {"Authorization": "Bearer " + str(token)}
         headers = kwargs.get("headers", def_headers)
         response = requests.request(headers=headers, **kwargs)
         if status is None:
             return response
         if response.status_code != status:
             raise DatabricksError(error=response.text, status_code=response.status_code)
-        return response
+        return response

cornflow/tests/integration/test_cornflowclient.py

@@ -1,6 +1,7 @@
 """
 
 """
+
 # Full imports
 import json
 import pulp
@@ -279,7 +280,7 @@ class TestCornflowClientOpen(TestCornflowClientBasic):
     def test_server_alive(self):
         data = self.client.is_alive()
         cf_status = data["cornflow_status"]
-        af_status = data["airflow_status"]
+        af_status = data["backend_status"]
         self.assertEqual(str, type(cf_status))
         self.assertEqual(str, type(af_status))
         self.assertEqual(cf_status, STATUS_HEALTHY)

cornflow/tests/unit/test_executions.py

@@ -62,13 +62,13 @@ class TestExecutionsListEndpoint(BaseTestCases.ListFilters):
     def test_new_execution(self):
         self.create_new_row(self.url, self.model, payload=self.payload)
 
-    @patch("cornflow.endpoints.execution.Airflow")
+    @patch("cornflow.endpoints.execution_databricks.Airflow")
     def test_new_execution_run(self, af_client_class):
         patch_af_client(af_client_class)
 
         self.create_new_row(EXECUTION_URL, self.model, payload=self.payload)
 
-    @patch("cornflow.endpoints.execution.Airflow")
+    @patch("cornflow.endpoints.execution_databricks.Airflow")
     def test_new_execution_bad_config(self, af_client_class):
         patch_af_client(af_client_class)
         response = self.create_new_row(
@@ -81,7 +81,7 @@ class TestExecutionsListEndpoint(BaseTestCases.ListFilters):
         self.assertIn("error", response)
         self.assertIn("jsonschema_errors", response)
 
-    @patch("cornflow.endpoints.execution.Airflow")
+    @patch("cornflow.endpoints.execution_databricks.Airflow")
     def test_new_execution_partial_config(self, af_client_class):
         patch_af_client(af_client_class)
         self.payload["config"].pop("solver")
@@ -91,7 +91,7 @@ class TestExecutionsListEndpoint(BaseTestCases.ListFilters):
         self.assertIn("solver", response["config"])
         self.assertEqual(response["config"]["solver"], "cbc")
 
-    @patch("cornflow.endpoints.execution.Airflow")
+    @patch("cornflow.endpoints.execution_databricks.Airflow")
     def test_new_execution_with_solution(self, af_client_class):
         patch_af_client(af_client_class)
         self.payload["data"] = self.solution
@@ -102,7 +102,7 @@ class TestExecutionsListEndpoint(BaseTestCases.ListFilters):
             check_payload=False,
         )
 
-    @patch("cornflow.endpoints.execution.Airflow")
+    @patch("cornflow.endpoints.execution_databricks.Airflow")
     def test_new_execution_with_solution_bad(self, af_client_class):
         patch_af_client(af_client_class)
         patch_af_client(af_client_class)
@@ -193,7 +193,7 @@ class TestExecutionRelaunchEndpoint(CustomTestCase):
         self.assertEqual(row["config"], self.payload["config"])
         self.assertIsNone(row["checks"])
 
-    @patch("cornflow.endpoints.execution.Airflow")
+    @patch("cornflow.endpoints.execution_databricks.Airflow")
     def test_relaunch_execution_run(self, af_client_class):
         patch_af_client(af_client_class)
 
@@ -372,7 +372,7 @@ class TestExecutionsDetailEndpoint(
 
         self.assertEqual(row.json["checks"], None)
 
-    @patch("cornflow.endpoints.execution.Airflow")
+    @patch("cornflow.endpoints.execution_databricks.Airflow")
     def test_stop_execution(self, af_client_class):
         patch_af_client(af_client_class)
 
@@ -479,7 +479,7 @@ class TestExecutionsStatusEndpoint(TestExecutionsDetailEndpointMock):
         self.response_items = {"id", "name", "status"}
         self.items_to_check = []
 
-    @patch("cornflow.endpoints.execution.Airflow")
+    @patch("cornflow.endpoints.execution_databricks.Airflow")
     def test_get_one_status(self, af_client_class):
         patch_af_client(af_client_class)
 
@@ -495,7 +495,7 @@ class TestExecutionsStatusEndpoint(TestExecutionsDetailEndpointMock):
         )
         self.assertEqual(data["state"], 1)
 
-    @patch("cornflow.endpoints.execution.Airflow")
+    @patch("cornflow.endpoints.execution_databricks.Airflow")
     def test_put_one_status(self, af_client_class):
         patch_af_client(af_client_class)
 

cornflow/tests/unit/test_health.py

@@ -28,7 +28,7 @@ class TestHealth(CustomTestCase):
         response = self.client.get(HEALTH_URL)
         self.assertEqual(200, response.status_code)
         cf_status = response.json["cornflow_status"]
-        af_status = response.json["airflow_status"]
+        backend_status = response.json["backend_status"]
         self.assertEqual(str, type(cf_status))
-        self.assertEqual(str, type(af_status))
+        self.assertEqual(str, type(backend_status))
         self.assertEqual(cf_status, STATUS_HEALTHY)

{cornflow-2.0.0a9.dist-info → cornflow-2.0.0a10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: cornflow
-Version: 2.0.0a9
+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.0a9.dist-info → cornflow-2.0.0a10.dist-info}/RECORD

@@ -6,7 +6,7 @@ airflow_config/plugins/XCom/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMp
 airflow_config/plugins/XCom/gce_xcom_backend.py,sha256=vCGvF2jbfZt5bOv-pk5Q_kUR6LomFUojIymimSJmj3o,1795
 cornflow/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cornflow/app.py,sha256=X73N64o8OGEqVIRWbC13e_4xb1lxzOH_BV3F3fzAmXE,7312
-cornflow/config.py,sha256=gcKhxmsC3UIbdl42SnrB8vOtrL8_7QCGj5yyYlmSOsI,6124
+cornflow/config.py,sha256=XRO0moHFw3JSIUsww8Id0JCj7JYz_2Bq7SMVfvyTkQ8,6126
 cornflow/gunicorn.py,sha256=uO-Yk7w7nvQSWh12iDxsVvlG-_2BiKIIjm2UiTk4P9E,480
 cornflow/orchestrator_constants.py,sha256=VO6EGcHhovH6nFpp3QGY_orYVJbjFrMavmemL-Gx_Vs,269
 cornflow/cli/__init__.py,sha256=5jBmSMpaE1S9rDaQjS8VHJ6x4FfJG8MhKzMzfw7G4Zc,743
@@ -17,7 +17,7 @@ cornflow/cli/migrations.py,sha256=Stc8H99rG8vgo3yRJcck11zBY_EA4WqyVybglfl8zJE,16
 cornflow/cli/permissions.py,sha256=4KXKysH4g8YYQIZcPuXFS2g0xEErp-e8I_FAqMGaV7U,1006
 cornflow/cli/roles.py,sha256=NFG__qrlyOT0h4L4nwo9FSV4DKjGtMVh3gwiJxwM37w,411
 cornflow/cli/schemas.py,sha256=sxuJOZf12SBZAXDiAYNPB-n9LSxzSwkB3xyhgS_4K9A,6086
-cornflow/cli/service.py,sha256=lCFOQXtBMYOm8sjntuyLhY_TshFOMMnEtysFy6zrgc8,10817
+cornflow/cli/service.py,sha256=h44X7U9QHdwHbACBzBNIkmCEhGYpE8QF9l3p3mKgbMU,10857
 cornflow/cli/users.py,sha256=nPnu8rQNLtwmeXLwYtJ_hjlsa_24XOnQLgBJRBP9bJw,2104
 cornflow/cli/utils.py,sha256=0tF41gTt6LL9XGOizTQg2GXuOXbqLg6gapCr-HWjJ0Q,733
 cornflow/cli/views.py,sha256=Xyx2l-Sm7panxQEfR3qksCIUoqF7woMKsYgZALkxUXM,636
@@ -33,7 +33,7 @@ cornflow/commands/access.py,sha256=NTZJFF9la8TDuMcD_ISQtJTj-wtM2p1dddokQJHtkj0,7
 cornflow/commands/actions.py,sha256=4AwgAmyI6VeaugkISvTlNGrIzMMU_-ZB3MhwDD_CIEA,1544
 cornflow/commands/cleanup.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cornflow/commands/dag.py,sha256=4c74pqtAg9y_qw_mOSQU1WULwdd9PnV0SmVqbY4T-uw,3721
-cornflow/commands/permissions.py,sha256=TtKJcX219RecAytm8PCdVctaVLZSA7ilhBlPnWxYuFQ,6934
+cornflow/commands/permissions.py,sha256=yvMB3kE2rrqRj-eAh53E0j8Zgk84rz8ZT438ylauHnk,7017
 cornflow/commands/roles.py,sha256=Oux-UkswkQ74zqaMEJYIEsZpQZGBcGaSahVzx9feAHU,1516
 cornflow/commands/schemas.py,sha256=QjLXLw5So3f8ZqTg5_uvXxwpo4vE0dMT4_gFMKZHGvQ,1828
 cornflow/commands/users.py,sha256=MEfqMm2ujso0NQgdUm-crOet-G0M43GNqVCx2Ls-2HY,2591
@@ -47,7 +47,7 @@ cornflow/endpoints/dag.py,sha256=x_F9gGc1_9-zoAErPGPFBhwQKqIgGEoCjDVMmRKgZRg,103
 cornflow/endpoints/data_check.py,sha256=FrsBCzwXD_E5qCYkXt0Byc_UPC0APBSxbj8EriPVKxI,16520
 cornflow/endpoints/example_data.py,sha256=e3Y_lZTKVQTkk_pzOnmd-VisuKo7kE-7IqhhLHV-6yw,4374
 cornflow/endpoints/execution.py,sha256=2tX4UJ23e-phAFfgRpjFASr54OWIROuA6a9Uum2Cg48,28010
-cornflow/endpoints/execution_databricks.py,sha256=-QiBp6sRWh-XD2DRe4cDUTiOrUNfqfWGdSQncT5_sik,34912
+cornflow/endpoints/execution_databricks.py,sha256=vsnvOgY5bBD5wDWsoF-z08DuVkKIG2Toz8g1GPTjBBs,33430
 cornflow/endpoints/health.py,sha256=K1l8YK7t5nfVRi0WXMn124l8-ezvYAYqHlCtitLU9AY,2558
 cornflow/endpoints/instance.py,sha256=YuB0TTs32eKFxd2GJc7WeVTVWZdv7tqNAtWLrOPsuXo,11618
 cornflow/endpoints/licenses.py,sha256=82hHWGYvVIiyw9mlwGtMwJMDJ-ShHOi9rvuM6KvfE4U,873
@@ -88,7 +88,7 @@ cornflow/models/base_data_model.py,sha256=mVMHJpEoJeH6Wly_ZIfzLfTPd39nSYpCgmtA_f
 cornflow/models/case.py,sha256=GEs-xeo0bJ5qJETDnIur-2q2IyR3NSj1K0jP3Arz4Xs,9572
 cornflow/models/dag.py,sha256=a3X_WSuLj9nNv-c2qziQb9RSOULprTNbWm4G3k5jGAo,3021
 cornflow/models/dag_permissions.py,sha256=LM2CacGyflwYbG8fbjRaaUry7pQDtvOXjfJpY9jj5EQ,1718
-cornflow/models/execution.py,sha256=63wECEMQKehozS0Fnh6-bxnwe2BFCZ0qLDljnRbVxTI,6140
+cornflow/models/execution.py,sha256=4zAnrkYyIJ4v0ulCgMDGyeyE8zmB8wc0PT8hL_HyagE,6098
 cornflow/models/instance.py,sha256=2E9kBKv1a8soaEAvG8X4qXQ4BVC-IWYD5WQcPmZQw00,3979
 cornflow/models/main_alarms.py,sha256=9S-Ohr2kYFFWB0HomrpSdDIoUr85Eu1rt90Om_Pa8VY,1748
 cornflow/models/meta_models.py,sha256=qeliGdpw0_q0GCeZzansF-09Ay5pueaT-QQPVPZ5aj4,12000
@@ -121,8 +121,8 @@ cornflow/schemas/user_role.py,sha256=e5y6RgdZZtLqD-h2B3sa5WokI5-pT78tWw85IG34I74
 cornflow/schemas/view.py,sha256=ctq9Y1TmjrWdyOqgDYeEx7qbbuNLKfSiNOlFTlXmpaw,429
 cornflow/shared/__init__.py,sha256=1ahcBwWOsSjGI4FEm77JBQjitBdBszOncKcEMjzwGYE,29
 cornflow/shared/compress.py,sha256=pohQaGs1xbH8CN6URIH6BAHA--pFq7Hmjz8oI3c3B5c,1347
-cornflow/shared/const.py,sha256=FQQQz-LekFN81YPh1KTHwyLaikddPYqQJHdzTHhvTfw,4034
-cornflow/shared/databricks.py,sha256=DYhF99DJCxAcTQjB3ih4w9qc_yb_3iG1bpx-eerumvE,4731
+cornflow/shared/const.py,sha256=DmeA1DL-9RXkRXgbGI0kykMuCFSuZ0LGFpwl2rjirOI,4029
+cornflow/shared/databricks.py,sha256=k8MSove2FMEsH5gYGgD2d9i0O4ug4v3lAGUGQOyzEp4,4847
 cornflow/shared/email.py,sha256=QNDDMv86LZObkevSCyUbLQeR2UD3zWScPIr82NDzYHQ,3437
 cornflow/shared/exceptions.py,sha256=bD_Eo1T2baSrgfGMKb6r9GZ1fCJbitykEEMy9wOnk4A,7033
 cornflow/shared/licenses.py,sha256=Lc71Jw2NxVTFWtoXdQ9wJX_o3BDfYg1xVoehDXvnCkQ,1328
@@ -141,7 +141,7 @@ cornflow/tests/custom_liveServer.py,sha256=I_0YNrcKIwVmRov3zCQMWwcCWkMe5V246Hpa4
 cornflow/tests/custom_test_case.py,sha256=X1j-cy9QKhF4W6_7jcJsTm-0Jn6lluq6gj-g126dFpQ,35945
 cornflow/tests/integration/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cornflow/tests/integration/test_commands.py,sha256=mGiMfqIqwvRx08Al6LcHXEKPgEQEJ33EoPIZhGcReX0,697
-cornflow/tests/integration/test_cornflowclient.py,sha256=ioAQmQKWW6mXVJhdF4LECZcGIOa_N0xPkFaGWGtxOO8,20963
+cornflow/tests/integration/test_cornflowclient.py,sha256=mwtlzn4MSjQNRhwV3yzhco3DHCS5BVFimTfmB8lB2Cc,20964
 cornflow/tests/ldap/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cornflow/tests/ldap/test_ldap_authentication.py,sha256=6Gu1WkF7MQmcV_10IJkpo2qEloZZ9zjpV18ANDD0HRw,4286
 cornflow/tests/unit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -155,9 +155,9 @@ cornflow/tests/unit/test_commands.py,sha256=HjeR7vytN_IQhkmKFJw5St4bSEOIXuTeLPuu
 cornflow/tests/unit/test_dags.py,sha256=Q8l0K1FoMI6Yn4tuzxfiZp2azCpIWclbUneQnkOattA,13403
 cornflow/tests/unit/test_data_checks.py,sha256=6s50d1iuRTUcAYn14oEcRS39ZZ6E9ussU4YpkpYhtC4,8612
 cornflow/tests/unit/test_example_data.py,sha256=rCj3wNW4OMHyLfuPdIKxBa7-eRuFsymij-9Nk27_Z3o,4128
-cornflow/tests/unit/test_executions.py,sha256=_hIaiZri7Blyx4DYhBDHh-0peU1HQh66RSPqQJFveE8,17501
+cornflow/tests/unit/test_executions.py,sha256=zSb2NzNpjQKFDX6t5-coqhROHvrevktG2FCbHQkdRKs,17600
 cornflow/tests/unit/test_generate_from_schema.py,sha256=L1EdnASbDJ8SjrX1V4WnUKKwV0sRTwVnNYnxSpyeSeQ,15376
-cornflow/tests/unit/test_health.py,sha256=0E0HXMb63_Z8drbLZdxnJwtTbQyaZS9ZEHut6qsDbh8,1033
+cornflow/tests/unit/test_health.py,sha256=er-CkoYI3Nb9lq0pGn-rIbhfIigxXI5wYNFNn-OsAoc,1043
 cornflow/tests/unit/test_instances.py,sha256=RaD9Tue2HODKThBNhciu6krdIvrauDLxOq4Y6a_z8DU,10573
 cornflow/tests/unit/test_instances_file.py,sha256=zXxSlOM_MMkFvpWNX-iatD40xoIAOGQkinCLf1txb0M,1986
 cornflow/tests/unit/test_licenses.py,sha256=jgnfE4UMFooGn44HK_KspJXIpmLjUpK_WgsBBeTO5eI,1534
@@ -172,8 +172,8 @@ cornflow/tests/unit/test_tables.py,sha256=dY55YgaCkyqwJnqn0LbZHNeXBoL4ZxXWwKkCoT
 cornflow/tests/unit/test_token.py,sha256=OEVPgG8swSMkUbuGJGfGF5Z27utMLICn1eIyma1cM9E,3760
 cornflow/tests/unit/test_users.py,sha256=WfaMcybPpR7rspXyvzHGgw25p751hMPAV0DOp_caSPM,22430
 cornflow/tests/unit/tools.py,sha256=BCAm_KGVgZO-CCb_rkaZlbK4SID_F2ab8FiBJzGwKtc,587
-cornflow-2.0.0a9.dist-info/METADATA,sha256=qRfLB20zSSHf-LfT1RzSGLEy8VwFxFhM75XdiYK3DIQ,9536
-cornflow-2.0.0a9.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-cornflow-2.0.0a9.dist-info/entry_points.txt,sha256=q9cPKAFBsmHkERCqQ2JcOTM-tVBLHTl-DGxwCXowAWM,46
-cornflow-2.0.0a9.dist-info/top_level.txt,sha256=Qj9kLFJW1PLb-ZV2s_aCkQ-Wi5W6KC6fFR-LTBrx-rU,24
-cornflow-2.0.0a9.dist-info/RECORD,,
+cornflow-2.0.0a10.dist-info/METADATA,sha256=KPriiqRbRG0MYnmQJyItpKiz8OF-G4ZL1Vu6H2LERhU,9537
+cornflow-2.0.0a10.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+cornflow-2.0.0a10.dist-info/entry_points.txt,sha256=q9cPKAFBsmHkERCqQ2JcOTM-tVBLHTl-DGxwCXowAWM,46
+cornflow-2.0.0a10.dist-info/top_level.txt,sha256=Qj9kLFJW1PLb-ZV2s_aCkQ-Wi5W6KC6fFR-LTBrx-rU,24
+cornflow-2.0.0a10.dist-info/RECORD,,