Subscription priming/ordinals

lib/absinthe.ex

@@ -44,9 +44,15 @@ defmodule Absinthe do
           %{message: String.t()}
           | %{message: String.t(), locations: [%{line: pos_integer, column: integer}]}
 
+  @type continuation_t :: nil | [Absinthe.Blueprint.Continuation.t()]
+
   @type result_t ::
-          %{data: nil | result_selection_t}
-          | %{data: nil | result_selection_t, errors: [result_error_t]}
+          %{
+            required(:data) => nil | result_selection_t,
+            optional(:ordinal) => term(),
+            optional(:continuation) => continuation_t,
+            optional(:errors) => [result_error_t]
+          }
           | %{errors: [result_error_t]}
 
   @doc """
@@ -95,7 +101,8 @@ defmodule Absinthe do
           max_complexity: non_neg_integer | :infinity
         ]
 
-  @type run_result :: {:ok, result_t} | {:error, String.t()}
+  @type run_result :: {:ok, result_t} | {:more, result_t} | {:error, String.t()}
+  @type continue_result :: run_result | :no_more_results
 
   @spec run(
           binary | Absinthe.Language.Source.t() | Absinthe.Language.Document.t(),
@@ -107,7 +114,26 @@ defmodule Absinthe do
       schema
       |> Absinthe.Pipeline.for_document(options)
 
-    case Absinthe.Pipeline.run(document, pipeline) do
+    document
+    |> Absinthe.Pipeline.run(pipeline)
+    |> build_result()
+  end
+
+  @spec continue([Absinthe.Blueprint.Continuation.t()]) :: continue_result
+  def continue(continuation) do
+    continuation
+    |> Absinthe.Pipeline.continue()
+    |> build_result()
+  end
+
+  defp build_result(output) do
+    case output do
+      {:ok, %{result: :no_more_results}, _phases} ->
+        :no_more_results
+
+      {:ok, %{result: %{continuation: c} = result}, _phases} when c != [] ->
+        {:more, result}
+
       {:ok, %{result: result}, _phases} ->
         {:ok, result}
 
@@ -131,6 +157,7 @@ defmodule Absinthe do
   def run!(input, schema, options \\ []) do
     case run(input, schema, options) do
       {:ok, result} -> result
+      {:more, result} -> result
       {:error, err} -> raise ExecutionError, message: err
     end
   end

lib/absinthe/blueprint/continuation.ex

@@ -0,0 +1,18 @@
+defmodule Absinthe.Blueprint.Continuation do
+  @moduledoc false
+
+  # Continuations allow further resolutions after the initial result is
+  # returned
+
+  alias Absinthe.Pipeline
+
+  defstruct [
+    :phase_input,
+    :pipeline
+  ]
+
+  @type t :: %__MODULE__{
+          phase_input: Pipeline.data_t(),
+          pipeline: Pipeline.t()
+        }
+end

lib/absinthe/blueprint/result/list.ex

@@ -9,14 +9,16 @@ defmodule Absinthe.Blueprint.Result.List do
     :values,
     errors: [],
     flags: %{},
-    extensions: %{}
+    extensions: %{},
+    continuations: []
   ]
 
   @type t :: %__MODULE__{
           emitter: Blueprint.Document.Field.t(),
           values: [Blueprint.Execution.node_t()],
           errors: [Phase.Error.t()],
           flags: Blueprint.flags_t(),
-          extensions: %{any => any}
+          extensions: %{any => any},
+          continuations: [Blueprint.Continuation.t()]
         }
 end

lib/absinthe/blueprint/result/object.ex

@@ -10,14 +10,16 @@ defmodule Absinthe.Blueprint.Result.Object do
     :fields,
     errors: [],
     flags: %{},
-    extensions: %{}
+    extensions: %{},
+    continuations: []
   ]
 
   @type t :: %__MODULE__{
           emitter: Blueprint.Document.Field.t(),
           fields: [Blueprint.Execution.node_t()],
           errors: [Phase.Error.t()],
           flags: Blueprint.flags_t(),
-          extensions: %{any => any}
+          extensions: %{any => any},
+          continuations: [Blueprint.Continuation.t()]
         }
 end

lib/absinthe/phase/document/result.ex

@@ -25,7 +25,9 @@ defmodule Absinthe.Phase.Document.Result do
           {:validation_failed, errors}
       end
 
-    format_result(result)
+    result
+    |> format_result()
+    |> maybe_add_continuations(blueprint.execution.result)
   end
 
   defp format_result({:ok, {data, []}}) do
@@ -134,4 +136,9 @@ defmodule Absinthe.Phase.Document.Result do
   end
 
   defp format_location(_), do: []
+
+  defp maybe_add_continuations(result, %{continuations: continuations}) when continuations != [],
+    do: Map.put(result, :continuation, continuations)
+
+  defp maybe_add_continuations(result, _), do: result
 end

lib/absinthe/phase/subscription/get_ordinal.ex

@@ -0,0 +1,42 @@
+defmodule Absinthe.Phase.Subscription.GetOrdinal do
+  use Absinthe.Phase
+
+  alias Absinthe.Phase.Subscription.SubscribeSelf
+
+  @moduledoc false
+
+  alias Absinthe.Blueprint
+
+  @spec run(any, Keyword.t()) :: {:ok, Blueprint.t()}
+  def run(blueprint, _options \\ []) do
+    op = Blueprint.current_operation(blueprint)
+
+    if op.type == :subscription do
+      {:ok,
+       %{blueprint | result: Map.put(blueprint.result, :ordinal, get_ordinal(op, blueprint))}}
+    else
+      {:ok, blueprint}
+    end
+  end
+
+  defp get_ordinal(op, blueprint) do
+    %{selections: [field]} = op
+    {:ok, config} = SubscribeSelf.get_config(field, blueprint.execution.context, blueprint)
+
+    case config[:ordinal] do
+      nil ->
+        nil
+
+      fun when is_function(fun, 1) ->
+        fun.(blueprint.execution.root_value)
+
+      _fun ->
+        IO.write(
+          :stderr,
+          "Ordinal function must be 1-arity"
+        )
+
+        nil
+    end
+  end
+end

lib/absinthe/phase/subscription/prime.ex

@@ -0,0 +1,46 @@
+defmodule Absinthe.Phase.Subscription.Prime do
+  @moduledoc false
+
+  alias Absinthe.Blueprint.Continuation
+  alias Absinthe.Phase
+
+  @spec run(any(), Keyword.t()) :: Absinthe.Phase.result_t()
+  def run(blueprint, prime_result: prime_result) do
+    {:ok, put_in(blueprint.execution.root_value, prime_result)}
+  end
+
+  def run(blueprint, prime_fun: prime_fun, resolution_options: options) do
+    {:ok, prime_results} = prime_fun.(blueprint.execution)
+
+    case prime_results do
+      [first | rest] ->
+        blueprint = put_in(blueprint.execution.root_value, first)
+        blueprint = maybe_add_continuations(blueprint, rest, options)
+        {:ok, blueprint}
+
+      [] ->
+        blueprint = put_in(blueprint.result, :no_more_results)
+        {:replace, blueprint, []}
+    end
+  end
+
+  defp maybe_add_continuations(blueprint, [], _options), do: blueprint
+
+  defp maybe_add_continuations(blueprint, remaining_results, options) do
+    continuations =
+      Enum.map(
+        remaining_results,
+        &%Continuation{
+          phase_input: blueprint,
+          pipeline: [
+            {__MODULE__, [prime_result: &1]},
+            {Phase.Document.Execution.Resolution, options},
+            Phase.Subscription.GetOrdinal,
+            Phase.Document.Result
+          ]
+        }
+      )
+
+    put_in(blueprint.result, %{continuation: continuations})
+  end
+end

lib/absinthe/phase/subscription/result.ex

@@ -5,10 +5,44 @@ defmodule Absinthe.Phase.Subscription.Result do
   # subscription
 
   alias Absinthe.Blueprint
+  alias Absinthe.Blueprint.Continuation
+  alias Absinthe.Phase
 
   @spec run(any, Keyword.t()) :: {:ok, Blueprint.t()}
-  def run(blueprint, topic: topic) do
+  def run(blueprint, options) do
+    topic = Keyword.fetch!(options, :topic)
+    prime = Keyword.get(options, :prime)
     result = %{"subscribed" => topic}
+
+    case prime do
+      nil ->
+        {:ok, put_in(blueprint.result, result)}
+
+      prime_fun when is_function(prime_fun, 1) ->
+        stash_prime(prime_fun, result, blueprint, options)
+
+      val ->
+        raise """
+        Invalid prime function. Must be a function of arity 1.
+
+        #{inspect(val)}
+        """
+    end
+  end
+
+  def stash_prime(prime_fun, base_result, blueprint, options) do
+    continuation = %Continuation{
+      phase_input: blueprint,
+      pipeline: [
+        {Phase.Subscription.Prime, [prime_fun: prime_fun, resolution_options: options]},
+        {Phase.Document.Execution.Resolution, options},
+        Phase.Subscription.GetOrdinal,
+        Phase.Document.Result
+      ]
+    }
+
+    result = Map.put(base_result, :continuation, [continuation])
+
     {:ok, put_in(blueprint.result, result)}
   end
 end

lib/absinthe/phase/subscription/subscribe_self.ex

@@ -30,7 +30,7 @@ defmodule Absinthe.Phase.Subscription.SubscribeSelf do
 
       {:replace, blueprint,
        [
-         {Phase.Subscription.Result, topic: subscription_id},
+         {Phase.Subscription.Result, topic: subscription_id, prime: config[:prime]},
          {Phase.Telemetry, Keyword.put(options, :event, [:execute, :operation, :stop])}
        ]}
     else
@@ -45,11 +45,11 @@ defmodule Absinthe.Phase.Subscription.SubscribeSelf do
     end
   end
 
-  defp get_config(
-         %{schema_node: schema_node, argument_data: argument_data} = field,
-         context,
-         blueprint
-       ) do
+  def get_config(
+        %{schema_node: schema_node, argument_data: argument_data} = field,
+        context,
+        blueprint
+      ) do
     name = schema_node.identifier
 
     config =

lib/absinthe/pipeline.ex

@@ -12,21 +12,42 @@ defmodule Absinthe.Pipeline do
   * See `Absinthe.Schema` on adjusting the schema pipeline for schema manipulation.
   """
 
