Add decoding to Ch.stream/4 (#277)

* add decoding to Ch.stream/4

* more info in changelog

* note CSVWithNames and other format are not affected

* fix formatting

* cleanup

* cleanup tests

* assume 0.5.7 is released

* update changelog

* cleanup

* shorter

* shorter

Author: ruslandoga

CHANGELOG.md

@@ -1,5 +1,49 @@
 # Changelog
 
+## Unreleased
+
+- added **automatic decoding** to `Ch.stream/4` when using `RowBinaryWithNamesAndTypes` format: https://github.com/plausible/ch/pull/277.
+
+    Previously, this function returned raw bytes.
+  
+    To **restore the previous behavior** (raw bytes/no automatic decoding), pass `decode: false` in the options (**fourth** argument).
+
+    **Example of required change to preserve the previous behavior**
+
+    ```elixir
+    # before, no decoding by default
+    DBConnection.run(pool, fn conn ->
+      conn
+      |> Ch.stream("select number from numbers(10)")
+      |> Enum.into([])
+    end)
+
+    # after, to keep the same behaviour add `decode: false` option
+    DBConnection.run(pool, fn conn ->
+      conn
+      |> Ch.stream("select number from numbers(10)", _params = %{}, decode: false)
+      |> Enum.into([])
+    end)
+  ```
+
+  Queries using other explicit formats like `CSVWithNames` are **unaffected** and can remain as they are.
+
+  **Examples of unaffected queries**
+
+  ```elixir
+  DBConnection.run(pool, fn conn ->
+    conn
+    |> Ch.stream("select number from numbers(10) format CSVWithNames")
+    |> Enum.into([])
+  end)
+
+  DBConnection.run(pool, fn conn ->
+    conn
+    |> Ch.stream("select number from numbers(10)", _params = %{}, format: "CSVWithNames")
+    |> Enum.into([])
+  end)
+  ```
+
 ## 0.5.7 (2025-11-26)
 
 - fix type decoding for strings containing newlines https://github.com/plausible/ch/pull/278

bench/stream.exs

@@ -32,26 +32,24 @@ Benchee.run(
         fn conn ->
           conn
           |> Ch.stream(
-            "SELECT number FROM system.numbers_mt LIMIT {limit:UInt64} FORMAT RowBinary",
-            %{"limit" => limit}
+            "SELECT number FROM system.numbers_mt LIMIT {limit:UInt64}",
+            %{"limit" => limit},
+            decode: false
           )
           |> Stream.run()
         end,
         timeout: :infinity
       )
     end,
-    "Ch.stream with manual RowBinary decoding" => fn %{pool: pool, limit: limit} ->
+    "Ch.stream with decoding" => fn %{pool: pool, limit: limit} ->
       DBConnection.run(
         pool,
         fn conn ->
           conn
           |> Ch.stream(
-            "SELECT number FROM system.numbers_mt LIMIT {limit:UInt64} FORMAT RowBinary",
+            "SELECT number FROM system.numbers_mt LIMIT {limit:UInt64}",
             %{"limit" => limit}
           )
-          |> Stream.each(fn %Ch.Result{data: data} ->
-            data |> IO.iodata_to_binary() |> Ch.RowBinary.decode_rows([:u64])
-          end)
           |> Stream.run()
         end,
         timeout: :infinity

lib/ch/connection.ex

@@ -2,7 +2,7 @@ defmodule Ch.Connection do
   @moduledoc false
   use DBConnection
   require Logger
-  alias Ch.{Error, Query, Result}
+  alias Ch.{Error, Query, Result, RowBinary}
   alias Mint.HTTP1, as: HTTP
 
   @user_agent "ch/" <> Mix.Project.config()[:version]
@@ -104,31 +104,42 @@ defmodule Ch.Connection do
   @impl true
   def handle_declare(query, params, opts, conn) do
     conn = maybe_reconnect(conn)
-    %Query{command: command} = query
+    %Query{command: command, decode: decode} = query
     {query_params, extra_headers, body} = params
 
     path = path(conn, query_params, opts)
     headers = headers(conn, extra_headers, opts)
+    timeout = timeout(conn, opts)
 
     with {:ok, conn, _ref} <- send_request(conn, "POST", path, headers, body),
-         {:ok, conn} <- eat_ok_status_and_headers(conn, timeout(conn, opts)) do
-      {:ok, query, %Result{command: command}, conn}
+         {:ok, conn, columns, headers, reader} <- recv_declare(conn, decode, timeout) do
+      result = %Result{
+        command: command,
+        columns: columns,
+        rows: [],
+        num_rows: 0,
+        headers: headers,
+        data: []
+      }
+
+      {:ok, query, result, {conn, reader}}
     end
   end
 
-  @spec eat_ok_status_and_headers(conn, timeout) ::
-          {:ok, %{conn: conn, buffer: [Mint.Types.response()]}}
-          | {:error, Ch.Error.t(), conn}
-          | {:disconnect, Mint.Types.error(), conn}
-  defp eat_ok_status_and_headers(conn, timeout) do
+  defp recv_declare(conn, decode, timeout) do
+    acc = %{decode: decode, step: :status, buffer: [], headers: []}
+    recv_declare_continue(conn, acc, timeout)
+  end
+
+  defp recv_declare_continue(conn, acc, timeout) do
     case HTTP.recv(conn, 0, timeout) do
       {:ok, conn, responses} ->
-        case eat_ok_status_and_headers(responses) do
-          {:ok, data} ->
-            {:ok, %{conn: conn, buffer: data}}
+        case handle_recv_declare(responses, acc) do
+          {:ok, columns, headers, reader} ->
+            {:ok, conn, columns, headers, reader}
 
-          :more ->
-            eat_ok_status_and_headers(conn, timeout)
+          {:more, acc} ->
+            recv_declare_continue(conn, acc, timeout)
 
           :error ->
             all_responses_result =
@@ -155,49 +166,102 @@ defmodule Ch.Connection do
     end
   end
 
-  defp eat_ok_status_and_headers([{:status, _ref, 200} | rest]) do
-    eat_ok_status_and_headers(rest)
+  defp handle_recv_declare([{:status, _ref, status} | responses], %{step: :status} = acc) do
+    case status do
+      200 -> handle_recv_declare(responses, %{acc | step: :headers})
+      _other -> :error
+    end
+  end
+
+  defp handle_recv_declare([{:headers, _ref, headers} | responses], %{step: :headers} = acc) do
+    with %{decode: true} <- acc,
+         "RowBinaryWithNamesAndTypes" <- get_header(headers, "x-clickhouse-format") do
+      handle_recv_declare(responses, %{acc | headers: headers, step: :columns})
+    else
+      _ ->
+        reader = %{decode: false, responses: responses}
+        {:ok, _columns = nil, headers, reader}
+    end
   end
 
-  defp eat_ok_status_and_headers([{:status, _ref, _status} | _rest]), do: :error
-  defp eat_ok_status_and_headers([{:headers, _ref, _headers} | data]), do: {:ok, data}
-  defp eat_ok_status_and_headers([]), do: :more
+  defp handle_recv_declare([{:data, _ref, data} | responses], %{step: :columns} = acc) do
+    buffer = maybe_concat_buffer(acc.buffer, data)
+
+    case RowBinary.decode_header(buffer) do
+      {:ok, names, types, buffer} ->
+        reader = %{buffer: buffer, types: types, state: nil, responses: responses}
+        {:ok, names, acc.headers, reader}
+
+      :more ->
+        handle_recv_declare(responses, %{acc | buffer: buffer})
+    end
+  end
+
+  defp handle_recv_declare([], acc), do: {:more, acc}
+
+  @compile inline: [maybe_concat_buffer: 2]
+  defp maybe_concat_buffer("", data), do: data
+  defp maybe_concat_buffer(buffer, data) when is_binary(buffer), do: buffer <> data
+  defp maybe_concat_buffer([], data), do: data
 
   @impl true
-  def handle_fetch(query, %Result{} = result, opts, %{conn: conn, buffer: buffer}) do
-    case buffer do
-      [] -> handle_fetch(query, result, opts, conn)
-      _not_empty -> {halt_or_cont(buffer), %{result | data: extract_data(buffer)}, conn}
+  def handle_fetch(query, %Result{} = result, opts, {conn, reader}) do
+    case reader do
+      %{responses: []} ->
+        handle_fetch_recv(query, result, opts, conn, reader)
+
+      %{decode: false, responses: responses} ->
+        case responses do
+          [{:data, _ref, data} | responses] ->
+            result = %Result{result | data: data}
+            reader = %{reader | responses: responses}
+            {:cont, result, {conn, reader}}
+
+          [{:done, _ref}] ->
+            reader = %{reader | responses: []}
+            {:halt, result, {conn, reader}}
+        end
+
+      %{buffer: buffer, types: types, state: state, responses: responses} ->
+        case responses do
+          [{:data, _ref, data} | responses] ->
+            buffer = maybe_concat_buffer(buffer, data)
+            {rows, buffer, state} = RowBinary.decode_rows_continue(buffer, types, state)
+            result = %Result{result | data: data, rows: rows, num_rows: length(rows)}
+            reader = %{reader | buffer: buffer, state: state, responses: responses}
+            {:cont, result, {conn, reader}}
+
+          [{:done, _ref}] ->
+            reader = %{reader | responses: []}
+            {:halt, result, {conn, reader}}
+        end
     end
   end
 
-  def handle_fetch(_query, %Result{} = result, opts, conn) do
-    case HTTP.recv(conn, 0, timeout(conn, opts)) do
+  defp handle_fetch_recv(query, result, opts, conn, reader) do
+    timeout = timeout(conn, opts)
+
+    case HTTP.recv(conn, 0, timeout) do
       {:ok, conn, responses} ->
-        {halt_or_cont(responses), %{result | data: extract_data(responses)}, conn}
+        reader = %{reader | responses: responses}
+        handle_fetch(query, result, opts, {conn, reader})
 
       {:error, conn, reason, _responses} ->
         {:disconnect, reason, conn}
     end
   end
 
-  defp halt_or_cont([{:done, _ref}]), do: :halt
-  defp halt_or_cont([_ | rest]), do: halt_or_cont(rest)
-  defp halt_or_cont([]), do: :cont
-
-  defp extract_data([{:data, _ref, data} | rest]), do: [data | extract_data(rest)]
-  defp extract_data([] = empty), do: empty
-  defp extract_data([{:done, _ref}]), do: []
-
   @impl true
-  def handle_deallocate(_query, %Result{} = result, _opts, conn) do
+  def handle_deallocate(_query, %Result{} = result, _opts, {conn, _reader}) do
     case HTTP.open_request_count(conn) do
       0 ->
-        # TODO data: [], anything else?
         {:ok, %{result | data: []}, conn}
 
       1 ->
-        {:disconnect, Error.exception("cannot stop stream before receiving full response"), conn}
+        error =
+          Error.exception("stopping stream before receiving full response by closing connection")
+
+        {:disconnect, error, conn}
     end
   end