fix(arrow): drain connection before release to avoid poisoning the pool

Author: zfarrell

hotdata/arrow.py

@@ -50,6 +50,22 @@ def __init__(self, status: str, result_id: str) -> None:
         )
 
 
+def _release_conn(response: Any) -> None:
+    """Return a urllib3 connection to the pool without poisoning it.
+
+    pyarrow stops reading at the Arrow end-of-stream marker, which can leave
+    bytes unconsumed at the HTTP layer (e.g. the terminating zero-length chunk
+    of a chunked response). ``release_conn`` on a partially-read connection
+    hands a tainted socket back to the pool — the next request to reuse it
+    fails with ``http.client.ResponseNotReady`` ("Request-sent"). Draining any
+    unread bytes first makes the connection safe to reuse.
+    """
+    try:
+        response.drain_conn()
+    finally:
+        response.release_conn()
+
+
 def _import_pyarrow() -> Any:
     try:
         import pyarrow.ipc as ipc  # type: ignore[import-untyped]
@@ -95,7 +111,7 @@ def get_result_arrow(
         try:
             return ipc.open_stream(response).read_all()
         finally:
-            response.release_conn()
+            _release_conn(response)
 
     @contextmanager
     def stream_result_arrow(
@@ -108,8 +124,10 @@ def stream_result_arrow(
     ) -> Iterator["pa.RecordBatchStreamReader"]:
         """Yield a :class:`pyarrow.RecordBatchStreamReader` for a ready result.
 
-        The HTTP connection is released when the context exits. Iterate the
-        reader to consume :class:`pyarrow.RecordBatch` messages, or call
+        The HTTP connection is drained and released when the context exits, so
+        exiting early (before the reader is exhausted) reads and discards the
+        remaining body to keep the connection reusable. Iterate the reader to
+        consume :class:`pyarrow.RecordBatch` messages, or call
         ``reader.read_all()`` for a full :class:`pyarrow.Table`.
 
         Example::
@@ -127,7 +145,7 @@ def stream_result_arrow(
         try:
             yield ipc.open_stream(response)
         finally:
-            response.release_conn()
+            _release_conn(response)
 
     def _call_arrow(
         self,

tests/test_arrow.py

@@ -45,7 +45,8 @@ class _FakeUrllib3Response(io.RawIOBase):
 
     pyarrow.ipc.open_stream wants a real file-like object (it checks
     ``closed`` and ``readable()``); the SDK's RESTResponse needs ``status``,
-    ``reason``, ``data``, and ``headers``. release_conn is recorded.
+    ``reason``, ``data``, and ``headers``. drain_conn / release_conn are
+    recorded so tests can assert the connection is drained before reuse.
     """
 
     def __init__(self, status: int, body: bytes, headers: Dict[str, str]):
@@ -56,6 +57,7 @@ def __init__(self, status: int, body: bytes, headers: Dict[str, str]):
         self._buf = io.BytesIO(body)
         self.headers = headers
         self.release_conn_called = False
+        self.drain_conn_called = False
 
     @property
     def data(self) -> bytes:
@@ -71,6 +73,9 @@ def read(self, amt: Optional[int] = -1) -> bytes:
     def readable(self) -> bool:
         return True
 
+    def drain_conn(self) -> None:
+        self.drain_conn_called = True
+
     def release_conn(self) -> None:
         self.release_conn_called = True
 
@@ -146,6 +151,9 @@ def test_get_result_arrow_returns_table(monkeypatch: pytest.MonkeyPatch) -> None
     got = results.get_result_arrow("res_123")
 
     assert got.equals(table)
+    # Connection is drained (not just released) so a partially-read body can't
+    # poison the pool and break the next request with ResponseNotReady.
+    assert fake.drain_conn_called
     assert fake.release_conn_called
 
     # Single request was made.
@@ -197,7 +205,8 @@ def test_stream_result_arrow_yields_reader(
         roundtrip = pa.Table.from_batches(batches, schema=reader.schema)
         assert roundtrip.equals(table)
 
-    # Connection is released after the context exits.
+    # Connection is drained and released after the context exits.
+    assert fake.drain_conn_called
     assert fake.release_conn_called