Handle late ClickHouse HTTP exceptions

lib/ch/connection.ex

@@ -343,6 +343,7 @@ defmodule Ch.Connection do
     case request(conn, "POST", path, headers, body, opts) do
       {:ok, conn, responses} -> {:ok, query, responses, conn}
       {:error, _reason, _conn} = client_error -> client_error
+      {:disconnect, %Error{} = reason, conn} -> {:disconnect, reason, conn}
       {:disconnect, reason, conn} -> {:disconnect_and_retry, reason, conn}
     end
   end
@@ -362,7 +363,7 @@ defmodule Ch.Connection do
   @spec request(conn, binary, binary, Mint.Types.headers(), iodata, [Ch.query_option()]) ::
           {:ok, conn, [response]}
           | {:error, Error.t(), conn}
-          | {:disconnect, Mint.Types.error(), conn}
+          | {:disconnect, Error.t() | Mint.Types.error(), conn}
   defp request(conn, method, path, headers, body, opts) do
     with {:ok, conn, _ref} <- send_request(conn, method, path, headers, body) do
       receive_full_response(conn, timeout(conn, opts))
@@ -372,7 +373,7 @@ defmodule Ch.Connection do
   @spec request_chunked(conn, binary, binary, Mint.Types.headers(), Enumerable.t(), Keyword.t()) ::
           {:ok, conn, [response]}
           | {:error, Error.t(), conn}
-          | {:disconnect, Mint.Types.error(), conn}
+          | {:disconnect, Error.t() | Mint.Types.error(), conn}
   def request_chunked(conn, method, path, headers, stream, opts) do
     with {:ok, conn, ref} <- send_request(conn, method, path, headers, :stream),
          {:ok, conn} <- stream_body(conn, ref, stream),
@@ -408,29 +409,71 @@ defmodule Ch.Connection do
   @spec receive_full_response(conn, timeout) ::
           {:ok, conn, [response]}
           | {:error, Error.t(), conn}
-          | {:disconnect, Mint.Types.error(), conn}
+          | {:disconnect, Error.t() | Mint.Types.error(), conn}
   defp receive_full_response(conn, timeout) do
-    with {:ok, conn, responses} <- recv_all(conn, [], timeout) do
-      case responses do
-        [200, headers | _rest] ->
-          conn = ensure_same_server(conn, headers)
-          {:ok, conn, responses}
-
-        [_status, headers | data] ->
-          message = IO.iodata_to_binary(data)
-
-          code =
-            if code = get_header(headers, "x-clickhouse-exception-code") do
-              String.to_integer(code)
+    case recv_all(conn, [], timeout) do
+      {:ok, conn, responses} ->
+        case responses do
+          [200, headers | _data] ->
+            conn = ensure_same_server(conn, headers)
+
+            case decode_row_binary_exception(responses) do
+              {:ok, error} -> {:error, error, conn}
+              :error -> {:ok, conn, responses}
             end
 
-          {:error, Error.exception(code: code, message: message), conn}
+          [_status, headers | data] ->
+            message = IO.iodata_to_binary(data)
+
+            code =
+              if code = get_header(headers, "x-clickhouse-exception-code") do
+                String.to_integer(code)
+              end
+
+            {:error, Error.exception(code: code, message: message), conn}
+        end
+
+      {:closed, conn, responses, reason} ->
+        case decode_row_binary_exception(responses) do
+          {:ok, error} -> {:disconnect, error, conn}
+          :error -> {:disconnect, reason, conn}
+        end
+
+      {:disconnect, _reason, _conn} = disconnect ->
+        disconnect
+    end
+  end
+
+  defp decode_row_binary_exception([200, headers | data]) do
+    try do
+      with "RowBinaryWithNamesAndTypes" <- get_header(headers, "x-clickhouse-format"),
+           nil <- get_header(headers, "content-encoding"),
+           tag when is_binary(tag) <- get_header(headers, "x-clickhouse-exception-tag"),
+           {:ok, message} <- RowBinary.decode_exception(data, tag) do
+        {:ok, Error.exception(code: exception_code(message), message: message)}
+      else
+        _ -> :error
       end
