feat(engine): Add workflow error handler to schema and API (#690)

TracecatHQ · Jan 2, 2025 · ba114d7 · ba114d7
1 parent 5fffb7b
commit ba114d7
Show file tree

Hide file tree

Showing 8 changed files with 100 additions and 0 deletions.
diff --git a/alembic/versions/f2e840c9fb2c_add_workflow_error_handler.py b/alembic/versions/f2e840c9fb2c_add_workflow_error_handler.py
@@ -0,0 +1,34 @@
+"""Add workflow error handler
+
+Revision ID: f2e840c9fb2c
+Revises: 8be6393b2ee0
+Create Date: 2024-12-31 15:39:57.045107
+
+"""
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+import sqlmodel.sql.sqltypes
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "f2e840c9fb2c"
+down_revision: str | None = "8be6393b2ee0"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "workflow",
+        sa.Column("error_handler", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("workflow", "error_handler")
+    # ### end Alembic commands ###
diff --git a/frontend/src/client/schemas.gen.ts b/frontend/src/client/schemas.gen.ts
@@ -671,6 +671,18 @@ export const $DSLInput = {
             ],
             title: 'Returns',
             description: 'The action ref or value to return.'
+        },
+        error_handler: {
+            anyOf: [
+                {
+                    type: 'string'
+                },
+                {
+                    type: 'null'
+                }
+            ],
+            title: 'Error Handler',
+            description: 'The action ref to handle errors.'
         }
     },
     type: 'object',
@@ -3672,6 +3684,17 @@ export const $WorkflowRead = {
                 }
             ],
             title: 'Alias'
+        },
+        error_handler: {
+            anyOf: [
+                {
+                    type: 'string'
+                },
+                {
+                    type: 'null'
+                }
+            ],
+            title: 'Error Handler'
         }
     },
     type: 'object',
@@ -3754,6 +3777,17 @@ export const $WorkflowReadMinimal = {
                 }
             ],
             title: 'Alias'
+        },
+        error_handler: {
+            anyOf: [
+                {
+                    type: 'string'
+                },
+                {
+                    type: 'null'
+                }
+            ],
+            title: 'Error Handler'
         }
     },
     type: 'object',
@@ -3908,6 +3942,17 @@ export const $WorkflowUpdate = {
                 }
             ],
             title: 'Alias'
+        },
+        error_handler: {
+            anyOf: [
+                {
+                    type: 'string'
+                },
+                {
+                    type: 'null'
+                }
+            ],
+            title: 'Error Handler'
         }
     },
     type: 'object',

diff --git a/frontend/src/client/types.gen.ts b/frontend/src/client/types.gen.ts
@@ -250,6 +250,10 @@ export type DSLInput = {
      * The action ref or value to return.
      */
     returns?: unknown | null;
+    /**
+     * The action ref to handle errors.
+     */
+    error_handler?: string | null;
 };
 
 export type DSLRunArgs = {
@@ -1128,6 +1132,7 @@ export type WorkflowRead = {
     returns: unknown;
     config: DSLConfig_Output | null;
     alias?: string | null;
+    error_handler?: string | null;
 };
 
 export type WorkflowReadMinimal = {
@@ -1141,6 +1146,7 @@ export type WorkflowReadMinimal = {
     version: number | null;
     tags?: Array<TagRead> | null;
     alias?: string | null;
+    error_handler?: string | null;
 };
 
 export type WorkflowTagCreate = {
@@ -1166,6 +1172,7 @@ export type WorkflowUpdate = {
     returns?: unknown | null;
     config?: DSLConfig_Input | null;
     alias?: string | null;
+    error_handler?: string | null;
 };
 
 export type WorkspaceMember = {

diff --git a/tracecat/db/schemas.py b/tracecat/db/schemas.py
@@ -313,6 +313,10 @@ class Workflow(Resource, table=True):
     alias: str | None = Field(
         default=None, description="Alias for the workflow", index=True
     )
+    error_handler: str | None = Field(
+        default=None,
+        description="Workflow alias or ID for the workflow to run when this fails.",
+    )
     icon_url: str | None = None
     # Owner
     owner_id: OwnerID = Field(

diff --git a/tracecat/dsl/common.py b/tracecat/dsl/common.py
@@ -79,6 +79,9 @@ class DSLInput(BaseModel):
     returns: Any | None = Field(
         default=None, description="The action ref or value to return."
     )
+    error_handler: str | None = Field(
+        default=None, description="The action ref to handle errors."
+    )
 
     @model_validator(mode="after")
     def validate_structure(self) -> Self:

diff --git a/tracecat/workflow/management/management.py b/tracecat/workflow/management/management.py
@@ -229,6 +229,7 @@ async def build_dsl_from_workflow(self, workflow: Workflow) -> DSLInput:
             inputs=workflow.static_inputs,
             config=DSLConfig(**workflow.config),
             returns=workflow.returns,
+            error_handler=workflow.error_handler,
         )
 
     async def create_workflow_from_external_definition(

diff --git a/tracecat/workflow/management/models.py b/tracecat/workflow/management/models.py
@@ -35,6 +35,7 @@ class WorkflowRead(BaseModel):
     returns: Any
     config: DSLConfig | None
     alias: str | None = None
+    error_handler: str | None = None
 
 
 class WorkflowReadMinimal(BaseModel):
@@ -48,6 +49,7 @@ class WorkflowReadMinimal(BaseModel):
     version: int | None
     tags: list[TagRead] | None = None
     alias: str | None = None
+    error_handler: str | None = None
 
 
 class WorkflowUpdate(BaseModel):
@@ -63,6 +65,7 @@ class WorkflowUpdate(BaseModel):
     returns: Any | None = None
     config: DSLConfig | None = None
     alias: str | None = None
+    error_handler: str | None = None
 
 
 class WorkflowCreate(BaseModel):

diff --git a/tracecat/workflow/management/router.py b/tracecat/workflow/management/router.py
@@ -74,6 +74,7 @@ async def list_workflows(
                 version=workflow.version,
                 tags=tags,
                 alias=workflow.alias,
+                error_handler=workflow.error_handler,
             )
         )
     return res
@@ -165,6 +166,7 @@ async def create_workflow(
         created_at=workflow.created_at,
         updated_at=workflow.updated_at,
         version=workflow.version,
+        error_handler=workflow.error_handler,
     )
 
 
@@ -205,6 +207,7 @@ async def get_workflow(
         webhook=WebhookResponse(**workflow.webhook.model_dump()),
         schedules=workflow.schedules or [],
         alias=workflow.alias,
+        error_handler=workflow.error_handler,
     )