Refactoring

lib/absinthe/subscription.ex

@@ -150,17 +150,13 @@ defmodule Absinthe.Subscription do
 
     storage_module = document_storage(pubsub)
     storage_process_name = document_storage_name(pubsub)
-
-    storage_module.put(storage_process_name, doc_id, doc_value)
-    storage_module.subscribe(storage_process_name, doc_id, field_keys)
+    storage_module.put(storage_process_name, doc_id, doc_value, field_keys)
   end
 
   @doc false
   def unsubscribe(pubsub, doc_id) do
     storage_module = document_storage(pubsub)
     storage_process_name = document_storage_name(pubsub)
-
-    storage_module.unsubscribe(storage_process_name, doc_id)
     storage_module.delete(storage_process_name, doc_id)
   end
 

lib/absinthe/subscription/default_document_storage.ex

@@ -12,33 +12,26 @@ defmodule Absinthe.Subscription.DefaultDocumentStorage do
   end
 
   @impl Absinthe.Subscription.DocumentStorage
-  def put(storage_process_name, doc_id, doc_value) do
-    {:ok, _} = Registry.register(storage_process_name, doc_id, doc_value)
-  end
-
-  @impl Absinthe.Subscription.DocumentStorage
-  def subscribe(storage_process_name, doc_id, field_keys) do
+  def put(storage_process_name, doc_id, doc_value, field_keys) do
     pdict_add_fields(doc_id, field_keys)
 
     for field_key <- field_keys do
       {:ok, _} = Registry.register(storage_process_name, field_key, doc_id)
     end
 
-    {:ok, field_keys}
+    {:ok, _} = Registry.register(storage_process_name, doc_id, doc_value)
   end
 
   @impl Absinthe.Subscription.DocumentStorage
   def delete(storage_process_name, doc_id) do
-    Registry.unregister(storage_process_name, doc_id)
-  end
-
-  @impl Absinthe.Subscription.DocumentStorage
-  def unsubscribe(storage_process_name, doc_id) do
     for field_key <- pdict_fields(doc_id) do
       Registry.unregister(storage_process_name, field_key)
     end
 
     pdict_delete_fields(doc_id)
+
+    Registry.unregister(storage_process_name, doc_id)
+
     :ok
   end
 

lib/absinthe/subscription/document_storage.ex

@@ -22,30 +22,16 @@ defmodule Absinthe.Subscription.DocumentStorage do
               doc :: %{
                 initial_phases: Absinthe.Subscription.PipelineSerializer.packed_pipeline(),
                 source: binary()
-              }
-            ) ::
-              {:ok, term} | {:error, :reason}
-
-  @doc """
-  Associates each `{field, key}` pair in `field_keys` to `doc_id`.
-  """
-  @callback subscribe(
-              storage_process_name :: atom,
-              doc_id :: term,
+              },
               field_keys :: [{field :: term, key :: term}]
             ) ::
               {:ok, term} | {:error, :reason}
 
   @doc """
-  Removes the document.
+  Removes the document. Along with any field_keys associated with it
   """
   @callback delete(storage_process_name :: atom, doc_id :: term) :: :ok
 
-  @doc """
-  Removes the field_keys associated with `doc_id`.
-  """
-  @callback unsubscribe(storage_process_name :: atom, doc_id :: term) :: :ok
-
   @doc """
   Get all docs associated with `field_key`
   """