From d8483895cb26ad3d00a0b49e756cedaa3e3b2565 Mon Sep 17 00:00:00 2001
From: Mat Trudel <mat@geeky.net>
Date: Sun, 22 Jun 2025 21:38:03 -0400
Subject: [PATCH] Allow `unallow_existing` as an opt to ownership_allow/4

---
 integration_test/ownership/manager_test.exs | 14 +++++++++
 lib/db_connection/ownership.ex              |  4 +++
 lib/db_connection/ownership/manager.ex      | 34 +++++++++++++++++++--
 3 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/integration_test/ownership/manager_test.exs b/integration_test/ownership/manager_test.exs
index d1ded2e..1496f9f 100644
--- a/integration_test/ownership/manager_test.exs
+++ b/integration_test/ownership/manager_test.exs
@@ -49,6 +49,7 @@ defmodule ManagerTest do
 
   test "connection may be shared with other processes" do
     {:ok, pool, _opts} = start_pool()
+    {:ok, pool_2, _opts} = start_pool()
     parent = self()
 
     Task.await(
@@ -63,6 +64,12 @@ defmodule ManagerTest do
     assert Ownership.ownership_allow(pool, self(), self(), []) ==
              {:already, :owner}
 
+    :ok = Ownership.ownership_checkout(pool_2, [])
+
+    assert Ownership.ownership_allow(pool_2, self(), self(), unallow_existing: true) == :ok
+
+    assert Ownership.ownership_allow(pool_2, self(), self(), []) == {:already, :allowed}
+
     Task.await(
       async_no_callers(fn ->
         assert Ownership.ownership_allow(pool, parent, self(), []) ==
@@ -73,6 +80,13 @@ defmodule ManagerTest do
 
         assert Ownership.ownership_checkin(pool, []) == :not_owner
 
+        # Test that we can switch our 'allow' to another pool and back again
+        assert Ownership.ownership_allow(pool_2, parent, self(), unallow_existing: true) == :ok
+
+        assert Ownership.ownership_allow(pool_2, parent, self(), []) == {:already, :allowed}
+
+        assert Ownership.ownership_allow(pool, parent, self(), unallow_existing: true) == :ok
+
         parent = self()
 
         Task.await(
diff --git a/lib/db_connection/ownership.ex b/lib/db_connection/ownership.ex
index f61a379..77fe394 100644
--- a/lib/db_connection/ownership.ex
+++ b/lib/db_connection/ownership.ex
@@ -144,6 +144,10 @@ defmodule DBConnection.Ownership do
   has a connection. `owner_or_allowed` may either be the owner or any
   other allowed process. Returns `:not_found` if the given process
   does not have any connection checked out.
+
+  Setting the `unallow_existing` option to `true` will remove the process given by `allow` from
+  any existing allowance it may have (this is necessary because a given process can only be
+  allowed on a single connection at a time).
   """
   @spec ownership_allow(GenServer.server(), owner_or_allowed :: pid, allow :: pid, Keyword.t()) ::
           :ok | {:already, :owner | :allowed} | :not_found
diff --git a/lib/db_connection/ownership/manager.ex b/lib/db_connection/ownership/manager.ex
index c8baeaa..c128933 100644
--- a/lib/db_connection/ownership/manager.ex
+++ b/lib/db_connection/ownership/manager.ex
@@ -56,7 +56,8 @@ defmodule DBConnection.Ownership.Manager do
           :ok | {:already, :owner | :allowed} | :not_found
   def allow(manager, parent, allow, opts) do
     timeout = Keyword.get(opts, :timeout, @timeout)
-    GenServer.call(manager, {:allow, parent, allow}, timeout)
+    passed_opts = Keyword.take(opts, [:unallow_existing])
+    GenServer.call(manager, {:allow, parent, allow, passed_opts}, timeout)
   end
 
   @spec get_connection_metrics(GenServer.server()) ::
@@ -170,15 +171,24 @@ defmodule DBConnection.Ownership.Manager do
     {:reply, reply, state}
   end
 
-  def handle_call({:allow, caller, allow}, _from, %{checkouts: checkouts} = state) do
-    if kind = already_checked_out(checkouts, allow) do
+  def handle_call({:allow, caller, allow, opts}, _from, %{checkouts: checkouts} = state) do
+    unallow_existing = Keyword.get(opts, :unallow_existing, false)
+    kind = already_checked_out(checkouts, allow)
+
+    if !unallow_existing && kind do
       {:reply, {:already, kind}, state}
     else
       case Map.get(checkouts, caller, :not_found) do
         {:owner, ref, proxy} ->
+          state =
+            if unallow_existing, do: owner_unallow(state, caller, allow, ref, proxy), else: state
+
           {:reply, :ok, owner_allow(state, caller, allow, ref, proxy)}
 
         {:allowed, ref, proxy} ->
+          state =
+            if unallow_existing, do: owner_unallow(state, caller, allow, ref, proxy), else: state
+
           {:reply, :ok, owner_allow(state, caller, allow, ref, proxy)}
 
         :not_found ->
@@ -310,6 +320,24 @@ defmodule DBConnection.Ownership.Manager do
     state
   end
 
+  defp owner_unallow(%{ets: ets, log: log} = state, caller, unallow, ref, proxy) do
+    if log do
+      Logger.log(log, fn ->
+        [inspect(unallow), " was unallowed by ", inspect(caller), " on proxy ", inspect(proxy)]
+      end)
+    end
+
+    state = update_in(state.checkouts, &Map.delete(&1, unallow))
+
+    state =
+      update_in(state.owners[ref], fn {proxy, caller, allowed} ->
+        {proxy, caller, List.delete(allowed, unallow)}
+      end)
+
+    ets && :ets.delete(ets, {unallow, proxy})
+    state
+  end
+
   defp owner_down(%{ets: ets, log: log} = state, ref) do
     case get_and_update_in(state.owners, &Map.pop(&1, ref)) do
       {{proxy, caller, allowed}, state} ->