Reuse prepared statements in prepare: :unnamed (#146)

Co-authored-by: Hans Krutzer <[email protected]>
Co-authored-by: José Valim <[email protected]>
Co-authored-by: Greg Rychlewski <[email protected]>

Author: 4 people

lib/myxql/connection.ex

@@ -13,7 +13,7 @@ defmodule MyXQL.Connection do
     prepare: :named,
     queries: nil,
     transaction_status: :idle,
-    last_ref: nil
+    last_query: nil
   ]
 
   @impl true
@@ -29,7 +29,7 @@ defmodule MyXQL.Connection do
           prepare: prepare,
           disconnect_on_error_codes: Keyword.fetch!(opts, :disconnect_on_error_codes),
           ping_timeout: ping_timeout,
-          queries: queries_new()
+          queries: queries_new(prepare)
         }
 
         {:ok, state}
@@ -64,7 +64,8 @@ defmodule MyXQL.Connection do
     query = rename_query(state, query)
 
     if cached_query = queries_get(state, query) do
-      {:ok, cached_query, %{state | last_ref: cached_query.ref}}
+      %{ref: ref, statement_id: statement_id} = cached_query
+      {:ok, cached_query, %{state | last_query: {ref, statement_id}}}
     else
       case prepare(query, state) do
         {:ok, _, _} = ok ->
@@ -469,14 +470,21 @@ defmodule MyXQL.Connection do
         ref = make_ref()
         query = %{query | num_params: num_params, statement_id: statement_id, ref: ref}
         queries_put(state, query)
-        {:ok, query, %{state | last_ref: ref}}
+        {:ok, query, %{state | last_query: {ref, statement_id}}}
 
       result ->
         result(result, query, state)
     end
   end
 
-  defp maybe_reprepare(%{ref: ref} = query, %{last_ref: ref} = state), do: {:ok, query, state}
+  defp maybe_reprepare(%{ref: ref} = query, %{last_query: {ref, _statement_id}} = state) do
+    {:ok, query, state}
+  end
+
+  defp maybe_reprepare(query, %{queries: nil, last_query: {_ref, statement_id}} = state) do
+    Client.com_stmt_close(state.client, statement_id)
+    prepare(query, state)
+  end
 
   defp maybe_reprepare(query, state) do
     if query_member?(state, query) do
@@ -490,12 +498,16 @@ defmodule MyXQL.Connection do
     %{state | cursors: Map.delete(state.cursors, cursor.ref)}
   end
 
-  # Close unnamed queries after executing them
+  # When prepare is :named, close unnamed queries after executing them.
+  # When prepare is :unnamed, queries will be closed the next time a
+  # query is prepared or a different query is executed. This allows us
+  # to re-execute the same unnamed query without preparing it again.
+  defp maybe_close(_query, %{queries: nil} = state), do: {:ok, state}
   defp maybe_close(%{name: ""} = query, state), do: close(query, state)
   defp maybe_close(_query, state), do: {:ok, state}
 
-  defp close(%{ref: ref} = query, %{last_ref: ref} = state) do
-    close(query, %{state | last_ref: nil})
+  defp close(%{ref: ref} = query, %{last_query: {ref, _statement_id}} = state) do
+    close(query, %{state | last_query: nil})
   end
 
   defp close(query, state) do
@@ -516,7 +528,8 @@ defmodule MyXQL.Connection do
 
   ## Cache query handling
 
-  defp queries_new(), do: :ets.new(__MODULE__, [:set, :public])
+  defp queries_new(:unnamed), do: nil
+  defp queries_new(_), do: :ets.new(__MODULE__, [:set, :public])
 
   defp queries_put(%{queries: nil}, _), do: :ok
   defp queries_put(_state, %{name: ""}), do: :ok
@@ -570,7 +583,11 @@ defmodule MyXQL.Connection do
     end
   end
 
-  defp queries_get(%{queries: nil}, _), do: nil
+  defp queries_get(%{queries: nil, last_query: {_ref, statement_id}} = state, _query) do
+    Client.com_stmt_close(state.client, statement_id)
+    nil
+  end
+
   defp queries_get(_state, %{name: ""}), do: nil
 
   defp queries_get(state, %{cache: :reference, name: name}) do

test/myxql/sync_test.exs

@@ -8,8 +8,36 @@ defmodule MyXQL.SyncTest do
     {:ok, conn} = MyXQL.start_link(@opts ++ [prepare: :unnamed])
     assert prepared_stmt_count() == 0
 
