add more cases

Author: greg-rychlewski

lib/postgrex/protocol.ex

@@ -376,9 +376,7 @@ defmodule Postgrex.Protocol do
           parse_describe_flush(s, status, query, comment)
       end
 
-    with {:disconnect, %{reason: :closed} = err, s} <- result do
-      {:disconnect_and_retry, err, s}
-    end
+    handle_disconnect_retry(result)
   end
 
   def handle_prepare(%Query{} = query, opts, %{queries: nil} = s) do
@@ -401,14 +399,8 @@ defmodule Postgrex.Protocol do
         end
 
       case result do
-        {:ok, query, s} ->
-          {:ok, query, %{s | messages: []}}
-
-        {:disconnect, %{reason: :closed} = err, s} ->
-          {:disconnect_and_retry, err, s}
-
-        other ->
-          other
+        {:ok, query, s} -> {:ok, query, %{s | messages: []}}
+        other -> handle_disconnect_retry(other)
       end
     end
   end
@@ -434,11 +426,15 @@ defmodule Postgrex.Protocol do
           | {:error, %ArgumentError{} | Postgrex.Error.t(), state}
           | {:error, %DBConnection.TransactionError{}, state}
           | {:disconnect, %RuntimeError{}, state}
-          | {:disconnect, %DBConnection.ConnectionError{}, state}
+          | {:disconnect | :disconnect_and_retry, %DBConnection.ConnectionError{}, state}
   def handle_execute(%Query{} = query, params, opts, s) do
     case Keyword.get(opts, :postgrex_copy, false) do
-      true -> handle_execute_copy(query, params, opts, s)
-      false -> handle_execute_result(query, params, opts, s)
+      true ->
+        handle_execute_copy(query, params, opts, s)
+
+      false ->
+        result = handle_execute_result(query, params, opts, s)
+        handle_disconnect_retry(result)
     end
   end
 
@@ -515,17 +511,19 @@ defmodule Postgrex.Protocol do
           {:ok, Postgrex.Result.t(), state}
           | {:error, %ArgumentError{} | Postgrex.Error.t(), state}
           | {:disconnect, %RuntimeError{}, state}
-          | {:disconnect, %DBConnection.ConnectionError{}, state}
+          | {:disconnect | :disconnect_and_retry, %DBConnection.ConnectionError{}, state}
   def handle_close(%Query{ref: ref} = query, opts, %{postgres: {_, ref}} = s) do
-    flushed_close(s, new_status(opts), query)
+    result = flushed_close(s, new_status(opts), query)
+    handle_disconnect_retry(result)
   end
 
   def handle_close(%Query{} = query, _, %{postgres: {_, _}} = s) do
     lock_error(s, :close, query)
   end
 
   def handle_close(%Query{} = query, opts, s) do
-    close(s, new_status(opts), query)
+    result = close(s, new_status(opts), query)
+    handle_disconnect_retry(result)
   end
 
   @impl true
@@ -604,11 +602,8 @@ defmodule Postgrex.Protocol do
     case Keyword.get(opts, :mode, :transaction) do
       :transaction when postgres == :idle ->
         statement = "BEGIN"
-
-        with {:disconnect, %{reason: :closed} = err, s} <-
-               handle_transaction(statement, opts, s) do
-          {:disconnect_and_retry, err, s}
-        end
+        result = handle_transaction(statement, opts, s)
+        handle_disconnect_retry(result)
 
       :savepoint when postgres == :transaction ->
         statement = "SAVEPOINT postgrex_savepoint"
@@ -2098,7 +2093,7 @@ defmodule Postgrex.Protocol do
         bind_execute_close(s, status, query, params)
 
       {error, _, _} = other when error in [:error, :disconnect] ->
-        other
+        handle_disconnect_retry(other)
     end
   end
 
@@ -2110,7 +2105,7 @@ defmodule Postgrex.Protocol do
         bind_execute(s, status, query, params)
 
       {error, _, _} = other when error in [:error, :disconnect] ->
-        other
+        handle_disconnect_retry(other)
     end
   end
 
