Move GenStage API to Ecto.Repo (via Ecto.Adapter.Stage)

Author: fishcakez

integration_test/pg/copy_test.exs

@@ -19,29 +19,4 @@ defmodule Ecto.Integration.CopyTest do
       assert TestRepo.all(Post) == [one, two]
     end
   end
-
-  test "stage copy to and from table" do
-    one = TestRepo.insert!(%Post{title: "one"})
-    two = TestRepo.insert!(%Post{title: "two"})
-
-    read = "COPY posts TO STDOUT"
-    {:ok, producer} = Ecto.Adapters.SQL.start_stage(TestRepo, read, [], [])
-
-    data =
-      [{producer, cancel: :transient}]
-      |> GenStage.stream()
-      |> Enum.to_list()
-
-    assert TestRepo.delete_all(Post) == {2, nil}
-
-    write = "COPY posts FROM STDIN"
-    {:ok, consumer} = Ecto.Adapters.SQL.start_stage(TestRepo, write, [], [stage_module: Postgrex.CopyConsumer])
-
-    {:ok, _} =
-      data
-      |> Flow.from_enumerable()
-      |> Flow.into_stages([consumer])
-
-    assert TestRepo.all(Post) == [one, two]
-  end
 end

integration_test/pg/test_helper.exs

@@ -34,10 +34,6 @@ defmodule Ecto.Integration.TestRepo do
   use Ecto.Integration.Repo, otp_app: :ecto
 end
 
-defmodule Ecto.Integration.TestStage do
-  use Ecto.Repo.Stage, repo: TestRepo
-end
-
 # Pool repo for non-async tests
 alias Ecto.Integration.PoolRepo
 

integration_test/sql/stage.exs

@@ -2,16 +2,15 @@ defmodule Ecto.Integration.StageTest do
   use Ecto.Integration.Case, async: true
 
   alias Ecto.Integration.TestRepo
-  alias Ecto.Integration.TestStage
   alias Ecto.Integration.Post
   alias Ecto.Integration.Comment
   import Ecto.Query
 
   test "stream empty" do
-    {:ok, stage} = TestStage.start_link(Post)
+    {:ok, stage} = TestRepo.start_producer(Post)
     assert to_list(stage) === []
 
-    {:ok, stage} = TestStage.start_link(from p in Post)
+    {:ok, stage} = TestRepo.start_producer(from p in Post)
     assert to_list(stage) == []
   end
 
@@ -20,7 +19,7 @@ defmodule Ecto.Integration.StageTest do
     %Post{} = TestRepo.insert!(%Post{title: "title2"})
 
     query = from(p in "posts", order_by: p.title, select: p.title)
-    {:ok, stage} = TestStage.start_link(query)
+    {:ok, stage} = TestRepo.start_producer(query)
 
     assert to_list(stage) == ["title1", "title2"]
   end
@@ -31,7 +30,7 @@ defmodule Ecto.Integration.StageTest do
     %Comment{id: cid1} = TestRepo.insert!(%Comment{text: "1", post_id: p1.id})
     %Comment{id: cid2} = TestRepo.insert!(%Comment{text: "2", post_id: p1.id})
 
-    {:ok, stage} = TestStage.start_link(Ecto.assoc(p1, :comments))
+    {:ok, stage} = TestRepo.start_producer(Ecto.assoc(p1, :comments))
     assert [c1, c2] = to_list(stage)
 
     assert c1.id == cid1
@@ -49,7 +48,7 @@ defmodule Ecto.Integration.StageTest do
     %Comment{id: cid4} = TestRepo.insert!(%Comment{text: "4", post_id: p2.id})
 
     query = from(p in Post, preload: [:comments], select: p)
-    {:ok, stage} = TestStage.start_link(query, [max_rows: 2])
+    {:ok, stage} = TestRepo.start_producer(query, [max_rows: 2])
 
     assert [p1, p2, p3] = stage |> to_list() |> sort_by_id()
 

lib/ecto/adapters/postgres/connection.ex

@@ -18,9 +18,8 @@ if Code.ensure_loaded?(Postgrex) do
       |> Postgrex.child_spec()
     end
 
-    def stage_spec(pool, statement, params, opts) do
-      stage_mod = Keyword.get(opts, :stage_module, Postgrex.Producer)
-      Supervisor.Spec.worker(stage_mod, [pool, statement, params, opts])
+    def start_producer(pool, statement, params, opts) do
+      Postgrex.Producer.start_link(pool, statement, params, opts)
     end
 
     def to_constraints(%Postgrex.Error{postgres: %{code: :unique_violation, constraint: constraint}}),

lib/ecto/adapters/sql.ex

@@ -17,6 +17,7 @@ defmodule Ecto.Adapters.SQL do
       @behaviour Ecto.Adapter
       @behaviour Ecto.Adapter.Migration
       @behaviour Ecto.Adapter.Transaction
+      @behaviour Ecto.Adapter.Stage
 
       @conn __MODULE__.Connection
       @adapter unquote(adapter)
@@ -139,13 +140,13 @@ defmodule Ecto.Adapters.SQL do
       ## GenStage
 
       @doc false
-      def stage_spec(repo, meta, query, params, process, flat_map, opts) do
-        Ecto.Adapters.SQL.stage_spec(repo, meta, query, params, process, flat_map, opts)
+      def start_producer(repo, meta, query, params, process, flat_map, opts) do
+        Ecto.Adapters.SQL.start_producer(repo, meta, query, params, process, flat_map, 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,
-                      stage_spec: 7]
+                      start_producer: 7]
     end
   end
 
@@ -640,41 +641,29 @@ defmodule Ecto.Adapters.SQL do
     end
   end
 
