From ad41c258c8beb4b59077c35b989cf5069ab902af Mon Sep 17 00:00:00 2001
From: "Brian J. Cardiff" <bcardiff@gmail.com>
Date: Fri, 10 Nov 2023 23:00:35 -0300
Subject: [PATCH] Allow statements to auto close when consumed if no cache

---
 src/db/connection.cr              | 4 +++-
 src/db/database.cr                | 4 ++++
 src/db/pool_prepared_statement.cr | 9 ++++++++-
 src/db/result_set.cr              | 2 +-
 src/db/statement.cr               | 9 +++++++++
 5 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/src/db/connection.cr b/src/db/connection.cr
index fb0baa8..0f68f75 100644
--- a/src/db/connection.cr
+++ b/src/db/connection.cr
@@ -55,7 +55,9 @@ module DB
       if @options.prepared_statements_cache
         @statements_cache.fetch(query) { build_prepared_statement(query) }
       else
-        build_prepared_statement(query)
+        stmt = build_prepared_statement(query)
+        stmt.auto_close = true
+        stmt
       end
     end
 
diff --git a/src/db/database.cr b/src/db/database.cr
index 47531a3..79741d5 100644
--- a/src/db/database.cr
+++ b/src/db/database.cr
@@ -59,6 +59,10 @@ module DB
       @connection_options.prepared_statements
     end
 
+    def prepared_statements_cache? : Bool
+      @connection_options.prepared_statements_cache
+    end
+
     # Run the specified block every time a new connection is established, yielding the new connection
     # to the block.
     #
diff --git a/src/db/pool_prepared_statement.cr b/src/db/pool_prepared_statement.cr
index cb598e2..8652cb9 100644
--- a/src/db/pool_prepared_statement.cr
+++ b/src/db/pool_prepared_statement.cr
@@ -15,7 +15,14 @@ module DB
       # otherwise the preparation is delayed until the first execution.
       # After the first initialization the connection must be released
       # it will be checked out when executing it.
-      statement_with_retry &.release_connection
+
+      # This only happens if the db is configured to use prepared statements cache.
+      # Without that there is no reference to the already prepared statement we can
+      # take advantage of.
+      if db.prepared_statements_cache?
+        statement_with_retry &.release_connection
+      end
+
       # TODO use a round-robin selection in the pool so multiple sequentially
       #      initialized statements are assigned to different connections.
     end
diff --git a/src/db/result_set.cr b/src/db/result_set.cr
index ccc58b4..b23be6b 100644
--- a/src/db/result_set.cr
+++ b/src/db/result_set.cr
@@ -29,7 +29,7 @@ module DB
     end
 
     protected def do_close
-      statement.release_connection
+      statement.release_from_result_set
     end
 
     # TODO add_next_result_set : Bool
diff --git a/src/db/statement.cr b/src/db/statement.cr
index 8ed9246..6602e6d 100644
--- a/src/db/statement.cr
+++ b/src/db/statement.cr
@@ -56,6 +56,15 @@ module DB
     def initialize(@connection : Connection, @command : String)
     end
 
+    # :nodoc:
+    property auto_close : Bool = false
+
+    # :nodoc:
+    def release_from_result_set
+      self.close if @auto_close
+      self.release_connection
+    end
+
     def release_connection
       @connection.release_from_statement
     end