feat: add support for multiple sql statements in single request (#40)

lib/avalanche/requests/statement_request.ex

@@ -143,7 +143,9 @@ defmodule Avalanche.StatementRequest do
       parameters: %{
         "TIME_OUTPUT_FORMAT" => "HH24:MI:SS",
         "TIMESTAMP_OUTPUT_FORMAT" => "YYYY-MM-DD HH24:MI:SS.FFTZH:TZM",
-        "TIMESTAMP_NTZ_OUTPUT_FORMAT" => "YYYY-MM-DD HH24:MI:SS.FF3"
+        "TIMESTAMP_NTZ_OUTPUT_FORMAT" => "YYYY-MM-DD HH24:MI:SS.FF3",
+        # variable number of SQL statements
+        "MULTI_STATEMENT_COUNT" => "0"
       },
       statement: statement,
       bindings: bindings
@@ -155,12 +157,20 @@ defmodule Avalanche.StatementRequest do
 
   defp handle_response(%Req.Response{status: 200, body: body}) do
     statement_handle = Map.fetch!(body, "statementHandle")
+    statement_handles = Map.get(body, "statementHandles")
     data = Map.fetch!(body, "data")
 
     metadata = Map.fetch!(body, "resultSetMetaData")
     num_rows = Map.fetch!(metadata, "numRows")
 
-    {:ok, %Result{status: :complete, statement_handle: statement_handle, num_rows: num_rows, rows: data}}
+    {:ok,
+     %Result{
+       status: :complete,
+       statement_handle: statement_handle,
+       statement_handles: statement_handles,
+       num_rows: num_rows,
+       rows: data
+     }}
   end
 
   defp handle_response(%Req.Response{status: 202, body: body}) do

lib/avalanche/result.ex

@@ -7,7 +7,11 @@ defmodule Avalanche.Result do
 
   * `:status` - the status of the statement being executed (`:running`, `:complete`)
 
-  * `:statement_handle` - the unique identifier for the statement being executed
+  * `:statement_handle` - the unique identifier for the statement being executed.
+  If multiple statements were specified in the request, this handle corresponds
+  to the set of those statements.
+
+  * `:statement_handles` - list of unique identifiers for the statements being executed for this request.
 
   * `:num_rows` - the number of fetched or affected rows
 
@@ -17,13 +21,14 @@ defmodule Avalanche.Result do
   """
 
   @enforce_keys [:status, :statement_handle]
-  defstruct status: :running, statement_handle: nil, num_rows: nil, rows: nil
+  defstruct status: :running, statement_handle: nil, statement_handles: nil, num_rows: nil, rows: nil
 
   @type result_status :: atom()
 
   @type t() :: %__MODULE__{
           status: result_status(),
           statement_handle: String.t() | nil,
+          statement_handles: list(String.t()) | nil,
           num_rows: non_neg_integer() | nil,
           rows: list(map()) | nil
         }

lib/avalanche/token_cache.ex

@@ -45,16 +45,18 @@ defmodule Avalanche.TokenCache do
   def fetch_token(options) do
     key = key_from_options(options)
 
-    Cachex.fetch(@cache, key, fn _key ->
-      case token_from_options(options) do
-        {:ok, token} ->
-          {:commit, token}
-
-        {:error, error} ->
-          {:ignore, {:error, error}}
-      end
-    end)
-    |> case do
+    result =
+      Cachex.fetch(@cache, key, fn _key ->
+        case token_from_options(options) do
+          {:ok, token} ->
+            {:commit, token}
+
+          {:error, error} ->
+            {:ignore, {:error, error}}
+        end
+      end)
+
+    case result do
       {:ok, token} ->
         token
 

test/avalanche_test.exs

@@ -55,6 +55,9 @@ defmodule AvalancheTest do
                %{"COLUMN1" => 1, "COLUMN2" => "one"},
                %{"COLUMN1" => 2, "COLUMN2" => "two"}
              ] = result.rows
+
+      # statement_handles should be nil when not a multi-statement query
+      assert {:ok, nil} = Map.fetch(result, :statement_handles)
     end
 
     @tag :capture_log
@@ -195,6 +198,28 @@ defmodule AvalancheTest do
     end
   end
 
+  describe "run/4 multi-statement" do
+    test "handles multi-statement requests", c do
+      expect_telemetry_mock_start()
+      expect_telemetry_mock_stop()
+
+      statement_handles = ["sh-one", "sh-two", "sh-three"]
+      result_set = result_set_fixture(%{"statementHandles" => statement_handles})
+
+      Bypass.expect(c.bypass, "POST", "/api/v2/statements", fn conn ->
+        conn
+        |> Plug.Conn.put_resp_header("content-type", "application/json")
+        |> Plug.Conn.send_resp(200, Jason.encode!(result_set))
+      end)
+
+      assert {:ok, %Avalanche.Result{statement_handles: result_handles} = result} =
+               Avalanche.run("begin transaction;select 1;commit", [], [], c.options)
+
+      assert result.status == :complete
+      assert ^statement_handles = result_handles
+    end
+  end
+
   describe "run/4 errors" do
     test "returns a bad request Error for 400 response code", c do
       expect_telemetry_mock_start()

test/integration_test.exs

@@ -71,7 +71,7 @@ defmodule AvalancheIntegrationTest do
     :ok
   end
 
-  describe "run/2 with OAuth token" do
+  describe "run/4 with OAuth token" do
     setup do
       options = test_options()
       [options: options]
@@ -83,7 +83,7 @@ defmodule AvalancheIntegrationTest do
     end
   end
 
-  describe "run/2 with Key Pair token" do
+  describe "run/4 with Key Pair token" do
     setup do
       options = test_key_pair_options()
       [options: options]
@@ -95,7 +95,7 @@ defmodule AvalancheIntegrationTest do
     end
   end
 
-  describe "run/2" do
+  describe "run/4" do
     setup do
       options = test_options()
       [options: options]
@@ -149,17 +149,8 @@ defmodule AvalancheIntegrationTest do
                | _rest
              ] = result1.rows
 
-      assert {:ok, %Avalanche.Result{} = result2} =
-               Avalanche.run(
-                 "SELECT * FROM SNOWFLAKE_SAMPLE_DATA.WEATHER.DAILY_14_TOTAL LIMIT 1",
-                 [],
-                 [],
-                 c.options
-               )
-
-      assert result2.num_rows == 1
-
-      assert [%{"T" => %NaiveDateTime{}, "V" => _stuff1}] = result2.rows
+      # statement_handles should be nil when not a multi-statement query
+      assert {:ok, nil} = Map.fetch(result1, :statement_handles)
     end
 
     test "auto loads partitions", c do
@@ -193,6 +184,32 @@ defmodule AvalancheIntegrationTest do
     # end
   end
 
+  describe "run/4 multi-statement" do
+    setup do
+      options = test_options()
+      [options: options]
+    end
+
+    test "handles multi-statement requests", c do
+      assert {:ok, %Avalanche.Result{statement_handles: statement_handles} = result} =
+               Avalanche.run("begin transaction;select ?;select ? as two;commit;", [1, 2], [], c.options)
+
+      assert result.status == :complete
+      assert length(statement_handles) == 4
+
+      [sh1, sh2, sh3, sh4] = statement_handles
+
+      assert {:ok, %Avalanche.Result{rows: [%{"status" => "Statement executed successfully."}]}} =
+               Avalanche.status(sh1, [], c.options)
+
+      assert {:ok, %Avalanche.Result{rows: [%{"?" => 1.0}]}} = Avalanche.status(sh2, [], c.options)
+      assert {:ok, %Avalanche.Result{rows: [%{"TWO" => 2.0}]}} = Avalanche.status(sh3, [], c.options)
+
+      assert {:ok, %Avalanche.Result{rows: [%{"status" => "Statement executed successfully."}]}} =
+               Avalanche.status(sh4, [], c.options)
+    end
+  end
+
   describe "decode_data/1 (integration)" do
     @describetag integration: true
 
@@ -227,18 +244,6 @@ defmodule AvalancheIntegrationTest do
                }
                | _rest
              ] = result1.rows
-
-      assert {:ok, %Avalanche.Result{} = result2} =
-               Avalanche.run(
-                 "SELECT * FROM SNOWFLAKE_SAMPLE_DATA.WEATHER.DAILY_14_TOTAL LIMIT 1",
-                 [],
-                 [],
-                 c.options
-               )
-
-      assert result2.num_rows == 1
-
-      assert [%{"T" => %NaiveDateTime{}, "V" => _stuff1}] = result2.rows
     end
   end