feat: Workflow refactoring, only async

.github/workflows/test-pr.yml

@@ -95,7 +95,7 @@ jobs:
           LOG_LEVEL: DEBUG
           SQLALCHEMY_WARN_20: 1
         run: |
-          poetry run coverage run --branch -m pytest --timeout 20 -n auto --non-integration --ignore=tests/e2e_tests/
+          poetry run coverage run --branch -m pytest -n auto --non-integration --ignore=tests/e2e_tests/
 
       - name: Run integration tests and report coverage
         run: |

docs/mint.json

@@ -121,6 +121,7 @@
             "providers/documentation/checkmk-provider",
             "providers/documentation/cilium-provider",
             "providers/documentation/clickhouse-provider",
+            "providers/documentation/clickhouse-http-provider",
             "providers/documentation/cloudwatch-provider",
             "providers/documentation/console-provider",
             "providers/documentation/coralogix-provider",

docs/providers/documentation/clickhouse-http-provider.mdx

@@ -0,0 +1,7 @@
+---
+title: 'ClickHouse HTTP'
+sidebarTitle: 'ClickHouse HTTP Provider'
+description: 'ClickHouse HTTP provider allows you to interact with ClickHouse database.'
+---
+
+This provider is an async (more performant) analog of [clickhouse-provider](providers/documentation/clickhouse-provider.mdx). It's using HTTP protocol to interact to the Clickhouse.

docs/providers/overview.mdx

@@ -156,6 +156,12 @@ By leveraging Keep Providers, users are able to deeply integrate Keep with the t
   }
 ></Card>
 
+<Card
+  title="ClickHouse HTTP"
+  href="/providers/documentation/clickhouse-http-provider"
+  icon={ <img src="https://img.logo.dev/clickhouse.com?token=pk_dfXfZBoKQMGDTIgqu7LvYg" /> }
+></Card>
+
 <Card
   title="ClickHouse"
   href="/providers/documentation/clickhouse-provider"

keep/api/bl/incidents_bl.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 import os
 import pathlib
@@ -153,7 +154,7 @@ def update_client_on_incident_change(self, incident_id: Optional[UUID] = None):
     def send_workflow_event(self, incident_dto: IncidentDto, action: str) -> None:
         try:
             workflow_manager = WorkflowManager.get_instance()
