cornflow 1.2.4__py3-none-any.whl → 1.3.0rc1__py3-none-any.whl

cornflow/endpoints/health.py

@@ -2,18 +2,28 @@
 Endpoint to check the health of the services.
 It performs a health check to airflow and a health check to cornflow database
 """
+
 import os
 
-# Import from internal modules
-from cornflow.endpoints.meta_resource import BaseMetaResource
-from cornflow.models import UserModel
-from cornflow.schemas.health import HealthResponse
-from cornflow.shared.const import STATUS_HEALTHY, STATUS_UNHEALTHY, CORNFLOW_VERSION
 # Import from libraries
 from cornflow_client.airflow.api import Airflow
 from flask import current_app
 from flask_apispec import marshal_with, doc
 
+# Import from internal modules
+from cornflow.endpoints.meta_resource import BaseMetaResource
+from cornflow.models import UserModel
+from cornflow.schemas.health import HealthResponse
+from cornflow.shared.const import (
+    AIRFLOW_BACKEND,
+    DATABRICKS_BACKEND,
+    STATUS_HEALTHY,
+    STATUS_UNHEALTHY,
+    CORNFLOW_VERSION
+)
+from cornflow_client.databricks.api import Databricks
+from cornflow.shared.exceptions import EndpointNotImplemented
+
 
 class HealthEndpoint(BaseMetaResource):
     @doc(description="Health check", tags=["Health"])
@@ -26,12 +36,9 @@ class HealthEndpoint(BaseMetaResource):
         :rtype: dict
         :doc-author: baobab soluciones
         """
-        af_client = Airflow.from_config(current_app.config)
-        airflow_status = STATUS_UNHEALTHY
+        backend_status = self.check_backend_status()
         cornflow_status = STATUS_UNHEALTHY
         cornflow_version = CORNFLOW_VERSION
