Merge pull request rails#51349 from Shopify/connection-optional-deprecation

Add `config.active_record.permanent_connection_checkout` setting

Author: byroot

activerecord/CHANGELOG.md

@@ -1,3 +1,18 @@
+*   Add `config.active_record.permanent_connection_checkout` setting.
+
+    Controls whether `ActiveRecord::Base.connection` raises an error, emits a deprecation warning, or neither.
+
+    `ActiveRecord::Base.connection` checkouts a database connection from the pool and keep it leased until the end of
+    the request or job. This behavior can be undesirable in environments that use many more threads or fibers than there
+    is available connections.
+
+    This configuration can be used to track down and eliminate code that calls `ActiveRecord::Base.connection` and
+    migrate it to use `ActiveRecord::Base.with_connection` instead.
+
+    The default behavior remains unchanged, and there is currently no plans to change the default.
+
+    *Jean Boussier*
+
 *   Add dirties option to uncached
 
     This adds a `dirties` option to `ActiveRecord::Base.uncached` and

activerecord/lib/active_record.rb

@@ -300,6 +300,17 @@ def self.global_executor_concurrency # :nodoc:
     @global_executor_concurrency ||= nil
   end
 
+  @permanent_connection_checkout = true
+  singleton_class.attr_reader :permanent_connection_checkout
+
+  # Defines whether +ActiveRecord::Base.connection+ is allowed, deprecated or entirely disallowed
+  def self.permanent_connection_checkout=(value)
+    unless [true, :deprecated, :disallowed].include?(value)
+      raise ArgumentError, "permanent_connection_checkout must be one of: `true`, `:deprecated` or `:disallowed`"
+    end
+    @permanent_connection_checkout = value
+  end
+
   singleton_class.attr_accessor :index_nested_attribute_errors
   self.index_nested_attribute_errors = false
 

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
@@ -290,14 +290,18 @@ def internal_metadata # :nodoc:
       # Retrieve the connection associated with the current thread, or call
       # #checkout to obtain one if necessary.
       #
-      # #connection can be called any number of times; the connection is
+      # #lease_connection can be called any number of times; the connection is
       # held in a cache keyed by a thread.
       def lease_connection
         lease = connection_lease
         lease.sticky = true
         lease.connection ||= checkout
       end
 
+      def permanent_lease? # :nodoc:
+        connection_lease.sticky.nil?
+      end
+
       def connection
         ActiveRecord.deprecator.warn(<<~MSG)
           ActiveRecord::ConnectionAdapters::ConnectionPool#connection is deprecated
@@ -348,18 +352,19 @@ def connection_class # :nodoc:
       # Returns true if there is an open connection being used for the current thread.
       #
       # This method only works for connections that have been obtained through
-      # #connection or #with_connection methods. Connections obtained through
+      # #lease_connection or #with_connection methods. Connections obtained through
       # #checkout will not be detected by #active_connection?
       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
       # and returns the connection to the pool.
       #
       # This method only works for connections that have been obtained through
-      # #connection or #with_connection methods, connections obtained through
+      # #lease_connection or #with_connection methods, connections obtained through
       # #checkout will not be automatically released.
       def release_connection(existing_lease = nil)
         if conn = connection_lease.release
@@ -373,19 +378,27 @@ def release_connection(existing_lease = nil)
       # is already checked out by the current thread, a connection will be checked
       # out from the pool, yielded to the block, and then returned to the pool when
       # the block is finished. If a connection has already been checked out on the
-      # current thread, such as via #connection or #with_connection, that existing
+      # current thread, such as via #lease_connection or #with_connection, that existing
       # connection will be the one yielded and it will not be returned to the pool
       # 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_adapters/abstract/query_cache.rb

@@ -93,7 +93,7 @@ def initialize(...)
             end
         end
 
-        def lease_connection
+        def lease_connection(**)
           connection = super
           connection.query_cache ||= query_cache
           connection

activerecord/lib/active_record/connection_handling.rb

@@ -256,18 +256,44 @@ def lease_connection
       connection_pool.lease_connection
     end
 
-    alias_method :connection, :lease_connection
+    # Soft deprecated. Use +#with_connection+ or +#lease_connection+ instead.
+    def connection
+      pool = connection_pool
+      if pool.permanent_lease?
+        case ActiveRecord.permanent_connection_checkout
+        when :deprecated
+          ActiveRecord.deprecator.warn <<~MESSAGE
+            Called deprecated `ActionRecord::Base.connection`method.
+
+            Either use `with_connection` or `lease_connection`.
+          MESSAGE
+        when :disallowed
+          raise ActiveRecordError, <<~MESSAGE
+            Called deprecated `ActionRecord::Base.connection`method.
+
+            Either use `with_connection` or `lease_connection`.
+          MESSAGE
+        end
+        pool.lease_connection
+      else
+        pool.active_connection
+      end
+    end
 
     # Return the currently leased connection into the pool
     def release_connection
-      connection.release_connection
+      connection_pool.release_connection
     end
 
     # Checkouts a connection from the pool, yield it and then check it back in.
-    # If a connection was already leased via #connection or a parent call to
+    # 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

@@ -4,9 +4,17 @@
 
 module ActiveRecord
   class ConnectionHandlingTest < ActiveRecord::TestCase
+    setup do
+      @_permanent_connection_checkout_was = ActiveRecord.permanent_connection_checkout
+    end
+
+    teardown do
+      ActiveRecord.permanent_connection_checkout = @_permanent_connection_checkout_was
+    end
+
     unless in_memory_db?
       test "#with_connection lease the connection for the duration of the block" do
