Split SQL stage start_link to 3 functions

Author: fishcakez

lib/ecto/adapters/sql.ex

@@ -628,7 +628,7 @@ defmodule Ecto.Adapters.SQL do
   end
 
   @doc false
-  def stage(repo, type, start, handle, stop, opts \\ []) do
+  def stage(fun, repo, start, handle, stop, opts) do
     {repo_mod, pool, default_opts} = lookup_pool(repo)
     default_opts =
       default_opts
@@ -649,7 +649,7 @@ defmodule Ecto.Adapters.SQL do
           delete_conn(pool)
         end
       end
-    DBConnection.Stage.start_link(pool, type, start, handle, stop, opts)
+    fun.(pool, start, handle, stop, opts)
   end
 
   ## Log

lib/ecto/adapters/sql/stage.ex

@@ -1,84 +1,137 @@
 defmodule Ecto.Adapters.SQL.Stage do
   @moduledoc """
   A `GenStage` process that encapsulates a SQL transaction.
+
+  ### Options
+
+    * `:name` - A name to register the started process (see the `:name` option
+    in `GenServer.start_link/3`)
+
+  See the "Shared options" section at the `Ecto.Repo` documentation. All options
+  are passed to the `GenStage` on init.
   """
 
   @doc """
-  Start link a `GenStage` process that will run a transaction for its duration.
+  Start link a `GenStage` producer that will run a transaction for its duration.
 
-  The first argument is the pool, the second argument is the `GenStage` type,
-  the third argument is the start function, the fourth argument is the handle
-  function, the fifth argument is the stop function and the optional sixth
-  argument are the options.
+  The first argument is the repo, the second argument is the start function,
+  the third argument is the handle demand function, the fourth argument is the
+  stop function and the optional fiftth argument are the options.
 
-  The start function is a o-arity anonymous function. This is called after the
-  transaction begins but before `start_link/6` returns. It should return the
-  `state` or call `MyRepo.rollback/1` to stop the `GenStage`.
+  The start function is a 0-arity anonymous function. This is called after the
+  transaction begins but before `producer/5` returns. It should return the
+  accumulator or call `repo.rollback/1` to stop the `GenStage`.
 
-  The handle function is a 2-arity anonymous function. If the `GenStage` type is
-  a `:producer`, then the first argument is the `demand` from a `GenStage`
-  `handle_demand` callback. Otherwise the first argument is the events from a
-  `GenStage` `handle_events` callback. The second argument is the state. This
-  function returns a 2-tuple, with first element as events (empty list for
-  `:consumer`) and second element as the `state`. This function can roll back
-  and stop the `GenStage` using `MyRepo.rollback/1`.
+  The handle demand function is a 2-arity anonymous function. The first argument
+  is the `demand`, and the second argument is the accumulator. This function
+  returns a 2-tuple, with first element as list of events to fulfil the demand
+  and second element as the accumulator. If the producer has emitted all events
+  (and so not fulfilled demand) it should call
+  `GenStage.async_notify(self(), {:producer, :done | :halted}` to signal to
+  consumers that it has finished. Also this function can rollback and stop the
+  `GenStage` using `repo.rollback/1`.
 
   The stop function is a 2-arity anonymous function. The first argument is the
-  terminate reason and the second argument is the `state`. This function will
+  terminate reason and the third argument is the accumulator. This function will
   only be called if connection is alive and the transaction has not been rolled
   back. If this function returns the transaction is committed. This function can
-  roll back and stop the `GenStage` using `MyRepo.rollback/1`.
+  rollback and stop the `GenStage` using `repo.rollback/1`.
+
+  For options see "Options" in the module documentation.
+
+  The `GenStage` process will behave like a `Flow` stage:
+
+    * It will stop with reason `:normal` when the last consumer cancels
+  """
+  @spec producer(module, start :: (() -> acc),
+    handle_demand :: ((demand :: pos_integer, acc) -> {[any], acc}),
+    stop :: ((reason :: any, acc) -> any), Keyword.t) ::
+    GenServer.on_start when acc: var
+  def producer(repo, start, handle_demand, stop, opts \\ []) do
+    fun = &DBConnection.Stage.producer/5
+    Ecto.Adapters.SQL.stage(fun, repo, start, handle_demand, stop, opts) 
+  end
+
+  @doc """
+  Start link a `GenStage` producer consumer that will run a transaction for its
+  duration.
+
+  The first argument is the repo, the second argument is the start function,
+  the third argument is the handle events function, the fourth argument is the
+  stop function and the optional fiftth argument are the options.
+
+  The start function is a 0-arity anonymous function. This is called after the
+  transaction begins but before `consumer_producer/5` returns. It should return
+  the accumulator or call `repo.rollback/1` to stop the `GenStage`.
+
+  The handle events function is a 2-arity anonymous function. The first argument
+  is a list of incoming events, and the second argument is the accumulator. This
+  function returns a 2-tuple, with first element as list of outgoing events and
+  second element as the accumulator. Also this function can rollback and stop
+  the `GenStage` using `repo.rollback/1`.
+
+  The stop function is a 2-arity anonymous function. The first argument is the
+  terminate reason and the second argument is the accumulator. This function
+  will only be called if connection is alive and the transaction has not been
+  rolled back. If this function returns the transaction is committed. This
+  function can rollback and stop the `GenStage` using `repo.rollback/1`.
+
+  For options see "Options" in the module documentation.
 
   The `GenStage` process will behave like a `Flow` stage:
 
     * It will stop with reason `:normal` when the last consumer cancels
     * It will notify consumers that it is done when all producers have cancelled
     or notified that they are done or halted
-    * It will cancel new and remaining producers when all producers have
-    notified that they are done or halted and it is a `:consumer`
     * It will not send demand to new producers when all producers have notified
-    that they are done or halted and it is a `:consumer_producer`
+    that they are done or halted
+  """
+  @spec producer_consumer(repo :: module, start :: (() -> acc),
+    handle_events :: ((events_in :: [any], acc) -> {events_out :: [any], acc}),
+    stop :: ((reason :: any, acc) -> any), Keyword.t) ::
+    GenServer.on_start when acc: var
+  def producer_consumer(repo, start, handle_events, stop, opts \\ []) do
+    fun = &DBConnection.Stage.producer_consumer/5
+    Ecto.Adapters.SQL.stage(fun, repo, start, handle_events, stop, opts)
+  end
 
