Merge pull request #51083 from Shopify/activerecord-with-connection

Add `ActiveRecord::Base.with_connection` as a shortcut

activerecord/CHANGELOG.md

@@ -1,3 +1,13 @@
+*   Add `ActiveRecord::Base.with_connection` as a shortcut for leasing a connection for a short duration.
+
+    The leased connection is yielded, and for the duration of the block, any call to `ActiveRecord::Base.connection`
+    will yield that same connection.
+
+    This is useful to perform a few database operations without causing a connection to be leased for the
+    entire duration of the request or job.
+
+    *Jean Boussier*
+
 *   Deprecate `config.active_record.warn_on_records_fetched_greater_than` now that `sql.active_record`
     notification includes `:row_count` field.
 

activerecord/lib/active_record/associations/builder/has_and_belongs_to_many.rb

@@ -42,8 +42,8 @@ def self.add_right_association(name, options)
           self.right_reflection = _reflect_on_association(rhs_name)
         end
 
-        def self.retrieve_connection
-          left_model.retrieve_connection
+        def self.connection_pool
+          left_model.connection_pool
         end
       }
 

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

@@ -182,7 +182,10 @@ def clear_active_connections!(role = nil)
           role = ActiveRecord::Base.current_role
         end
 
-        each_connection_pool(role).each(&:release_connection)
+        each_connection_pool(role).each do |pool|
+          pool.release_connection
+          pool.disable_query_cache!
+        end
       end
 
       # Clears the cache which maps classes.
@@ -223,20 +226,7 @@ def flush_idle_connections!(role = nil)
       # opened and set as the active connection for the class it was defined
       # for (not necessarily the current class).
       def retrieve_connection(connection_name, role: ActiveRecord::Base.current_role, shard: ActiveRecord::Base.current_shard) # :nodoc:
-        pool = retrieve_connection_pool(connection_name, role: role, shard: shard)
-
-        unless pool
-          if shard != ActiveRecord::Base.default_shard
-            message = "No connection pool for '#{connection_name}' found for the '#{shard}' shard."
-          elsif role != ActiveRecord::Base.default_role
-            message = "No connection pool for '#{connection_name}' found for the '#{role}' role."
-          else
-            message = "No connection pool for '#{connection_name}' found."
-          end
-
-          raise ConnectionNotEstablished, message
-        end
-
+        pool = retrieve_connection_pool(connection_name, role: role, shard: shard, strict: true)
         pool.connection
       end
 
@@ -256,9 +246,22 @@ def remove_connection_pool(connection_name, role: ActiveRecord::Base.current_rol
       # Retrieving the connection pool happens a lot, so we cache it in @connection_name_to_pool_manager.
       # This makes retrieving the connection pool O(1) once the process is warm.
       # When a connection is established or removed, we invalidate the cache.
-      def retrieve_connection_pool(connection_name, role: ActiveRecord::Base.current_role, shard: ActiveRecord::Base.current_shard)
-        pool_config = get_pool_manager(connection_name)&.get_pool_config(role, shard)
-        pool_config&.pool
+      def retrieve_connection_pool(connection_name, role: ActiveRecord::Base.current_role, shard: ActiveRecord::Base.current_shard, strict: false)
+        pool = get_pool_manager(connection_name)&.get_pool_config(role, shard)&.pool
+
+        if strict && !pool
+          if shard != ActiveRecord::Base.default_shard
+            message = "No connection pool for '#{connection_name}' found for the '#{shard}' shard."
+          elsif role != ActiveRecord::Base.default_role
+            message = "No connection pool for '#{connection_name}' found for the '#{role}' role."
+          else
+            message = "No connection pool for '#{connection_name}' found."
+          end
+
+          raise ConnectionNotEstablished, message
+        end
+
+        pool
       end
 
       private

activerecord/lib/active_record/connection_handling.rb

@@ -243,15 +243,22 @@ def connected_to?(role:, shard: ActiveRecord::Base.default_shard)
     # Clears the query cache for all connections associated with the current thread.
     def clear_query_caches_for_current_thread
       connection_handler.each_connection_pool do |pool|
-        pool.connection.clear_query_cache if pool.active_connection?
+        pool.connection.clear_query_cache
       end
     end
 
     # Returns the connection currently associated with the class. This can
     # also be used to "borrow" the connection to do database work unrelated
     # to any of the specific Active Records.
     def connection
-      retrieve_connection
+      connection_pool.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
+    # #with_connection, that same connection is yieled.
+    def with_connection(&block) # :nodoc:
+      connection_pool.with_connection(&block)
     end
 
     attr_writer :connection_specification_name
@@ -280,7 +287,7 @@ def connection_db_config
     end
 
     def connection_pool
-      connection_handler.retrieve_connection_pool(connection_specification_name, role: current_role, shard: current_shard) || raise(ConnectionNotEstablished)
+      connection_handler.retrieve_connection_pool(connection_specification_name, role: current_role, shard: current_shard, strict: true)
     end
 
     def retrieve_connection

activerecord/lib/active_record/transactions.rb

@@ -209,7 +209,9 @@ module Transactions
     module ClassMethods
       # See the ConnectionAdapters::DatabaseStatements#transaction API docs.
       def transaction(**options, &block)
-        connection.transaction(**options, &block)
+        with_connection do |connection|
+          connection.transaction(**options, &block)
+        end
       end
 
       def before_commit(*args, &block) # :nodoc:

activerecord/test/cases/connection_handling_test.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+require "cases/helper"
+
+module ActiveRecord
+  class ConnectionHandlingTest < ActiveRecord::TestCase
+    unless in_memory_db?
+      test "#with_connection lease the connection for the duration of the block" do
+        ActiveRecord::Base.connection_pool.release_connection
+        assert_not_predicate ActiveRecord::Base.connection_pool, :active_connection?
+
+        ActiveRecord::Base.with_connection do |connection|
+          assert_predicate ActiveRecord::Base.connection_pool, :active_connection?
+          assert_same connection, ActiveRecord::Base.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.connection
+        assert_predicate ActiveRecord::Base.connection_pool, :active_connection?
+
+        ActiveRecord::Base.with_connection do |connection|
+          assert_same leased_connection, connection
+          assert_same ActiveRecord::Base.connection, connection
+        end
+
+        assert_predicate ActiveRecord::Base.connection_pool, :active_connection?
+        assert_same ActiveRecord::Base.connection, leased_connection
+      end
+
+      test "#with_connection is reentrant" do
+        leased_connection = ActiveRecord::Base.connection
+        assert_predicate ActiveRecord::Base.connection_pool, :active_connection?
+
+        ActiveRecord::Base.with_connection do |connection|
+          assert_same leased_connection, connection
+          assert_same ActiveRecord::Base.connection, connection
+
+          ActiveRecord::Base.with_connection do |connection2|
+            assert_same leased_connection, connection
+            assert_same ActiveRecord::Base.connection, connection
+          end
+        end
+
+        assert_predicate ActiveRecord::Base.connection_pool, :active_connection?
+        assert_same ActiveRecord::Base.connection, leased_connection
+      end
+    end
+  end
+end