activerecord 7.1.6 → 7.2.3

data/lib/active_record/connection_adapters/abstract/database_statements.rb

@@ -14,7 +14,7 @@ module ActiveRecord
         sql
       end
 
-      def to_sql_and_binds(arel_or_sql_string, binds = [], preparable = nil) # :nodoc:
+      def to_sql_and_binds(arel_or_sql_string, binds = [], preparable = nil, allow_retry = false) # :nodoc:
         # Arel::TreeManager -> Arel::Node
         if arel_or_sql_string.respond_to?(:ast)
           arel_or_sql_string = arel_or_sql_string.ast
@@ -27,6 +27,7 @@ module ActiveRecord
           end
 
           collector = collector()
+          collector.retryable = true
 
           if prepared_statements
             collector.preparable = true
@@ -41,10 +42,11 @@ module ActiveRecord
           else
             sql = visitor.compile(arel_or_sql_string, collector)
           end
-          [sql.freeze, binds, preparable]
+          allow_retry = collector.retryable
+          [sql.freeze, binds, preparable, allow_retry]
         else
           arel_or_sql_string = arel_or_sql_string.dup.freeze unless arel_or_sql_string.frozen?
-          [arel_or_sql_string, binds, preparable]
+          [arel_or_sql_string, binds, preparable, allow_retry]
         end
       end
       private :to_sql_and_binds
@@ -64,11 +66,15 @@ module ActiveRecord
       end
 
       # Returns an ActiveRecord::Result instance.
-      def select_all(arel, name = nil, binds = [], preparable: nil, async: false)
+      def select_all(arel, name = nil, binds = [], preparable: nil, async: false, allow_retry: false)
         arel = arel_from_relation(arel)
-        sql, binds, preparable = to_sql_and_binds(arel, binds, preparable)
+        sql, binds, preparable, allow_retry = to_sql_and_binds(arel, binds, preparable, allow_retry)
 
-        select(sql, name, binds, prepare: prepared_statements && preparable, async: async && FutureResult::SelectAll)
+        select(sql, name, binds,
+          prepare: prepared_statements && preparable,
+          async: async && FutureResult::SelectAll,
+          allow_retry: allow_retry
+        )
       rescue ::RangeError
         ActiveRecord::Result.empty(async: async)
       end
@@ -214,7 +220,7 @@ module ActiveRecord
       end
 
       def truncate_tables(*table_names) # :nodoc:
-        table_names -= [schema_migration.table_name, internal_metadata.table_name]
+        table_names -= [pool.schema_migration.table_name, pool.internal_metadata.table_name]
 
         return if table_names.empty?
 
@@ -229,6 +235,17 @@ module ActiveRecord
       # Runs the given block in a database transaction, and returns the result
       # of the block.
       #
+      # == Transaction callbacks
+      #
+      # #transaction yields an ActiveRecord::Transaction object on which it is
+      # possible to register callback:
+      #
+      #   ActiveRecord::Base.transaction do |transaction|
+      #     transaction.before_commit { puts "before commit!" }
+      #     transaction.after_commit { puts "after commit!" }
+      #     transaction.after_rollback { puts "after rollback!" }
+      #   end
+      #
       # == Nested transactions support
       #
       # #transaction calls can be nested. By default, this makes all database
@@ -296,9 +313,9 @@ module ActiveRecord
       # #transaction will raise exceptions when it tries to release the
       # already-automatically-released savepoints:
       #
-      #   Model.connection.transaction do  # BEGIN
-      #     Model.connection.transaction(requires_new: true) do  # CREATE SAVEPOINT active_record_1
-      #       Model.connection.create_table(...)
+      #   Model.lease_connection.transaction do  # BEGIN
+      #     Model.lease_connection.transaction(requires_new: true) do  # CREATE SAVEPOINT active_record_1
+      #       Model.lease_connection.create_table(...)
       #       # active_record_1 now automatically released
       #     end  # RELEASE SAVEPOINT active_record_1  <--- BOOM! database error!
       #   end
@@ -339,7 +356,7 @@ module ActiveRecord
           if isolation
             raise ActiveRecord::TransactionIsolationError, "cannot set isolation when joining a transaction"
           end
