Add subscription ordinal

Author: bernardd

lib/absinthe.ex

@@ -44,12 +44,15 @@ defmodule Absinthe do
           %{message: String.t()}
           | %{message: String.t(), locations: [%{line: pos_integer, column: integer}]}
 
-  @type continuation_t :: nil | [Continuation.t()]
+  @type continuation_t :: nil | [Absinthe.Blueprint.Continuation.t()]
 
   @type result_t ::
-          %{required(:data) => nil | result_selection_t,
+          %{
+            required(:data) => nil | result_selection_t,
+            optional(:ordinal) => term(),
             optional(:continuation) => continuation_t,
-            optional(:errors) => [result_error_t]}
+            optional(:errors) => [result_error_t]
+          }
           | %{errors: [result_error_t]}
 
   @doc """
@@ -115,7 +118,7 @@ defmodule Absinthe do
     |> build_result()
   end
 
-  @spec continue([Continuation.t()]) :: run_result()
+  @spec continue([Absinthe.Blueprint.Continuation.t()]) :: run_result()
   def continue(continuation) do
     continuation
     |> Absinthe.Pipeline.continue()

lib/absinthe/blueprint/continuation.ex

@@ -12,8 +12,7 @@ defmodule Absinthe.Blueprint.Continuation do
   ]
 
   @type t :: %__MODULE__{
-    phase_input: Pipeline.data_t,
-    pipeline: Pipeline.t()
-  }
-
+          phase_input: Pipeline.data_t(),
+          pipeline: Pipeline.t()
+        }
 end

lib/absinthe/blueprint/result/list.ex

@@ -19,6 +19,6 @@ defmodule Absinthe.Blueprint.Result.List do
           errors: [Phase.Error.t()],
           flags: Blueprint.flags_t(),
           extensions: %{any => any},
-          continuations: [Continuation.t()]
+          continuations: [Blueprint.Continuation.t()]
         }
 end

lib/absinthe/blueprint/result/object.ex

@@ -20,6 +20,6 @@ defmodule Absinthe.Blueprint.Result.Object do
           errors: [Phase.Error.t()],
           flags: Blueprint.flags_t(),
           extensions: %{any => any},
-          continuations: [Continuation.t()]
+          continuations: [Blueprint.Continuation.t()]
         }
 end

lib/absinthe/phase/document/result.ex

@@ -138,7 +138,7 @@ defmodule Absinthe.Phase.Document.Result do
   defp format_location(_), do: []
 
   defp maybe_add_continuations(result, %{continuations: continuations}) when continuations != [],
-  do: Map.put(result, :continuation, 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

@@ -1,8 +1,8 @@
 defmodule Absinthe.Phase.Subscription.Prime do
   @moduledoc false
 
-  @spec run(any(), Keyword.t()) :: Phase.result_t()
-  def run(blueprint, [prime_result: cr]) do
+  @spec run(any(), Keyword.t()) :: Absinthe.Phase.result_t()
+  def run(blueprint, prime_result: cr) do
     {:ok, put_in(blueprint.execution.root_value, cr)}
   end
 end

lib/absinthe/phase/subscription/result.ex

@@ -6,46 +6,53 @@ defmodule Absinthe.Phase.Subscription.Result do
 
   alias Absinthe.Blueprint
   alias Absinthe.Blueprint.Continuation
+  alias Absinthe.Phase
 
   @spec run(any, Keyword.t()) :: {:ok, Blueprint.t()}
   def run(blueprint, options) do
-    topic = Keyword.get(options, :topic)
+    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, 0) ->
-        {:ok, prime_results} = prime_fun.()
-
-        result =
-          if prime_results != [] do
-            continuations =
-              Enum.map(prime_results, fn cr ->
-                %Continuation{
-                  phase_input: blueprint,
-                  pipeline: [
-                    {Absinthe.Phase.Subscription.Prime, [prime_result: cr]},
-                    {Absinthe.Phase.Document.Execution.Resolution, options},
-                    Absinthe.Phase.Document.Result
-                  ]
-                }
-              end)
-
-            Map.put(result, :continuation, continuations)
-          else
-            result
-          end
-
-        {:ok, put_in(blueprint.result, result)}
+      prime_fun when is_function(prime_fun, 1) ->
+        do_prime(prime_fun, result, blueprint, options)
 
       val ->
         raise """
-        Invalid prime function. Must be a function of arity 0.
+        Invalid prime function. Must be a function of arity 1.
 
         #{inspect(val)}
         """
     end
   end
+
+  def do_prime(prime_fun, base_result, blueprint, options) do
+    {:ok, prime_results} = prime_fun.(blueprint.execution)
+
+    result =
+      if prime_results != [] do
+        continuations =
+          Enum.map(prime_results, fn cr ->
+            %Continuation{
+              phase_input: blueprint,
+              pipeline: [
+                {Phase.Subscription.Prime, [prime_result: cr]},
+                {Phase.Document.Execution.Resolution, options},
+                Phase.Subscription.GetOrdinal,
+                Phase.Document.Result
+              ]
+            }
+          end)
+
+        Map.put(base_result, :continuation, continuations)
+      else
+        base_result
+      end
+
+    {:ok, put_in(blueprint.result, result)}
+  end
 end

lib/absinthe/phase/subscription/subscribe_self.ex

@@ -22,15 +22,15 @@ defmodule Absinthe.Phase.Subscription.SubscribeSelf do
     %{selections: [field]} = op
 
     with {:ok, config} <- get_config(field, context, blueprint) do
-      {field_keys, prime} = get_field_keys(field, config)
+      field_keys = get_field_keys(field, config)
       subscription_id = get_subscription_id(config, blueprint, options)
 
       for field_key <- field_keys,
           do: Absinthe.Subscription.subscribe(pubsub, field_key, subscription_id, blueprint)
 
       {:replace, blueprint,
        [
-         {Phase.Subscription.Result, topic: subscription_id, prime: prime},
+         {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 =
@@ -96,9 +96,8 @@ defmodule Absinthe.Phase.Subscription.SubscribeSelf do
   defp get_field_keys(%{schema_node: schema_node} = _field, config) do
     name = schema_node.identifier
 
-    {keys, prime} = find_field_keys!(config)
-    field_keys = Enum.map(keys, fn key -> {name, key} end)
-    {field_keys, prime}
+    find_field_keys!(config)
+    |> Enum.map(fn key -> {name, key} end)
   end
 
   defp ensure_pubsub!(context) do
@@ -133,12 +132,8 @@ defmodule Absinthe.Phase.Subscription.SubscribeSelf do
         """
 
       val ->
-        topics = List.wrap(val)
-                 |> Enum.map(&to_string/1)
-
-        prime = config[:prime] || nil
-
-        {topics, prime}
+        List.wrap(val)
+        |> Enum.map(&to_string/1)
     end
   end
 

lib/absinthe/pipeline.ex

@@ -37,11 +37,14 @@ defmodule Absinthe.Pipeline do
     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
+
+      error ->
+        error
     end
   end
 
@@ -132,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])}

lib/absinthe/subscription.ex

@@ -47,7 +47,7 @@ defmodule Absinthe.Subscription do
 
   @type subscription_field_spec :: {atom, term | (term -> term)}
 
-  @type prime_fun :: (-> {:ok, [map()]})
+  @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)