@@ -2322,7 +2317,9 @@ defmodule Postgrex.Protocol do
         recv_execute(s, status, query, rows, buffer)
 
       {:disconnect, _, _} = dis ->
-        dis
+        with {_, %{reason: :closed} = err, s} <- dis do
+          {:disconnect, %{err | reason: :execute_closed}, s}
+        end
     end
   end
 
@@ -3443,6 +3440,14 @@ defmodule Postgrex.Protocol do
     {:disconnect, err, %{s | buffer: buffer}}
   end
 
+  defp handle_disconnect_retry({:disconnect, %{reason: :closed} = err, s}),
+    do: {:disconnect_and_retry, err, s}
+
+  defp handle_disconnect_retry({:disconnect, err, s}),
+    do: {:disconnect, err, s}
+
+  defp handle_disconnect_retry(other), do: other
+
   defp sync_recv(s, status, buffer) do
     %{postgres: postgres, transactions: transactions} = s
 

test/query_test.exs

@@ -1966,23 +1966,69 @@ defmodule QueryTest do
     end
   end
 
-  test "disconnect_and_retry", context do
+  test "disconnect_and_retry with prepare" do
     # Start new connection so we can retry on disconnect
     opts = [database: "postgrex_test", backoff_min: 1, backoff_max: 1]
     {:ok, pid} = P.start_link(opts)
 
-    # Have to get socket by knowing DBConnection specifics
-    fun = fn conn ->
-      {:pool_ref, _, _, _, holder, _} = conn.pool_ref
-      [{:conn, _, _, state, _, _, _, _}] = :ets.lookup(holder, :conn)
-      {:gen_tcp, sock} = state.sock
-      sock
-    end
+    # Drop socket
+    disconnect(pid)
+
+    # Assert preparation happens instead of returning error
+    assert {:ok, _} = P.prepare(pid, "42", "SELECT 42")
+  end
+
+  test "disconnect_and_retry with transaction" do
+    # Start new connection so we can retry on disconnect
+    opts = [database: "postgrex_test", backoff_min: 1, backoff_max: 1]
+    {:ok, pid} = P.start_link(opts)
+
+    # Drop socket
+    disconnect(pid)
+
+    # Assert transaction happens instead of returning error
+    assert {:ok, _} = P.transaction(pid, fn conn -> P.query(conn, "SELECT 1", []) end)
+  end
+
+  test "disconnect_and_retry with closing prepared statement" do
+    # Start new connection so we can retry on disconnect
+    opts = [database: "postgrex_test", backoff_min: 1, backoff_max: 1]
+    {:ok, pid} = P.start_link(opts)
 
-    sock = DBConnection.run(pid, fun)
+    # Prepare query that we wil try to close after disconnecting
+    {:ok, query} = P.prepare(pid, "42", "SELECT 42")
+
+    # Drop socket
+    disconnect(pid)
+
+    # Assert close happens instead of returning error
+    assert :ok = P.close(pid, query)
+  end
+
+  test "disconnect_and_retry on attempting execution of prepared statement" do
+    # Start new connection so we can retry on disconnect
+    opts = [database: "postgrex_test", backoff_min: 1, backoff_max: 1]
+    {:ok, pid} = P.start_link(opts)
+
+    # Prepare query that we wil try to close after disconnecting
+    {:ok, query} = P.prepare(pid, "42", "SELECT 42")
+
+    # Drop socket
+    disconnect(pid)
+
+    # Assert execute happens instead of returning error
+    assert {:ok, _, _} = P.execute(pid, query, [])
+  end
+
+  defp disconnect(pid) do
+    sock = DBConnection.run(pid, &get_socket/1)
     :gen_tcp.shutdown(sock, :read_write)
+  end
 
-    # Assert preparation happens instead of returning error
-    assert {:ok, %Postgrex.Query{}} = P.prepare(pid, "42", "SELECT 42")
+  defp get_socket(conn) do
+    {:pool_ref, _, _, _, holder, _} = conn.pool_ref
+    [{:conn, _, _, state, _, _, _, _}] = :ets.lookup(holder, :conn)
+    {:gen_tcp, sock} = state.sock
+    sock
   end
 end