-  @doc """
-  Return child specification for `GenStage` process running query.
-  """
-  def stage_spec(repo, statement, params, opts) do
-    stage_spec(repo, statement, params, fn x -> x end, nil, opts)
-  end
-
-  @doc """
-  Start link `GenStage` process running query.
-  """
-  def start_stage(repo, statement, params, opts) do
-    case stage_spec(repo, statement, params, opts) do
-      {_, {mod, fun, args}, _, _, _, _} ->
-        apply(mod, fun, args)
-      %{start: {mod, fun, args}} ->
-        apply(mod, fun, args)
-    end
-  end
-
   @doc false
-  def stage_spec(repo, meta, prepared, params, mapper, flat_map, opts) do
+  def start_producer(repo, meta, prepared, params, mapper, flat_map, opts) do
     stream = stream(repo, meta, prepared, params, mapper, flat_map, opts)
     %Ecto.Adapters.SQL.Stream{repo: repo, statement: statement, params: params,
       mapper: mapper, flat_map: flat_map, opts: opts} = stream
-    stage_spec(repo, statement, params, mapper, flat_map, opts)
+    start_producer(repo, statement, params, mapper, flat_map, opts)
+  end
+
+  @doc """
+  Start link a `GenStage` producer that streams the result of a query.
+  """
+  def start_producer(repo, statement, params, opts) do
+    start_producer(repo, statement, params, fn x -> x end, nil, opts)
   end
 
-  defp stage_spec(repo, statement, params, mapper, flat_map, opts) do
+  defp start_producer(repo, statement, params, mapper, flat_map, opts) do
     {repo_mod, pool, default_opts} = lookup_pool(repo)
     stage_opts = Keyword.delete(default_opts, :name)
     stream_mapper = flat_map && {__MODULE__, :stage_mapper, [pool, flat_map]}
     map_opts = [decode_mapper: mapper, stream_mapper: stream_mapper]
     opts = map_opts ++ with_log(repo, params, opts ++ stage_opts)
     opts = Keyword.put_new(opts, :caller, self())
-    apply(repo_mod.__sql__, :stage_spec, [pool, statement, params, opts])
+    repo_mod.__sql__.start_producer(pool, statement, params, opts)
   end
 
   @doc false

lib/ecto/adapters/sql/connection.ex

@@ -15,13 +15,6 @@ defmodule Ecto.Adapters.SQL.Connection do
   """
   @callback child_spec(options :: Keyword.t) :: {module, Keyword.t}
 
-  @doc """
-  Receives `pool`, `statement`, `params` and `options` and returns `GenStage`
-  child specification
-  """
-  @callback stage_spec(pool :: GenServer.server, statement :: prepared, params :: [term], options :: Keyword.t) ::
-            Supervisor.spec()
-
   @doc """
   Prepares and executes the given query with `DBConnection`.
   """
@@ -101,4 +94,12 @@ defmodule Ecto.Adapters.SQL.Connection do
   Receives a DDL command and returns a query that executes it.
   """
   @callback execute_ddl(command :: Ecto.Adapter.Migration.command) :: String.t | [iodata]
+
+  ## GenStage
+
+  @doc """
+  Start and link a `GenStage` producer that streams the result of a query.
+  """
+  @callback start_producer(pool :: GenServer.server, statement :: prepared, params :: [term], options :: Keyword.t) ::
+            GenServer.on_start
 end

lib/ecto/repo.ex

@@ -229,6 +229,12 @@ defmodule Ecto.Repo do
       def load(schema_or_types, data) do
         Ecto.Repo.Schema.load(@adapter, schema_or_types, data)
       end
+
+      if function_exported?(@adapter, :start_producer, 7) do
+        def start_producer(queryable, opts \\ []) do
+          Ecto.Repo.Queryable.start_producer(__MODULE__, @adapter, queryable, opts)
+        end
+      end
     end
   end
 
@@ -1009,4 +1015,40 @@ defmodule Ecto.Repo do
 
   """
   @callback load(Ecto.Schema.t | map(), map() | Keyword.t | {list, list}) :: Ecto.Schema.t | map()
+
+  @doc """
+  Starts and links to a `GenStage` producer that executes a query.
+
+
+  Returns a `GenStage` that produces all entries from the data store to its
+  consumer(s). SQL adapters, such as Postgres and MySQL, will wrap the query
+  inside a transaction.
+
+  May raise `Ecto.QueryError` if query validation fails.
+
+  ## Options
+
+    * `:prefix` - The prefix to run the query on (such as the schema path
+      in Postgres or the database in MySQL). This overrides the prefix set
+      in the query
+
+    * `: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 module documentation.
+
+  ## Example
+
+      # Fetch all post titles
+      query = from p in Post,
+           select: p.title
+      {:ok, producer} = MyRepo.start_producer(query)
+      [{producer, cancel: :transient}]
+      |> Flow.from_stages()
+      |> Flow.map(&IO.inspect/1)
+      |> Flow.run()
+  """
+  @callback start_producer(queryable :: Ecto.Query.t, opts :: Keyword.t) ::
+            GenServer.on_start
+  @optional_callbacks [start_producer: 2]
 end

lib/ecto/repo/queryable.ex

@@ -109,13 +109,13 @@ defmodule Ecto.Repo.Queryable do
     execute(:delete_all, repo, adapter, query, opts)
   end
 
-  def stage_spec(repo, adapter, queryable, opts) when is_list(opts) do
+  def start_producer(repo, adapter, queryable, opts) when is_list(opts) do
     query =
       queryable
       |> Ecto.Queryable.to_query
       |> Ecto.Query.Planner.returning(true)
       |> attach_prefix(opts)
-    lazy(&adapter.stage_spec/7, :all, repo, adapter, query, opts)
+    lazy(&adapter.start_producer/7, :all, repo, adapter, query, opts)
   end
 
   ## Helpers