-          yield
+          yield current_transaction.user_transaction
         else
           transaction_manager.within_new_transaction(isolation: isolation, joinable: joinable, &block)
         end
@@ -457,8 +474,8 @@ module ActiveRecord
         statements = table_deletes + fixture_inserts
 
         with_multi_statements do
-          disable_referential_integrity do
-            transaction(requires_new: true) do
+          transaction(requires_new: true) do
+            disable_referential_integrity do
               execute_batch(statements, "Fixtures Load")
             end
           end
@@ -495,7 +512,7 @@ module ActiveRecord
       end
 
       # This is a safe default, even if not high precision on all databases
-      HIGH_PRECISION_CURRENT_TIMESTAMP = Arel.sql("CURRENT_TIMESTAMP").freeze # :nodoc:
+      HIGH_PRECISION_CURRENT_TIMESTAMP = Arel.sql("CURRENT_TIMESTAMP", retryable: true).freeze # :nodoc:
       private_constant :HIGH_PRECISION_CURRENT_TIMESTAMP
 
       # Returns an Arel SQL literal for the CURRENT_TIMESTAMP for usage with
@@ -507,7 +524,7 @@ module ActiveRecord
         HIGH_PRECISION_CURRENT_TIMESTAMP
       end
 
-      def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false) # :nodoc:
+      def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false, allow_retry: false) # :nodoc:
         raise NotImplementedError
       end
 
@@ -606,7 +623,7 @@ module ActiveRecord
         end
 
         # Returns an ActiveRecord::Result instance.
-        def select(sql, name = nil, binds = [], prepare: false, async: false)
+        def select(sql, name = nil, binds = [], prepare: false, async: false, allow_retry: false)
           if async && async_enabled?
             if current_transaction.joinable?
               raise AsynchronousQueryInsideTransactionError, "Asynchronous queries are not allowed inside transactions"
@@ -627,7 +644,7 @@ module ActiveRecord
             return future_result
           end
 
-          result = internal_exec_query(sql, name, binds, prepare: prepare)
+          result = internal_exec_query(sql, name, binds, prepare: prepare, allow_retry: allow_retry)
           if async
             FutureResult.wrap(result)
           else

data/lib/active_record/connection_adapters/abstract/query_cache.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "concurrent/map"
+require "concurrent/atomic/atomic_fixnum"
 
 module ActiveRecord
   module ConnectionAdapters # :nodoc:
@@ -13,15 +14,16 @@ module ActiveRecord
             :truncate_tables, :rollback_to_savepoint, :rollback_db_transaction, :restart_db_transaction,
             :exec_insert_all
 
-          base.set_callback :checkout, :after, :configure_query_cache!
-          base.set_callback :checkin, :after, :disable_query_cache!
+          base.set_callback :checkin, :after, :unset_query_cache!
         end
 
         def dirties_query_cache(base, *method_names)
           method_names.each do |method_name|
             base.class_eval <<-end_code, __FILE__, __LINE__ + 1
               def #{method_name}(...)
-                ActiveRecord::Base.clear_query_caches_for_current_thread
+                if pool.dirties_query_cache
+                  ActiveRecord::Base.clear_query_caches_for_current_thread
+                end
                 super
               end
             end_code
@@ -29,60 +31,207 @@ module ActiveRecord
         end
       end
 