+  alias Absinthe.Blueprint.Continuation
   alias Absinthe.Phase
 
   @type data_t :: any
 
+  @type run_result_t :: {:ok, data_t, [Phase.t()]} | {:error, String.t(), [Phase.t()]}
+
   @type phase_config_t :: Phase.t() | {Phase.t(), Keyword.t()}
 
   @type t :: [phase_config_t | [phase_config_t]]
 
-  @spec run(data_t, t) :: {:ok, data_t, [Phase.t()]} | {:error, String.t(), [Phase.t()]}
+  @spec run(data_t, t) :: run_result_t
   def run(input, pipeline) do
     pipeline
     |> List.flatten()
     |> run_phase(input)
   end
 
+  @spec continue([Continuation.t()]) :: run_result_t
+  def continue([continuation | rest]) do
+    result = run_phase(continuation.pipeline, continuation.phase_input)
+
+    case result do
+      {:ok, blueprint, phases} when rest == [] ->
+        {:ok, blueprint, phases}
+
+      {:ok, blueprint, phases} ->
+        bp_result = Map.put(blueprint.result, :continuation, rest)
+        blueprint = Map.put(blueprint, :result, bp_result)
+        {:ok, blueprint, phases}
+
+      error ->
+        error
+    end
+  end
+
   @defaults [
     adapter: Absinthe.Adapter.LanguageConventions,
     operation_name: nil,
@@ -114,6 +135,7 @@ defmodule Absinthe.Pipeline do
       # Execution
       {Phase.Subscription.SubscribeSelf, options},
       {Phase.Document.Execution.Resolution, options},
+      Phase.Subscription.GetOrdinal,
       # Format Result
       Phase.Document.Result,
       {Phase.Telemetry, Keyword.put(options, :event, [:execute, :operation, :stop])}
@@ -388,8 +410,8 @@ defmodule Absinthe.Pipeline do
     end)
   end
 
-  @spec run_phase(t, data_t, [Phase.t()]) ::
-          {:ok, data_t, [Phase.t()]} | {:error, String.t(), [Phase.t()]}
+  @spec run_phase(t, data_t, [Phase.t()]) :: run_result_t
+
   def run_phase(pipeline, input, done \\ [])
 
   def run_phase([], input, done) do

lib/absinthe/subscription.ex

@@ -47,6 +47,8 @@ defmodule Absinthe.Subscription do
 
   @type subscription_field_spec :: {atom, term | (term -> term)}
 
+  @type prime_fun :: (Absinthe.Resolution.t() -> {:ok, [map()]})
+
   @doc """
   Publish a mutation
 

lib/absinthe/subscription/local.ex

@@ -46,7 +46,10 @@ defmodule Absinthe.Subscription.Local do
           |> Pipeline.without(Phase.Subscription.SubscribeSelf)
           |> Pipeline.insert_before(
             Phase.Document.Execution.Resolution,
-            {Phase.Document.OverrideRoot, root_value: mutation_result}
+            [
+              {Phase.Document.OverrideRoot, root_value: mutation_result},
+              Phase.Subscription.GetOrdinal
+            ]
           )
           |> Pipeline.upto(Phase.Document.Execution.Resolution)