Implement with_connection(prevent_permanent_checkout: true) option

When `.connection` is called inside a `.with_connection` block
it cause the lease to become permanent (or sticky). This is to
ensure that legacy code that may hold onto this connection
won't cause thread safety issues.

However if you autidted the code that calls `.connection` and
know for sure it won't hold into the returned connection,
you can wrap it with `with_connection(prevent_permanent_checkout: true)`
to prevent the lease from becoming permanent and ensure the connection
is released at the end of the block.

Author: byroot

activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb

@@ -123,20 +123,20 @@ class Lease # :nodoc:
 
         def initialize
           @connection = nil
-          @sticky = false
+          @sticky = nil
         end
 
         def release
           conn = @connection
           @connection = nil
-          @sticky = false
+          @sticky = nil
           conn
         end
 
         def clear(connection)
           if @connection == connection
             @connection = nil
-            @sticky = false
+            @sticky = nil
             true
           else
             false
@@ -298,8 +298,8 @@ def lease_connection
         lease.connection ||= checkout
       end
 
-      def sticky_lease? # :nodoc:
-        connection_lease.sticky
+      def permanent_lease? # :nodoc:
+        connection_lease.sticky.nil?
       end
 
       def connection
@@ -357,6 +357,7 @@ def connection_class # :nodoc:
       def active_connection?
         connection_lease.connection
       end
+      alias_method :active_connection, :active_connection? # :nodoc:
 
       # Signal that the thread is finished with the current connection.
       # #release_connection releases the connection-thread association
@@ -382,14 +383,22 @@ def release_connection(existing_lease = nil)
       # automatically at the end of the block; it is expected that such an existing
       # connection will be properly returned to the pool by the code that checked
       # it out.
-      def with_connection
+      def with_connection(prevent_permanent_checkout: false)
         lease = connection_lease
+        sticky_was = lease.sticky
+        lease.sticky = false if prevent_permanent_checkout
+
         if lease.connection
-          yield lease.connection
+          begin
+            yield lease.connection
+          ensure
+            lease.sticky = sticky_was if prevent_permanent_checkout && !sticky_was
+          end
         else
           begin
             yield lease.connection = checkout
           ensure
+            lease.sticky = sticky_was if prevent_permanent_checkout && !sticky_was
             release_connection(lease) unless lease.sticky
           end
         end

activerecord/lib/active_record/connection_handling.rb

@@ -259,7 +259,7 @@ def lease_connection
     # Soft deprecated. Use +#with_connection+ or +#lease_connection+ instead.
     def connection
       pool = connection_pool
-      unless pool.sticky_lease?
+      if pool.permanent_lease?
         case ActiveRecord.permanent_connection_checkout
         when :deprecated
           ActiveRecord.deprecator.warn <<~MESSAGE
@@ -274,8 +274,10 @@ def connection
             Either use `with_connection` or `lease_connection`.
           MESSAGE
         end
+        pool.lease_connection
+      else
+        pool.active_connection
       end
-      pool.lease_connection
     end
 
     # Return the currently leased connection into the pool
@@ -286,8 +288,12 @@ def release_connection
     # Checkouts a connection from the pool, yield it and then check it back in.
     # If a connection was already leased via #lease_connection or a parent call to
     # #with_connection, that same connection is yieled.
-    def with_connection(&block)
-      connection_pool.with_connection(&block)
+    # If #lease_connection is called inside the block, the connection won't be checked
+    # back in.
+    # If #connection is called inside the block, the connection won't be checked back in
+    # unless the +prevent_permanent_checkout+ argument is set to +true+.
+    def with_connection(prevent_permanent_checkout: false, &block)
+      connection_pool.with_connection(prevent_permanent_checkout: prevent_permanent_checkout, &block)
     end
 
     attr_writer :connection_specification_name

activerecord/test/cases/connection_handling_test.rb

@@ -41,6 +41,16 @@ class ConnectionHandlingTest < ActiveRecord::TestCase
         assert_same conn, ActiveRecord::Base.lease_connection
       end
 
+      test "#lease_connection makes the lease permanent even inside #with_connection(prevent_permanent_checkout: true)" do
+        ActiveRecord::Base.release_connection
+
+        ActiveRecord::Base.with_connection(prevent_permanent_checkout: true) do |connection|
+          assert_same connection, ActiveRecord::Base.lease_connection
+        end
+
+        assert_not_predicate ActiveRecord::Base.connection_pool, :active_connection?
+      end
+
       test "#with_connection use the already leased connection if available" do
         leased_connection = ActiveRecord::Base.lease_connection
         assert_predicate ActiveRecord::Base.connection_pool, :active_connection?
@@ -142,6 +152,18 @@ class ConnectionHandlingTest < ActiveRecord::TestCase
           ActiveRecord::Base.connection
         end
       end
+
+      test "#connection doesn't make the lease permanent if inside #with_connection(prevent_permanent_checkout: true)" do
+        ActiveRecord.permanent_connection_checkout = :disallowed
+
+        ActiveRecord::Base.release_connection
+
+        ActiveRecord::Base.with_connection(prevent_permanent_checkout: true) do |connection|
+          assert_same connection, ActiveRecord::Base.connection
+        end
+
+        assert_not_predicate ActiveRecord::Base.connection_pool, :active_connection?
+      end
     end
   end
 end