Search scroll API

lib/snap/scroll.ex

@@ -0,0 +1,96 @@
+defmodule Snap.Scroll do
+  @moduledoc """
+  Streams documents from a search query using the [scroll API](https://www.elastic.co/guide/en/elasticsearch/reference/current/paginate-search-results.html#scroll-search-results).
+
+  Returns a lazy `Stream` that yields `Snap.Hit` structs one at a time,
+  transparently fetching subsequent batches via `Snap.Search.scroll/6` as
+  needed, then clearing the scroll cursor via `Snap.Search.clear_scroll/5`
+  when the stream is exhausted, terminated early, or raises.
+
+  Unlike most functions in `Snap`, errors are not returned as `{:error, _}`
+  tuples — lazy enumerables can't surface errors mid-iteration. Instead, any
+  underlying request failure raises the corresponding exception (typically
+  `Snap.ResponseError`) when the consumer pulls the next element.
+  """
+  alias Snap.Hits
+  alias Snap.Search
+  alias Snap.SearchResponse
+
+  @spec stream(
+          cluster :: module(),
+          index_or_alias :: String.t(),
+          query :: map(),
+          opts :: Keyword.t()
+        ) :: Enumerable.t()
+  @doc """
+  Returns a lazy stream of `Snap.Hit` structs matching `query` in `index_or_alias`.
+
+  ## Options
+
+  - `:scroll` — TTL string for the server-side cursor, refreshed on each
+    continuation. Defaults to `"1m"`.
+  - `:params` — extra query params merged into the initial search and each
+    scroll continuation.
+  - `:headers` — passed through to the underlying requests.
+  - `:opts` — request opts passed through to the underlying HTTP client.
+
+  ## Examples
+
+      query = %{"query" => %{"match_all" => %{}}, "size" => 100}
+
+      Snap.Scroll.stream(Cluster, "products", query)
+      |> Stream.map(& &1.source)
+      |> Enum.each(&IO.inspect/1)
+  """
+  def stream(cluster, index_or_alias, query, opts \\ []) do
+    scroll = Keyword.get(opts, :scroll, "1m")
+    params = Keyword.get(opts, :params, [])
+    headers = Keyword.get(opts, :headers, [])
+    request_opts = Keyword.get(opts, :opts, [])
+
+    start_fun = fn -> {:start, query} end
+
+    next_fun = fn
+      {:start, query} ->
+        initial_params = Keyword.put(params, :scroll, scroll)
+
+        case Search.search(cluster, index_or_alias, query, initial_params, headers, request_opts) do
+          {:ok, %SearchResponse{hits: %Hits{hits: []}, scroll_id: sid}} ->
+            {:halt, sid}
+
+          {:ok, %SearchResponse{hits: %Hits{hits: hits}, scroll_id: sid}} ->
+            {hits, {:continue, sid}}
+
+          {:error, exc} ->
+            raise exc
+        end
+
+      {:continue, scroll_id} ->
+        case Search.scroll(cluster, scroll_id, scroll, params, headers, request_opts) do
+          {:ok, %SearchResponse{hits: %Hits{hits: []}, scroll_id: sid}} ->
+            {:halt, sid || scroll_id}
+
+          {:ok, %SearchResponse{hits: %Hits{hits: hits}, scroll_id: sid}} ->
+            {hits, {:continue, sid || scroll_id}}
+
+          {:error, exc} ->
+            raise exc
+        end
+    end
+
+    after_fun = fn
+      {:continue, scroll_id} when is_binary(scroll_id) ->
+        Search.clear_scroll(cluster, scroll_id, params, headers, request_opts)
+        :ok
+
+      scroll_id when is_binary(scroll_id) ->
+        Search.clear_scroll(cluster, scroll_id, params, headers, request_opts)
+        :ok
+
+      _ ->
+        :ok
+    end
+
+    Stream.resource(start_fun, next_fun, after_fun)
+  end
+end

