Allow statements to auto close when consumed if no cache (#198)

spec/dummy_driver.cr

@@ -34,14 +34,20 @@ class DummyDriver < DB::Driver
     end
 
     def build_prepared_statement(query) : DB::Statement
+      assert_not_closed!
+
       DummyStatement.new(self, query, true)
     end
 
     def build_unprepared_statement(query) : DB::Statement
+      assert_not_closed!
+
       DummyStatement.new(self, query, false)
     end
 
     def last_insert_id : Int64
+      assert_not_closed!
+
       0
     end
 
@@ -54,12 +60,18 @@ class DummyDriver < DB::Driver
     end
 
     def create_transaction
+      assert_not_closed!
+
       DummyTransaction.new(self)
     end
 
     protected def do_close
       super
     end
+
+    private def assert_not_closed!
+      raise "Statement is closed" if closed?
+    end
   end
 
   class DummyTransaction < DB::TopLevelTransaction
@@ -114,6 +126,8 @@ class DummyDriver < DB::Driver
     end
 
     protected def perform_query(args : Enumerable) : DB::ResultSet
+      assert_not_closed!
+
       Fiber.yield
       @connection.as(DummyConnection).check
       set_params args
@@ -121,6 +135,8 @@ class DummyDriver < DB::Driver
     end
 
     protected def perform_exec(args : Enumerable) : DB::ExecResult
+      assert_not_closed!
+
       @connection.as(DummyConnection).check
       set_params args
       raise DB::Error.new("forced exception due to query") if command == "raise"
@@ -153,6 +169,10 @@ class DummyDriver < DB::Driver
     protected def do_close
       super
     end
+
+    private def assert_not_closed!
+      raise "Statement is closed" if closed?
+    end
   end
 
   class DummyResultSet < DB::ResultSet

spec/statement_spec.cr

@@ -43,6 +43,17 @@ describe DB::Statement do
       end
     end
 
+    it "should leave statements open to be reused if true" do
+      with_dummy_connection("prepared_statements=true&prepared_statements_cache=true") do |cnn|
+        rs = cnn.query("the query")
+        # do not close while iterating
+        rs.statement.closed?.should be_false
+        rs.close
+        # do not close to be reused
+        rs.statement.closed?.should be_false
+      end
+    end
+
     it "should not reuse prepared statements if false" do
       with_dummy_connection("prepared_statements=true&prepared_statements_cache=false") do |cnn|
         stmt1 = cnn.query("the query").statement
@@ -50,6 +61,31 @@ describe DB::Statement do
         stmt1.object_id.should_not eq(stmt2.object_id)
       end
     end
+
+    it "should close statements if false" do
+      with_dummy_connection("prepared_statements=true&prepared_statements_cache=false") do |cnn|
+        rs = cnn.query("the query")
+        # do not close while iterating
+        rs.statement.closed?.should be_false
+        rs.close
+        # do close after iterating
+        rs.statement.closed?.should be_true
+      end
+    end
+
+    it "should not close statements if false and created explicitly" do
+      with_dummy_connection("prepared_statements=true&prepared_statements_cache=false") do |cnn|
+        stmt = cnn.prepared("the query")
+
+        rs = stmt.query
+        # do not close while iterating
+        stmt.closed?.should be_false
+        rs.close
+
+        # do not close after iterating
+        stmt.closed?.should be_false
+      end
+    end
   end
 
   it "should initialize positional params in query" do

src/db/connection.cr

@@ -50,6 +50,10 @@ module DB
       @options.prepared_statements
     end
 
+    def prepared_statements_cache? : Bool
+      @options.prepared_statements_cache
+    end
+
     # :nodoc:
     def fetch_or_build_prepared_statement(query) : Statement
       if @options.prepared_statements_cache

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.
     #

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.
+
+      # 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
@@ -46,7 +53,7 @@ module DB
         conn.release
         raise ex
       end
-      unless existing
+      if !existing && @db.prepared_statements_cache?
         @mutex.synchronize do
           @connections << WeakRef.new(conn)
         end
@@ -55,6 +62,8 @@ module DB
     end
 
     private def clean_connections
+      return unless @db.prepared_statements_cache?
+
       @mutex.synchronize do
         # remove disposed or closed connections
         @connections.each do |ref|

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

src/db/session_methods.cr

@@ -14,13 +14,27 @@ module DB
     # be prepared or not.
     abstract def prepared_statements? : Bool
 
+    abstract def prepared_statements_cache? : Bool
+
     abstract def fetch_or_build_prepared_statement(query) : Stmt
 
     abstract def build_unprepared_statement(query) : Stmt
 
     def build(query) : Stmt
       if prepared_statements?
-        fetch_or_build_prepared_statement(query)
+        stmt = fetch_or_build_prepared_statement(query)
+
+        # #build is a :nodoc: method used on QueryMethods where
+        # the statements are not exposed. As such if the cache
+        # is disabled we should auto_close the statement.
+        # When the statements are build explicitly the #prepared
+        # and #unprepared methods are used. In that case the
+        # statement is closed by the user explicitly also.
+        if !prepared_statements_cache?
+          stmt.auto_close = true if stmt.responds_to?(:auto_close=)
+        end
+
+        stmt
       else
         build_unprepared_statement(query)
       end

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