Add optional ordinal comparison fun

Author: bernardd

lib/absinthe.ex

@@ -46,10 +46,15 @@ defmodule Absinthe do
 
   @type continuations_t :: nil | [Absinthe.Blueprint.Continuation.t()]
 
+  @type ordinal_fun :: (term() -> term())
+
+  @type ordinal_compare_fun :: (term(), term() -> {boolean(), term()})
+
   @type result_t ::
           %{
             required(:data) => nil | result_selection_t,
-            optional(:ordinal) => term(),
+            optional(:ordinal_fun) => ordinal_fun(),
+            optional(:ordinal_compare_fun) => ordinal_compare_fun(),
             optional(:continuation) => continuations_t,
             optional(:errors) => [result_error_t]
           }

lib/absinthe/phase/subscription/get_ordinal.ex

@@ -11,20 +11,44 @@ defmodule Absinthe.Phase.Subscription.GetOrdinal do
   def run(blueprint, _options \\ []) do
     with %{type: :subscription, selections: [field]} <- Blueprint.current_operation(blueprint),
          {:ok, config} = SubscribeSelf.get_config(field, blueprint.execution.context, blueprint),
-         ordinal_fun when is_function(ordinal_fun, 1) <- config[:ordinal] do
-      result = ordinal_fun.(blueprint.execution.root_value)
-      {:ok, %{blueprint | result: Map.put(blueprint.result, :ordinal, result)}}
+         {_, ordinal_fun} when is_function(ordinal_fun, 1) <- {:ordinal_fun, config[:ordinal]},
+         {_, ordinal_compare_fun} when is_function(ordinal_compare_fun, 2) <-
+           {:ordinal_compare_fun,
+            Keyword.get(config, :ordinal_compare, &default_ordinal_compare/2)} do
+      ordinal = ordinal_fun.(blueprint.execution.root_value)
+
+      result =
+        blueprint.result
+        |> Map.put(:ordinal, ordinal)
+        |> Map.put(:ordinal_compare_fun, ordinal_compare_fun)
+
+      {:ok, %{blueprint | result: result}}
     else
-      f when is_function(f) ->
+      {:ordinal_fun, f} when is_function(f) ->
         IO.write(
           :stderr,
           "Ordinal function must be 1-arity"
         )
 
         {:ok, blueprint}
 
+      {:ordinal_compare_fun, f} when is_function(f) ->
+        IO.write(
+          :stderr,
+          "Ordinal compare function must be 2-arity"
+        )
+
+        {:ok, blueprint}
+
       _ ->
         {:ok, blueprint}
     end
   end
+
+  defp default_ordinal_compare(nil, new_ordinal), do: {true, new_ordinal}
+
+  defp default_ordinal_compare(old_ordinal, new_ordinal) when old_ordinal < new_ordinal,
+    do: {true, new_ordinal}
+
+  defp default_ordinal_compare(old_ordinal, _new_ordinal), do: {false, old_ordinal}
 end

lib/absinthe/phase/subscription/result.ex

@@ -12,25 +12,15 @@ defmodule Absinthe.Phase.Subscription.Result 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)}
+    result = maybe_add_prime(%{"subscribed" => topic}, prime, blueprint, options)
 
-      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
+    {:ok, put_in(blueprint.result, result)}
   end
 
-  def stash_prime(prime_fun, base_result, blueprint, options) do
+  def maybe_add_prime(result, nil, _blueprint, _options), do: result
+
+  def maybe_add_prime(result, prime_fun, blueprint, options) when is_function(prime_fun, 1) do
     continuation = %Continuation{
       phase_input: blueprint,
       pipeline: [
@@ -41,8 +31,14 @@ defmodule Absinthe.Phase.Subscription.Result do
       ]
     }
 
-    result = Map.put(base_result, :continuations, [continuation])
+    Map.put(result, :continuations, [continuation])
+  end
 
-    {:ok, put_in(blueprint.result, result)}
+  def maybe_add_prime(_result, prime_fun, _blueprint, _options) do
+    raise """
+        Invalid prime function. Must be a function of arity 1.
+
+    #{inspect(prime_fun)}
+    """
   end
 end

test/absinthe/execution/subscription_test.exs

@@ -174,6 +174,23 @@ defmodule Absinthe.Execution.SubscriptionTest do
           }
         end
       end
+
+      field :prime_ordinal_with_compare, :user do
+        arg :client_id, non_null(:id)
+        arg :prime_data, list_of(:string)
+
+        config fn args, _ ->
+          {
+            :ok,
+            topic: args.client_id,
+            prime: fn _ ->
+              {:ok, [%{name: "first_user", version: 1}, %{name: "second_user", version: 2}]}
+            end,
+            ordinal: fn %{version: version} -> version end,
+            ordinal_compare: &custom_ordinal_compare/2
+          }
+        end
+      end
     end
 
     mutation do
@@ -185,6 +202,10 @@ defmodule Absinthe.Execution.SubscriptionTest do
         end
       end
     end
+
+    def custom_ordinal_compare(a, b) do
+      {a > b, b + 0.5}
+    end
   end
 
   setup_all do
@@ -761,6 +782,39 @@ defmodule Absinthe.Execution.SubscriptionTest do
              Absinthe.continue(continuations)
   end
 
+  @query """
+  subscription ($clientId: ID!) {
+    primeOrdinalWithCompare(clientId: $clientId) {
+      name
+    }
+  }
+  """
+  test "subscription with priming, ordinals, and custom ordinal compare function" do
+    client_id = "abc"
+
+    assert {:more, %{"subscribed" => _topic, continuations: continuations}} =
+             run_subscription(
+               @query,
+               Schema,
+               variables: %{
+                 "clientId" => client_id
+               }
+             )
+
+    assert {:more,
+            %{
+              data: %{"primeOrdinalWithCompare" => %{"name" => "first_user"}},
+              ordinal: 1,
+              ordinal_compare_fun: custom_ordinal_compare,
+              continuations: continuations
+            }} = Absinthe.continue(continuations)
+
+    assert custom_ordinal_compare.(1, 2) == {false, 2.5}
+
+    assert {:ok, %{data: %{"primeOrdinalWithCompare" => %{"name" => "second_user"}}, ordinal: 2}} =
+             Absinthe.continue(continuations)
+  end
+
   defp run_subscription(query, schema, opts \\ []) do
     opts = Keyword.update(opts, :context, %{pubsub: PubSub}, &Map.put(&1, :pubsub, PubSub))