py: wait for a failed pipeline to shutdown (feldera#2435)

abhizer · web-flow · commit 60390d35c230 · 2024-09-07T00:58:13.000+05:45
Wait for 15 seconds for a pipeline to shutdown. If the pipeline doesn't
shutdown even after the timeout, resend the shutdown request, and wait a
little longer.

Signed-off-by: Abhinav Gyawali &lt;22275402+abhizer@users.noreply.github.com&gt;
diff --git a/python/feldera/enums.py b/python/feldera/enums.py
@@ -181,3 +181,6 @@ def from_str(value):
             if member.name.lower() == value.lower():
                 return member
         raise ValueError(f"Unknown value '{value}' for enum {PipelineStatus.__name__}")
+
+    def __eq__(self, other):
+        return self.value == other.value
diff --git a/python/feldera/rest/feldera_client.py b/python/feldera/rest/feldera_client.py
@@ -225,18 +225,38 @@ def shutdown_pipeline(self, pipeline_name: str):
             path=f"/pipelines/{pipeline_name}/shutdown",
         )
 
-        while True:
+        start = time.time()
+        timeout = 15
+
+        while time.time() - start < timeout:
             status = self.get_pipeline(pipeline_name).deployment_status
 
             if status == "Shutdown":
-                break
-            elif status == "Failed":
-                raise RuntimeError(f"Failed to shutdown pipeline")
+                return
 
             logging.debug("still shutting down %s, waiting for 100 more milliseconds", pipeline_name)
             time.sleep(0.1)
 
-    # TODO: better name for this method
+        # retry sending shutdown request as the pipline hasn't shutdown yet
+        logging.debug("pipeline %s hasn't shutdown after %s s, retrying", pipeline_name, timeout)
+        self.http.post(
+            path=f"/pipelines/{pipeline_name}/shutdown",
+        )
+
+        start = time.time()
+        timeout = 5
+
+        while time.time() - start < timeout:
+            status = self.get_pipeline(pipeline_name).deployment_status
+
+            if status == "Shutdown":
+                return
+
+            logging.debug("still shutting down %s, waiting for 100 more milliseconds", pipeline_name)
+            time.sleep(0.1)
+
+        raise RuntimeError(f"Failed to shutdown pipeline {pipeline_name}")
+
     def push_to_pipeline(
             self,
             pipeline_name: str,
diff --git a/python/tests/requirements.txt b/python/tests/requirements.txt
@@ -1,2 +1,2 @@
-kafka-python==2.0.2
+kafka-python-ng==2.2.2
 pytest
diff --git a/python/tests/test_pipeline_builder.py b/python/tests/test_pipeline_builder.py
@@ -1,9 +1,11 @@
 import os
+import time
 import unittest
 import pandas as pd
 from kafka import KafkaProducer, KafkaConsumer
 from kafka.admin import KafkaAdminClient, NewTopic
 
+from build.lib.feldera.enums import PipelineStatus
 from feldera import PipelineBuilder, Pipeline
 from tests import TEST_CLIENT
 
@@ -699,6 +701,7 @@ def test_pandas_binary(self):
         got = out.to_dict()
 
         assert expected_data == got
+        pipeline.delete()
 
     def test_pandas_decimal(self):
         from decimal import Decimal
@@ -721,6 +724,7 @@ def test_pandas_decimal(self):
         got = out.to_dict()
 
         assert expected == got
+        pipeline.delete()
 
     def test_pandas_array(self):
         sql = f"""
@@ -742,6 +746,7 @@ def test_pandas_array(self):
             datum.update({"insert_delete": 1})
 
         assert got == data
+        pipeline.delete()
 
     def test_pandas_struct(self):
         sql = f"""
@@ -765,6 +770,7 @@ def test_pandas_struct(self):
             datum.update({"insert_delete": 1})
 
         assert data == got
+        pipeline.delete()
 
     def test_pandas_date_time_timestamp(self):
         from pandas import Timestamp, Timedelta
@@ -785,6 +791,7 @@ def test_pandas_date_time_timestamp(self):
         got = out.to_dict()
 
         assert expected == got
+        pipeline.delete()
 
     def test_pandas_simple(self):
         sql = f"""
@@ -805,6 +812,7 @@ def test_pandas_simple(self):
             datum.update({"insert_delete": 1})
 
         assert data == got
+        pipeline.delete()
 
     def test_pandas_map(self):
         sql = f"""
@@ -823,6 +831,27 @@ def test_pandas_map(self):
         got = out.to_dict()
 
         assert expected == got
+        pipeline.delete()
+
+    def test_failed_pipeline_shutdown(self):
+        sql = f"""
+            CREATE TABLE t0 (c1 TINYINT);
+            CREATE VIEW v0 AS SELECT c1 + 127::TINYINT FROM t0;"""
+
+        pipeline = PipelineBuilder(TEST_CLIENT, name="test_failed_pipeline_shutdown", sql=sql).create_or_replace()
+        pipeline.start()
+        data = [{"c1": 127}]
+        pipeline.input_json("t0", data)
+
+        while True:
+            status = pipeline.status()
+            expected = PipelineStatus.FAILED
+            if status == expected:
+                break
+            time.sleep(1)
+
+        pipeline.shutdown()
+        pipeline.delete()
 
 
 if __name__ == '__main__':

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`		`-kafka-python==2.0.2`
	`1`	`+kafka-python-ng==2.2.2`
`2`	`2`	`pytest`