Add fire-and-forget async statement submission

CHANGELOG.md

@@ -7,6 +7,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Enhancements
+
+- Add `Snowflex.submit_async/4`, `Snowflex.statement_status/3`, `Snowflex.fetch_result/3` and `Snowflex.cancel_statement/3` for fire-and-forget statement submission. `submit_async/4` sends Snowflake's `async=true` query parameter, so the API acknowledges with a statement handle as soon as it accepts the statement instead of holding the response open until the statement finishes. A connection is checked out only for the submission round-trip, so long-running statements no longer occupy a pool slot for their full duration. `fetch_result/3` returns the statement's rows once it completes (and `{:ok, :running}` before that, so it can be used to poll), decoded exactly as `Ecto.Repo.query/4` decodes them. Submitted statements are bounded only by Snowflake's statement timeout, since nothing on the client side waits on them. ([#197](https://github.com/pepsico-ecommerce/snowflex/issues/197))
+- `Snowflex.Transport` gains four optional callbacks — `submit_async/4`, `statement_status/3`, `fetch_result/3` and `cancel_statement/3` — implemented by `Snowflex.Transport.Http`. They are optional, so existing custom transports continue to satisfy the behaviour without change; calling an async function on a transport that does not implement them returns a descriptive `Snowflex.Error` rather than raising `UndefinedFunctionError`.
+- `Snowflex.Transport.Http` now caches its signed JWT for the configured `:token_lifetime` instead of re-signing on every request. Each request previously performed a PEM decode and an RSA private-key operation; this is most noticeable when polling an async statement, where one workflow issues many requests.
+
+### Bug Fixes
+
+- `Snowflex.Transport.Http` sent the `timeout` field of a statement request in milliseconds, but Snowflake defines it in seconds. The default 45-second timeout therefore asked Snowflake to wait 45,000 seconds (about 12.5 hours), effectively disabling the server-side statement timeout; a configured `:timeout` is now converted to seconds. Statements that previously ran unbounded will now be cancelled by Snowflake at the configured timeout.
+- JWT `exp` claims were computed by adding a millisecond `:token_lifetime` to a second-precision timestamp, so tokens advertised an expiry roughly a week away instead of the configured lifetime (10 minutes by default).
+
 ### Breaking Changes
 
 - The minimum supported Req version is now 0.7 (`{:req, "~> 0.7"}`, previously `"~> 0.5"`). Consumers pinned to an older Req should stay on snowflex 1.6.0 until they can upgrade.

README.md