-      module ConnectionPoolConfiguration
-        def initialize(*)
+      class Store # :nodoc:
+        attr_accessor :enabled, :dirties
+        alias_method :enabled?, :enabled
+        alias_method :dirties?, :dirties
+
+        def initialize(version, max_size)
+          @version = version
+          @current_version = version.value
+          @map = {}
+          @max_size = max_size
+          @enabled = false
+          @dirties = true
+        end
+
+        def size
+          check_version
+          @map.size
+        end
+
+        def empty?
+          check_version
+          @map.empty?
+        end
+
+        def [](key)
+          check_version
+          return unless @enabled
+
+          if entry = @map.delete(key)
+            @map[key] = entry
+          end
+        end
+
+        def compute_if_absent(key)
+          check_version
+
+          return yield unless @enabled
+
+          if entry = @map.delete(key)
+            return @map[key] = entry
+          end
+
+          if @max_size && @map.size >= @max_size
+            @map.shift # evict the oldest entry
+          end
+
+          @map[key] ||= yield
+        end
+
+        def clear
+          @map.clear
+          self
+        end
+
+        private
+          def check_version
+            if @current_version != @version.value
+              @map.clear
+              @current_version = @version.value
+            end
+          end
+      end
+
+      class QueryCacheRegistry # :nodoc:
+        def initialize
+          @mutex = Mutex.new
+          @map = ConnectionPool::WeakThreadKeyMap.new
+        end
+
+        def compute_if_absent(context)
+          @map[context] || @mutex.synchronize do
+            @map[context] ||= yield
+          end
+        end
+
+        def clear
+          @map.synchronize do
+            @map.clear
+          end
+        end
+      end
+
+      module ConnectionPoolConfiguration # :nodoc:
+        def initialize(...)
+          super
+          @query_cache_version = Concurrent::AtomicFixnum.new
+          @thread_query_caches = QueryCacheRegistry.new
+          @query_cache_max_size = \
+            case query_cache = db_config&.query_cache
+            when 0, false
+              nil
+            when Integer
+              query_cache
+            when nil
+              DEFAULT_SIZE
+            end
+        end
+
+        def checkout_and_verify(connection)
           super
-          @query_cache_enabled = Concurrent::Map.new { false }
+          connection.query_cache ||= query_cache
+          connection
+        end
+
+        # Disable the query cache within the block.
+        def disable_query_cache(dirties: true)
+          cache = query_cache
+          old_enabled, cache.enabled, old_dirties, cache.dirties = cache.enabled, false, cache.dirties, dirties
+          begin
+            yield
+          ensure
+            cache.enabled, cache.dirties = old_enabled, old_dirties
+          end
+        end
+
+        def enable_query_cache
+          cache = query_cache
+          old_enabled, cache.enabled, old_dirties, cache.dirties = cache.enabled, true, cache.dirties, true
+          begin
+            yield
+          ensure
+            cache.enabled, cache.dirties = old_enabled, old_dirties
+          end
         end
 
         def enable_query_cache!
-          @query_cache_enabled[connection_cache_key(current_thread)] = true
-          connection.enable_query_cache! if active_connection?
+          query_cache.enabled = true
+          query_cache.dirties = true
         end
 
         def disable_query_cache!
-          @query_cache_enabled.delete connection_cache_key(current_thread)
-          connection.disable_query_cache! if active_connection?
+          query_cache.enabled = false
+          query_cache.dirties = true
         end
 
         def query_cache_enabled
-          @query_cache_enabled[connection_cache_key(current_thread)]
+          query_cache.enabled
+        end
+
+        def dirties_query_cache
+          query_cache.dirties
         end
-      end
 
-      attr_reader :query_cache, :query_cache_enabled
+        def clear_query_cache
+          if @pinned_connection
+            # With transactional fixtures, and especially systems test
+            # another thread may use the same connection, but with a different
+            # query cache. So we must clear them all.
+            @query_cache_version.increment
+          end
+          query_cache.clear
+        end
+
+        def query_cache
+          @thread_query_caches.compute_if_absent(ActiveSupport::IsolatedExecutionState.context) do
+            Store.new(@query_cache_version, @query_cache_max_size)
+          end
+        end
+      end
 
       def initialize(*)
         super
-        @query_cache         = {}
-        @query_cache_enabled = false
-        @query_cache_max_size = nil
+        @query_cache = nil
+      end
+
+      attr_writer :query_cache
+
+      def query_cache
+        if @pinned && @owner != ActiveSupport::IsolatedExecutionState.context
+          # With transactional tests, if the connection is pinned, any thread
+          # other than the one that pinned the connection need to go through the
+          # query cache pool, so each thread get a different cache.
+          pool.query_cache
+        else
+          @query_cache
+        end
+      end
+
+      def query_cache_enabled
+        query_cache&.enabled?
       end
 
       # Enable the query cache within the block.
-      def cache
-        old, @query_cache_enabled = @query_cache_enabled, true
-        yield
-      ensure
-        @query_cache_enabled = old
-        clear_query_cache unless @query_cache_enabled
+      def cache(&block)
+        pool.enable_query_cache(&block)
       end
 
       def enable_query_cache!
-        @query_cache_enabled = true
+        pool.enable_query_cache!
       end
 
