defmodule Mongo.Stream do alias Mongo.Session alias Mongo.Error import Record, only: [defrecordp: 2] defstruct [:topology_pid, :session, :cursor, :coll, :docs, :opts] alias Mongo.Session def new(topology_pid, cmd, opts) do ## check, if retryable reads are enabled opts = Mongo.retryable_reads(opts) with {:ok, session} <- Session.start_implicit_session(topology_pid, :read, opts), {:ok, %{"ok" => ok, "cursor" => %{ "id" => cursor_id, "ns" => coll, "firstBatch" => docs}}} when ok == 1 <- Mongo.exec_command_session(session, cmd, opts) do %Mongo.Stream{topology_pid: topology_pid, session: session, cursor: cursor_id, coll: coll, docs: docs, opts: opts} else {:error, error} -> case Error.should_retry_read(error, cmd, opts) do true -> new(topology_pid, cmd, Keyword.put(opts, :read_counter, 2)) false -> {:error, error} end end end defimpl Enumerable do defrecordp :state, [:topology_pid, :session, :cursor, :coll, :docs] def reduce(%Mongo.Stream{topology_pid: topology_pid, session: session, cursor: cursor_id, coll: coll, docs: docs, opts: opts}, acc, reduce_fun) do start_fun = fn -> state(topology_pid: topology_pid, session: session, cursor: cursor_id, coll: coll, docs: docs) end next_fun = next_fun(opts) after_fun = after_fun(opts) Stream.resource(start_fun, next_fun, after_fun).(acc, reduce_fun) end defp next_fun(opts) do fn state(docs: [], cursor: 0) = state -> {:halt, state} # this is a regular cursor state(docs: [], topology_pid: topology_pid, session: session, cursor: cursor, coll: coll) = state -> case get_more(topology_pid, session, only_coll(coll), cursor, nil, opts) do {:ok, %{cursor_id: cursor_id, docs: []}} -> {:halt, state(state, cursor: cursor_id)} {:ok, %{cursor_id: cursor_id, docs: docs}} -> {docs, state(state, cursor: cursor_id)} {:error, error} -> raise error end state(docs: docs) = state -> {docs, state(state, docs: [])} {:error, error} -> raise error ## In case of an error, we should raise the error end end @doc """ Calls the GetCore-Command See https://github.com/mongodb/specifications/blob/master/source/find_getmore_killcursors_commands.rst """ def get_more(_topology_pid, session, coll, cursor, nil, opts) do cmd = [ getMore: %BSON.LongNumber{value: cursor}, collection: coll, batchSize: opts[:batch_size], maxTimeMS: opts[:max_time] ] |> filter_nils() with {:ok, %{"cursor" => %{ "id" => cursor_id, "nextBatch" => docs}, "ok" => ok}} when ok == 1 <- Mongo.exec_command_session(session, cmd, opts) do {:ok, %{cursor_id: cursor_id, docs: docs}} end end @doc""" Calls the KillCursors-Command See https://github.com/mongodb/specifications/blob/master/source/find_getmore_killcursors_commands.rst """ def kill_cursors(session, coll, cursor_ids, opts) do cmd = [ killCursors: coll, cursors: cursor_ids |> Enum.map(fn id -> %BSON.LongNumber{value: id} end) ] |> filter_nils() with {:ok, %{"cursorsAlive" => [], "cursorsNotFound" => [], "cursorsUnknown" => [], "ok" => ok}} when ok == 1 <- Mongo.exec_command_session(session, cmd, opts) do :ok end end defp filter_nils(keyword) when is_list(keyword) do Enum.reject(keyword, fn {_key, value} -> is_nil(value) end) end defp after_fun(opts) do fn state(topology_pid: topology_pid, session: session, cursor: 0) -> Session.end_implict_session(topology_pid, session) state(topology_pid: topology_pid, session: session, cursor: cursor, coll: coll) -> with :ok <- kill_cursors(session, only_coll(coll), [cursor], opts) do Session.end_implict_session(topology_pid, session) end error -> error end end defp only_coll(coll) do [_db, coll] = String.split(coll, ".", parts: 2) coll end # we cannot determinstically slice, so tell Enumerable to # fall back on brute force def slice(_cursor), do: { :error, __MODULE__ } def count(_stream), do: {:error, __MODULE__} def member?(_stream, _term), do: {:error, __MODULE__} end end