+    # Multiple queries do not increase statement count
     MyXQL.query!(conn, "SELECT 42", [], cache_statement: "42")
+    assert prepared_stmt_count() == 1
+
+    MyXQL.query!(conn, "SELECT 43", [], cache_statement: "43")
+    assert prepared_stmt_count() == 1
+
+    # Multiple preparations don't increase statement count
+    {:ok, _query} = MyXQL.prepare(conn, "1", "SELECT 1")
+    assert prepared_stmt_count() == 1
+
+    {:ok, _query} = MyXQL.prepare(conn, "2", "SELECT 2")
+    assert prepared_stmt_count() == 1
+  end
+
+  test "do not leak statements with prepare: :named" do
+    {:ok, conn} = MyXQL.start_link(@opts)
     assert prepared_stmt_count() == 0
+
+    MyXQL.query!(conn, "SELECT 42", [], cache_statement: "42")
+    assert prepared_stmt_count() == 1
+
+    MyXQL.query!(conn, "SELECT 1337", [], cache_statement: "1337")
+    assert prepared_stmt_count() == 2
+
+    MyXQL.query!(conn, "SELECT 42", [], cache_statement: "42")
+    assert prepared_stmt_count() == 2
+
+    MyXQL.query!(conn, "SELECT 43", [])
+    assert prepared_stmt_count() == 2
   end
 
   test "do not leak statements with rebound :cache_statement" do
@@ -23,15 +51,31 @@ defmodule MyXQL.SyncTest do
     assert prepared_stmt_count() == 1
   end
 
-  test "do not leak statements with insert and failed insert" do
+  test "do not leak statements with insert and failed insert with prepare: :named" do
     {:ok, conn} = MyXQL.start_link(@opts)
     assert prepared_stmt_count() == 0
+
     {:ok, _} = MyXQL.query(conn, "INSERT INTO uniques(a) VALUES (1)")
     assert prepared_stmt_count() == 0
+
     {:error, _} = MyXQL.query(conn, "INSERT INTO uniques(a) VALUES (1)")
     assert prepared_stmt_count() == 0
   end
 
+  test "do not leak statements with insert and failed insert with prepare: :unnamed" do
+    {:ok, conn} = MyXQL.start_link(@opts ++ [prepare: :unnamed])
+    assert prepared_stmt_count() == 0
+
+    {:ok, _} = MyXQL.query(conn, "INSERT INTO uniques(a) VALUES (2)")
+    assert prepared_stmt_count() == 1
+
+    {:error, _} = MyXQL.query(conn, "INSERT INTO uniques(a) VALUES (2)")
+    assert prepared_stmt_count() == 1
+
+    MyXQL.query!(conn, "SELECT 123", [], cache_statement: "123")
+    assert prepared_stmt_count() == 1
+  end
+
   test "do not leak statements on multiple executions of the same name in prepare_execute" do
     {:ok, conn} = MyXQL.start_link(@opts)
     {:ok, _, _} = MyXQL.prepare_execute(conn, "foo", "SELECT 42")

test/myxql_test.exs

@@ -205,14 +205,26 @@ defmodule MyXQLTest do
       assert query == query3
     end
 
-    test ":unnamed" do
+    test ":unnamed re-executes last query without preparing" do
       {:ok, pid} = MyXQL.start_link(@opts ++ [prepare: :unnamed])
       {:ok, query} = MyXQL.prepare(pid, "1", "SELECT 1")
       {:ok, query2, _} = MyXQL.execute(pid, query, [])
       assert query == query2
       {:ok, query3, _} = MyXQL.execute(pid, query, [])
-      assert query2.ref != query3.ref
-      assert query2.statement_id != query3.statement_id
+      assert query2 == query3
+    end
+
+    test ":unnamed re-prepares if last query is not the same" do
+      {:ok, pid} = MyXQL.start_link(@opts ++ [prepare: :unnamed])
+
+      {:ok, query1} = MyXQL.prepare(pid, "1", "SELECT 1")
+      {:ok, query2} = MyXQL.prepare(pid, "2", "SELECT 2")
+
+      {:ok, query1b, _} = MyXQL.execute(pid, query1, [])
+      {:ok, query2b, _} = MyXQL.execute(pid, query2, [])
+
+      assert query1 != query1b
+      assert query2 != query2b
     end
 
     test ":force_named" do