lib/snap/search.ex

@@ -42,6 +42,55 @@ defmodule Snap.Search do
     end
   end
 
+  @spec scroll(
+          cluster :: module(),
+          scroll_id :: String.t(),
+          scroll :: String.t(),
+          params :: Keyword.t(),
+          headers :: Keyword.t(),
+          opts :: Keyword.t()
+        ) :: {:ok, SearchResponse.t()} | Snap.Cluster.error()
+  @doc """
+  Continues a scroll search, retrieving the next batch of hits.
+
+  Given a `scroll_id` returned from a previous `search/6` (with a `scroll`
+  parameter) or `scroll/6` call, fetches the next batch of hits and returns a
+  fresh `Snap.SearchResponse` with an updated `scroll_id`.
+
+  The `scroll` argument is the lifetime of the scroll cursor on the server,
+  refreshed on each call. Defaults to `"1m"`.
+
+  This endpoint is global on the cluster — index namespaces do not apply.
+  """
+  def scroll(cluster, scroll_id, scroll \\ "1m", params \\ [], headers \\ [], opts \\ []) do
+    body = %{"scroll" => scroll, "scroll_id" => scroll_id}
+
+    case cluster.post("/_search/scroll", body, params, headers, opts) do
+      {:ok, response} -> {:ok, SearchResponse.new(response)}
+      err -> err
+    end
+  end
+
+  @spec clear_scroll(
+          cluster :: module(),
+          scroll_id :: String.t(),
+          params :: Keyword.t(),
+          headers :: Keyword.t(),
+          opts :: Keyword.t()
+        ) :: {:ok, map()} | Snap.Cluster.error()
+  @doc """
+  Releases a scroll cursor on the server, freeing the resources it holds.
+
+  Once a scroll is exhausted (or no longer needed), call this to clear the
+  underlying search context. Scrolls also expire naturally based on the TTL
+  passed to `search/6` or `scroll/6`, but explicit cleanup is preferred.
+  """
+  def clear_scroll(cluster, scroll_id, params \\ [], headers \\ [], opts \\ []) do
+    encoded = Request.encode_segment(scroll_id)
+
+    cluster.delete("/_search/scroll/#{encoded}", params, headers, opts)
+  end
+
   @doc """
   Runs a count of the documents in an index, using an optional query.
   """

test/scroll_test.exs

@@ -0,0 +1,71 @@
+defmodule Snap.ScrollTest do
+  @moduledoc false
+  use Snap.IntegrationCase, async: true
+
+  alias Snap.Bulk.Action
+  alias Snap.Scroll
+  alias Snap.Test.Cluster
+
+  @test_index "scroll"
+
+  test "stream/4 yields every hit across multiple scroll batches" do
+    {:ok, _} = Snap.Indexes.create(Cluster, @test_index, %{})
+
+    1..7
+    |> Enum.map(fn i -> %Action.Index{id: i, doc: %{"n" => i}} end)
+    |> Snap.Bulk.perform(Cluster, @test_index, refresh: true)
+
+    query = %{"query" => %{"match_all" => %{}}, "size" => 2, "sort" => ["_doc"]}
+
+    hits = Scroll.stream(Cluster, @test_index, query) |> Enum.to_list()
+
+    assert length(hits) == 7
+    assert Enum.map(hits, & &1.id) == ~w(1 2 3 4 5 6 7)
+  end
+
+  test "stream/4 calls clear_scroll when the consumer halts early" do
+    {:ok, _} = Snap.Indexes.create(Cluster, @test_index, %{})
+
+    1..10
+    |> Enum.map(fn i -> %Action.Index{id: i, doc: %{"n" => i}} end)
+    |> Snap.Bulk.perform(Cluster, @test_index, refresh: true)
+
+    parent = self()
+    handler_id = {__MODULE__, :clear_scroll_observer, make_ref()}
+
+    :telemetry.attach(
+      handler_id,
+      [:snap, :snap, :request],
+      fn _event, _measurements, metadata, _config ->
+        send(parent, {:request, metadata.method, metadata.path})
+      end,
+      nil
+    )
+
+    on_exit(fn -> :telemetry.detach(handler_id) end)
+
+    query = %{"query" => %{"match_all" => %{}}, "size" => 2, "sort" => ["_doc"]}
+
+    hits = Scroll.stream(Cluster, @test_index, query) |> Enum.take(3)
+
+    assert length(hits) == 3
+    assert_receive {:request, "DELETE", "/_search/scroll/" <> _}, 1_000
+  end
+
+  test "stream/4 returns an empty stream when nothing matches" do
+    {:ok, _} = Snap.Indexes.create(Cluster, @test_index, %{})
+    :ok = Snap.Indexes.refresh(Cluster, @test_index)
+
+    query = %{"query" => %{"match_all" => %{}}, "size" => 2}
+
+    assert [] = Scroll.stream(Cluster, @test_index, query) |> Enum.to_list()
+  end
+
+  test "stream/4 raises if the initial search fails" do
+    query = %{"query" => %{"this_is_not_a_real_query" => %{}}}
+
+    assert_raise Snap.ResponseError, fn ->
+      Scroll.stream(Cluster, @test_index, query) |> Enum.to_list()
+    end
+  end
+end

