Allow client to retry when specified (#337)

Author: josevalim

integration_test/cases/execute_test.exs

@@ -348,6 +348,66 @@ defmodule ExecuteTest do
            ] = A.record(agent)
   end
 
+  test "execute disconnect_and_retry succeeds" do
+    err = RuntimeError.exception("oops")
+
+    stack = [
+      {:ok, :state},
+      {:disconnect_and_retry, err, :state},
+      :ok,
+      fn opts ->
+        send(opts[:parent], :reconnected)
+        {:ok, :new_state}
+      end,
+      {:ok, %Q{}, %R{}, :new_state}
+    ]
+
+    {:ok, agent} = A.start_link(stack)
+
+    opts = [agent: agent, parent: self()]
+    {:ok, pool} = P.start_link(opts)
+    assert P.execute(pool, %Q{}, [:param]) == {:ok, %Q{}, %R{}}
+
+    assert_receive :reconnected
+
+    assert [
+             connect: [opts2],
+             handle_execute: [%Q{}, [:param], _, :state],
+             disconnect: [^err, :state],
+             connect: [opts2],
+             handle_execute: [%Q{}, [:param], _, :new_state]
+           ] = A.record(agent)
+  end
+
+  test "execute disconnect_and_retry errors if there are no retries" do
+    err = RuntimeError.exception("oops")
+
+    stack = [
+      {:ok, :state},
+      {:disconnect_and_retry, err, :new_state},
+      :ok,
+      fn opts ->
+        send(opts[:parent], :reconnected)
+        {:ok, :state}
+      end
+    ]
+
+    {:ok, agent} = A.start_link(stack)
+
+    opts = [agent: agent, parent: self()]
+    {:ok, pool} = P.start_link(opts)
+    assert P.execute(pool, %Q{}, [:param], checkout_retries: 0) == {:error, err}
+
+    assert_receive :reconnected
+
+    assert [
+             connect: [opts2],
+             handle_execute: [%Q{}, [:param], _, :state],
+             disconnect: [^err, :new_state],
+             connect: [opts2]
+           ] = A.record(agent)
+  end
+
   test "execute bad return raises DBConnection.ConnectionError and stops" do
     stack = [
       fn opts ->

integration_test/cases/prepare_execute_test.exs

@@ -496,6 +496,39 @@ defmodule PrepareExecuteTest do
            ] = A.record(agent)
   end
 
+  test "prepare_execute execute disconnect_and_retry succeeds" do
+    err = RuntimeError.exception("oops")
+
+    stack = [
+      {:ok, :state},
+      {:disconnect_and_retry, err, :state},
+      :ok,
+      fn opts ->
+        send(opts[:parent], :reconnected)
+        {:ok, :new_state}
+      end,
+      {:ok, %Q{}, :newer_state},
+      {:ok, %Q{state: :executed}, %R{}, :newest_state}
+    ]
+
+    {:ok, agent} = A.start_link(stack)
+
+    opts = [agent: agent, parent: self()]
+    {:ok, pool} = P.start_link(opts)
+    assert P.prepare_execute(pool, %Q{}, [:param]) == {:ok, %Q{state: :executed}, %R{}}
+
+    assert_receive :reconnected
+
+    assert [
+             connect: [opts2],
+             handle_prepare: [%Q{}, _, :state],
+             disconnect: [^err, :state],
+             connect: [opts2],
+             handle_prepare: [%Q{}, _, :new_state],
+             handle_execute: [%Q{}, [:param], _, :newer_state]
+           ] = A.record(agent)
+  end
+
   test "prepare_execute describe or encode raises and closes query" do
     stack = [
       {:ok, :state},

integration_test/cases/transaction_test.exs

@@ -435,6 +435,70 @@ defmodule TransactionTest do
            ] = A.record(agent)
   end
 
+  test "transaction begin disconnect_and_retry succeeds" do
+    err = RuntimeError.exception("oops")
+
+    stack = [
+      {:ok, :state},
+      {:disconnect_and_retry, err, :state},
+      :ok,
+      fn opts ->
+        send(opts[:parent], :reconnected)
+        {:ok, :new_state}
+      end,
+      {:ok, :began, :newer_state},
+      {:ok, :committed, :newest_state}
+    ]
+
+    {:ok, agent} = A.start_link(stack)
+
+    opts = [agent: agent, parent: self()]
+    {:ok, pool} = P.start_link(opts)
+    assert P.transaction(pool, fn _ -> :ok end) == {:ok, :ok}
+    assert_receive :reconnected
+
+    assert [
+             connect: [_],
+             handle_begin: [_, :state],
+             disconnect: [_, :state],
+             connect: [_],
+             handle_begin: [_, :new_state],
+             handle_commit: [_, :newer_state]
+           ] = A.record(agent)
+  end
+
+  test "transaction begin disconnect_and_retry errors if there are no retries" do
+    err = RuntimeError.exception("oops")
+
+    stack = [
+      {:ok, :state},
+      {:disconnect_and_retry, err, :new_state},
+      :ok,
+      fn opts ->
+        send(opts[:parent], :reconnected)
+        {:ok, :newest_state}
+      end
+    ]
+
+    {:ok, agent} = A.start_link(stack)
+
+    opts = [agent: agent, parent: self()]
+    {:ok, pool} = P.start_link(opts)
+
+    assert_raise RuntimeError, "oops", fn ->
+      P.transaction(pool, fn _ -> flunk("transaction ran") end, checkout_retries: 0)
+    end
+
+    assert_receive :reconnected
+
+    assert [
+             connect: [_],
+             handle_begin: [_, :state],
+             disconnect: [_, :new_state],
+             connect: [_]
+           ] = A.record(agent)
+  end
+
   test "transaction begin bad return raises and stops connection" do
     stack = [
       fn opts ->