-  ### Options
+  @doc """
+  Start link a `GenStage` consumer that will run a transaction for its duration.
 
-    * `:name` - A name to register the started process (see the `:name` option
-    in `GenServer.start_link/3`)
+  The first argument is the repo, the second argument is the start function,
+  the third argument is the handle events function, the fourth argument is the
+  stop function and the optional fiftth argument are the options.
 
-  See the "Shared options" section at the `Ecto.Repo` documentation. All options
-  are passed to the `GenStage` on init.
+  The start function is a 0-arity anonymous function. This is called after the
+  transaction begins but before `consumer/5` returns. It should return the
+  accumulator or call `repo.rollback/1` to stop the `GenStage`.
 
-  ### Example
-
-      start = fn() -> Post end
-      handle =
-        fn(entries, schema) ->
-          MyRepo.insert_all(schema, entries)
-          {[], schema}
-        end
-      stop =
-        fn
-          :normal, _ -> :ok
-          reason,  _ -> MyRepo.rollback(reason)
-        end
-      Ecto.Adapters.SQL.Stage.start_link(MyRepo, :consumer, start, handle, stop)
+  The handle events function is a 2-arity anonymous function. The first argument
+  is the list of events, and the second argument is the accumulator. This
+  function returns a 2-tuple, with first element is an empty list (as no
+  outgoing events) and second element as the accumulator. Also this function can
+  rollback and stop the `GenStage` using `repo.rollback/1`.
+
+  The stop function is a 2-arity anonymous function. The first argument is the
+  terminate reason and the second argument is the accumulator. This function
+  will only be called if connection is alive and the transaction has not been
+  rolled back. If this function returns the transaction is committed. This
+  function can rollback and stop the `GenStage` using `repo.rollback/1`.
+
+  See the "Shared options" section at the `Ecto.Repo` documentation.
+
+  The `GenStage` process will behave like a `Flow` stage:
+
+    * It will cancel new and remaining producers when all producers have
+    notified that they are done or halted and it is a `:consumer`
   """
