diff --git a/lib/ch/query.ex b/lib/ch/query.ex index b142621..f321376 100644 --- a/lib/ch/query.ex +++ b/lib/ch/query.ex @@ -75,14 +75,14 @@ defimpl DBConnection.Query, for: Ch.Query do query_params: [{String.t(), String.t()}], body: iodata | Enumerable.t() - def encode(%Query{command: :insert, encode: false, statement: statement}, data, _opts) do + def encode(%Query{command: :insert, encode: false, statement: statement}, data, opts) do body = case data do _ when is_list(data) or is_binary(data) -> [statement, ?\n | data] _ -> Stream.concat([[statement, ?\n]], data) end - {_query_params = [], _extra_headers = [], body} + {_query_params = [], headers(opts), body} end def encode(%Query{command: :insert, statement: statement}, params, opts) do @@ -91,15 +91,15 @@ defimpl DBConnection.Query, for: Ch.Query do types = Keyword.fetch!(opts, :types) header = RowBinary.encode_names_and_types(names, types) data = RowBinary.encode_rows(params, types) - {_query_params = [], _extra_headers = [], [statement, ?\n, header | data]} + {_query_params = [], headers(opts), [statement, ?\n, header | data]} format_row_binary?(statement) -> types = Keyword.fetch!(opts, :types) data = RowBinary.encode_rows(params, types) - {_query_params = [], _extra_headers = [], [statement, ?\n | data]} + {_query_params = [], headers(opts), [statement, ?\n | data]} true -> - {query_params(params), _extra_headers = [], statement} + {query_params(params), headers(opts), statement} end end @@ -107,7 +107,7 @@ defimpl DBConnection.Query, for: Ch.Query do types = Keyword.get(opts, :types) default_format = if types, do: "RowBinary", else: "RowBinaryWithNamesAndTypes" format = Keyword.get(opts, :format) || default_format - {query_params(params), [{"x-clickhouse-format", format}], statement} + {query_params(params), [{"x-clickhouse-format", format} | headers(opts)], statement} end defp format_row_binary?(statement) when is_binary(statement) do @@ -266,6 +266,9 @@ defimpl DBConnection.Query, for: Ch.Query do end defp escape_param([], param), do: param + + @spec headers(Keyword.t()) :: Mint.Types.headers() + defp headers(opts), do: Keyword.get(opts, :headers, []) end defimpl String.Chars, for: Ch.Query do diff --git a/test/ch/headers_test.exs b/test/ch/headers_test.exs new file mode 100644 index 0000000..52aa7c3 --- /dev/null +++ b/test/ch/headers_test.exs @@ -0,0 +1,56 @@ +defmodule Ch.HeadersTest do + use ExUnit.Case, async: true + + setup do + {:ok, conn} = Ch.start_link() + {:ok, conn: conn} + end + + test "can request gzipped response through headers", %{conn: conn} do + assert {:ok, %{rows: rows, headers: headers}} = + Ch.query(conn, "select number from system.numbers limit 100", [], + decode: false, + settings: [enable_http_compression: 1], + headers: [{"accept-encoding", "gzip"}] + ) + + assert :proplists.get_value("content-type", headers) == "application/octet-stream" + assert :proplists.get_value("content-encoding", headers) == "gzip" + assert :proplists.get_value("x-clickhouse-format", headers) == "RowBinaryWithNamesAndTypes" + + # https://en.wikipedia.org/wiki/Gzip + assert <<0x1F, 0x8B, _rest::bytes>> = IO.iodata_to_binary(rows) + end + + test "can request lz4 response through headers", %{conn: conn} do + assert {:ok, %{rows: rows, headers: headers}} = + Ch.query(conn, "select number from system.numbers limit 100", [], + decode: false, + settings: [enable_http_compression: 1], + headers: [{"accept-encoding", "lz4"}] + ) + + assert :proplists.get_value("content-type", headers) == "application/octet-stream" + assert :proplists.get_value("content-encoding", headers) == "lz4" + assert :proplists.get_value("x-clickhouse-format", headers) == "RowBinaryWithNamesAndTypes" + + # https://en.wikipedia.org/wiki/LZ4_(compression_algorithm) + assert <<0x04, 0x22, 0x4D, 0x18, _rest::bytes>> = IO.iodata_to_binary(rows) + end + + test "can request zstd response through headers", %{conn: conn} do + assert {:ok, %{rows: rows, headers: headers}} = + Ch.query(conn, "select number from system.numbers limit 100", [], + decode: false, + settings: [enable_http_compression: 1], + headers: [{"accept-encoding", "zstd"}] + ) + + assert :proplists.get_value("content-type", headers) == "application/octet-stream" + assert :proplists.get_value("content-encoding", headers) == "zstd" + assert :proplists.get_value("x-clickhouse-format", headers) == "RowBinaryWithNamesAndTypes" + + # https://en.wikipedia.org/wiki/LZ4_(compression_algorithm) + assert <<0x28, 0xB5, 0x2F, 0xFD, _rest::bytes>> = IO.iodata_to_binary(rows) + end +end