From ee49e7c1bc430ec558a5f03815c4f3fc8521a3ea Mon Sep 17 00:00:00 2001
From: Shahar Glazner <shaharglazner@gmail.com>
Date: Mon, 6 Jan 2025 15:56:56 +0200
Subject: [PATCH] fix(api): workflowdb handler (#2988)

---
 keep/api/logging.py              | 31 ++++++++++---------------------
 pyproject.toml                   |  2 +-
 tests/test_workflow_execution.py |  4 ++--
 3 files changed, 13 insertions(+), 24 deletions(-)

diff --git a/keep/api/logging.py b/keep/api/logging.py
index 88be1bf08..36c13d85f 100644
--- a/keep/api/logging.py
+++ b/keep/api/logging.py
@@ -11,7 +11,6 @@
 import urllib3
 from sqlmodel import Session
 
-from keep.api.consts import RUNNING_IN_CLOUD_RUN
 from keep.api.core.db import get_session, push_logs_to_db
 from keep.api.models.db.provider import ProviderExecutionLog
 
@@ -140,26 +139,16 @@ def process(self, msg, kwargs):
 
     def dump(self):
         self.logger.info("Dumping workflow logs")
-        # TODO - this is a POC level code.
-        # TODO - we should:
-        # TODO - 1. find the right handler to push the logs to the DB
-        # TODO - 2. find a better way to push the logs async (maybe another service)
-        workflow_db_handler = next(
-            iter(
-                [
-                    handler
-                    for handler in (
-                        # tb: for some reason, when running in cloud run, the handler is nested in another handler
-                        #   this needs to be handled in a better way
-                        self.logger.parent.parent.handlers
-                        if RUNNING_IN_CLOUD_RUN
-                        else self.logger.parent.handlers
-                    )
-                    if isinstance(handler, WorkflowDBHandler)
-                ]
-            ),
-            None,
-        )
+        root_logger = logging.getLogger()
+        handlers = root_logger.handlers
+        workflow_db_handler = None
+
+        for handler in handlers:
+            # should be always the second
+            if isinstance(handler, WorkflowDBHandler):
+                workflow_db_handler = handler
+                break
+
         if workflow_db_handler:
             self.logger.info("Pushing logs to DB")
             workflow_db_handler.push_logs_to_db()
diff --git a/pyproject.toml b/pyproject.toml
index f8d595770..47fc30ea5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "keep"
-version = "0.33.9"
+version = "0.34.0"
 description = "Alerting. for developers, by developers."
 authors = ["Keep Alerting LTD"]
 packages = [{include = "keep"}]
diff --git a/tests/test_workflow_execution.py b/tests/test_workflow_execution.py
index 0272f30cd..f662a4bcf 100644
--- a/tests/test_workflow_execution.py
+++ b/tests/test_workflow_execution.py
@@ -931,7 +931,7 @@ def test_workflow_execution_logs(
     with patch(
         "keep.contextmanager.contextmanager.WorkflowLoggerAdapter",
         side_effect=fake_workflow_adapter,
-    ), patch("keep.api.logging.RUNNING_IN_CLOUD_RUN", value=True):
+    ):
         base_time = datetime.now(tz=pytz.utc)
 
         # Create alerts with specified statuses and timestamps
@@ -1009,7 +1009,7 @@ def test_workflow_execution_logs_log_level_debug_console_provider(
         with patch(
             "keep.contextmanager.contextmanager.WorkflowLoggerAdapter",
             side_effect=fake_workflow_adapter,
-        ), patch("keep.api.logging.RUNNING_IN_CLOUD_RUN", value=True):
+        ):
             base_time = datetime.now(tz=pytz.utc)
 
             # Create alerts with specified statuses and timestamps