From 399f83a75b84e4be94d2171f730f4cc92ab1be8a Mon Sep 17 00:00:00 2001
From: Vladimir Filonov <vladimir@filonov.dev>
Date: Wed, 11 Dec 2024 16:28:35 +0400
Subject: [PATCH] Make __update_client_on_incident_change and __run_workflows
 public

---
 keep/api/bl/incidents_bl.py | 42 ++++++++++---------------------------
 keep/api/core/db.py         |  4 ++--
 2 files changed, 13 insertions(+), 33 deletions(-)

diff --git a/keep/api/bl/incidents_bl.py b/keep/api/bl/incidents_bl.py
index 5bc3d9216..30d11d097 100644
--- a/keep/api/bl/incidents_bl.py
+++ b/keep/api/bl/incidents_bl.py
@@ -80,12 +80,12 @@ def create_incident(
             "Incident DTO created",
             extra={"incident_id": new_incident_dto.id, "tenant_id": self.tenant_id},
         )
-        self.__update_client_on_incident_change()
+        self.update_client_on_incident_change()
         self.logger.info(
             "Client updated on incident change",
             extra={"incident_id": new_incident_dto.id, "tenant_id": self.tenant_id},
         )
-        self.__run_workflows(new_incident_dto, "created")
+        self.send_workflow_event(new_incident_dto, "created")
         self.logger.info(
             "Workflows run on incident",
             extra={"incident_id": new_incident_dto.id, "tenant_id": self.tenant_id},
@@ -113,13 +113,13 @@ async def add_alerts_to_incident(
             "Alerts pushed to elastic",
             extra={"incident_id": incident_id, "alert_fingerprints": alert_fingerprints},
         )
-        self.__update_client_on_incident_change(incident_id)
+        self.update_client_on_incident_change(incident_id)
         self.logger.info(
             "Client updated on incident change",
             extra={"incident_id": incident_id, "alert_fingerprints": alert_fingerprints},
         )
         incident_dto = IncidentDto.from_db_incident(incident)
-        self.__run_workflows(incident_dto, "updated")
+        self.send_workflow_event(incident_dto, "updated")
         self.logger.info(
             "Workflows run on incident",
             extra={"incident_id": incident_id, "alert_fingerprints": alert_fingerprints},
@@ -146,7 +146,7 @@ def __update_elastic(self, incident_id: UUID, alert_fingerprints: List[str]):
         except Exception:
             self.logger.exception("Failed to push alert to elasticsearch")
 
-    def __update_client_on_incident_change(self, incident_id: Optional[UUID] = None):
+    def update_client_on_incident_change(self, incident_id: Optional[UUID] = None):
         if self.pusher_client is not None:
             self.logger.info(
                 "Pushing incident change to client",
@@ -162,7 +162,7 @@ def __update_client_on_incident_change(self, incident_id: Optional[UUID] = None)
                 extra={"incident_id": incident_id, "tenant_id": self.tenant_id},
             )
 
-    def __run_workflows(self, incident_dto: IncidentDto, action: str):
+    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)
@@ -238,17 +238,8 @@ def delete_incident(self, incident_id: UUID) -> None:
         )
         if not deleted:
             raise HTTPException(status_code=404, detail="Incident not found")
-        self.__update_client_on_incident_change()
-        try:
-            workflow_manager = WorkflowManager.get_instance()
-            self.logger.info("Adding incident to the workflow manager queue")
-            workflow_manager.insert_incident(self.tenant_id, incident_dto, "deleted")
-            self.logger.info("Added incident to the workflow manager queue")
-        except Exception:
-            self.logger.exception(
-                "Failed to run workflows based on incident",
-                extra={"incident_id": incident_dto.id, "tenant_id": self.tenant_id},
-            )
+        self.update_client_on_incident_change()
+        self.send_workflow_event(incident_dto, "deleted")
 
     def update_incident(
         self,
@@ -269,17 +260,6 @@ def update_incident(
         if not incident:
             raise HTTPException(status_code=404, detail="Incident not found")
 
-        new_incident_dto = IncidentDto.from_db_incident(incident)
-        try:
-            workflow_manager = WorkflowManager.get_instance()
-            self.logger.info("Adding incident to the workflow manager queue")
-            workflow_manager.insert_incident(
-                self.tenant_id, new_incident_dto, "updated"
-            )
-            self.logger.info("Added incident to the workflow manager queue")
-        except Exception:
-            self.logger.exception(
-                "Failed to run workflows based on incident",
-                extra={"incident_id": new_incident_dto.id, "tenant_id": self.tenant_id},
-            )
-        return new_incident_dto
+        incident_dto = IncidentDto.from_db_incident(incident)
+        self.send_workflow_event(incident_dto, "updated")
+        return incident_dto
diff --git a/keep/api/core/db.py b/keep/api/core/db.py
index 59631d4b7..8990e7bcb 100644
--- a/keep/api/core/db.py
+++ b/keep/api/core/db.py
@@ -3291,7 +3291,7 @@ def create_incident_from_dict(
 
 def update_incident_from_dto_by_id(
     tenant_id: str,
-    incident_id: str,
+    incident_id: str | UUID,
     updated_incident_dto: IncidentDtoIn | IncidentDto,
     generated_by_ai: bool = False,
 ) -> Optional[Incident]:
@@ -3650,7 +3650,7 @@ def add_alerts_to_incident(
             return incident
 
 
-def get_incident_unique_fingerprint_count(tenant_id: str, incident_id: str) -> int:
+def get_incident_unique_fingerprint_count(tenant_id: str, incident_id: str | UUID) -> int:
     with Session(engine) as session:
         return session.execute(
             select(func.count(1))