+    rescue
+      _ -> :error
+    end
+  end
+
+  defp decode_row_binary_exception(_responses), do: :error
+
+  defp exception_code(<<"Code: ", rest::binary>>) do
+    case Integer.parse(rest) do
+      {code, <<".", _rest::binary>>} -> code
+      _ -> nil
     end
   end
 
+  defp exception_code(_message), do: nil
+
   @spec recv_all(conn, [response], timeout()) ::
-          {:ok, conn, [response]} | {:disconnect, Mint.Types.error(), conn}
+          {:ok, conn, [response]}
+          | {:closed, conn, [response], Mint.Types.error()}
+          | {:disconnect, Mint.Types.error(), conn}
   defp recv_all(conn, acc, timeout) do
     case HTTP.recv(conn, 0, timeout) do
       {:ok, conn, responses} ->
@@ -439,6 +482,9 @@ defmodule Ch.Connection do
           {:more, acc} -> recv_all(conn, acc, timeout)
         end
 
+      {:error, conn, %Mint.TransportError{reason: :closed} = reason, _responses} ->
+        {:closed, conn, :lists.reverse(acc), reason}
+
       {:error, conn, reason, _responses} ->
         {:disconnect, reason, conn}
     end

lib/ch/row_binary.ex

@@ -8,6 +8,10 @@ defmodule Ch.RowBinary do
 
   @epoch_gregorian_seconds 62_167_219_200
   @epoch_gregorian_days 719_528
+  @exception_marker "__exception__"
+  @exception_tag_length 16
+  @max_exception_size 16 * 1024
+  @max_exception_length_digits 8
 
   @doc false
   def encode_names_and_types(names, types) do
@@ -660,6 +664,90 @@ defmodule Ch.RowBinary do
     end
   end
 
