py: fix race condition in listen()

abhizer · abhizer · commit 4a103bd49e40 · 2025-10-23T10:14:13.000Z
Fixes: #4898 Signed-off-by: Abhinav Gyawali <22275402+abhizer@users.noreply.github.com>
diff --git a/python/feldera/_callback_runner.py b/python/feldera/_callback_runner.py
@@ -1,10 +1,13 @@
-from threading import Thread
-from typing import Callable, Optional
+from threading import Thread, Event
+from typing import Callable, List, Optional, Mapping, Any
 
 import pandas as pd
 from feldera import FelderaClient
 from feldera._helpers import dataframe_from_response
 from feldera.enums import PipelineFieldSelector
+from feldera.rest.sql_table import SQLTable
+from feldera.rest.sql_view import SQLView
+from feldera.rest.pipeline import Pipeline
 
 
 class CallbackRunner(Thread):
@@ -15,6 +18,7 @@ def __init__(
         view_name: str,
         callback: Callable[[pd.DataFrame, int], None],
         exception_callback: Callable[[BaseException], None],
+        event: Event,
     ):
         super().__init__()
         self.daemon = True
@@ -23,7 +27,32 @@ def __init__(
         self.view_name: str = view_name
         self.callback: Callable[[pd.DataFrame, int], None] = callback
         self.exception_callback: Callable[[BaseException], None] = exception_callback
-        self.schema: Optional[dict] = None
+        self.event: Event = event
+
+        self.pipeline: Pipeline = self.client.get_pipeline(
+            self.pipeline_name, PipelineFieldSelector.ALL
+        )
+
+        view_schema = None
+
+        schemas: List[SQLTable | SQLView] = self.pipeline.tables + self.pipeline.views
+        for schema in schemas:
+            if schema.name == self.view_name:
+                view_schema = schema
+                break
+
+        if view_schema is None:
+            raise ValueError(
+                f"Table or View {self.view_name} not found in the pipeline schema."
+            )
+
+        self.schema: SQLTable | SQLView = view_schema
+
+    def process_chunk(self, chunk: Mapping[str, Any]):
+        data: Optional[list[Mapping[str, Any]]] = chunk.get("json_data")
+        seq_no: Optional[int] = chunk.get("sequence_number")
+        if data is not None and seq_no is not None:
+            self.callback(dataframe_from_response([data], self.schema.fields), seq_no)
 
     def run(self):
         """
@@ -33,21 +62,6 @@ def run(self):
         """
 
         try:
-            pipeline = self.client.get_pipeline(
-                self.pipeline_name, PipelineFieldSelector.ALL
-            )
-
-            schemas = pipeline.tables + pipeline.views
-            for schema in schemas:
-                if schema.name == self.view_name:
-                    self.schema = schema
-                    break
-
-            if self.schema is None:
-                raise ValueError(
-                    f"Table or View {self.view_name} not found in the pipeline schema."
-                )
-
             gen_obj = self.client.listen_to_pipeline(
                 self.pipeline_name,
                 self.view_name,
@@ -57,13 +71,16 @@ def run(self):
 
             iterator = gen_obj()
 
+            # Trigger the HTTP call
+            chunk = next(iterator)
+
+            # Unblock the main thread
+            self.event.set()
+
+            self.process_chunk(chunk)
+
             for chunk in iterator:
-                chunk: dict = chunk
-                data: Optional[list[dict]] = chunk.get("json_data")
-                seq_no: Optional[int] = chunk.get("sequence_number")
-                if data is not None and seq_no is not None:
-                    self.callback(
-                        dataframe_from_response([data], self.schema.fields), seq_no
-                    )
+                self.process_chunk(chunk)
+
         except BaseException as e:
             self.exception_callback(e)
diff --git a/python/feldera/_helpers.py b/python/feldera/_helpers.py
@@ -2,6 +2,7 @@
 
 import pandas as pd
 from decimal import Decimal
+from typing import Mapping, Any
 
 
 def sql_type_to_pandas_type(sql_type: str):
@@ -60,7 +61,9 @@ def ensure_dataframe_has_columns(df: pd.DataFrame):
         )
 
 
-def dataframe_from_response(buffer: list[list[dict]], fields: list[dict]):
+def dataframe_from_response(
+    buffer: list[list[Mapping[str, Any]]], fields: list[Mapping[str, Any]]
+):
     """
     Converts the response from Feldera to a pandas DataFrame.
     """
diff --git a/python/feldera/output_handler.py b/python/feldera/output_handler.py
@@ -1,6 +1,7 @@
 import pandas as pd
 
 from typing import Optional
+from threading import Event
 
 from feldera import FelderaClient
 from feldera._callback_runner import CallbackRunner
@@ -23,6 +24,7 @@ def __init__(
         self.view_name: str = view_name
         self.buffer: list[pd.DataFrame] = []
         self.exception: Optional[BaseException] = None
+        self.event = Event()
 
         # the callback that is passed to the `CallbackRunner`
         def callback(df: pd.DataFrame, _: int):
@@ -39,6 +41,7 @@ def exception_callback(exception: BaseException):
             self.view_name,
             callback,
             exception_callback,
+            self.event,
         )
 
     def start(self):
@@ -47,6 +50,7 @@ def start(self):
         """
 
         self.handler.start()
+        _ = self.event.wait()
 
     def to_pandas(self, clear_buffer: bool = True):
         """
diff --git a/python/feldera/pipeline.py b/python/feldera/pipeline.py
@@ -7,6 +7,7 @@
 from uuid import UUID
 
 from typing import List, Dict, Callable, Optional, Generator, Mapping, Any
+from threading import Event
 from collections import deque
 
 from feldera.rest.errors import FelderaAPIError
@@ -294,10 +295,12 @@ def foreach_chunk(
         if self.status() not in [PipelineStatus.RUNNING, PipelineStatus.PAUSED]:
             raise RuntimeError("Pipeline must be running or paused to listen to output")
 
+        event = Event()
         handler = CallbackRunner(
-            self.client, self.name, view_name, callback, lambda exception: None
+            self.client, self.name, view_name, callback, lambda exception: None, event
         )
         handler.start()
+        event.wait()
 
     def wait_for_completion(
         self, force_stop: bool = False, timeout_s: float | None = None