-        ActiveRecord::Base.connection_pool.release_connection
+        ActiveRecord::Base.release_connection
         assert_not_predicate ActiveRecord::Base.connection_pool, :active_connection?
 
         ActiveRecord::Base.with_connection do |connection|
@@ -16,8 +24,8 @@ class ConnectionHandlingTest < ActiveRecord::TestCase
         assert_not_predicate ActiveRecord::Base.connection_pool, :active_connection?
       end
 
-      test "#connection makes the lease permanent even inside #with_connection" do
-        ActiveRecord::Base.connection_pool.release_connection
+      test "#lease_connection makes the lease permanent even inside #with_connection" do
+        ActiveRecord::Base.release_connection
         assert_not_predicate ActiveRecord::Base.connection_pool, :active_connection?
 
         conn = nil
@@ -33,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?
@@ -63,6 +81,89 @@ class ConnectionHandlingTest < ActiveRecord::TestCase
         assert_predicate ActiveRecord::Base.connection_pool, :active_connection?
         assert_same ActiveRecord::Base.lease_connection, leased_connection
       end
+
+      test "#connection is a soft-deprecated alias to #lease_connection" do
+        ActiveRecord.permanent_connection_checkout = true
+
+        ActiveRecord::Base.release_connection
+        assert_not_predicate ActiveRecord::Base.connection_pool, :active_connection?
+
+        conn = nil
+        ActiveRecord::Base.with_connection do |connection|
+          conn = connection
+          assert_predicate ActiveRecord::Base.connection_pool, :active_connection?
+          2.times do
+            assert_same connection, ActiveRecord::Base.connection
+          end
+        end
+
+        assert_predicate ActiveRecord::Base.connection_pool, :active_connection?
+        assert_same conn, ActiveRecord::Base.connection
+
+        ActiveRecord::Base.release_connection
+      end
+
+      test "#connection emits a deprecation warning if ActiveRecord.permanent_connection_checkout == :deprecated" do
+        ActiveRecord.permanent_connection_checkout = :deprecated
+
+        ActiveRecord::Base.release_connection
+
+        assert_deprecated(ActiveRecord.deprecator) do
+          ActiveRecord::Base.connection
+        end
+
+        assert_not_deprecated(ActiveRecord.deprecator) do
+          ActiveRecord::Base.connection
+        end
+
+        ActiveRecord::Base.release_connection
+
+        assert_deprecated(ActiveRecord.deprecator) do
+          ActiveRecord::Base.connection
+        end
+
+        ActiveRecord::Base.release_connection
+
+        ActiveRecord::Base.with_connection do
+          assert_deprecated(ActiveRecord.deprecator) do
+            ActiveRecord::Base.connection
+          end
+        end
+      end
+
+      test "#connection raises an error if ActiveRecord.permanent_connection_checkout == :disallowed" do
+        ActiveRecord.permanent_connection_checkout = :disallowed
+
+        ActiveRecord::Base.release_connection
+
+        assert_raises(ActiveRecordError) do
+          ActiveRecord::Base.connection
+        end
+
+        ActiveRecord::Base.with_connection do
+          assert_raises(ActiveRecordError) do
+            ActiveRecord::Base.connection
+          end
+        end
+
+        ActiveRecord::Base.lease_connection
+
+        assert_nothing_raised do
+          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

activerecord/test/cases/helper.rb

@@ -22,9 +22,9 @@
 # Show backtraces for deprecated behavior for quicker cleanup.
 ActiveRecord.deprecator.debug = true
 
-# ActiveRecord::Base.connection is only soft deprecated but we remove it
-# in the test suite to ensure we're not using it internally.
-ActiveRecord::ConnectionHandling.remove_method(:connection)
+# ActiveRecord::Base.connection is only soft deprecated but we ban it from the test suite
+# to ensure it's not used internally.
+ActiveRecord.permanent_connection_checkout = :disallowed
 
 # Disable available locale checks to avoid warnings running the test suite.
 I18n.enforce_available_locales = false

guides/source/configuring.md

@@ -1526,6 +1526,25 @@ record.token # => "fwZcXX6SkJBJRogzMdciS7wf"
 | (original)            | `:create`            |
 | 7.1                   | `:initialize`        |
 
+
+#### `config.active_record.permanent_connection_checkout`
+
+Controls whether `ActiveRecord::Base.connection` raises an error, emits a deprecation warning, or neither.
+
+`ActiveRecord::Base.connection` checkouts a database connection from the pool and keep it leased until the end of
+the request or job. This behavior can be undesirable in environments that use many more threads or fibers than there
+is available connections.
+
+This configuration can be used to track down and eliminate code that calls `ActiveRecord::Base.connection` and
+migrate it to use `ActiveRecord::Base.with_connection` instead.
+
+The value can be set to `:disallowed`, `:deprecated` or `true` to respectively raise an error, emit a deprecation
+warning, or neither.
+
+| Starting with version | The default value is |
+| --------------------- | -------------------- |
+| (original)            | `true`               |
+
 #### `ActiveRecord::ConnectionAdapters::Mysql2Adapter.emulate_booleans` and `ActiveRecord::ConnectionAdapters::TrilogyAdapter.emulate_booleans`
 
 Controls whether the Active Record MySQL adapter will consider all `tinyint(1)` columns as booleans. Defaults to `true`.