@@ -199,6 +199,33 @@ end)
 
 When streaming with `Snowflex.Transport.Http`, keep in mind that [Snowflake dictates the number of partitions returned](https://docs.snowflake.com/en/developer-guide/sql-api/handling-responses#retrieving-additional-partitions). This is different than a normal TCP protocol like `Postgrex`, where the stream iterates one row at a time: partitions are fetched whole, so memory usage is bounded by the partition size rather than by a row count. Pass a `:timeout` that covers statement execution plus however long consumption takes, and prefer `Ecto.Repo.all/2` when you want the entire result set anyway — the eager path fetches partitions in parallel and is faster for full materialization.
 
+### Asynchronous Statements
+
+Ordinary queries occupy a connection for as long as the statement runs. When you only need to *start* a statement — a stored procedure, a long-running `MERGE`, a warehouse maintenance task — `Snowflex.submit_async/4` submits it and returns Snowflake's statement handle as soon as the submission round-trip completes, so no pool slot is held while the statement runs:
+
+``` elixir
+{:ok, handle} = Snowflex.submit_async(MyRepo, "CALL rebuild_scenarios(?)", [scenario_id])
+```
+
+Later, retrieve the result with the handle. `fetch_result/3` answers `{:ok, :running}` until the statement finishes, so you can poll with it alone:
+
+``` elixir
+case Snowflex.fetch_result(MyRepo, handle) do
+  {:ok, :running} -> :still_running
+  {:ok, %Snowflex.Result{rows: rows}} -> rows
+  {:error, error} -> Logger.error("statement failed: #{Exception.message(error)}")
+end
+```
+
+Rows decode exactly as they would from `MyRepo.query/3`. If you only care whether the statement finished and not about its rows, `statement_status/3` returns `{:ok, :running | :succeeded}`. To stop a statement, `cancel_statement/3` returns `:ok`:
+
+``` elixir
+{:ok, :running} = Snowflex.statement_status(MyRepo, handle)
+:ok = Snowflex.cancel_statement(MyRepo, handle)
+```
+
+Three things to keep in mind. Nothing on the Elixir side waits on a submitted statement, so the only thing bounding it is Snowflake's own statement timeout — prefer statements that are safe to leave running unattended. Snowflake retains results for a limited window, so a handle fetched long after completion returns an error rather than rows. And `fetch_result/3` materializes the entire result set (fetching partitions in parallel), so stream large results with `Snowflex.stream_query/5` instead.
+
 ### Migrations
 
 Migrations are not currently supported by Snowflex.

lib/snowflex.ex

@@ -19,11 +19,14 @@ defmodule Snowflex do
   alias Ecto.Adapters.SQL
   alias Ecto.UUID
   alias Snowflex.Ecto.Adapter.Stream, as: AdapterStream
+  alias Snowflex.Error
   alias Snowflex.Query
   alias Snowflex.Result
   alias Snowflex.VariantField
   alias String.Chars
 
+  import Snowflex.Query, only: [is_statement: 1]
+
   require Logger
 
   @conn __MODULE__.Ecto.Adapter.Connection
@@ -334,6 +337,160 @@ defmodule Snowflex do
     end
   end
 
+  ## Asynchronous statements
+
+  @doc """
+  Submits `statement` for execution and returns its Snowflake statement handle
+  without waiting for the statement to finish.
+
+  A connection is checked out only for the submission round-trip, so a
+  long-running statement does not occupy a pool slot while it runs. Use
+  `statement_status/3` to check on it and `cancel_statement/3` to stop it.
+
+  Snowflake bounds the statement itself with `STATEMENT_TIMEOUT_IN_SECONDS`;
+  nothing on this side will cancel it, so prefer statements that are safe to
+  leave running unattended.
+
+  ## Options
+
+  Options are merged over the repo's configured connection options.
+
+    * `:timeout` - bounds the submission round-trip only, *not* the statement's
+      own execution.
+    * `:query_tag` - sets Snowflake's `QUERY_TAG`, useful for correlating the
+      statement later.
+
+  ## Examples
+
+      {:ok, handle} = Snowflex.submit_async(MyRepo, "CALL long_running_proc()")
+
+      case Snowflex.statement_status(MyRepo, handle) do
+        {:ok, :running} -> :still_going
+        {:ok, :succeeded} -> :done
+        {:error, error} -> Logger.error(Exception.message(error))
+      end
+
+  """
+  @spec submit_async(
+          repo :: Ecto.Repo.t() | pid(),
+          statement :: String.t(),
+          params :: list(),
+          opts :: Keyword.t()
+        ) :: {:ok, String.t()} | {:error, Error.t()}
+  def submit_async(repo, statement, params \\ [], opts \\ [])
+      when (is_atom(repo) or is_pid(repo)) and is_statement(statement) do
+    case run_op(repo, :submit_async, statement, params, opts) do
+      {:ok, %Result{query_id: handle}} -> {:ok, handle}
+      {:error, error} -> {:error, error}
+    end
+  end
+
+  @doc """
+  Reports the status of a statement previously submitted with `submit_async/4`.
+
+  Returns `{:ok, :running}` while the statement is still executing and
+  `{:ok, :succeeded}` once it has completed. A statement that failed is reported
+  as `{:error, error}` carrying Snowflake's error code and message.
+
+  Note that Snowflake only retains a statement's result for a limited period; a
+  handle queried long after completion may report an error rather than success.
+  """
+  @spec statement_status(
+          repo :: Ecto.Repo.t() | pid(),
+          handle :: String.t(),
+          opts :: Keyword.t()
+        ) :: {:ok, :running | :succeeded} | {:error, Error.t()}
+  def statement_status(repo, handle, opts \\ [])
+      when (is_atom(repo) or is_pid(repo)) and is_binary(handle) do
+    run_op(repo, :status, handle, [], opts)
+  end
+
+  @doc """
+  Fetches the result of a statement previously submitted with `submit_async/4`.
+
+  Returns `{:ok, :running}` while the statement is still executing, so you can
+  poll with this function alone instead of pairing it with `statement_status/3`.
+  Once the statement has finished, the full `t:Snowflex.Result.t/0` is returned,
+  with rows decoded exactly as `Ecto.Repo.query/4` decodes them.
+
+  The entire result set is materialized (partitions are fetched in parallel), so
+  prefer this for statements returning a bounded number of rows; stream large
+  result sets with `stream_query/5` instead.
+
+  Snowflake only retains a statement's result for a limited period; a handle
+  fetched long after completion returns an error rather than rows.
+
+  ## Examples
+
+      {:ok, handle} = Snowflex.submit_async(MyRepo, "CALL build_report()")
+
+      case Snowflex.fetch_result(MyRepo, handle) do
+        {:ok, :running} -> :not_finished_yet
+        {:ok, %Snowflex.Result{rows: rows}} -> rows
+        {:error, error} -> Logger.error(Exception.message(error))
+      end
+
+  """
+  @spec fetch_result(
+          repo :: Ecto.Repo.t() | pid(),
+          handle :: String.t(),
+          opts :: Keyword.t()
+        ) :: {:ok, Result.t() | :running} | {:error, Error.t()}
+  def fetch_result(repo, handle, opts \\ [])
+      when (is_atom(repo) or is_pid(repo)) and is_binary(handle) do
+    run_op(repo, :fetch_result, handle, [], opts)
+  end
+
+  @doc """
+  Requests cancellation of a statement previously submitted with `submit_async/4`.
+
+  Cancellation is a request to Snowflake, not a guarantee: a statement that has
+  already finished (or that Snowflake no longer knows about) returns an error.
+  """
+  @spec cancel_statement(
+          repo :: Ecto.Repo.t() | pid(),
+          handle :: String.t(),
+          opts :: Keyword.t()
+        ) :: :ok | {:error, Error.t()}
+  def cancel_statement(repo, handle, opts \\ [])
+      when (is_atom(repo) or is_pid(repo)) and is_binary(handle) do
+    case run_op(repo, :cancel, handle, [], opts) do
+      {:ok, %Result{}} -> :ok
+      {:error, error} -> {:error, error}
+    end
+  end
+
+  # Runs a non-:execute transport operation over the pool, reusing the ordinary
+  # execute path so these calls get the same logging, telemetry and error
+  # enrichment as a normal query.
+  #
+  # Runs on the connection held by an enclosing Ecto.Repo.checkout/2 when there
+  # is one: asking the pool for a second connection there would queue behind the
+  # one the caller already holds, which at pool_size: 1 exhausts the queue
+  # timeout and disconnects the outer checkout.
+  defp run_op(repo, op, statement, params, opts) do
+    %{pid: pool, telemetry: telemetry, opts: default_opts} = Adapter.lookup_meta(repo)
+    opts = with_log(telemetry, params, opts ++ default_opts)
+    # Query.new/1 normalizes iodata statements to a binary; building the struct
+    # directly would send an iolist to Snowflake as a JSON array.
+    query = Query.new(statement: statement, op: op)
+
+    pool
+    |> checked_out_conn()
+    |> DBConnection.execute(query, params, opts)
+    |> case do
+      {:ok, _query, result} -> {:ok, result}
+      {:error, error} -> {:error, error}
+    end
+  end
+
+  defp checked_out_conn(pool) do
+    case Process.get({SQL, pool}) do
+      %DBConnection{} = conn -> conn
+      nil -> pool
+    end
+  end
+
   @doc false
   @spec reduce(
           adapter_meta :: map(),

lib/snowflex/connection.ex

@@ -76,7 +76,7 @@ defmodule Snowflex.Connection do
     # Set base metadata at the start so it's available even if DBConnection times out
     set_base_metadata(state, query)
 
-    case transport.execute_statement(state.pid, query.statement, params, opts) do
+    case dispatch(transport, state.pid, query, params, opts) do
       {:ok, result} ->
         {:ok, query, result, state}
 
@@ -86,6 +86,49 @@ defmodule Snowflex.Connection do
     end
   end
 
+  defp dispatch(transport, pid, %{op: :submit_async, statement: statement}, params, opts) do
+    with :ok <- ensure_supported(transport, :submit_async, 4) do
+      transport.submit_async(pid, statement, params, opts)
+    end
+  end
+
+  defp dispatch(transport, pid, %{op: :status, statement: handle}, _params, opts) do
+    with :ok <- ensure_supported(transport, :statement_status, 3) do
+      transport.statement_status(pid, handle, opts)
+    end
+  end
+
+  defp dispatch(transport, pid, %{op: :fetch_result, statement: handle}, _params, opts) do
+    with :ok <- ensure_supported(transport, :fetch_result, 3) do
+      transport.fetch_result(pid, handle, opts)
+    end
+  end
+
+  defp dispatch(transport, pid, %{op: :cancel, statement: handle}, _params, opts) do
+    with :ok <- ensure_supported(transport, :cancel_statement, 3) do
+      transport.cancel_statement(pid, handle, opts)
+    end
+  end
+
+  defp dispatch(transport, pid, query, params, opts) do
+    transport.execute_statement(pid, query.statement, params, opts)
+  end
+
+  # The async callbacks are optional on Snowflex.Transport, so a custom
+  # transport may not implement them. Report that as a Snowflex.Error instead of
+  # letting it surface as an UndefinedFunctionError.
+  defp ensure_supported(transport, fun, arity) do
+    if Code.ensure_loaded?(transport) and function_exported?(transport, fun, arity) do
+      :ok
+    else
+      {:error,
+       Error.exception(
+         "transport #{inspect(transport)} does not implement #{fun}/#{arity}, " <>
+           "which #{inspect(transport)} must provide to support this operation"
+       )}
+    end
+  end
+
   @impl DBConnection
   def handle_close(_query, _opts, state) do
     {:ok, %Result{}, state}

lib/snowflex/query.ex

@@ -1,6 +1,12 @@
 defmodule Snowflex.Query do
   @moduledoc """
   Snowflake implementation of `DBConnection.Query`.
+
+  `:op` selects which transport operation the query runs. It is `:execute` for
+  ordinary statements. For the async operations (`:submit_async`, `:status`,
+  `:fetch_result` and `:cancel`) `:statement` holds a statement handle rather
+  than SQL, so that the handle reaches query logging and telemetry the same way
+  a statement does.
   """
 
   alias String.Chars
@@ -9,7 +15,8 @@ defmodule Snowflex.Query do
     :statement,
     :transport,
     name: "",
-    cache: :reference
+    cache: :reference,
+    op: :execute
   ]
 
   defguard is_statement(statement) when is_list(statement) or is_binary(statement)

lib/snowflex/result.ex

@@ -7,7 +7,8 @@ defmodule Snowflex.Result do
   * `:columns` - List of column names returned in the result set
   * `:rows` - List of tuples containing the row data
   * `:num_rows` - Number of rows in the result set
-  * `:metadata` - Additional metadata about the query execution
+  * `:metadata` - Snowflake's `resultSetMetaData` for the statement, as a map
+    (it defaults to `[]` when no statement has produced metadata yet)
   * `:messages` - Any messages returned by Snowflake during query execution
   * `:query` - The SQL statement that was executed
   * `:query_id` - The ID of the query that was executed
@@ -28,7 +29,7 @@ defmodule Snowflex.Result do
           columns: [String.t()] | nil,
           rows: [tuple()] | nil,
           num_rows: integer(),
-          metadata: [map()],
+          metadata: map() | [map()],
           messages: [map()],
           query: Snowflex.Query.t() | nil,
           query_id: String.t() | nil,

lib/snowflex/transport.ex

@@ -9,6 +9,7 @@ defmodule Snowflex.Transport do
   alias Snowflex.Result
 
   @type query_result :: {:ok, Result.t()} | {:error, Error.t()}
+  @type status_result :: {:ok, :running | :succeeded} | {:error, Error.t()}
   @type connection_opts :: Keyword.t()
   @type cursor :: term()
   @type fetch_result :: {:ok, Result.t()} | {:halt, Result.t()} | {:error, Error.t()}
@@ -44,4 +45,47 @@ defmodule Snowflex.Transport do
   Periodic pings to the server, default is once per second. See `c:DBConnection.ping/1` for more information.
   """
   @callback ping(pid()) :: query_result()
+
+  @doc """
+  Submit a statement for execution and return as soon as the server acknowledges
+  it, without waiting for the statement to finish.
+
+  The returned `t:Snowflex.Result.t/0` carries the server's statement handle in
+  `:query_id`; it has no rows. Use `c:statement_status/3` to check on the
+  statement later and `c:cancel_statement/3` to stop it.
+  """
+  @callback submit_async(pid(), String.t(), any(), Keyword.t()) :: query_result()
+
+  @doc """
+  Report the status of a previously submitted statement.
+
+  Returns `{:ok, :running}` while the statement is executing and
+  `{:ok, :succeeded}` once it has finished successfully. A statement that
+  finished with an error is reported as `{:error, t:Snowflex.Error.t/0}`.
+  """
+  @callback statement_status(pid(), String.t(), Keyword.t()) :: status_result()
+
+  @doc """
+  Fetch the result of a previously submitted statement.
+
+  Returns `{:ok, :running}` if the statement has not finished yet, so a caller
+  can poll with this function alone. Once the statement has completed the full
+  result set is returned as `{:ok, t:Snowflex.Result.t/0}`.
+  """
+  @callback fetch_result(pid(), String.t(), Keyword.t()) ::
+              {:ok, Result.t() | :running} | {:error, Error.t()}
+
+  @doc """
+  Request cancellation of a previously submitted statement.
+  """
+  @callback cancel_statement(pid(), String.t(), Keyword.t()) :: query_result()
+
+  # Optional so that transports predating async submission (and any custom
+  # transport that cannot support it) still satisfy the behaviour. Callers reach
+  # these through Snowflex.Connection, which checks for them and raises a
+  # descriptive error rather than UndefinedFunctionError.
+  @optional_callbacks submit_async: 4,
+                      statement_status: 3,
+                      fetch_result: 3,
+                      cancel_statement: 3
 end