-            workflow_manager.insert_incident(self.tenant_id, incident_dto, action)
+            asyncio.run(workflow_manager.insert_incident(self.tenant_id, incident_dto, action))
         except Exception:
             self.logger.exception(
                 "Failed to run workflows based on incident",
@@ -231,6 +232,7 @@ def delete_incident(self, incident_id: UUID) -> None:
         self.update_client_on_incident_change()
         self.send_workflow_event(incident_dto, "deleted")
 
+
     def update_incident(
         self,
         incident_id: UUID,

keep/api/core/db.py

@@ -39,6 +39,7 @@
 from sqlalchemy.exc import IntegrityError, OperationalError
 from sqlalchemy.orm import joinedload, selectinload, subqueryload
 from sqlalchemy.sql import exists, expression
+from sqlmodel.ext.asyncio.session import AsyncSession
 from sqlmodel import Session, SQLModel, col, or_, select, text
 
 from keep.api.consts import STATIC_PRESETS
@@ -82,6 +83,8 @@
 
 
 engine = create_db_engine()
+engine_async = create_db_engine(_async=True)
+
 SQLAlchemyInstrumentor().instrument(enable_commenter=True, engine=engine)
 
 
@@ -145,7 +148,7 @@ def __convert_to_uuid(value: str) -> UUID | None:
         return None
 
 
-def create_workflow_execution(
+async def create_workflow_execution(
     workflow_id: str,
     tenant_id: str,
     triggered_by: str,
@@ -155,7 +158,7 @@ def create_workflow_execution(
     execution_id: str = None,
     event_type: str = "alert",
 ) -> str:
-    with Session(engine) as session:
+    async with AsyncSession(engine_async) as session:
         try:
             if len(triggered_by) > 255:
                 triggered_by = triggered_by[:255]
@@ -170,7 +173,7 @@ def create_workflow_execution(
             )
             session.add(workflow_execution)
             # Ensure the object has an id
-            session.flush()
+            await session.flush()
             execution_id = workflow_execution.id
             if KEEP_AUDIT_EVENTS_ENABLED:
                 if fingerprint and event_type == "alert":
@@ -188,10 +191,10 @@ def create_workflow_execution(
                     )
                     session.add(workflow_to_incident_execution)
 
-            session.commit()
+            await session.commit()
             return execution_id
         except IntegrityError:
-            session.rollback()
+            await session.rollback()
             logger.debug(
                 f"Failed to create a new execution for workflow {workflow_id}. Constraint is met."
             )
@@ -226,12 +229,12 @@ def get_last_completed_execution(
     ).first()
 
 
-def get_workflows_that_should_run():
-    with Session(engine) as session:
+async def get_workflows_that_should_run():
+    async with AsyncSession(engine_async) as session:
         logger.debug("Checking for workflows that should run")
         workflows_with_interval = []
         try:
-            result = session.exec(
+            result = await session.exec(
                 select(Workflow)
                 .filter(Workflow.is_deleted == False)
                 .filter(Workflow.is_disabled == False)
@@ -252,7 +255,7 @@ def get_workflows_that_should_run():
             if not last_execution:
                 try:
                     # try to get the lock
-                    workflow_execution_id = create_workflow_execution(
+                    workflow_execution_id = await create_workflow_execution(
                         workflow.id, workflow.tenant_id, "scheduler"
                     )
                     # we succeed to get the lock on this execution number :)
@@ -274,7 +277,7 @@ def get_workflows_that_should_run():
             ):
                 try:
                     # try to get the lock with execution_number + 1
-                    workflow_execution_id = create_workflow_execution(
+                    workflow_execution_id = await create_workflow_execution(
                         workflow.id,
                         workflow.tenant_id,
                         "scheduler",
@@ -294,10 +297,10 @@ def get_workflows_that_should_run():
                 # some other thread/instance has already started to work on it
                 except IntegrityError:
                     # we need to verify the locking is still valid and not timeouted
-                    session.rollback()
+                    await session.rollback()
                     pass
                 # get the ongoing execution
-                ongoing_execution = session.exec(
+                ongoing_execution = await session.exec(
                     select(WorkflowExecution)
                     .where(WorkflowExecution.workflow_id == workflow.id)
                     .where(
@@ -319,10 +322,10 @@ def get_workflows_that_should_run():
                 # if the ongoing execution runs more than 60 minutes, than its timeout
                 elif ongoing_execution.started + timedelta(minutes=60) <= current_time:
                     ongoing_execution.status = "timeout"
-                    session.commit()
+                    await session.commit()
                     # re-create the execution and try to get the lock
                     try:
-                        workflow_execution_id = create_workflow_execution(
+                        workflow_execution_id = await create_workflow_execution(
                             workflow.id,
                             workflow.tenant_id,
                             "scheduler",
@@ -479,25 +482,29 @@ def get_last_workflow_workflow_to_alert_executions(
     return latest_workflow_to_alert_executions
 
 
-def get_last_workflow_execution_by_workflow_id(
+async def get_last_workflow_execution_by_workflow_id(
     tenant_id: str, workflow_id: str, status: str = None
 ) -> Optional[WorkflowExecution]:
-    with Session(engine) as session:
-        query = (
-            session.query(WorkflowExecution)
-            .filter(WorkflowExecution.workflow_id == workflow_id)
-            .filter(WorkflowExecution.tenant_id == tenant_id)
-            .filter(WorkflowExecution.started >= datetime.now() - timedelta(days=1))
-            .order_by(WorkflowExecution.started.desc())
+    async with AsyncSession(engine_async) as session:
+        await session.flush()
+        q = select(WorkflowExecution).filter(
+            WorkflowExecution.workflow_id == workflow_id
+        ).filter(WorkflowExecution.tenant_id == tenant_id).filter(
+            WorkflowExecution.started >= datetime.now() - timedelta(days=1)
+        ).order_by(
+            WorkflowExecution.started.desc()
         )
 
-        if status:
-            query = query.filter(WorkflowExecution.status == status)
+        if status is not None:
+            q = q.filter(WorkflowExecution.status == status)
 
-        workflow_execution = query.first()
+        workflow_execution = (
+            (await session.exec(q)).first()
+        )
     return workflow_execution
 
 
+
 def get_workflows_with_last_execution(tenant_id: str) -> List[dict]:
     with Session(engine) as session:
         latest_execution_cte = (
@@ -582,30 +589,32 @@ def get_all_workflows_yamls(tenant_id: str) -> List[str]:
     return workflows
 
 
-def get_workflow(tenant_id: str, workflow_id: str) -> Workflow:
-    with Session(engine) as session:
+async def get_workflow(tenant_id: str, workflow_id: str) -> Workflow:
+    async with AsyncSession(engine_async) as session:
         # if the workflow id is uuid:
         if validators.uuid(workflow_id):
-            workflow = session.exec(
+            workflow = await session.exec(
                 select(Workflow)
                 .where(Workflow.tenant_id == tenant_id)
                 .where(Workflow.id == workflow_id)
                 .where(Workflow.is_deleted == False)
-            ).first()
+            )
+            workflow = workflow.first()
         else:
-            workflow = session.exec(
+            workflow = await session.exec(
                 select(Workflow)
                 .where(Workflow.tenant_id == tenant_id)
                 .where(Workflow.name == workflow_id)
                 .where(Workflow.is_deleted == False)
-            ).first()
+            )
+            workflow = workflow.first()
     if not workflow:
         return None
     return workflow
 
 
-def get_raw_workflow(tenant_id: str, workflow_id: str) -> str:
-    workflow = get_workflow(tenant_id, workflow_id)
+async def get_raw_workflow(tenant_id: str, workflow_id: str) -> str:
+    workflow = await get_workflow(tenant_id, workflow_id)
     if not workflow:
         return None
     return workflow.workflow_raw
@@ -653,33 +662,36 @@ def get_consumer_providers() -> List[Provider]:
     return providers
 
 
-def finish_workflow_execution(tenant_id, workflow_id, execution_id, status, error):
-    with Session(engine) as session:
-        workflow_execution = session.exec(
+async def finish_workflow_execution(tenant_id, workflow_id, execution_id, status, error):
+    async with AsyncSession(engine_async) as session:
+        random_number = random.randint(1, 2147483647 - 1)  # max int
+
+        workflow_execution_old = (await session.exec(
             select(WorkflowExecution).where(WorkflowExecution.id == execution_id)
-        ).first()
-        # some random number to avoid collisions
-        if not workflow_execution:
+        )).first()
+
+        # Perform the update query
+        result = await session.exec(
+            update(WorkflowExecution)
+            .where(WorkflowExecution.id == execution_id)
+            .values(
+                is_running=random_number,
+                status=status,
+                error=error[:255] if error else None,
+                execution_time=(datetime.utcnow() - workflow_execution_old.started).total_seconds()
+            )
+        )
+
+        # Check if the update affected any rows
+        if result.rowcount == 0:
             logger.warning(
-                f"Failed to finish workflow execution {execution_id} for workflow {workflow_id}. Execution not found.",
-                extra={
-                    "tenant_id": tenant_id,
-                    "workflow_id": workflow_id,
-                    "execution_id": execution_id,
-                },
+                f"Failed to finish workflow execution {execution_id} for workflow {workflow_id}. Execution not found."
             )
             raise ValueError("Execution not found")
-        workflow_execution.is_running = random.randint(1, 2147483647 - 1)  # max int
-        workflow_execution.status = status
-        # TODO: we had a bug with the error field, it was too short so some customers may fail over it.
-        #   we need to fix it in the future, create a migration that increases the size of the error field
-        #   and then we can remove the [:511] from here
-        workflow_execution.error = error[:511] if error else None
-        workflow_execution.execution_time = (
-            datetime.utcnow() - workflow_execution.started
-        ).total_seconds()
-        # TODO: logs
-        session.commit()
+
+        # Commit the transaction
+        await session.commit()
+        await session.flush()
 
 
 def get_workflow_executions(
@@ -787,14 +799,14 @@ def delete_workflow_by_provisioned_file(tenant_id, provisioned_file):
             session.commit()
 
 
-def get_workflow_id(tenant_id, workflow_name):
-    with Session(engine) as session:
-        workflow = session.exec(
+async def get_workflow_id(tenant_id, workflow_name):
+    async with AsyncSession(engine_async) as session:
+        workflow = (await session.exec(
             select(Workflow)
             .where(Workflow.tenant_id == tenant_id)
             .where(Workflow.name == workflow_name)
             .where(Workflow.is_deleted == False)
-        ).first()
+        )).first()
 
         if workflow:
             return workflow.id
@@ -1606,16 +1618,16 @@ def update_user_role(tenant_id, username, role):
     return user
 
 
-def save_workflow_results(tenant_id, workflow_execution_id, workflow_results):
-    with Session(engine) as session:
-        workflow_execution = session.exec(
-            select(WorkflowExecution)
+async def save_workflow_results(tenant_id, workflow_execution_id, workflow_results):
+    async with AsyncSession(engine_async) as session:
+        await session.exec(
+            update(WorkflowExecution)
             .where(WorkflowExecution.tenant_id == tenant_id)
             .where(WorkflowExecution.id == workflow_execution_id)
-        ).one()
-
-        workflow_execution.results = workflow_results
-        session.commit()
+            .values(results=workflow_results)
+        )
+        await session.commit()
+        await session.flush()
 
 
 def get_workflow_by_name(tenant_id, workflow_name):
@@ -1629,10 +1641,10 @@ def get_workflow_by_name(tenant_id, workflow_name):
 
         return workflow
 
-
-def get_previous_execution_id(tenant_id, workflow_id, workflow_execution_id):
-    with Session(engine) as session:
-        previous_execution = session.exec(
+        
+async def get_previous_execution_id(tenant_id, workflow_id, workflow_execution_id):
+    async with AsyncSession(engine_async) as session:
+        previous_execution = (await session.exec(
             select(WorkflowExecution)
             .where(WorkflowExecution.tenant_id == tenant_id)
             .where(WorkflowExecution.workflow_id == workflow_id)
@@ -1642,13 +1654,14 @@ def get_previous_execution_id(tenant_id, workflow_id, workflow_execution_id):
             )  # no need to check more than 1 day ago
             .order_by(WorkflowExecution.started.desc())
             .limit(1)
-        ).first()
+        )).first()
         if previous_execution:
             return previous_execution
         else:
             return None
 
 
+
 def create_rule(
     tenant_id,
     name,