test/search_test.exs

@@ -150,6 +150,51 @@ defmodule Snap.SearchTest do
     assert {:ok, 1} = Snap.Search.count(Cluster, @test_index, %{query: %{term: %{foo: "bar"}}})
   end
 
+  test "scroll/2 retrieves subsequent batches" do
+    {:ok, _} = Snap.Indexes.create(Cluster, @test_index, %{})
+
+    1..5
+    |> Enum.map(fn i -> %Action.Index{id: i, doc: %{"title" => "Document #{i}"}} end)
+    |> Snap.Bulk.perform(Cluster, @test_index, refresh: true)
+
+    query = %{"query" => %{"match_all" => %{}}, "sort" => ["_doc"]}
+
+    {:ok, first} = Search.search(Cluster, @test_index, query, scroll: "1m", size: 2)
+    assert Enum.count(first) == 2
+    assert is_binary(first.scroll_id)
+
+    {:ok, second} = Search.scroll(Cluster, first.scroll_id)
+    assert Enum.count(second) == 2
+    assert Enum.at(second, 0).id == "3"
+
+    {:ok, third} = Search.scroll(Cluster, second.scroll_id, "30s")
+    assert Enum.count(third) == 1
+    assert Enum.at(third, 0).id == "5"
+  end
+
+  test "scroll/2 returns an error for an unknown scroll_id" do
+    assert {:error, %Snap.ResponseError{type: type}} =
+             Search.scroll(Cluster, "not-a-real-scroll-id")
+
+    assert type in ["search_context_missing_exception", "illegal_argument_exception"]
+  end
+
+  test "clear_scroll/2 releases the cursor" do
+    {:ok, _} = Snap.Indexes.create(Cluster, @test_index, %{})
+
+    1..3
+    |> Enum.map(fn i -> %Action.Index{id: i, doc: %{"title" => "Document #{i}"}} end)
+    |> Snap.Bulk.perform(Cluster, @test_index, refresh: true)
+
+    query = %{"query" => %{"match_all" => %{}}, "sort" => ["_doc"]}
+    {:ok, first} = Search.search(Cluster, @test_index, query, scroll: "1m", size: 1)
+
+    assert {:ok, %{"succeeded" => true, "num_freed" => freed}} =
+             Search.clear_scroll(Cluster, first.scroll_id)
+
+    assert freed >= 1
+  end
+
   test "delete_by_query" do
     {:ok, _} = Snap.Indexes.create(Cluster, @test_index, %{})
     {:ok, _} = Snap.Document.add(Cluster, @test_index, %{foo: "bar"})