add workflow definitions

Author: andrii-i

jupyter_scheduler/extension.py

@@ -7,6 +7,8 @@
 
 from jupyter_scheduler.orm import create_tables
 from jupyter_scheduler.workflows import (
+    WorkflowDefinitionsHandler,
+    WorkflowDefinitionsTasksHandler,
     WorkflowsHandler,
     WorkflowsRunHandler,
     WorkflowsTasksHandler,
@@ -25,6 +27,7 @@
 
 JOB_DEFINITION_ID_REGEX = r"(?P<job_definition_id>\w+(?:-\w+)+)"
 JOB_ID_REGEX = r"(?P<job_id>\w+(?:-\w+)+)"
+WORKFLOW_DEFINITION_ID_REGEX = r"(?P<workflow_definition_id>\w+(?:-\w+)+)"
 WORKFLOW_ID_REGEX = r"(?P<workflow_id>\w+(?:-\w+)+)"
 
 
@@ -51,6 +54,18 @@ class SchedulerApp(ExtensionApp):
             rf"scheduler/worklows/{WORKFLOW_ID_REGEX}/tasks",
             WorkflowsTasksHandler,
         ),
+        (
+            rf"scheduler/worklow_definitions/{WORKFLOW_DEFINITION_ID_REGEX}",
+            WorkflowDefinitionsHandler,
+        ),
+        (
+            rf"scheduler/worklows/{WORKFLOW_DEFINITION_ID_REGEX}/run",
+            WorkflowDefinitionsHandler,
+        ),
+        (
+            rf"scheduler/worklows/{WORKFLOW_ID_REGEX}/tasks",
+            WorkflowDefinitionsTasksHandler,
+        ),
     ]
 
     drop_tables = Bool(False, config=True, help="Drop the database tables before starting.")

jupyter_scheduler/orm.py

@@ -120,6 +120,19 @@ class Workflow(Base):
     # Any default values specified for new columns will be ignored during the migration process.
 
 
+class WorkflowDefinition(Base):
+    __tablename__ = "workflow_definitions"
+    __table_args__ = {"extend_existing": True}
+    workflow_id = Column(String(36), primary_key=True, default=generate_uuid)
+    tasks = Column(JsonType)
+    status = Column(String(64), default=Status.CREATED)
+    active = Column(Boolean, default=False)
+    schedule = Column(String(256))
+    timezone = Column(String(36))
+    # All new columns added to this table must be nullable to ensure compatibility during database migrations.
+    # Any default values specified for new columns will be ignored during the migration process.
+
+
 class JobDefinition(CommonColumns, Base):
     __tablename__ = "job_definitions"
     __table_args__ = {"extend_existing": True}

jupyter_scheduler/workflows.py

@@ -148,6 +148,79 @@ async def post(self, workflow_id: str):
             self.finish(json.dumps(dict(workflow_id=workflow_id)))
 
 
+class WorkflowDefinitionsHandler(ExtensionHandlerMixin, JobHandlersMixin, APIHandler):
+    @authenticated
+    async def post(self):
+        payload = self.get_json_body() or {}
+        try:
+            workflow_id = await ensure_async(
+                self.scheduler.create_workflow(CreateWorkflow(**payload))
+            )
+        except ValidationError as e:
+            self.log.exception(e)
+            raise HTTPError(500, str(e)) from e
+        except InputUriError as e:
+            self.log.exception(e)
+            raise HTTPError(500, str(e)) from e
+        except IdempotencyTokenError as e:
+            self.log.exception(e)
+            raise HTTPError(409, str(e)) from e
+        except SchedulerError as e:
+            self.log.exception(e)
+            raise HTTPError(500, str(e)) from e
+        except Exception as e:
+            self.log.exception(e)
+            raise HTTPError(500, "Unexpected error occurred during creation of a workflow.") from e
+        else:
+            self.finish(json.dumps(dict(workflow_id=workflow_id)))
+
+    @authenticated
+    async def get(self, workflow_id: str = None):
+        if not workflow_id:
+            raise HTTPError(400, "Missing workflow_id in the request URL.")
+        try:
+            workflow = await ensure_async(self.scheduler.get_workflow(workflow_id))
+        except SchedulerError as e:
+            self.log.exception(e)
+            raise HTTPError(500, str(e)) from e
+        except Exception as e:
+            self.log.exception(e)
+            raise HTTPError(500, "Unexpected error occurred while getting workflow details.") from e
+        else:
+            self.finish(workflow.json())
+
+
+class WorkflowDefinitionsTasksHandler(ExtensionHandlerMixin, JobHandlersMixin, APIHandler):
+    @authenticated
+    async def post(self, workflow_id: str):
+        payload = self.get_json_body()
+        try:
+            task_id = await ensure_async(
+                self.scheduler.create_workflow_task(
+                    workflow_id=workflow_id, model=CreateJob(**payload)
+                )
+            )
+        except ValidationError as e:
+            self.log.exception(e)
+            raise HTTPError(500, str(e)) from e
+        except InputUriError as e:
+            self.log.exception(e)
+            raise HTTPError(500, str(e)) from e
+        except IdempotencyTokenError as e:
+            self.log.exception(e)
+            raise HTTPError(409, str(e)) from e
+        except SchedulerError as e:
+            self.log.exception(e)
+            raise HTTPError(500, str(e)) from e
+        except Exception as e:
+            self.log.exception(e)
+            raise HTTPError(
+                500, "Unexpected error occurred during creation of workflow job."
+            ) from e
+        else:
+            self.finish(json.dumps(dict(task_id=task_id)))
+
+
 class CreateWorkflow(BaseModel):
     tasks: List[str] = []
 
@@ -169,3 +242,35 @@ class UpdateWorkflow(BaseModel):
 
     class Config:
         orm_mode = True
+
+
+class CreateWorkflowDefinition(BaseModel):
+    tasks: List[str] = []
+    # any field added to CreateWorkflow should also be added to this model as well
+    schedule: Optional[str] = None
+    timezone: Optional[str] = None
+
+    class Config:
+        orm_mode = True
+
+
+class DescribeWorkflowDefinition(BaseModel):
+    workflow_definition_id: str
+    tasks: List[str] = None
+    schedule: Optional[str] = None
+    timezone: Optional[str] = None
+    status: Status = Status.CREATED
+    active: Optional[bool] = None
+
+    class Config:
+        orm_mode = True
+
+
+class UpdateWorkflowDefinition(BaseModel):
+    tasks: Optional[List[str]] = None
+    schedule: Optional[str] = None
+    timezone: Optional[str] = None
+    active: Optional[bool] = None
+
+    class Config:
+        orm_mode = True