Update DBConnection/Postgrex and introduce start_consumer

Author: fishcakez

integration_test/sql/stage.exs

@@ -57,6 +57,20 @@ defmodule Ecto.Integration.StageTest do
     assert [] = p3.comments
   end
 
+  test "insert entries" do
+    {:ok, stage} = TestRepo.start_consumer(Post)
+
+    mon = Process.monitor(stage)
+
+    [[title: "hello"], [title: "world"]]
+    |> Flow.from_enumerable()
+    |> Flow.into_stages([{stage, [cancel: :transient, max_demand: 1]}])
+
+    assert_receive {:DOWN, ^mon, _, _, :normal}
+
+    assert [%Post{title: "hello"}, %Post{title: "world"}] = TestRepo.all(Post)
+  end
+
   defp to_list(stage) do
     stage
     |> Flow.from_stage()

lib/ecto/adapter/stage.ex

@@ -0,0 +1,32 @@
+defmodule Ecto.Adapter.Stage do
+  @moduledoc """
+  Specifies the adapter `GenStage` API.
+  """
+
+  @type flat_map :: ({integer, list | nil} -> list)
+  @type insert_all :: (list -> integer)
+  @type options :: Keyword.t
+
+  @doc """
+  Starts and links to a `GenStage` producer.
+
+  The producer executes a query and returns the result to consumer(s).
+
+  See `Ecto.Repo.start_producer/2`.
+  """
+  @callback start_producer(repo :: Ecto.Repo.t, Ecto.Adapter.query_meta, query, params :: list(), Ecto.Adapter.process | nil, flat_map, options) ::
+            GenServer.on_start when
+              query: {:nocache, Ecto.Adapter.prepared} |
+                     {:cached, (Ecto.Adapter.prepared -> :ok), Ecto.Adapter.cached} |
+                     {:cache, (Ecto.Adapter.cached -> :ok), Ecto.Adapter.prepared}
+
+  @doc """
+  Starts and links to a `GenStage` consumer.
+
+  The consumers inserts the entries it receives for the schema.
+
+  See `Ecto.Repo.start_consumer/2`
+  """
+  @callback start_consumer(repo :: Ecto.Repo.t, insert_all, options) ::
+            GenServer.on_start
+end

lib/ecto/adapters/postgres/connection.ex

@@ -22,6 +22,10 @@ if Code.ensure_loaded?(Postgrex) do
       Postgrex.Producer.start_link(pool, statement, params, opts)
     end
 
+    def start_consumer(pool, fun, opts) do
+      Postgrex.Consumer.start_link(pool, fun, opts)
+    end
+
     def to_constraints(%Postgrex.Error{postgres: %{code: :unique_violation, constraint: constraint}}),
       do: [unique: constraint]
     def to_constraints(%Postgrex.Error{postgres: %{code: :foreign_key_violation, constraint: constraint}}),

lib/ecto/adapters/sql.ex

@@ -144,9 +144,14 @@ defmodule Ecto.Adapters.SQL do
         Ecto.Adapters.SQL.start_producer(repo, meta, query, params, process, flat_map, opts)
       end
 
+      @doc false
+      def start_consumer(repo, insert, opts) do
+        Ecto.Adapters.SQL.start_consumer(repo, insert, opts)
+      end
+
       defoverridable [prepare: 2, execute: 6, insert: 6, update: 6, delete: 4, insert_all: 7,
                       execute_ddl: 3, loaders: 2, dumpers: 2, autogenerate: 1, ensure_all_started: 2,
-                      start_producer: 7]
+                      start_producer: 7, start_consumer: 3]
     end
   end
 
@@ -652,7 +657,7 @@ defmodule Ecto.Adapters.SQL do
   @doc """
   Start link a `GenStage` producer that streams the result of a query.
   """
-  def start_producer(repo, statement, params, opts) do
+  def start_producer(repo, statement, params, opts \\ []) do
     start_producer(repo, statement, params, fn x -> x end, nil, opts)
   end
 
@@ -681,6 +686,25 @@ defmodule Ecto.Adapters.SQL do
     end
   end
 