-        if af_client.is_alive():
-            airflow_status = STATUS_HEALTHY
 
         if (
             UserModel.get_one_user_by_username(os.getenv("CORNFLOW_SERVICE_USER"))
@@ -40,6 +47,18 @@ class HealthEndpoint(BaseMetaResource):
             cornflow_status = STATUS_HEALTHY
 
         current_app.logger.info(
-            f"Health check: cornflow {cornflow_status}, airflow {airflow_status}"
+            f"Health check: cornflow {cornflow_status}, backend {backend_status}"
         )
-        return {"cornflow_status": cornflow_status, "airflow_status": airflow_status, "cornflow_version":cornflow_version}
+        return {"cornflow_status": cornflow_status, "backend_status": backend_status,
+                "cornflow_version": cornflow_version}
+
+    def check_backend_status(self):  
+        if current_app.config["CORNFLOW_BACKEND"] == AIRFLOW_BACKEND:  
+            client = Airflow.from_config(current_app.config)  
+        elif current_app.config["CORNFLOW_BACKEND"] == DATABRICKS_BACKEND:  
+            client = Databricks.from_config(current_app.config)  
+        else:  
+            raise EndpointNotImplemented()  
+        if client.is_alive():  
+            return STATUS_HEALTHY  
+        return STATUS_UNHEALTHY  

cornflow/endpoints/instance.py

@@ -16,7 +16,7 @@ from werkzeug.utils import secure_filename
 
 # Import from internal modules
 from cornflow.endpoints.meta_resource import BaseMetaResource
-from cornflow.models import InstanceModel, DeployedDAG
+from cornflow.models import InstanceModel, DeployedWorkflow
 from cornflow.schemas.instance import (
     InstanceSchema,
     InstanceEndpointResponse,
@@ -91,7 +91,7 @@ class InstanceEndpoint(BaseMetaResource):
         # We validate the instance data
         config = current_app.config
 
-        instance_schema = DeployedDAG.get_one_schema(
+        instance_schema = DeployedWorkflow.get_one_schema(
             config, data_schema, INSTANCE_SCHEMA
         )
         instance_errors = json_schema_validate_as_string(
@@ -166,7 +166,7 @@ class InstanceDetailsEndpoint(InstanceDetailsEndpointBase):
 
             config = current_app.config
 
-            instance_schema = DeployedDAG.get_one_schema(
+            instance_schema = DeployedWorkflow.get_one_schema(
                 config, schema, INSTANCE_SCHEMA
             )
             instance_errors = json_schema_validate_as_string(

cornflow/endpoints/login.py

@@ -188,8 +188,15 @@ class LoginBaseEndpoint(BaseMetaResource):
                 )
 
                 email = decoded_token.get("email", f"{username}@cornflow.org")
-                first_name = decoded_token.get("given_name", "")
-                last_name = decoded_token.get("family_name", "")
+
+                email_names = email.split("@")[0]
+                if "." in email_names:
+                    default_first_name, default_last_name = email_names.split(".")[0:2]
+                else:
+                    default_first_name, default_last_name = email_names, ""
+
+                first_name = decoded_token.get("given_name", default_first_name)
+                last_name = decoded_token.get("family_name", default_last_name)
 
                 data = {
                     "username": username,

cornflow/endpoints/schemas.py

@@ -8,7 +8,7 @@ from flask_apispec import marshal_with, doc
 
 # Import from internal modules
 from cornflow.endpoints.meta_resource import BaseMetaResource
-from cornflow.models import PermissionsDAG, DeployedDAG
+from cornflow.models import PermissionsDAG, DeployedWorkflow
 from cornflow.schemas.schemas import SchemaOneApp, SchemaListApp
 from cornflow.shared.authentication import Auth, authenticate
 from cornflow.shared.const import ALL_DEFAULT_ROLES
@@ -63,7 +63,7 @@ class SchemaDetailsEndpoint(BaseMetaResource):
         )
 
         if permission:
-            deployed_dag = DeployedDAG.get_one_object(dag_name)
+            deployed_dag = DeployedWorkflow.get_one_object(dag_name)
             current_app.logger.info("User gets schema {}".format(dag_name))
             return {
                 "instance": deployed_dag.instance_schema,
@@ -71,7 +71,7 @@ class SchemaDetailsEndpoint(BaseMetaResource):
                 "instance_checks": deployed_dag.instance_checks_schema,
                 "solution_checks": deployed_dag.solution_checks_schema,
                 "config": deployed_dag.config_schema,
-                "name": dag_name
+                "name": dag_name,
             }, 200
         else:
             err = "User does not have permission to access this dag"

cornflow/migrations/versions/999b98e24225.py

@@ -0,0 +1,34 @@
+"""
+Rename dag_run_id column to run_id in executions table
+
+Revision ID: 999b98e24225
+Revises: 991b98e24225
+Create Date: 2024-04-08 12:00:00.000000
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "abc123456789"
+down_revision = "991b98e24225"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("executions", schema=None) as batch_op:
+        batch_op.alter_column("dag_run_id", new_column_name="run_id")
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("executions", schema=None) as batch_op:
+        batch_op.alter_column("run_id", new_column_name="dag_run_id")
+
+    # ### end Alembic commands ###

cornflow/migrations/versions/cef1df240b27_.py

@@ -0,0 +1,34 @@
+"""empty message
+
+Revision ID: cef1df240b27
+Revises: abc123456789
+Create Date: 2025-07-22 19:05:11.146449
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = 'cef1df240b27'
+down_revision = 'abc123456789'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+
+    op.rename_table("deployed_dags", "deployed_workflows")
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+
+    # Rename table back from deployed_workflows to deployed_dags
+    op.rename_table("deployed_workflows", "deployed_dags")
+
+    # ### end Alembic commands ###

cornflow/models/__init__.py

@@ -1,10 +1,11 @@
 """
 Initialization file for the models module
 """
+
 from .action import ActionModel
 from .alarms import AlarmsModel
 from .case import CaseModel
-from .dag import DeployedDAG
+from .dag import DeployedWorkflow
 from .dag_permissions import PermissionsDAG
 from .execution import ExecutionModel
 from .instance import InstanceModel

cornflow/models/dag.py

@@ -1,13 +1,12 @@
-"""
+""" """
 
-"""
 # Import from libraries
 from cornflow_client.airflow.api import Airflow
 from cornflow_client.constants import (
     INSTANCE_SCHEMA,
     SOLUTION_SCHEMA,
     INSTANCE_CHECKS_SCHEMA,
-    SOLUTION_CHECKS_SCHEMA
+    SOLUTION_CHECKS_SCHEMA,
 )
 from sqlalchemy.dialects.postgresql import TEXT, JSON
 
@@ -17,12 +16,12 @@ from cornflow.shared import db
 from cornflow.shared.exceptions import ObjectDoesNotExist
 
 
-class DeployedDAG(TraceAttributesModel):
+class DeployedWorkflow(TraceAttributesModel):
     """
     This model contains the registry of the DAGs that are deployed on the corresponding Airflow server
     """
 
-    __tablename__ = "deployed_dags"
+    __tablename__ = "deployed_workflows"
     id = db.Column(db.String(128), primary_key=True)
     description = db.Column(TEXT, nullable=True)
     instance_schema = db.Column(JSON, nullable=True)
@@ -34,8 +33,8 @@ class DeployedDAG(TraceAttributesModel):
     dag_permissions = db.relationship(
         "PermissionsDAG",
         cascade="all,delete",
-        backref="deployed_dags",
-        primaryjoin="and_(DeployedDAG.id==PermissionsDAG.dag_id)",
+        backref="deployed_workflows",
+        primaryjoin="and_(DeployedWorkflow.id==PermissionsDAG.dag_id)",
     )
 
     def __init__(self, data):
@@ -53,14 +52,14 @@ class DeployedDAG(TraceAttributesModel):
 
     @staticmethod
     def get_one_schema(config, dag_name, schema=INSTANCE_SCHEMA):
-        item = DeployedDAG.get_one_object(dag_name)
+        item = DeployedWorkflow.get_one_object(dag_name)
 
         if item is None:
             err = f"The DAG {dag_name} does not exist in the database."
             raise ObjectDoesNotExist(
                 err,
                 log_txt=f"Error while user tries to register data for DAG {dag_name} "
-                            f"from instance and execution. " + err
+                f"from instance and execution. " + err,
             )
 
         if schema == INSTANCE_SCHEMA:

cornflow/models/dag_permissions.py

@@ -1,4 +1,4 @@
-from cornflow.models.dag import DeployedDAG
+from cornflow.models.dag import DeployedWorkflow
 from cornflow.models.meta_models import TraceAttributesModel
 from cornflow.shared import db
 
@@ -10,7 +10,7 @@ class PermissionsDAG(TraceAttributesModel):
     id = db.Column(db.Integer, primary_key=True, autoincrement=True)
 
     dag_id = db.Column(
-        db.String(128), db.ForeignKey("deployed_dags.id"), nullable=False
+        db.String(128), db.ForeignKey("deployed_workflows.id"), nullable=False
     )
     user_id = db.Column(db.Integer, db.ForeignKey("users.id"), nullable=False)
     user = db.relationship("UserModel", viewonly=True)
@@ -29,7 +29,7 @@ class PermissionsDAG(TraceAttributesModel):
 
     @staticmethod
     def add_all_permissions_to_user(user_id):
-        dags = DeployedDAG.get_all_objects()
+        dags = DeployedWorkflow.get_all_objects()
         permissions = [
             PermissionsDAG({"dag_id": dag.id, "user_id": user_id}) for dag in dags
         ]

cornflow/models/execution.py

@@ -55,7 +55,7 @@ class ExecutionModel(BaseDataModel):
         db.String(256), db.ForeignKey("instances.id"), nullable=False
     )
     config = db.Column(JSON, nullable=False)
-    dag_run_id = db.Column(db.String(256), nullable=True)
+    run_id = db.Column(db.String(256), nullable=True)
     log_text = db.Column(TEXT, nullable=True)
     log_json = db.Column(JSON, nullable=True)
     state = db.Column(db.SmallInteger, default=DEFAULT_EXECUTION_CODE, nullable=False)
@@ -78,8 +78,7 @@ class ExecutionModel(BaseDataModel):
                 + str(self.instance_id)
             ).encode()
         ).hexdigest()
-
-        self.dag_run_id = data.get("dag_run_id")
+        self.run_id = data.get("run_id")
         self.state = data.get("state", DEFAULT_EXECUTION_CODE)
         self.state_message = EXECUTION_STATE_MESSAGE_DICT[self.state]
         self.config = data.get("config")

cornflow/models/permissions.py

@@ -1,6 +1,7 @@
 """
 This file contains the PermissionViewRoleModel
 """
+
 # Imports from internal modules
 from cornflow.models.meta_models import TraceAttributesModel
 from cornflow.shared import db

cornflow/models/user.py

@@ -5,7 +5,7 @@ This file contains the UserModel
 # Imports from external libraries
 import random
 import string
-from datetime import datetime, timezone, timedelta
+from datetime import datetime, timezone
 
 # Imports from internal modules
 from cornflow.models.meta_models import TraceAttributesModel

cornflow/schemas/execution.py

@@ -43,7 +43,7 @@ class ExecutionSchema(Schema):
     name = fields.Str()
     description = fields.Str()
     schema = fields.Str(required=False)
-    dag_run_id = fields.Str(required=False, dump_only=True)
+    run_id = fields.Str(required=False, dump_only=True)
     config = fields.Nested(ConfigSchema, required=True)
     data = fields.Raw(dump_only=True)
     checks = fields.Raw(required=False, allow_none=True)
@@ -78,6 +78,7 @@ class ExecutionEditRequest(Schema):
     name = fields.Str()
     description = fields.Str()
     data = fields.Raw()
+    instance_id = fields.Str(required=False)
 
 
 class ExecutionDagRequest(Schema):
@@ -119,7 +120,19 @@ class ExecutionDetailsEndpointWithIndicatorsResponse(ExecutionDetailsEndpointRes
             return obj.user.username
         return None
 
+    def get_first_name(self, obj):
+        if hasattr(obj, "user") and obj.user is not None:
+            return obj.user.first_name
+        return None
+
+    def get_last_name(self, obj):
+        if hasattr(obj, "user") and obj.user is not None:
+            return obj.user.last_name
+        return None
+
     username = fields.Method("get_username")
+    first_name = fields.Method("get_first_name")
+    last_name = fields.Method("get_last_name")
 
 
 class ExecutionDetailsWithIndicatorsAndLogResponse(

cornflow/schemas/health.py

@@ -4,5 +4,5 @@ from marshmallow import fields, Schema
 class HealthResponse(Schema):
 
     cornflow_status = fields.Str()
-    airflow_status = fields.Str()
+    backend_status = fields.Str()
     cornflow_version = fields.Str()

cornflow/shared/authentication/auth.py

@@ -25,6 +25,7 @@ from cornflow.shared.const import (
     AUTH_OID,
     PERMISSION_METHOD_MAP,
     INTERNAL_TOKEN_ISSUER,
+    OID_PROVIDER_AZURE,
 )
 from cornflow.shared.exceptions import (
     CommunicationError,
@@ -300,7 +301,19 @@ class Auth:
         :rtype: dict
         """
         # Fetch keys from provider
-        jwks_url = f"{provider_url.rstrip('/')}/.well-known/jwks.json"
+        # For Azure AD, we need to use the discovery endpoint to get the jwks_uri
+        oid_provider = current_app.config["OID_PROVIDER"]
+        if oid_provider == OID_PROVIDER_AZURE:
+            tenant_id = provider_url.split("/")[
+                3
+            ]  # Extract tenant ID from provider URL
+            jwks_url = (
+                f"https://login.microsoftonline.com/{tenant_id}/discovery/v2.0/keys"
+            )
+        else:
+            # For other providers, use the standard well-known endpoint
+            jwks_url = f"{provider_url.rstrip('/')}/.well-known/jwks.json"
+
         try:
             response = requests.get(jwks_url)
             response.raise_for_status()

cornflow/shared/authentication/decorators.py

@@ -6,7 +6,6 @@ from functools import wraps
 from .auth import Auth
 from cornflow.shared.exceptions import InvalidCredentials
 from flask import current_app
-import os
 
 
 def authenticate(

cornflow/shared/const.py

@@ -2,7 +2,12 @@
 In this file we import the values for different constants on cornflow server
 """
 
-CORNFLOW_VERSION = "1.2.4"
+# CORNFLOW BACKEND
+AIRFLOW_BACKEND = 1
+DATABRICKS_BACKEND = 2
+
+
+CORNFLOW_VERSION = "1.3.0rc1"
 INTERNAL_TOKEN_ISSUER = "cornflow"
 
 # endpoints responses for health check
@@ -53,6 +58,24 @@ NO_SIGNUP = 0
 SIGNUP_WITH_NO_AUTH = 1
 SIGNUP_WITH_AUTH = 2
 
+DATABRICKS_TO_STATE_MAP = dict(
+    BLOCKED=EXEC_STATE_QUEUED,
+    PENDING=EXEC_STATE_QUEUED,
+    QUEUED=EXEC_STATE_QUEUED,
+    RUNNING=EXEC_STATE_RUNNING,
+    TERMINATING=EXEC_STATE_RUNNING,
+    SUCCESS=EXEC_STATE_CORRECT,
+    USER_CANCELED=EXEC_STATE_STOPPED,
+    OTHER_FINISH_ERROR=EXEC_STATE_ERROR,
+    RUN_EXECUTION_ERROR=EXEC_STATE_ERROR,
+)
+
+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
 # as this library that is the base of airflow
 AUTH_DB = 1
@@ -60,6 +83,11 @@ AUTH_LDAP = 2
 AUTH_OAUTH = 4
 AUTH_OID = 0
 
+# OID possible providers
+OID_PROVIDER_AWS = 0
+OID_PROVIDER_AZURE = 1
+OID_OTHER = 2
+
 GET_ACTION = 1
 PATCH_ACTION = 2
 POST_ACTION = 3
@@ -135,3 +163,18 @@ CONDITIONAL_ENDPOINTS = {
     "signup": "/signup/",
     "login": "/login/",
 }
+
+
+# Orchestrator constants
+config_orchestrator = {
+    "airflow": {
+        "name": "Airflow",
+        "def_schema": "solve_model_dag",
+        "run_id": "dag_run_id",
+    },
+    "databricks": {
+        "name": "Databricks",
+        "def_schema": "979073949072767",
+        "run_id": "run_id",
+    },
+}

cornflow/shared/exceptions.py

@@ -5,7 +5,7 @@ on a flask REST API server
 
 from flask import jsonify
 from webargs.flaskparser import parser
-from cornflow_client.constants import AirflowError
+from cornflow_client.constants import AirflowError, DatabricksError
 from werkzeug.exceptions import HTTPException
 import traceback
 
@@ -148,6 +148,7 @@ def initialize_errorhandlers(app):
     @app.errorhandler(InvalidCredentials)
     @app.errorhandler(EndpointNotImplemented)
     @app.errorhandler(AirflowError)
+    @app.errorhandler(DatabricksError)
     @app.errorhandler(InvalidData)
     @app.errorhandler(InvalidPatch)
     @app.errorhandler(ConfigurationError)