-      def disable_query_cache!
-        @query_cache_enabled = false
-        clear_query_cache
+      # Disable the query cache within the block.
+      #
+      # Set <tt>dirties: false</tt> to prevent query caches on all connections from being cleared by write operations.
+      # (By default, write operations dirty all connections' query caches in case they are replicas whose cache would now be outdated.)
+      def uncached(dirties: true, &block)
+        pool.disable_query_cache(dirties: dirties, &block)
       end
 
-      # Disable the query cache within the block.
-      def uncached
-        old, @query_cache_enabled = @query_cache_enabled, false
-        yield
-      ensure
-        @query_cache_enabled = old
+      def disable_query_cache!
+        pool.disable_query_cache!
       end
 
       # Clears the query cache.
@@ -92,24 +241,22 @@ module ActiveRecord
       # the same SQL query and repeatedly return the same result each time, silently
       # undermining the randomness you were expecting.
       def clear_query_cache
-        @lock.synchronize do
-          @query_cache.clear
-        end
+        pool.clear_query_cache
       end
 
-      def select_all(arel, name = nil, binds = [], preparable: nil, async: false) # :nodoc:
+      def select_all(arel, name = nil, binds = [], preparable: nil, async: false, allow_retry: false) # :nodoc:
         arel = arel_from_relation(arel)
 
         # If arel is locked this is a SELECT ... FOR UPDATE or somesuch.
         # Such queries should not be cached.
-        if @query_cache_enabled && !(arel.respond_to?(:locked) && arel.locked)
-          sql, binds, preparable = to_sql_and_binds(arel, binds, preparable)
+        if query_cache_enabled && !(arel.respond_to?(:locked) && arel.locked)
+          sql, binds, preparable, allow_retry = to_sql_and_binds(arel, binds, preparable)
 
           if async
-            result = lookup_sql_cache(sql, name, binds) || super(sql, name, binds, preparable: preparable, async: async)
+            result = lookup_sql_cache(sql, name, binds) || super(sql, name, binds, preparable: preparable, async: async, allow_retry: allow_retry)
             FutureResult.wrap(result)
           else
-            cache_sql(sql, name, binds) { super(sql, name, binds, preparable: preparable, async: async) }
+            cache_sql(sql, name, binds) { super(sql, name, binds, preparable: preparable, async: async, allow_retry: allow_retry) }
           end
         else
           super
@@ -117,42 +264,37 @@ module ActiveRecord
       end
 
       private
+        def unset_query_cache!
+          @query_cache = nil
+        end
+
         def lookup_sql_cache(sql, name, binds)
           key = binds.empty? ? sql : [sql, binds]
-          hit = false
-          result = nil
 
+          result = nil
           @lock.synchronize do
-            if (result = @query_cache.delete(key))
-              hit = true
-              @query_cache[key] = result
-            end
+            result = query_cache[key]
           end
 
-          if hit
+          if result
             ActiveSupport::Notifications.instrument(
               "sql.active_record",
               cache_notification_info(sql, name, binds)
             )
-
-            result
           end
+
+          result
         end
 
         def cache_sql(sql, name, binds)
           key = binds.empty? ? sql : [sql, binds]
           result = nil
-          hit = false
+          hit = true
 
           @lock.synchronize do
-            if (result = @query_cache.delete(key))
-              hit = true
-              @query_cache[key] = result
-            else
-              result = @query_cache[key] = yield
-              if @query_cache_max_size && @query_cache.size > @query_cache_max_size
-                @query_cache.shift
-              end
+            result = query_cache.compute_if_absent(key) do
+              hit = false
+              yield
             end
           end
 
@@ -175,26 +317,10 @@ module ActiveRecord
             type_casted_binds: -> { type_casted_binds(binds) },
             name: name,
             connection: self,
+            transaction: current_transaction.user_transaction.presence,
             cached: true
           }
         end
-
-        def configure_query_cache!
-          case query_cache = pool.db_config.query_cache
-          when 0, false
-            return
-          when Integer
-            @query_cache_max_size = query_cache
-          when nil
-            @query_cache_max_size = DEFAULT_SIZE
-          else
-            @query_cache_max_size = nil # no limit
-          end
-
-          if pool.query_cache_enabled
-            enable_query_cache!
-          end
-        end
     end
   end
 end