+  @doc """
+  Start link a `GenStage` consumers that runs a `transaction` for every batch of
+  events and maintains a state.
+  """
+  def start_consumer(repo, fun, opts \\ []) when is_function(fun, 1) do
+    {repo_mod, pool, default_opts} = lookup_pool(repo)
+
+    transaction =
+      fn(conn, rows) ->
+        put_conn(pool, conn)
+        fun.(rows)
+      end
+
+    stage_opts = Keyword.delete(default_opts, :name)
+    opts = with_log(repo, [], opts ++ stage_opts)
+    opts = Keyword.put_new(opts, :caller, self())
+    repo_mod.__sql__.start_consumer(pool, transaction, opts)
+  end
+
   ## Log
 
   defp with_log(repo, params, opts) do

lib/ecto/adapters/sql/connection.ex

@@ -9,6 +9,9 @@ defmodule Ecto.Adapters.SQL.Connection do
   @typedoc "The cache query which is a DBConnection Query"
   @type cached :: map
 
+  @typedoc "The DBConection transaction fun"
+  @type transaction :: (DBConnection.t, list -> term)
+
   @doc """
   Receives options and returns `DBConnection` supervisor child
   specification.
@@ -102,4 +105,10 @@ defmodule Ecto.Adapters.SQL.Connection do
   """
   @callback start_producer(pool :: GenServer.server, statement :: prepared, params :: [term], options :: Keyword.t) ::
             GenServer.on_start
+
+  @doc """
+  Start and link a `GenStage` consumer that runs a transaction for the events.
+  """
+  @callback start_consumer(pool :: GenServer.server, transaction, options :: Keyword.t) ::
+            GenServer.on_start
 end

lib/ecto/repo.ex

@@ -234,6 +234,10 @@ defmodule Ecto.Repo do
         def start_producer(queryable, opts \\ []) do
           Ecto.Repo.Queryable.start_producer(__MODULE__, @adapter, queryable, opts)
         end
+
+        def start_consumer(schema_or_source, opts \\ []) do
+          Ecto.Repo.Schema.start_consumer(__MODULE__, @adapter, schema_or_source, opts)
+        end
       end
     end
   end
@@ -1051,4 +1055,27 @@ defmodule Ecto.Repo do
   @callback start_producer(queryable :: Ecto.Query.t, opts :: Keyword.t) ::
             GenServer.on_start
   @optional_callbacks [start_producer: 2]
+
+  @doc """
+  Starts and links to a `GenStage` consumer that inserts all entries it
+  consumes.
+
+  Returns a `GenStage` consumer that uses `c:insert_all/3` to insert all
+  consumed entries in batches to the database. SQL adapters, such as Postgres
+  and MySQL, will wrap all the inserts inside a single transaction.
+
+  See `c:insert_all/3` for options.
+
+  ## Examples
+
+      {:ok, consumer} = MyRepo.start_consumer(Post)
+      "posts"
+      |> File.stream!()
+      |> Flow.from_enumerable(stages: 1)
+      |> Flow.map(&MyDecoder.decode/1)
+      |> Flow.into_stages([{consumer, cancel: :transient}])
+  """
+  @callback start_consumer(schema_or_source :: binary | {binary, Ecto.Schema.t} | Ecto.Schema.t, opts :: Keyword.t) ::
+            GenServer.on_start
+  @optional_callbacks [start_consumer: 2]
 end

lib/ecto/repo/schema.ex

@@ -50,6 +50,45 @@ defmodule Ecto.Repo.Schema do
     {count, postprocess(rows, fields, adapter, schema, source)}
   end
 
