Add streaming support

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,

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

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

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")
 

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

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