Support durable request identity for async statements

lib/snowflex.ex

@@ -367,6 +367,15 @@ defmodule Snowflex do
       this to the maximum the statement itself may take.
     * `:query_tag` - sets Snowflake's `QUERY_TAG`, useful for correlating the
       statement later.
+    * `:request_id` - a caller-generated UUID string sent as the SQL API
+      `requestId`.
+    * `:retry` - boolean, default `false`. Set to `true` only when resubmitting
+      the original request with its original `:request_id`. Missing/invalid UUIDs
+      and non-boolean retry values return an error before submission.
+
+  ## Durable submission and recovery
+
+  Snowflake supports idempotent requests using the `request_id`.  For more information, see [Snowflake's resubmission contract](https://docs.snowflake.com/en/developer-guide/sql-api/submitting-requests#resubmitting-a-request-to-execute-sql-statements).
 
   ## Examples
 

lib/snowflex/transport/http.ex

@@ -212,7 +212,51 @@ defmodule Snowflex.Transport.Http do
   @impl Snowflex.Transport
   def submit_async(pid, statement, params, opts) do
     opts = add_default_timeout(opts)
-    call(pid, {:submit_async, statement, params, opts}, statement, opts)
+
+    with :ok <- validate_submission_identity(opts) do
+      pid
+      |> call({:submit_async, statement, params, opts}, statement, opts)
+      |> put_submission_identity(opts[:request_id])
+    end
+  end
+
+  # Validate before dispatch so a malformed identity cannot submit a statement.
+  defp validate_submission_identity(opts) do
+    request_id = opts[:request_id]
+    retry = Keyword.get(opts, :retry, false)
+
+    cond do
+      Keyword.has_key?(opts, :request_id) and not valid_request_id?(request_id) ->
+        {:error,
+         Error.exception(
+           ":request_id must be a UUID string (xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx)"
+         )}
+
+      not is_boolean(retry) ->
+        {:error, Error.exception(":retry must be a boolean")}
+
+      retry and is_nil(request_id) ->
+        {:error, Error.exception(":retry requires the original :request_id")}
+
+      true ->
+        :ok
+    end
+  end
+
+  defp valid_request_id?(id) when is_binary(id) do
+    Regex.match?(~r/\A[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\z/i, id)
+  end
+
+  defp valid_request_id?(_id), do: false
+
+  defp put_submission_identity(reply, nil), do: reply
+
+  defp put_submission_identity({:ok, %Result{} = result}, request_id) do
+    {:ok, %{result | request_id: request_id}}
+  end
+
+  defp put_submission_identity({:error, %Error{} = error}, request_id) do
+    {:error, %{error | metadata: Map.put(error.metadata || %{}, :request_id, request_id)}}
   end
 
   @impl Snowflex.Transport
@@ -953,7 +997,11 @@ defmodule Snowflex.Transport.Http do
     # finishes (or until its own ~45s cutoff). Req treats an empty *list* as
     # "no params" and leaves the URL untouched; an empty map would take the
     # encode path instead.
-    query_params = if call_opts[:async?], do: %{async: true}, else: []
+    query_params = if call_opts[:async?], do: submission_query_params(opts), else: []
+
+    # Recovery owns resubmission. Do not let Req repeat an async POST (including
+    # one configured with retry: :transient) without an explicit caller retry.
+    req_client = if call_opts[:async?], do: Req.merge(req_client, retry: false), else: req_client
 
     case Req.post(req_client,
            url: url,
@@ -1003,6 +1051,12 @@ defmodule Snowflex.Transport.Http do
     end
   end
 
+  defp submission_query_params(opts) do
+    params = [async: true]
+    params = if opts[:request_id], do: params ++ [requestId: opts[:request_id]], else: params
+    if opts[:retry], do: params ++ [retry: true], else: params
+  end
+
   defp fetch_partition(state, statement_handle, partition_index, opts) do
     url = "/api/v2/statements/#{statement_handle}"
     params = %{partition: partition_index}

test/snowflex/async_statement_test.exs

@@ -3,10 +3,13 @@ defmodule Snowflex.AsyncStatementTest do
 
   import Req.Test, only: [set_req_test_to_shared: 1]
 
+  alias Ecto.Adapter
   alias Plug.Conn
   alias Req.Test, as: ReqTest
+  alias Snowflex.Query
 
   @handle "01b7e043-0206-7a43-0008-8b8300073d86"
+  @request_id "ea7b46ed-bdc1-8c32-d593-764fcad64e83"
 
   setup :set_req_test_to_shared
 
@@ -74,6 +77,126 @@ defmodule Snowflex.AsyncStatementTest do
       refute_received {:polled, _path}
     end
 
+    test "resubmits a lost response only on explicit retry with the original identity and body" do
+      test_pid = self()
+
+      ReqTest.stub(MockAsyncHttp, fn conn ->
+        if health_check?(conn) do
+          ReqTest.json(conn, %{})
+        else
+          query = URI.decode_query(conn.query_string)
+          send(test_pid, {:submission, query, conn.body_params})
+
+          if query["retry"] == "true" do
+            json(conn, 202, %{"statementHandle" => @handle})
+          else
+            # The remote accepted the request, but its acknowledgement was lost.
+            ReqTest.transport_error(conn, :closed)
+          end
+        end
+      end)
+
+      start_repo(req_options: [plug: {Req.Test, MockAsyncHttp}, retry: :transient])
+      opts = [request_id: @request_id, query_tag: "operation-1", statement_timeout: 60_000]
+
+      assert {:error, error} = Snowflex.submit_async(AsyncRepo, "CALL proc(?)", ["arg"], opts)
+      assert error.code == "HTTP_ERROR"
+      assert error.metadata.request_id == @request_id
+      assert_received {:submission, initial_query, initial_body}
+      assert initial_query == %{"async" => "true", "requestId" => @request_id}
+      refute_received {:submission, _, _}
+
+      assert {:ok, @handle} =
+               Snowflex.submit_async(AsyncRepo, "CALL proc(?)", ["arg"], [retry: true] ++ opts)
+
+      assert_received {:submission, retry_query, retry_body}
+      assert retry_query == Map.put(initial_query, "retry", "true")
+      assert retry_body == initial_body
+      refute_received {:submission, _, _}
+    end
+
+    test "rejects invalid identity and retry options without submitting" do
+      test_pid = self()
+
+      ReqTest.stub(MockAsyncHttp, fn conn ->
+        unless health_check?(conn), do: send(test_pid, :unexpected_submission)
+        ReqTest.json(conn, %{})
+      end)
+
+      start_repo()
+
+      for id <- [
+            nil,
+            123,
+            "",
+            "not-a-uuid",
+            String.replace(@request_id, "-", ""),
+            @request_id <> "\n",
+            "gg7b46ed-bdc1-8c32-d593-764fcad64e83"
+          ] do
+        assert {:error, error} =
+                 Snowflex.submit_async(AsyncRepo, "CALL proc()", [], request_id: id)
+
+        assert error.message =~ ":request_id must be a UUID"
+      end
+
+      assert {:error, error} = Snowflex.submit_async(AsyncRepo, "CALL proc()", [], retry: true)
+      assert error.message =~ "requires the original :request_id"
+
+      for retry <- [nil, "true", 1, :transient] do
+        assert {:error, error} =
+                 Snowflex.submit_async(AsyncRepo, "CALL proc()", [],
+                   request_id: @request_id,
+                   retry: retry
+                 )
+
+        assert error.message =~ ":retry must be a boolean"
+      end
+
+      refute_received :unexpected_submission
+    end
+
+    test "retains caller identity in transport acknowledgement when response omits requestId" do
+      test_pid = self()
+      request_id = String.upcase(@request_id)
+
+      ReqTest.stub(MockAsyncHttp, fn conn ->
+        if health_check?(conn) do
+          ReqTest.json(conn, %{})
+        else
+          send(test_pid, {:query, URI.decode_query(conn.query_string)})
+          json(conn, 200, %{"statementHandle" => @handle, "sqlState" => "00000"})
+        end
+      end)
+
+      start_repo()
+      %{pid: pool} = Adapter.lookup_meta(AsyncRepo)
+      query = Query.new(statement: "CALL proc()", op: :submit_async)
+
+      assert {:ok, _, result} =
+               DBConnection.execute(pool, query, [], request_id: request_id, retry: false)
+
+      assert result.query_id == @handle
+      assert result.request_id == request_id
+      assert result.sql_state == "00000"
+      assert_received {:query, %{"async" => "true", "requestId" => ^request_id} = params}
+      refute Map.has_key?(params, "retry")
+    end
+
+    test "preserves request identity when acknowledgement has no usable handle" do
+      ReqTest.stub(MockAsyncHttp, fn conn ->
+        if health_check?(conn), do: ReqTest.json(conn, %{}), else: json(conn, 202, %{})
+      end)
+
+      start_repo()
+
+      assert {:error, error} =
+               Snowflex.submit_async(AsyncRepo, "CALL proc()", [], request_id: @request_id)
+
+      assert error.metadata.request_id == @request_id
+      assert error.message =~ "no usable statementHandle"
+    end
+
     test "passes bindings and query_tag through" do
       test_pid = self()