+  @doc """
+  Implementation for `Ecto.Repo.start_consumer/2`.
+  """
+  def start_consumer(repo, adapter, schema, opts) when is_atom(schema) do
+    start_consumer(repo, adapter, schema, schema.__schema__(:prefix),
+                      schema.__schema__(:source), opts)
+  end
+
+  def start_consumer(repo, adapter, table, opts) when is_binary(table) do
+    start_consumer(repo, adapter, nil, nil, table, opts)
+  end
+
+  def start_consumer(repo, adapter, {source, schema}, opts) when is_atom(schema) do
+    start_consumer(repo, adapter, schema, schema.__schema__(:prefix), source, opts)
+  end
+
+  defp start_consumer(repo, adapter, schema, prefix, source, opts) do
+    returning = opts[:returning] || false
+    autogen   = schema && schema.__schema__(:autogenerate_id)
+    source    = {Keyword.get(opts, :prefix, prefix), source}
+    fields    = preprocess(returning, schema)
+
+    insert_all = fn(rows) ->
+        {rows, header} = extract_header_and_fields(rows, schema, autogen, adapter)
+        counter = fn -> Enum.reduce(rows, 0, &length(&1) + &2) end
+        metadata = %{source: source, context: nil, schema: schema, autogenerate_id: autogen}
+
+        {on_conflict, opts} = Keyword.pop(opts, :on_conflict, :raise)
+        {conflict_target, opts} = Keyword.pop(opts, :conflict_target, [])
+        on_conflict = on_conflict(on_conflict, conflict_target, metadata, counter, adapter)
+
+        {count, _} =
+        adapter.insert_all(repo, metadata, Map.keys(header), rows, on_conflict, fields || [], opts)
+        count
+    end
+
+    adapter.start_consumer(repo, insert_all, opts)
+  end
+
   defp preprocess([_|_] = fields, _schema),
     do: fields
   defp preprocess([], _schema),

mix.lock

@@ -1,6 +1,6 @@
 %{"backoff": {:hex, :backoff, "1.1.1"},
   "connection": {:hex, :connection, "1.0.4", "a1cae72211f0eef17705aaededacac3eb30e6625b04a6117c1b2db6ace7d5976", [], [], "hexpm"},
-  "db_connection": {:git, "https://github.com/elixir-ecto/db_connection.git", "9c6e4945112554bf40c240244c033f6dcc9c4a13", [branch: "jf-stream_stage"]},
+  "db_connection": {:git, "https://github.com/elixir-ecto/db_connection.git", "845b46c234a3723bc4c1af1059a86bc9e9d85af0", [branch: "jf-stream_stage"]},
   "decimal": {:hex, :decimal, "1.3.1", "157b3cedb2bfcb5359372a7766dd7a41091ad34578296e951f58a946fcab49c6", [:mix], [], "hexpm"},
   "earmark": {:hex, :earmark, "1.0.3", "89bdbaf2aca8bbb5c97d8b3b55c5dd0cff517ecc78d417e87f1d0982e514557b", [:mix], [], "hexpm"},
   "ex_doc": {:hex, :ex_doc, "0.14.5", "c0433c8117e948404d93ca69411dd575ec6be39b47802e81ca8d91017a0cf83c", [:mix], [{:earmark, "~> 1.0", [hex: :earmark, repo: "hexpm", optional: false]}], "hexpm"},
@@ -10,5 +10,5 @@
   "mariaex": {:hex, :mariaex, "0.8.2", "a9ee64a02fd72579f844934b4cbecca9566593e499125edf6032c58f9d50b5f9", [:mix], [{:db_connection, "~> 1.1", [hex: :db_connection, repo: "hexpm", optional: false]}, {:decimal, "~> 1.0", [hex: :decimal, repo: "hexpm", optional: false]}], "hexpm"},
   "poison": {:hex, :poison, "3.1.0", "d9eb636610e096f86f25d9a46f35a9facac35609a7591b3be3326e99a0484665", [:mix], [], "hexpm"},
   "poolboy": {:hex, :poolboy, "1.5.1", "6b46163901cfd0a1b43d692657ed9d7e599853b3b21b95ae5ae0a777cf9b6ca8", [], [], "hexpm"},
-  "postgrex": {:git, "https://github.com/elixir-ecto/postgrex.git", "f9f2ae1ca374d462ff86622951fbd216f28d74f0", [branch: "jf-stage"]},
+  "postgrex": {:git, "https://github.com/elixir-ecto/postgrex.git", "b825546e8954f9623e256f2423483bf42a2c6596", [branch: "jf-stage"]},
   "sbroker": {:hex, :sbroker, "1.0.0", "28ff1b5e58887c5098539f236307b36fe1d3edaa2acff9d6a3d17c2dcafebbd0", [:rebar3], [], "hexpm"}}