Merge pull request rails#51353 from Shopify/get-rid-lease-connection

Eliminate remaining uses of `lease_connection` inside Active Record

Author: byroot

activerecord/lib/active_record/associations/join_dependency/join_association.rb

@@ -37,39 +37,41 @@ def join_constraints(foreign_table, foreign_klass, join_type, alias_tracker)
             chain << [reflection, table]
           end
 
-          # The chain starts with the target table, but we want to end with it here (makes
-          # more sense in this context), so we reverse
-          chain.reverse_each do |reflection, table|
-            klass = reflection.klass
+          base_klass.with_connection do |connection|
+            # The chain starts with the target table, but we want to end with it here (makes
+            # more sense in this context), so we reverse
+            chain.reverse_each do |reflection, table|
+              klass = reflection.klass
 
-            scope = reflection.join_scope(table, foreign_table, foreign_klass)
+              scope = reflection.join_scope(table, foreign_table, foreign_klass)
 
-            unless scope.references_values.empty?
-              associations = scope.eager_load_values | scope.includes_values
+              unless scope.references_values.empty?
+                associations = scope.eager_load_values | scope.includes_values
 
-              unless associations.empty?
-                scope.joins! scope.construct_join_dependency(associations, Arel::Nodes::OuterJoin)
+                unless associations.empty?
+                  scope.joins! scope.construct_join_dependency(associations, Arel::Nodes::OuterJoin)
+                end
               end
-            end
 
-            arel = scope.arel(alias_tracker.aliases)
-            nodes = arel.constraints.first
+              arel = scope.arel(alias_tracker.aliases)
+              nodes = arel.constraints.first
 
-            if nodes.is_a?(Arel::Nodes::And)
-              others = nodes.children.extract! do |node|
-                !Arel.fetch_attribute(node) { |attr| attr.relation.name == table.name }
+              if nodes.is_a?(Arel::Nodes::And)
+                others = nodes.children.extract! do |node|
+                  !Arel.fetch_attribute(node) { |attr| attr.relation.name == table.name }
+                end
               end
-            end
 
-            joins << join_type.new(table, Arel::Nodes::On.new(nodes))
+              joins << join_type.new(table, Arel::Nodes::On.new(nodes))
 
-            if others && !others.empty?
-              joins.concat arel.join_sources
-              append_constraints(joins.last, others)
-            end
+              if others && !others.empty?
+                joins.concat arel.join_sources
+                append_constraints(connection, joins.last, others)
+              end
 
-            # The current table in this iteration becomes the foreign table in the next
-            foreign_table, foreign_klass = table, klass
+              # The current table in this iteration becomes the foreign table in the next
+              foreign_table, foreign_klass = table, klass
+            end
           end
 
           joins
@@ -88,10 +90,10 @@ def strict_loading?
         end
 
         private
-          def append_constraints(join, constraints)
+          def append_constraints(connection, join, constraints)
             if join.is_a?(Arel::Nodes::StringJoin)
               join_string = Arel::Nodes::And.new(constraints.unshift join.left)
-              join.left = Arel.sql(base_klass.lease_connection.visitor.compile(join_string))
+              join.left = Arel.sql(connection.visitor.compile(join_string))
             else
               right = join.right
               right.expr = Arel::Nodes::And.new(constraints.unshift right.expr)

activerecord/lib/active_record/integration.rb

@@ -178,7 +178,10 @@ def collection_cache_key(collection = all, timestamp_column = :updated_at) # :no
       def can_use_fast_cache_version?(timestamp)
         timestamp.is_a?(String) &&
           cache_timestamp_format == :usec &&
-          self.class.lease_connection.default_timezone == :utc &&
+          # FIXME: checking out a connection for this is wasteful
+          # we should store/cache this information in the schema cache
+          # or similar.
+          self.class.with_connection(&:default_timezone) == :utc &&
           !updated_at_came_from_user?
       end
 

activerecord/lib/active_record/table_metadata.rb

@@ -44,7 +44,7 @@ def associated_table(table_name)
         arel_table = arel_table.alias(table_name) if arel_table.name != table_name
         TableMetadata.new(association_klass, arel_table, reflection)
       else
-        type_caster = TypeCaster::Connection.new(klass, table_name)
+        type_caster = TypeCaster::Connection.new(klass)
         arel_table = Arel::Table.new(table_name, type_caster: type_caster)
         TableMetadata.new(nil, arel_table, reflection)
       end

activerecord/lib/active_record/transactions.rb

@@ -393,18 +393,19 @@ def rolledback!(force_restore_state: false, should_run_callbacks: true) # :nodoc
     # This method is available within the context of an ActiveRecord::Base
     # instance.
     def with_transaction_returning_status
