support passing a function to Writer.transaction/3

which applies the transaction and returns the txid

Author: magnetised

lib/phoenix/sync/writer.ex

@@ -1160,17 +1160,28 @@ defmodule Phoenix.Sync.Writer do
 
   @txn_name {:__phoenix_sync__, :txn}
   @txid_name {:__phoenix_sync__, :txid}
+  @txid_query "SELECT txid_current() as txid"
 
   defp start_multi(txn) do
     Ecto.Multi.new()
+    |> txid_step()
     |> Ecto.Multi.put(@txn_name, txn)
-    |> Ecto.Multi.run(@txid_name, fn repo, _ ->
-      with {:ok, %{rows: [[txid]]}} <- repo.query("SELECT txid_current() as txid") do
+  end
+
+  defp txid_step(multi \\ Ecto.Multi.new()) do
+    Ecto.Multi.run(multi, @txid_name, fn repo, _ ->
+      with {:ok, %{rows: [[txid]]}} <- repo.query(@txid_query) do
         {:ok, txid}
       end
     end)
   end
 
+  defp has_txid_step?(multi) do
+    multi
+    |> Ecto.Multi.to_list()
+    |> Enum.any?(fn {name, _} -> name == @txid_name end)
+  end
+
   defp apply_change(multi, %Operation{} = op, %__MODULE__{} = writer) do
     with {:ok, actions} <- mutation_actions(op, writer),
          {:ok, action} <- Map.fetch(actions, op.operation) do
@@ -1455,18 +1466,53 @@ defmodule Phoenix.Sync.Writer do
               end)
             Plug.Conn.send_resp(conn, 400, Jason.encode!(error))
           end
+
+  Also supports normal fun/0 or fun/1 style transactions much like
+  `Ecto.Repo.transaction/2`, returning the txid of the operation:
+
+      {:ok, txid, todo} =
+        Phoenix.Sync.Writer.transaction(fn ->
+          Repo.insert!(changeset)
+        end, Repo)
   """
   @spec transaction(Ecto.Multi.t(), Ecto.Repo.t(), keyword()) ::
           {:ok, txid(), Ecto.Multi.changes()} | Ecto.Multi.failure()
   def transaction(multi, repo, opts \\ [])
 
   def transaction(%Ecto.Multi{} = multi, repo, opts) when is_atom(repo) do
-    with {:ok, changes} <- repo.transaction(multi, opts) do
+    wrapped_multi =
+      if has_txid_step?(multi) do
+        multi
+      else
+        Ecto.Multi.prepend(multi, txid_step())
+      end
+
+    with {:ok, changes} <- repo.transaction(wrapped_multi, opts) do
       {txid, changes} = Map.pop!(changes, @txid_name)
       {:ok, txid, changes}
     end
   end
 
+  def transaction(fun, repo, opts) when is_function(fun) and is_atom(repo) do
+    with {:ok, {txid, result}} <-
+           repo.transaction(
+             fn ->
+               %{rows: [[txid]]} = repo.query!(@txid_query)
+
+               result =
+                 case fun do
+                   fun0 when is_function(fun0, 0) -> fun.()
+                   fun1 when is_function(fun1, 1) -> fun1.(repo)
+                 end
+
+               {txid, result}
+             end,
+             opts
+           ) do
+      {:ok, txid, result}
+    end
+  end
+
   @doc """
   Extract the transaction id from changes returned from `Repo.transaction`.
 

test/phoenix/sync/writer_test.exs

@@ -929,6 +929,38 @@ defmodule Phoenix.Sync.WriterTest do
     end
   end
 
+  describe "transaction/3" do
+    setup [:with_repo, :with_todos]
+
+    test "supports a function and returns the {:ok, txid, result} tuple" do
+      changeset =
+        Ecto.Changeset.change(%Support.Todo{}, %{id: 1111, title: "my todo", completed: false})
+
+      assert {:ok, txid, %Support.Todo{title: "my todo"}} =
+               Writer.transaction(
+                 fn repo ->
+                   assert repo == Repo
+                   repo.insert!(changeset)
+                 end,
+                 Repo
+               )
+
+      assert is_integer(txid)
+    end
+
+    test "supports any Ecto.Multi argument and returns the txid" do
+      changeset =
+        Ecto.Changeset.change(%Support.Todo{}, %{id: 1111, title: "my todo", completed: false})
+
+      multi = Ecto.Multi.insert(Ecto.Multi.new(), :todo, changeset)
+
+      assert {:ok, txid, %{todo: %Support.Todo{title: "my todo"}}} =
+               Writer.transaction(multi, Repo)
+
+      assert is_integer(txid)
+    end
+  end
+
   def parse_transaction(m) when is_list(m) do
     with {:ok, operations} <-
            Writer.Transaction.parse_operations(m, fn op ->