Detect in-band server exceptions on Arrow streaming query paths

CHANGELOG.md

@@ -32,6 +32,7 @@
 - DB-API `Cursor.description` now reports the result type's top-level nullability instead of hardcoding `null_ok=True`, including the implicit null values supported by `Variant`, `Dynamic`, and `SimpleAggregateFunction` over a nullable element type. Existing `type_code` values are unchanged, and types whose nullability is unknown report `None`. The empty-result metadata probe also recognizes leading ClickHouse comments, including nested block comments, and is best effort, so a failed probe leaves `description` empty instead of raising after the original query succeeded. Closes [#902](https://github.com/ClickHouse/clickhouse-connect/issues/902), [#907](https://github.com/ClickHouse/clickhouse-connect/issues/907), and [#909](https://github.com/ClickHouse/clickhouse-connect/issues/909).
 - Compound values stored in JSON shared data, such as arrays of objects, heterogeneous arrays, and nested arrays, are now decoded to Python objects instead of being returned as raw bytes. `Date`, `DateTime`, and `DateTime64` values in shared data, both as scalars and inside arrays, now decode as well. Closes [#897](https://github.com/ClickHouse/clickhouse-connect/issues/897).
 - `AsyncClient` no longer tears down the aiohttp response from the parser's executor thread when a query fails mid-stream. The synchronous cleanup cancelled the producer task and closed the response directly, which raced with the event loop handling the server's connection abort and could surface an `AttributeError` from asyncio's SSL shutdown on TLS connections instead of the real `StreamFailureError`. Cleanup is now scheduled onto the event loop with `call_soon_threadsafe`.
+- Arrow streaming query methods now report a mid-stream server error instead of a truncated stream. When a query fails after ClickHouse has already committed a `200 OK` and started streaming rows, the server appends its exception to the response body and drops the connection. The native read path detects that and raises a clean error, but the Arrow paths handed the raw bytes to `pyarrow` and never consulted it, so the failure surfaced as a `pyarrow` parse error or a raw transport error such as `urllib3.ProtocolError` or `aiohttp.ClientPayloadError`, with the actual server message lost. `Client.query_arrow_stream`, `Client.query_df_arrow_stream`, `AsyncClient.query_arrow`, `AsyncClient.query_arrow_stream`, and `AsyncClient.query_df_arrow_stream` now raise `StreamFailureError` carrying the ClickHouse error, matching the native streaming path and honoring `show_clickhouse_errors`. Detection reads a bounded tail of the response, so a successful stream is unaffected. Both the tagged exception block used by newer servers and the untagged error text used by older ones are recognized; because an Arrow payload is binary, the untagged fallback requires the `Code: ` marker rather than treating any decodable tail as an error. `Client.query_arrow` is buffered and already reported these errors. Closes [#913](https://github.com/ClickHouse/clickhouse-connect/issues/913).
 
 ## 1.6.0, 2026-07-23
 

clickhouse_connect/driver/asyncclient.py

@@ -877,7 +877,7 @@ async def query_arrow(
         streaming_source = await start_streaming_response(response, encoding=encoding, exception_tag=exception_tag)
 
         def parse_arrow_stream():
-            file_adapter = StreamingFileAdapter(streaming_source)
+            file_adapter = StreamingFileAdapter(streaming_source, self.show_clickhouse_errors)
             reader = options.arrow.ipc.open_stream(file_adapter)
             table = reader.read_all()
             return _apply_arrow_tz_policy(table, self.tz_mode)
@@ -925,7 +925,7 @@ def _arrow_batch_stream(self, streaming_source: StreamingResponseSource, convert
         queued = QueuedStreamSource(streaming_source)
 
         def batches():
-            file_adapter = StreamingFileAdapter(streaming_source)
+            file_adapter = StreamingFileAdapter(streaming_source, self.show_clickhouse_errors)
             reader = options.arrow.ipc.open_stream(file_adapter)
             for batch in reader:
                 yield converter(batch)

clickhouse_connect/driver/client.py

@@ -18,6 +18,7 @@
 from clickhouse_connect.datatypes.base import ClickHouseType
 from clickhouse_connect.datatypes.registry import get_from_name
 from clickhouse_connect.driver import options, tzutil
+from clickhouse_connect.driver._backend.httpcommon import ex_tag_header
 from clickhouse_connect.driver._backend.models import ClientConfig, QueryRuntime
 from clickhouse_connect.driver._backend.operations import CommandOp, Operation, QueryOp, RawQueryOp
 from clickhouse_connect.driver._backend.orchestration import (
@@ -41,6 +42,7 @@
     ProgrammingError,
 )
 from clickhouse_connect.driver.external import ExternalData
+from clickhouse_connect.driver.httputil import ResponseSource
 from clickhouse_connect.driver.insert import InsertContext
 from clickhouse_connect.driver.models import SettingDef, SettingStatus, setting_status
 from clickhouse_connect.driver.options import (
@@ -60,6 +62,7 @@
     to_arrow,
     to_arrow_batches,
 )
+from clickhouse_connect.driver.streaming import StreamingFileAdapter
 from clickhouse_connect.driver.summary import QuerySummary
 from clickhouse_connect.driver.types import Closable
 
@@ -854,19 +857,16 @@ def query_arrow_stream(
         check_arrow()
         self._add_integration_tag("arrow")
         settings = self._update_arrow_settings(settings, use_strings)
-        return to_arrow_batches(
-            cast(
-                io.IOBase,
-                self.raw_stream(
-                    query,
-                    parameters,
-                    settings,
-                    fmt="ArrowStream",
-                    external_data=external_data,
-                    transport_settings=transport_settings,
-                ),
-            )
+        stream = self.raw_stream(
+            query,
+            parameters,
+            settings,
+            fmt="ArrowStream",
+            external_data=external_data,
+            transport_settings=transport_settings,
         )
+        buffer, source = self._arrow_stream_source(stream)
+        return to_arrow_batches(buffer, source)
 
     def query_df_arrow(
         self,
@@ -968,13 +968,31 @@ def converter(table: pyarrow.Table) -> polars.DataFrame:  # type: ignore[misc]
         raw_stream = self.raw_stream(
             query, parameters, settings, fmt="ArrowStream", external_data=external_data, transport_settings=transport_settings
         )
-        reader = options.arrow.ipc.open_stream(raw_stream)
+        buffer, source = self._arrow_stream_source(raw_stream)
+        reader = options.arrow.ipc.open_stream(buffer)
 
         def df_generator():
             for batch in reader:
                 yield converter(batch)
 
-        return StreamContext(cast(Closable, raw_stream), df_generator())
+        return StreamContext(source, df_generator())
+
+    def _arrow_stream_source(self, stream: Any) -> tuple[io.IOBase, Closable]:
+        """Prepare a raw ArrowStream response for PyArrow, returning the file to parse and the
+        object that owns the underlying connection.
+
+        Over HTTP the server can report a failure after it has already committed a 200 and started
+        streaming, by appending an exception block to the body. Reading the response directly loses
+        that block when the connection drops, so the bytes go through ResponseSource, which keeps
+        the trailing chunk, and then through the file adapter, which turns it into the same error
+        the native path raises. Backends that return a plain file object report their own errors and
+        are passed through untouched.
+        """
+        headers = getattr(stream, "headers", None)
+        if headers is None:
+            return cast(io.IOBase, stream), cast(Closable, stream)
+        source = ResponseSource(stream, exception_tag=headers.get(ex_tag_header))
+        return cast(io.IOBase, StreamingFileAdapter(source, self.show_clickhouse_errors)), cast(Closable, source)
 
     def _update_arrow_settings(self, settings: dict[str, Any] | None, use_strings: bool | None) -> dict[str, Any]:
         settings = dict_copy(settings)

clickhouse_connect/driver/query.py

@@ -439,10 +439,12 @@ def to_arrow(content: bytes):
     return reader.read_all()
 
 
-def to_arrow_batches(buffer: IOBase) -> StreamContext:
+def to_arrow_batches(buffer: IOBase, source: Closable | None = None) -> StreamContext:
     pyarrow = check_arrow()
     reader = pyarrow.ipc.open_stream(buffer)
-    return StreamContext(buffer, reader)
+    #  When the buffer is an adapter over a separate response source, that source owns the
+    #  connection and is what has to be closed when the stream context exits.
+    return StreamContext(source if source is not None else buffer, reader)
 
 
 def arrow_buffer(table, compression: str | None = None) -> tuple[Sequence[str], bytes | BinaryIO]:

clickhouse_connect/driver/streaming.py

@@ -2,13 +2,16 @@
 import logging
 import threading
 import zlib
+from collections import deque
 from collections.abc import Callable, Iterable, Iterator
 
 import lz4.frame
 
 from clickhouse_connect.driver.asyncqueue import EOF_SENTINEL, AsyncSyncQueue
+from clickhouse_connect.driver.common import ShowClickHouseErrors
 from clickhouse_connect.driver.compression import _zstd_decompressor, available_compression
-from clickhouse_connect.driver.exceptions import OperationalError
+from clickhouse_connect.driver.exceptions import OperationalError, StreamFailureError
+from clickhouse_connect.driver.transform import extract_stream_error, format_stream_error
 from clickhouse_connect.driver.types import Closable
 
 logger = logging.getLogger(__name__)
@@ -283,15 +286,75 @@ def close(self):
         self.source.close()
 
 
+#  ClickHouse appends an in-band exception block at the very end of the response body, so only a
+#  bounded tail of the stream can ever contain it. 64KB comfortably holds a server error message
+#  plus its tag markers without retaining an unbounded amount of a large result.
+_EXCEPTION_TAIL_SIZE = 1 << 16
+
+
+class _ExceptionTail:
+    """The trailing bytes of a response body, bounded by ``_EXCEPTION_TAIL_SIZE``.
+
+    Chunks are retained by reference and only joined when an error is actually suspected, so a
+    successful stream pays one deque append per chunk and never copies its payload.
+    """
+
+    __slots__ = ("_chunks", "_size")
+
+    def __init__(self):
+        self._chunks: deque[bytes] = deque()
+        self._size = 0
+
+    def add(self, chunk: bytes) -> None:
+        self._chunks.append(chunk)
+        self._size += len(chunk)
+        #  Keep the oldest chunk only while the newer ones cannot cover the tail on their own
+        while self._chunks and self._size - len(self._chunks[0]) >= _EXCEPTION_TAIL_SIZE:
+            self._size -= len(self._chunks.popleft())
+
+    @property
+    def message(self) -> bytes:
+        if len(self._chunks) == 1:
+            return self._chunks[0]
+        return b"".join(self._chunks)
+
+
 class StreamingFileAdapter:
-    """File-like adapter for PyArrow streaming."""
+    """File-like adapter for PyArrow streaming.
 
-    def __init__(self, streaming_source):
+    PyArrow parses the response bytes directly, so a mid-stream server error would otherwise reach
+    the caller as a truncated Arrow stream or a raw transport error. This adapter keeps a bounded
+    tail of the body and consults it when the stream ends, whether that end is a clean EOF or a
+    dropped connection, so the caller sees the same ClickHouse error the native path reports.
+    """
+
+    def __init__(self, streaming_source, show_clickhouse_errors: ShowClickHouseErrors = True):
         self.streaming_source = streaming_source
         self.gen = streaming_source.gen
         self.buffer = b""
         self.closed = False
         self.eof = False
+        self._show_clickhouse_errors = show_clickhouse_errors
+        self._exception_tag = getattr(streaming_source, "exception_tag", None)
+        self._tail = _ExceptionTail()
+
+    def _server_error(self, tagged_only: bool) -> str | None:
+        """The ClickHouse error carried in the response tail, if there is one."""
+        tail = self._tail.message
+        #  An Arrow payload is binary, so unlike the native reader this cannot treat an arbitrary
+        #  decodable tail as an error message. Without a tagged block, the `Code: ` marker is the
+        #  only trustworthy evidence that the tail holds an error at all.
+        if not tagged_only and b"Code: " not in tail[-1024:]:
+            tagged_only = True
+        error_msg = extract_stream_error(
+            tail,
+            self._exception_tag,
+            self._show_clickhouse_errors,
+            tagged_only=tagged_only,
+        )
+        if error_msg is None:
+            return None
+        return format_stream_error(error_msg, self._show_clickhouse_errors)
 
     def read(self, size: int = -1) -> bytes:
         """Read up to size bytes from stream"""
@@ -311,14 +374,32 @@ def read(self, size: int = -1) -> bytes:
             try:
                 chunk = next(self.gen)
                 if chunk:
+                    self._tail.add(chunk)
                     chunks.append(chunk)
                     current_len += len(chunk)
                 else:
                     self.eof = True
                     break
             except StopIteration:
                 self.eof = True
+                #  The body ended cleanly, so a truncated-stream guess would be unsafe here: trust
+                #  only a tagged block, which the server writes exclusively to report a failure.
+                error_msg = self._server_error(tagged_only=True)
+                if error_msg:
+                    raise StreamFailureError(error_msg) from None
                 break
+            except Exception as ex:
+                self.eof = True
+                #  A read failure partway through the stream: OperationalError from the sync reader,
+                #  ClientPayloadError from aiohttp. ClickHouse may have written the real error into
+                #  the response body before the connection dropped, so prefer that over the
+                #  transport error, exactly as the native path does.
+                if isinstance(ex, OperationalError) or ex.__class__.__name__ == "ClientPayloadError":
+                    error_msg = self._server_error(tagged_only=False)
+                    if error_msg:
+                        raise StreamFailureError(error_msg) from None
+                    raise StreamFailureError("Stream failed during read (connection closed by server)") from ex
+                raise
 
         full_data = b"".join(chunks)
 

clickhouse_connect/driver/transform.py

@@ -1,7 +1,7 @@
 import logging
 
 from clickhouse_connect.datatypes import registry
-from clickhouse_connect.driver.common import write_leb128
+from clickhouse_connect.driver.common import ShowClickHouseErrors, write_leb128
 from clickhouse_connect.driver.compression import get_compressor
 from clickhouse_connect.driver.exceptions import (
     GENERIC_CLICKHOUSE_ERROR,
@@ -29,26 +29,13 @@ def parse_response(source: ByteSource, context: QueryContext = _EMPTY_CTX) -> Nu
         renamer = context.column_renamer
         show_clickhouse_errors = context.show_clickhouse_errors
 
-        def format_stream_error(error_msg: str) -> str:
-            if show_clickhouse_errors is False:
-                return GENERIC_CLICKHOUSE_ERROR
-            if show_clickhouse_errors == "scrub":
-                return scrub_error_details(error_msg)
-            return error_msg
-
         def extract_source_error(tagged_only: bool = False) -> str | None:
-            if not source.last_message:
-                return None
-            exception_tag = getattr(source, "exception_tag", None)
-            if exception_tag:
-                error_msg = extract_exception_with_tag(source.last_message, exception_tag)
-                if error_msg:
-                    return error_msg
-            if tagged_only:
-                return None
-            if show_clickhouse_errors is not True and b"Code: " not in source.last_message[-1024:]:
-                return None
-            return extract_error_message(source.last_message)
+            return extract_stream_error(
+                source.last_message,
+                getattr(source, "exception_tag", None),
+                show_clickhouse_errors,
+                tagged_only=tagged_only,
+            )
 
         def get_block():
             nonlocal block_num
@@ -61,7 +48,7 @@ def get_block():
                 except StreamCompleteException:
                     error_msg = extract_source_error(tagged_only=True)
                     if error_msg:
-                        raise StreamFailureError(format_stream_error(error_msg)) from None
+                        raise StreamFailureError(format_stream_error(error_msg, show_clickhouse_errors)) from None
                     return None
                 num_rows = source.read_leb128()
                 for col_num in range(num_cols):
@@ -87,7 +74,7 @@ def get_block():
                     # in the response
                     error_msg = extract_source_error()
                     if error_msg:
-                        raise StreamFailureError(format_stream_error(error_msg)) from None
+                        raise StreamFailureError(format_stream_error(error_msg, show_clickhouse_errors)) from None
                     raise StreamFailureError("Stream ended unexpectedly (connection closed by server)") from ex
 
                 # A read failure partway through the stream: OperationalError from the sync reader,
@@ -96,7 +83,7 @@ def get_block():
                 if isinstance(ex, OperationalError) or ex.__class__.__name__ == "ClientPayloadError":
                     error_msg = extract_source_error()
                     if error_msg:
-                        raise StreamFailureError(format_stream_error(error_msg)) from None
+                        raise StreamFailureError(format_stream_error(error_msg, show_clickhouse_errors)) from None
                     raise StreamFailureError("Stream failed during read (connection closed by server)") from ex
 
                 raise
@@ -229,3 +216,37 @@ def extract_error_message(message: bytes) -> str:
     except UnicodeError:
         message_str = f"unrecognized data found in stream: `{message.hex()[128:]}`"
     return message_str
+
+
+def format_stream_error(error_msg: str, show_clickhouse_errors: ShowClickHouseErrors) -> str:
+    if show_clickhouse_errors is False:
+        return GENERIC_CLICKHOUSE_ERROR
+    if show_clickhouse_errors == "scrub":
+        return scrub_error_details(error_msg)
+    return error_msg
+
+
+def extract_stream_error(
+    last_message: bytes | None,
+    exception_tag: str | None,
+    show_clickhouse_errors: ShowClickHouseErrors,
+    tagged_only: bool = False,
+) -> str | None:
+    """Recover an error the server wrote into the response body after it had already committed a 200
+    and started streaming results.
+
+    A tagged block is unambiguous, so it is trusted anywhere. The untagged fallback only looks for a
+    `Code: ` prefix, which is a guess against binary payloads, so callers that cannot distinguish a
+    truncated stream from a complete one pass ``tagged_only``.
+    """
+    if not last_message:
+        return None
+    if exception_tag:
+        error_msg = extract_exception_with_tag(last_message, exception_tag)
+        if error_msg:
+            return error_msg
+    if tagged_only:
+        return None
+    if show_clickhouse_errors is not True and b"Code: " not in last_message[-1024:]:
+        return None
+    return extract_error_message(last_message)