-      status = nil
-      connection = self.class.lease_connection
-      ensure_finalize = !connection.transaction_open?
+      self.class.with_connection do |connection|
+        status = nil
+        ensure_finalize = !connection.transaction_open?
 
-      connection.transaction do
-        add_to_transaction(ensure_finalize || has_transactional_callbacks?)
-        remember_transaction_record_state
+        connection.transaction do
+          add_to_transaction(ensure_finalize || has_transactional_callbacks?)
+          remember_transaction_record_state
 
-        status = yield
-        raise ActiveRecord::Rollback unless status
+          status = yield
+          raise ActiveRecord::Rollback unless status
+        end
+        status
       end
-      status
     end
 
     def trigger_transactional_callbacks? # :nodoc:
@@ -496,7 +497,9 @@ def transaction_include_any_action?(actions)
       # Add the record to the current transaction so that the #after_rollback and #after_commit
       # callbacks can be called.
       def add_to_transaction(ensure_finalize = true)
-        self.class.lease_connection.add_transaction_record(self, ensure_finalize)
+        self.class.with_connection do |connection|
+          connection.add_transaction_record(self, ensure_finalize)
+        end
       end
 
       def has_transactional_callbacks?

activerecord/lib/active_record/type_caster/connection.rb

@@ -3,9 +3,8 @@
 module ActiveRecord
   module TypeCaster
     class Connection # :nodoc:
-      def initialize(klass, table_name)
+      def initialize(klass)
         @klass = klass
-        @table_name = table_name
       end
 
       def type_cast_for_database(attr_name, value)
@@ -14,20 +13,8 @@ def type_cast_for_database(attr_name, value)
       end
 
       def type_for_attribute(attr_name)
-        schema_cache = @klass.schema_cache
-
-        if schema_cache.data_source_exists?(table_name)
-          column = schema_cache.columns_hash(table_name)[attr_name.to_s]
-          if column
-            type = @klass.lease_connection.lookup_cast_type_from_column(column)
-          end
-        end
-
-        type || Type.default_value
+        @klass.type_for_attribute(attr_name) || Type.default_value
       end
-
-      private
-        attr_reader :table_name
     end
   end
 end

activerecord/lib/active_record/validations/uniqueness.rb

@@ -110,16 +110,20 @@ def resolve_attributes(record, attributes)
 
       def build_relation(klass, attribute, value)
         relation = klass.unscoped
-        comparison = relation.bind_attribute(attribute, value) do |attr, bind|
-          return relation.none! if bind.unboundable?
+        # TODO: Add case-sensitive / case-insensitive operators to Arel
+        # to no longer need to checkout a connection here.
+        comparison = klass.with_connection do |connection|
+          relation.bind_attribute(attribute, value) do |attr, bind|
+            return relation.none! if bind.unboundable?
 
-          if !options.key?(:case_sensitive) || bind.nil?
-            klass.lease_connection.default_uniqueness_comparison(attr, bind)
-          elsif options[:case_sensitive]
-            klass.lease_connection.case_sensitive_comparison(attr, bind)
-          else
-            # will use SQL LOWER function before comparison, unless it detects a case insensitive collation
-            klass.lease_connection.case_insensitive_comparison(attr, bind)
+            if !options.key?(:case_sensitive) || bind.nil?
+              connection.default_uniqueness_comparison(attr, bind)
+            elsif options[:case_sensitive]
+              connection.case_sensitive_comparison(attr, bind)
+            else
+              # will use SQL LOWER function before comparison, unless it detects a case insensitive collation
+              connection.case_insensitive_comparison(attr, bind)
+            end
           end
         end
 

activerecord/lib/arel/nodes/node.rb

@@ -147,8 +147,9 @@ def invert
       # Maybe we should just use `Table.engine`?  :'(
       def to_sql(engine = Table.engine)
         collector = Arel::Collectors::SQLString.new
-        collector = engine.lease_connection.visitor.accept self, collector
-        collector.value
+        engine.with_connection do |connection|
+          connection.visitor.accept(self, collector).value
+        end
       end
 
       def fetch_attribute

activerecord/lib/arel/tree_manager.rb

@@ -52,8 +52,9 @@ def to_dot
 
     def to_sql(engine = Table.engine)
       collector = Arel::Collectors::SQLString.new
-      collector = engine.lease_connection.visitor.accept @ast, collector
-      collector.value
+      engine.with_connection do |connection|
+        engine.lease_connection.visitor.accept(@ast, collector).value
+      end
     end
 
     def initialize_copy(other)

activerecord/test/cases/arel/support/fake_record.rb

@@ -129,6 +129,10 @@ def initialize
       @connection_pool = ConnectionPool.new
     end
 
+    def with_connection(...)
+      connection_pool.with_connection(...)
+    end
+
     def lease_connection
       connection_pool.lease_connection
     end