[py] bug: fix http request error validation for streaming responses

abhizer · abhizer · commit 6268d6761487 · 2024-10-02T13:17:24.000Z
Instead of entirely skipping validation checks for streaming responses,
raise any HTTP errors first, then return the response.

Signed-off-by: Abhinav Gyawali &lt;22275402+abhizer@users.noreply.github.com&gt;
diff --git a/python/feldera/pipeline.py b/python/feldera/pipeline.py
@@ -353,11 +353,6 @@ def query(self, query: str) -> Generator[Mapping[str, Any], None, None]:
         :return: A generator that yields the rows of the result as Python dictionaries.
         """
 
-        if self.status() not in [
-            PipelineStatus.RUNNING,
-            PipelineStatus.PAUSED,
-        ]:
-            raise RuntimeError("Pipeline must be running or paused to run a query")
         return self.client.query_as_json(self.name, query)
 
     def query_parquet(self, query: str, path: str):
@@ -369,11 +364,6 @@ def query_parquet(self, query: str, path: str):
         :param path: The path of the parquet file.
         """
 
-        if self.status() not in [
-            PipelineStatus.RUNNING,
-            PipelineStatus.PAUSED,
-        ]:
-            raise RuntimeError("Pipeline must be running or paused to run a query")
         self.client.query_as_parquet(self.name, query, path)
 
     def query_tabular(self, query: str) -> Generator[str, None, None]:
@@ -384,9 +374,4 @@ def query_tabular(self, query: str) -> Generator[str, None, None]:
         :return: A generator that yields a string representing the query result in a human-readable, tabular format.
         """
 
-        if self.status() not in [
-            PipelineStatus.RUNNING,
-            PipelineStatus.PAUSED,
-        ]:
-            raise RuntimeError("Pipeline must be running or paused to run a query")
         return self.client.query_as_text(self.name, query)
diff --git a/python/feldera/rest/_httprequests.py b/python/feldera/rest/_httprequests.py
@@ -83,14 +83,7 @@ def send_request(
                     stream=stream,
                 )
 
-            if stream:
-                return request
-            if request.headers.get("content-type") == "text/plain":
-                return request.text
-            elif request.headers.get("content-type") == "application/octet-stream":
-                return request.content
-
-            resp = self.__validate(request)
+            resp = self.__validate(request, stream=stream)
             logging.debug("got response: %s", str(resp))
             return resp
 
@@ -164,9 +157,17 @@ def __to_json(request: requests.Response) -> Any:
         return request.json()
 
     @staticmethod
-    def __validate(request: requests.Response) -> Any:
+    def __validate(request: requests.Response, stream=False) -> Any:
         try:
             request.raise_for_status()
+
+            if stream:
+                return request
+            if request.headers.get("content-type") == "text/plain":
+                return request.text
+            elif request.headers.get("content-type") == "application/octet-stream":
+                return request.content
+
             resp = HttpRequests.__to_json(request)
             return resp
         except requests.exceptions.HTTPError as err:
diff --git a/python/tests/test_pipeline.py b/python/tests/test_pipeline.py
@@ -161,7 +161,7 @@ def test_adhoc_query_text(self):
         TEST_CLIENT.start_pipeline(name)
 
         TEST_CLIENT.push_to_pipeline(name, "tbl", "csv", data)
-        resp = TEST_CLIENT.query(pipeline.name, "SELECT * FROM tbl", "text")
+        resp = TEST_CLIENT.query_as_text(pipeline.name, "SELECT * FROM tbl")
         expected = """+----+
 | id |
 +----+
@@ -214,7 +214,7 @@ def test_adhoc_query_json(self):
         TEST_CLIENT.start_pipeline(name)
 
         TEST_CLIENT.push_to_pipeline(name, "tbl", "csv", data)
-        resp = TEST_CLIENT.query(pipeline.name, "SELECT * FROM tbl", "json")
+        resp = TEST_CLIENT.query_as_json(pipeline.name, "SELECT * FROM tbl")
         expected = [{"id": 2}, {"id": 1}]
         got = list(resp)
 
diff --git a/python/uv.lock b/python/uv.lock
