From 4f341269843c1aa494b52b982aef91c288d88b3d Mon Sep 17 00:00:00 2001
From: Jason Stiebs <jason@fly.io>
Date: Thu, 2 Nov 2023 13:28:01 -0500
Subject: [PATCH 1/3] Adding before_disconnect hook

---
 lib/exqlite/connection.ex        | 18 ++++++++++++++----
 test/exqlite/connection_test.exs | 13 +++++++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/lib/exqlite/connection.ex b/lib/exqlite/connection.ex
index 09624821..9977262c 100644
--- a/lib/exqlite/connection.ex
+++ b/lib/exqlite/connection.ex
@@ -36,7 +36,8 @@ defmodule Exqlite.Connection do
     :path,
     :transaction_status,
     :status,
-    :chunk_size
+    :chunk_size,
+    :before_disconnect
   ]
 
   @type t() :: %__MODULE__{
@@ -44,7 +45,9 @@ defmodule Exqlite.Connection do
           directory: String.t() | nil,
           path: String.t(),
           transaction_status: :idle | :transaction,
-          status: :idle | :busy
+          status: :idle | :busy,
+          chunk_size: integer(),
+          before_disconnect: (t -> any) | {module, atom, [any]} | nil
         }
 
   @type journal_mode() :: :delete | :truncate | :persist | :memory | :wal | :off
@@ -73,6 +76,7 @@ defmodule Exqlite.Connection do
           | {:hard_heap_limit, integer()}
           | {:key, String.t()}
           | {:custom_pragmas, [{keyword(), integer() | boolean() | String.t()}]}
+          | {:before_disconnect, (t -> any) | {module, atom, [any]} | nil}
 
   @impl true
   @doc """
@@ -155,6 +159,7 @@ defmodule Exqlite.Connection do
             "./priv/sqlite/\#{arch_dir}/vss0"
           ]
       ```
+    * `:before_disconnect` - A function to run before disconnect, either a 2-arity fun or `{module, function, args}` with the close reason and `t:Exqlite.Connection.t/0` prepended to `args` or `nil` (default: `nil`)
 
   For more information about the options above, see [sqlite documentation][1]
 
@@ -190,7 +195,11 @@ defmodule Exqlite.Connection do
   end
 
   @impl true
-  def disconnect(_err, %__MODULE__{db: db}) do
+  def disconnect(err, %__MODULE__{db: db} = state) do
+    if state.before_disconnect != nil do
+      apply(state.before_disconnect, [err, state])
+    end
+
     case Sqlite3.close(db) do
       :ok -> :ok
       {:error, reason} -> {:error, %Error{message: to_string(reason)}}
@@ -539,7 +548,8 @@ defmodule Exqlite.Connection do
         path: database,
         transaction_status: :idle,
         status: :idle,
-        chunk_size: Keyword.get(options, :chunk_size)
+        chunk_size: Keyword.get(options, :chunk_size),
+        before_disconnect: Keyword.get(options, :before_disconnect, nil)
       }
 
       {:ok, state}
diff --git a/test/exqlite/connection_test.exs b/test/exqlite/connection_test.exs
index 288cf022..4bda34f0 100644
--- a/test/exqlite/connection_test.exs
+++ b/test/exqlite/connection_test.exs
@@ -184,6 +184,19 @@ defmodule Exqlite.ConnectionTest do
 
       assert :ok == Connection.disconnect(nil, conn)
     end
+
+    test "executes before_disconnect before disconnecting" do
+      {:ok, conn} =
+        Connection.connect(
+          database: :memory,
+          before_disconnect: fn err, db ->
+            assert err
+            assert is_function(db.before_disconnect)
+          end
+        )
+
+      assert :ok == Connection.disconnect(true, conn)
+    end
   end
 
   describe ".handle_execute/4" do

From 40b3658e2533ac69d6e8413aa2829f937b8121c0 Mon Sep 17 00:00:00 2001
From: Jason <jason@fly.io>
Date: Thu, 2 Nov 2023 14:02:46 -0500
Subject: [PATCH 2/3] Update lib/exqlite/connection.ex

Co-authored-by: Matthew Johnston <warmwaffles@gmail.com>
---
 lib/exqlite/connection.ex | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/lib/exqlite/connection.ex b/lib/exqlite/connection.ex
index 9977262c..b4957d8f 100644
--- a/lib/exqlite/connection.ex
+++ b/lib/exqlite/connection.ex
@@ -159,7 +159,9 @@ defmodule Exqlite.Connection do
             "./priv/sqlite/\#{arch_dir}/vss0"
           ]
       ```
-    * `:before_disconnect` - A function to run before disconnect, either a 2-arity fun or `{module, function, args}` with the close reason and `t:Exqlite.Connection.t/0` prepended to `args` or `nil` (default: `nil`)
+    * `:before_disconnect` - A function to run before disconnect, either a 
+      2-arity fun or `{module, function, args}` with the close reason and 
+      `t:Exqlite.Connection.t/0` prepended to `args` or `nil` (default: `nil`)
 
   For more information about the options above, see [sqlite documentation][1]
 

From f00718f7f32d3f34bed8cb6bba51b8f192e601de Mon Sep 17 00:00:00 2001
From: Jason Stiebs <jason@fly.io>
Date: Thu, 2 Nov 2023 14:06:43 -0500
Subject: [PATCH 3/3] Fixing test

---
 test/exqlite/connection_test.exs | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/test/exqlite/connection_test.exs b/test/exqlite/connection_test.exs
index 4bda34f0..36693805 100644
--- a/test/exqlite/connection_test.exs
+++ b/test/exqlite/connection_test.exs
@@ -186,16 +186,20 @@ defmodule Exqlite.ConnectionTest do
     end
 
     test "executes before_disconnect before disconnecting" do
+      {:ok, pid} = Agent.start_link(fn -> 0 end)
+
       {:ok, conn} =
         Connection.connect(
           database: :memory,
           before_disconnect: fn err, db ->
-            assert err
-            assert is_function(db.before_disconnect)
+            Agent.update(pid, fn count -> count + 1 end)
+            assert err == true
+            assert db
           end
         )
 
       assert :ok == Connection.disconnect(true, conn)
+      assert Agent.get(pid, &Function.identity/1) == 1
     end
   end