+  @doc false
+  def decode_exception(body, tag)
+      when is_list(body) and is_binary(tag) and byte_size(tag) == @exception_tag_length do
+    closing = " #{tag}\r\n#{@exception_marker}\r\n"
+
+    if exception_tail(body, byte_size(closing)) == closing do
+      body
+      |> exception_tail(@max_exception_size)
+      |> decode_exception(tag)
+    else
+      :error
+    end
+  end
+
+  def decode_exception(body, tag)
+      when is_binary(tag) and byte_size(tag) == @exception_tag_length do
+    opening = "\r\n#{@exception_marker}\r\n#{tag}\r\n"
+    closing = " #{tag}\r\n#{@exception_marker}\r\n"
+    body_size = byte_size(body)
+    closing_size = byte_size(closing)
+    closing_start = body_size - closing_size
+
+    with true <- closing_start >= 0,
+         ^closing <- binary_part(body, closing_start, closing_size),
+         {:ok, length_start} <- trailing_decimal_start(body, closing_start),
+         {message_length, ""} <-
+           body |> binary_part(length_start, closing_start - length_start) |> Integer.parse(),
+         message_start = length_start - message_length,
+         opening_start = message_start - byte_size(opening),
+         true <- opening_start >= 0,
+         true <- body_size - opening_start <= @max_exception_size,
+         ^opening <- binary_part(body, opening_start, byte_size(opening)),
+         message <- binary_part(body, message_start, message_length),
+         true <- String.ends_with?(message, "\n") do
+      {:ok, message}
+    else
+      _ -> :error
+    end
+  end
+
+  def decode_exception(_body, _tag), do: :error
+
+  defp exception_tail(chunks, limit) do
+    chunks
+    |> :lists.reverse()
+    |> exception_tail([], 0, limit)
+    |> IO.iodata_to_binary()
+  end
+
+  defp exception_tail(_chunks, acc, size, limit) when size == limit, do: acc
+
+  defp exception_tail([chunk | chunks], acc, size, limit) do
+    remaining = limit - size
+
+    if byte_size(chunk) <= remaining do
+      exception_tail(chunks, [chunk | acc], size + byte_size(chunk), limit)
+    else
+      [binary_part(chunk, byte_size(chunk) - remaining, remaining) | acc]
+    end
+  end
+
+  defp exception_tail([], acc, _size, _limit), do: acc
+
+  defp trailing_decimal_start(body, index), do: trailing_decimal_start(body, index, 0)
+
+  defp trailing_decimal_start(body, index, digits) when index > 0 do
+    case :binary.at(body, index - 1) do
+      digit when digit in ?0..?9 and digits < @max_exception_length_digits ->
+        trailing_decimal_start(body, index - 1, digits + 1)
+
+      digit when digit in ?0..?9 ->
+        :error
+
+      _other when digits > 0 ->
+        {:ok, index}
+
+      _other ->
+        :error
+    end
+  end
+
+  defp trailing_decimal_start(_body, 0, digits) when digits > 0, do: {:ok, 0}
+  defp trailing_decimal_start(_body, 0, 0), do: :error
+
   @doc """
   Decodes [RowBinary](https://clickhouse.com/docs/en/interfaces/formats/RowBinary) into rows.
 

test/ch/query_test.exs

@@ -492,6 +492,34 @@ defmodule Ch.QueryTest do
       assert {:error, %Ch.Error{code: 62}} = Ch.query(conn, "wat", [], query_options)
     end
 
+    @tag :exception
+    test "detects __exception__ block in HTTP 200 response", %{
+      conn: conn,
+      query_options: query_options
+    } do
+      assert {:error, %Ch.Error{code: 395, message: message}} =
+               Ch.query(
+                 conn,
+                 """
+                 SELECT number, throwIf(number = 2, 'late exception')
+                 FROM numbers(5)
+                 """,
+                 %{},
+                 Keyword.merge(query_options,
+                   settings: [
+                     wait_end_of_query: 0,
+                     http_response_buffer_size: 1,
+                     output_format_parallel_formatting: 0,
+                     max_threads: 1,
+                     max_block_size: 1
+                   ]
+                 )
+               )
+
+      assert message =~ "Code: 395"
+      assert message =~ "late exception"
+    end
+
     test "connection works after failure in execute", %{conn: conn, query_options: query_options} do
       assert {:error, %Ch.Error{}} = Ch.query(conn, "wat", [], query_options)
       assert [[42]] = Ch.query!(conn, "SELECT 42", [], query_options).rows

test/ch/row_binary_test.exs

@@ -312,6 +312,36 @@ defmodule Ch.RowBinaryTest do
     end
   end
 
+  describe "decode_exception/2" do
+    test "decodes a tagged exception split across response chunks" do
+      tag = "0123456789abcdef"
+      message = "Code: 395. DB::Exception: late exception\n"
+
+      trailer =
+        "\r\n__exception__\r\n#{tag}\r\n#{message}#{byte_size(message)} #{tag}\r\n__exception__\r\n"
+
+      assert {:ok, ^message} =
+               decode_exception(
+                 [
+                   <<1, 1, "x", 5, "UInt8", 42>>,
+                   binary_part(trailer, 0, 13),
+                   binary_part(trailer, 13, byte_size(trailer) - 13)
+                 ],
+                 tag
+               )
+    end
+
+    test "rejects a trailer with a different tag" do
+      tag = "0123456789abcdef"
+      message = "Code: 395. DB::Exception: late exception\n"
+
+      trailer =
+        "\r\n__exception__\r\n#{tag}\r\n#{message}#{byte_size(message)} fedcba9876543210\r\n__exception__\r\n"
+
+      assert :error = decode_exception(trailer, tag)
+    end
+  end
+
   describe "decode_rows/2" do
     test "empty" do
       assert decode_rows(<<>>, [:u8, :string]) == []

test/test_helper.exs

@@ -45,7 +45,7 @@ extra_exclude =
 
     true ->
       # Time, Variant, JSON, and Dynamic types are not supported in older ClickHouse versions we have in the CI
-      [:time, :variant, :json, :dynamic]
+      [:time, :variant, :json, :dynamic, :exception]
   end
 
 ExUnit.start(exclude: [:slow | extra_exclude])