diff --git a/lib/avalanche.ex b/lib/avalanche.ex index 448d268..2189bb8 100644 --- a/lib/avalanche.ex +++ b/lib/avalanche.ex @@ -109,6 +109,11 @@ defmodule Avalanche do ) @run_options_schema NimbleOptions.new!( + streaming: [ + type: :boolean, + default: false, + doc: "Set to true to stream the result of the statement." + ], async: [ type: :boolean, default: false, diff --git a/lib/avalanche/requests/statement_request.ex b/lib/avalanche/requests/statement_request.ex index 5d80a18..e837348 100644 --- a/lib/avalanche/requests/statement_request.ex +++ b/lib/avalanche/requests/statement_request.ex @@ -84,7 +84,10 @@ defmodule Avalanche.StatementRequest do end defp build_pipeline(request, opts) do - disable_polling = Keyword.fetch!(opts, :async) + async? = Keyword.fetch!(opts, :async) + streaming? = Keyword.get(opts, :streaming, false) + # Only disable polling if async is true AND we're not streaming + disable_polling = async? and not streaming? params = build_params(opts) req_options = @@ -116,11 +119,17 @@ defmodule Avalanche.StatementRequest do decode_data_options = Keyword.get(request.options, :decode_data, []) get_partitions_options = Keyword.get(request.options, :get_partitions, []) - req_options - |> Req.new() - |> Steps.Poll.attach(disable_polling, poll_options) - |> Steps.DecodeData.attach(decode_data_options) - |> Steps.GetPartitions.attach(get_partitions_options) + base_pipeline = + req_options + |> Req.new() + |> Steps.Poll.attach(disable_polling, poll_options) + |> Steps.DecodeData.attach(decode_data_options) + + if streaming? do + Steps.StreamPartitions.attach(base_pipeline, get_partitions_options) + else + Steps.GetPartitions.attach(base_pipeline, get_partitions_options) + end end defp build_params(opts) do diff --git a/lib/avalanche/result.ex b/lib/avalanche/result.ex index 2c00110..7f6146f 100644 --- a/lib/avalanche/result.ex +++ b/lib/avalanche/result.ex @@ -30,6 +30,6 @@ defmodule Avalanche.Result do statement_handle: String.t() | nil, statement_handles: list(String.t()) | nil, num_rows: non_neg_integer() | nil, - rows: list(map()) | nil + rows: list(map()) | function() | nil } end diff --git a/lib/avalanche/steps/decode_data.ex b/lib/avalanche/steps/decode_data.ex index 1a011c0..eb9638a 100644 --- a/lib/avalanche/steps/decode_data.ex +++ b/lib/avalanche/steps/decode_data.ex @@ -47,7 +47,7 @@ defmodule Avalanche.Steps.DecodeData do def decode_data(request_response), do: request_response - defp decode_data_rows(types, data, downcase_column_names) do + defp decode_data_rows(types, data, downcase_column_names) when is_list(data) do Enum.map(data, fn row -> Enum.zip_reduce(types, row, %{}, fn type, value, result -> column_name = maybe_downcased_column_name(type, downcase_column_names) @@ -57,6 +57,20 @@ defmodule Avalanche.Steps.DecodeData do end) end + defp decode_data_rows(types, %Stream{} = data, downcase_column_names) do + Stream.map(data, fn row -> + Enum.zip_reduce(types, row, %{}, fn type, value, result -> + column_name = maybe_downcased_column_name(type, downcase_column_names) + column_value = decode(type, value) + Map.put(result, column_name, column_value) + end) + end) + end + + defp decode_data_rows(types, data, downcase_column_names) do + decode_data_rows(types, Stream.map(data, & &1), downcase_column_names) + end + defp maybe_downcased_column_name(type, true), do: type |> Map.fetch!("name") |> String.downcase() defp maybe_downcased_column_name(type, false), do: Map.fetch!(type, "name") diff --git a/lib/avalanche/steps/stream_partitions.ex b/lib/avalanche/steps/stream_partitions.ex new file mode 100644 index 0000000..1a6da9c --- /dev/null +++ b/lib/avalanche/steps/stream_partitions.ex @@ -0,0 +1,137 @@ +defmodule Avalanche.Steps.StreamPartitions do + @moduledoc """ + A custom `Req` pipeline step to stream partitions of data from a statement execution. + + This module implements streaming of partitioned data responses from Snowflake using the SQL API. + It ensures partitions are retrieved and processed in the correct order. + + See: https://docs.snowflake.com/en/developer-guide/sql-api/handling-responses + """ + + require Logger + + @doc """ + Attach the streaming step to a Req pipeline. + + ## Options + + * `:max_concurrency` - sets the maximum number of tasks to run at the same time. + Defaults to `System.schedulers_online/0`. + + * `:timeout` - the maximum amount of time to wait (in milliseconds) for each partition. + Defaults to 2 minutes. + """ + def attach(%Req.Request{} = request, options \\ []) do + request + |> Req.Request.register_options([:max_concurrency, :timeout, :params]) + |> Req.Request.merge_options(options) + |> Req.Request.append_response_steps(stream_partitions: &stream_partitions/1) + end + + def stream_partitions(request_response) + + def stream_partitions({request, %{status: 200, body: %{"resultSetMetaData" => metadata} = body} = response}) do + options = request.options || %{} + max_concurrency = Map.get(options, :max_concurrency, System.schedulers_online()) + timeout = Map.get(options, :timeout, 120_000) # Default to 2 minutes + + path = Map.get(body, "statementStatusUrl") + data = Map.get(body, "data", []) + + row_types = Map.get(metadata, "rowType", []) + partitions = Map.get(metadata, "partitionInfo", []) + + partition_stream = + case {path, partitions} do + {nil, _} -> + [] + + {_, []} -> + [] + + {_, "0"} -> + [] + + {path, [_head | rest]} when is_binary(path) -> + rest + |> Stream.with_index(1) + |> Stream.map(fn {_info, partition} -> + build_status_request(request, path, partition, row_types) + end) + |> Task.async_stream( + fn req -> + Req.Request.run_request(req) + end, + ordered: true, + max_concurrency: max_concurrency, + timeout: timeout, + on_timeout: :kill_task + ) + |> Stream.map(fn + {:ok, {:ok, value}} -> value + {:ok, error} -> error_response(error) + {:exit, reason} -> error_response(reason) + end) + |> Stream.map(&handle_partition_response/1) + |> Stream.filter(&(&1.status == 200)) + |> Stream.flat_map(&Map.get(&1.body, "data", [])) + end + + # Create a stream of the initial data and partition data + stream = Stream.concat(Stream.map([data], & &1), partition_stream) + {request, %{response | body: Map.put(body, "data", stream)}} + end + + def stream_partitions({request, %{status: 200, body: ""} = response}) do + {request, response} + end + + def stream_partitions({request, %{status: 202} = response}) do + {request, response} + end + + def stream_partitions({request, %{status: status} = response}) when status >= 400 do + {request, response} + end + + def stream_partitions(request_response), do: request_response + + # Private helpers + + defp build_status_request(request, path, partition, row_types) do + # Start with a new request but copy over the relevant options from the original + # Handle cases where headers might be nil + headers = request.headers || %{} + options = request.options || %{} + + Req.new( + method: :get, + url: path, + params: [partition: partition], + auth: options[:auth], # Get auth from options + headers: headers, + receive_timeout: options[:timeout] # Map our timeout to Req's receive_timeout + ) + |> Req.Request.put_private(:avalanche_row_types, row_types) + end + + defp handle_partition_response(response) do + case response do + {:ok, {_request, %Req.Response{} = response}} -> + response + + {:ok, {_request, exception}} -> + error_response(exception) + + {:exit, reason} -> + error_response(reason) + + other -> + error_response(other) + end + end + + defp error_response(reason) do + %{status: 500, body: %{"message" => inspect(reason)}} + end +end diff --git a/test/avalanche/steps/stream_partitions_test.exs b/test/avalanche/steps/stream_partitions_test.exs new file mode 100644 index 0000000..4f32418 --- /dev/null +++ b/test/avalanche/steps/stream_partitions_test.exs @@ -0,0 +1,95 @@ +defmodule Avalanche.Steps.StreamPartitionsTest do + use ExUnit.Case, async: true + alias Avalanche.Steps.StreamPartitions + + setup do + # Set up Task.Supervisor for async streaming + start_supervised({Task.Supervisor, name: Avalanche.TaskSupervisor}) + + # Mock response data + base_response = %Req.Response{ + status: 200, + body: %{ + "resultSetMetaData" => %{ + "rowType" => [ + %{"name" => "col1", "type" => "text"} + ], + "partitionInfo" => [ + %{"rowCount" => 2}, + %{"rowCount" => 3} + ] + }, + "statementStatusUrl" => "https://example.snowflakecomputing.com/status", + "data" => [["first"]] + } + } + + # Mock request with all necessary fields + request = + Req.new( + url: "https://example.snowflakecomputing.com", + auth: {:bearer, "test-token"}, + headers: %{"Content-Type" => "application/json"}, + receive_timeout: 120_000 + ) + |> Req.Request.register_options([:max_concurrency, :timeout]) # Register our custom options + |> Req.Request.merge_options( + timeout: 120_000, + max_concurrency: System.schedulers_online() + ) + + {:ok, request: request, base_response: base_response} + end + + describe "attach/2" do + test "attaches streaming step to request pipeline" do + request = %Req.Request{} + result = StreamPartitions.attach(request) + + assert Keyword.has_key?(result.response_steps, :stream_partitions) + end + end + + describe "stream_partitions/1" do + test "handles empty response body", %{request: request} do + empty_response = %Req.Response{body: Stream.into([], []), status: 200} + assert {^request, ^empty_response} = StreamPartitions.stream_partitions({request, empty_response}) + end + + test "handles response with no partitions", %{request: request} do + response = %Req.Response{ + status: 200, + body: %{ + "resultSetMetaData" => %{ + "rowType" => [], + "partitionInfo" => [] + }, + "data" => "" + } + } + + {_req, result} = StreamPartitions.stream_partitions({request, response}) + assert Enum.to_list(result.body["data"]) == [""] + end + + test "handles response with missing fields", %{request: request} do + response = %Req.Response{ + status: 200, + body: %{ + "resultSetMetaData" => %{}, + "data" => [] + } + } + + {_req, result} = StreamPartitions.stream_partitions({request, response}) + assert Enum.to_list(result.body["data"]) == [[]] + end + + test "processes partitioned response", %{request: request, base_response: response} do + {_req, result} = StreamPartitions.stream_partitions({request, response}) + + assert is_function(result.body["data"]) + assert length(Enum.to_list(result.body["data"])) >= 1 + end + end +end