-  @spec start_link(repo :: module, :producer,
-    start :: (() -> state),
-    handle_demand :: ((demand :: pos_integer, state) -> {[any], state}),
-    stop :: ((reason :: any, state) -> any), opts :: Keyword.t) ::
-    GenServer.on_start when state: var
-  @spec start_link(repo :: module, :producer_consumer,
-    start :: (() -> state),
-    handle_events :: (([any], state) -> {[any], state}),
-    stop :: ((reason :: any, state) -> any), opts :: Keyword.t) ::
-    GenServer.on_start when state: var
-  @spec start_link(repo :: module, :consumer,
-    start :: (() -> state),
-    handle_events :: (([any], state) -> {[], state}),
-    stop :: ((reason :: any, state) -> any), opts :: Keyword.t) ::
-    GenServer.on_start when state: var
-  def start_link(repo, type, start, handle, stop, opts \\ []) do
-    Ecto.Adapters.SQL.stage(repo, type, start, handle, stop, opts)
+  @spec consumer(repo :: module, start :: (() -> acc),
+    handle_events :: ((events_in :: [any], acc) -> {[], acc}),
+    stop :: ((reason :: any, acc) -> any), Keyword.t) ::
+    GenServer.on_start when acc: var
+  def consumer(pool, start, handle_events, stop, opts \\ []) do
+    fun = &DBConnection.Stage.consumer/5
+    Ecto.Adapters.SQL.stage(fun, pool, start, handle_events, stop, opts)
   end
 
   @doc """
@@ -97,7 +150,7 @@ defmodule Ecto.Adapters.SQL.Stage do
     * `:max_rows` - The number of rows to load from the database as we stream.
       It is supported at least by Postgres and MySQL and defaults to 500.
 
-  See the "Shared options" section at the `Ecto.Repo` documentation.
+  For more options see "Options" in the module documentation.
 
   ## Example
 
@@ -110,8 +163,7 @@ defmodule Ecto.Adapters.SQL.Stage do
       |> Flow.each(&IO.inspect/1)
       |> Flow.start_link()
   """
-
-  @callback stream(repo :: module, queryable :: Ecto.Query.t, opts :: Keyword.t) ::
+  @spec stream(repo :: module, queryable :: Ecto.Query.t, opts :: Keyword.t) ::
     GenServer.on_start()
   def stream(repo, queryable, opts \\ []) do
     stream = apply(repo, :stream, [queryable, opts])
@@ -121,7 +173,7 @@ defmodule Ecto.Adapters.SQL.Stage do
         {:suspended, _, cont} = Enumerable.reduce(stream, acc, &stream_reduce/2)
         {repo, :cont, cont}
       end
-    start_link(repo, :producer, start, &stream_handle/2, &stream_stop/2, opts)
+    producer(repo, start, &stream_handle/2, &stream_stop/2, opts)
   end
 
   ## Helpers

mix.lock

@@ -1,6 +1,6 @@
 %{"backoff": {:hex, :backoff, "1.1.1"},
   "connection": {:hex, :connection, "1.0.4", "a1cae72211f0eef17705aaededacac3eb30e6625b04a6117c1b2db6ace7d5976", [:mix], []},
-  "db_connection": {:git, "https://github.com/elixir-ecto/db_connection.git", "392affacc4a730d748bf95d7261de9e8a9709fe2", [branch: "jf-stream_stage"]},
+  "db_connection": {:git, "https://github.com/elixir-ecto/db_connection.git", "73238baf9778b1263e108a0df9198aade2921bea", [branch: "jf-stream_stage"]},
   "decimal": {:hex, :decimal, "1.3.1", "157b3cedb2bfcb5359372a7766dd7a41091ad34578296e951f58a946fcab49c6", [:mix], []},
   "earmark": {:hex, :earmark, "1.0.3", "89bdbaf2aca8bbb5c97d8b3b55c5dd0cff517ecc78d417e87f1d0982e514557b", [:mix], []},
   "ex_doc": {:hex, :ex_doc, "0.14.5", "c0433c8117e948404d93ca69411dd575ec6be39b47802e81ca8d91017a0cf83c", [:mix], [{:earmark, "~> 1.0", [hex: :earmark, optional: false]}]},