activerecord 7.2.3 → 8.1.3

data/lib/active_record/encryption/config.rb

@@ -8,7 +8,8 @@ module ActiveRecord
     class Config
       attr_accessor :primary_key, :deterministic_key, :store_key_references, :key_derivation_salt, :hash_digest_class,
                     :support_unencrypted_data, :encrypt_fixtures, :validate_column_size, :add_to_filter_parameters,
-                    :excluded_from_filter_parameters, :extend_queries, :previous_schemes, :forced_encoding_for_deterministic_encryption
+                    :excluded_from_filter_parameters, :extend_queries, :previous_schemes, :forced_encoding_for_deterministic_encryption,
+                    :compressor
 
       def initialize
         set_defaults
@@ -55,6 +56,7 @@ module ActiveRecord
           self.previous_schemes = []
           self.forced_encoding_for_deterministic_encryption = Encoding::UTF_8
           self.hash_digest_class = OpenSSL::Digest::SHA1
+          self.compressor = Zlib
 
           # TODO: Setting to false for now as the implementation is a bit experimental
           self.extend_queries = false

data/lib/active_record/encryption/encryptable_record.rb

@@ -30,10 +30,10 @@ module ActiveRecord
         #   will use the oldest encryption scheme to encrypt new data by default. You can change this by setting
         #   <tt>deterministic: { fixed: false }</tt>. That will make it use the newest encryption scheme for encrypting new
         #   data.
-        # * <tt>:support_unencrypted_data</tt> - If `config.active_record.encryption.support_unencrypted_data` is +true+,
-        #   you can set this to +false+ to opt out of unencrypted data support for this attribute. This is useful for
-        #   scenarios where you encrypt one column, and want to disable support for unencrypted data without having to tweak
-        #   the global setting.
+        # * <tt>:support_unencrypted_data</tt> - When true, unencrypted data can be read normally. When false, it will raise errors.
+        #   Falls back to +config.active_record.encryption.support_unencrypted_data+ if no value is provided.
+        #   This is useful for scenarios where you encrypt one column, and want to disable support for unencrypted data
+        #   without having to tweak the global setting.
         # * <tt>:downcase</tt> - When true, it converts the encrypted content to downcase automatically. This allows to
         #   effectively ignore case when querying data. Notice that the case is lost. Use +:ignore_case+ if you are interested
         #   in preserving it.
@@ -46,11 +46,11 @@ module ActiveRecord
         # * <tt>:previous</tt> - List of previous encryption schemes. When provided, they will be used in order when trying to read
         #   the attribute. Each entry of the list can contain the properties supported by #encrypts. Also, when deterministic
         #   encryption is used, they will be used to generate additional ciphertexts to check in the queries.
-        def encrypts(*names, key_provider: nil, key: nil, deterministic: false, support_unencrypted_data: nil, downcase: false, ignore_case: false, previous: [], **context_properties)
+        def encrypts(*names, key_provider: nil, key: nil, deterministic: false, support_unencrypted_data: nil, downcase: false, ignore_case: false, previous: [], compress: true, compressor: nil, **context_properties)
           self.encrypted_attributes ||= Set.new # not using :default because the instance would be shared across classes
 
           names.each do |name|
-            encrypt_attribute name, key_provider: key_provider, key: key, deterministic: deterministic, support_unencrypted_data: support_unencrypted_data, downcase: downcase, ignore_case: ignore_case, previous: previous, **context_properties
+            encrypt_attribute name, key_provider: key_provider, key: key, deterministic: deterministic, support_unencrypted_data: support_unencrypted_data, downcase: downcase, ignore_case: ignore_case, previous: previous, compress: compress, compressor: compressor, **context_properties
           end
         end
 
@@ -81,12 +81,12 @@ module ActiveRecord
             end
           end
 
-          def encrypt_attribute(name, key_provider: nil, key: nil, deterministic: false, support_unencrypted_data: nil, downcase: false, ignore_case: false, previous: [], **context_properties)
+          def encrypt_attribute(name, key_provider: nil, key: nil, deterministic: false, support_unencrypted_data: nil, downcase: false, ignore_case: false, previous: [], compress: true, compressor: nil, **context_properties)
             encrypted_attributes << name.to_sym
 
             decorate_attributes([name]) do |name, cast_type|
               scheme = scheme_for key_provider: key_provider, key: key, deterministic: deterministic, support_unencrypted_data: support_unencrypted_data, \
-                downcase: downcase, ignore_case: ignore_case, previous: previous, **context_properties
+                downcase: downcase, ignore_case: ignore_case, previous: previous, compress: compress, compressor: compressor, **context_properties
 
               ActiveRecord::Encryption::EncryptedAttributeType.new(scheme: scheme, cast_type: cast_type, default: columns_hash[name.to_s]&.default)
             end

data/lib/active_record/encryption/encrypted_attribute_type.rb

@@ -59,7 +59,7 @@ module ActiveRecord
       end
 
       def support_unencrypted_data?
-        ActiveRecord::Encryption.config.support_unencrypted_data && scheme.support_unencrypted_data? && !previous_type?
+        scheme.support_unencrypted_data? && !previous_type?
       end
 
       private
@@ -100,7 +100,7 @@ module ActiveRecord
         end
 
         def decrypt(value)
-          text_to_database_type decrypt_as_text(value)
+          text_to_database_type decrypt_as_text(database_type_to_text(value))
         end
 
         def try_to_deserialize_with_previous_encrypted_types(value)
@@ -170,6 +170,15 @@ module ActiveRecord
             value
           end
         end
+
+        def database_type_to_text(value)
+          if value && cast_type.binary?
+            binary_cast_type = cast_type.serialized? ? cast_type.subtype : cast_type
+            binary_cast_type.deserialize(value)
+          else
+            value
+          end
+        end
     end
   end
 end

data/lib/active_record/encryption/encryptor.rb

@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
 require "openssl"
-require "zlib"
 require "active_support/core_ext/numeric"
 
 module ActiveRecord
@@ -12,13 +11,22 @@ module ActiveRecord
     # It interacts with a KeyProvider for getting the keys, and delegate to
     # ActiveRecord::Encryption::Cipher the actual encryption algorithm.
     class Encryptor
+      # The compressor to use for compressing the payload.
+      attr_reader :compressor
+
       # ==== Options
       #
       # [+:compress+]
       #   Boolean indicating whether records should be compressed before
       #   encryption. Defaults to +true+.
-      def initialize(compress: true)
+      #
+      # [+:compressor+]
+      #   The compressor to use. It must respond to +deflate+ and +inflate+.
+      #   If not provided, will default to +ActiveRecord::Encryption.config.compressor+,
+      #   which itself defaults to +Zlib+.
+      def initialize(compress: true, compressor: nil)
         @compress = compress
+        @compressor = compressor || ActiveRecord::Encryption.config.compressor
       end
 
       # Encrypts +clean_text+ and returns the encrypted result.
@@ -79,9 +87,25 @@ module ActiveRecord
         serializer.binary?
       end
 
+      def compress? # :nodoc:
+        @compress
+      end
+
       private
         DECRYPT_ERRORS = [OpenSSL::Cipher::CipherError, Errors::EncryptedContentIntegrity, Errors::Decryption]
         ENCODING_ERRORS = [EncodingError, Errors::Encoding]
+
+        # This threshold cannot be changed.
+        #
+        # Users can search for attributes encrypted with `deterministic: true`.
+        # That is possible because we are able to generate the message for the
+        # given clear text deterministically, and with that perform a regular
+        # string lookup in SQL.
+        #
+        # Problem is, messages may have a "c" header that is present or not
+        # depending on whether compression was applied on encryption. If this
+        # threshold was modified, the message generated for lookup could vary
+        # for the same clear text, and searches on exisiting data could fail.
         THRESHOLD_TO_JUSTIFY_COMPRESSION = 140.bytes
 
         def default_key_provider
@@ -131,12 +155,8 @@ module ActiveRecord
           end
         end
 
-        def compress?
-          @compress
-        end
-
         def compress(data)
-          Zlib::Deflate.deflate(data).tap do |compressed_data|
+          @compressor.deflate(data).tap do |compressed_data|
             compressed_data.force_encoding(data.encoding)
           end
         end
@@ -150,7 +170,7 @@ module ActiveRecord
         end
 
         def uncompress(data)
-          Zlib::Inflate.inflate(data).tap do |uncompressed_data|
+          @compressor.inflate(data).tap do |uncompressed_data|
             uncompressed_data.force_encoding(data.encoding)
           end
         end

data/lib/active_record/encryption/extended_deterministic_queries.rb

@@ -41,6 +41,8 @@ module ActiveRecord
       module EncryptedQuery # :nodoc:
         class << self
           def process_arguments(owner, args, check_for_additional_values)
+            owner = owner.model if owner.is_a?(Relation)
+
             return args if owner.deterministic_encrypted_attributes&.empty?
 
             if args.is_a?(Array) && (options = args.first).is_a?(Hash)
@@ -102,12 +104,12 @@ module ActiveRecord
         end
 
         def scope_for_create
-          return super unless klass.deterministic_encrypted_attributes&.any?
+          return super unless model.deterministic_encrypted_attributes&.any?
 
           scope_attributes = super
           wheres = where_values_hash
 
-          klass.deterministic_encrypted_attributes.each do |attribute_name|
+          model.deterministic_encrypted_attributes.each do |attribute_name|
             attribute_name = attribute_name.to_s
             values = wheres[attribute_name]
             if values.is_a?(Array) && values[1..].all?(AdditionalValue)

data/lib/active_record/encryption/scheme.rb

@@ -11,7 +11,7 @@ module ActiveRecord
       attr_accessor :previous_schemes
 
       def initialize(key_provider: nil, key: nil, deterministic: nil, support_unencrypted_data: nil, downcase: nil, ignore_case: nil,
-                     previous_schemes: nil, **context_properties)
+                     previous_schemes: nil, compress: true, compressor: nil, **context_properties)
         # Initializing all attributes to +nil+ as we want to allow a "not set" semantics so that we
         # can merge schemes without overriding values with defaults. See +#merge+
 
@@ -24,8 +24,13 @@ module ActiveRecord
         @previous_schemes_param = previous_schemes
         @previous_schemes = Array.wrap(previous_schemes)
         @context_properties = context_properties
+        @compress = compress
+        @compressor = compressor
 
         validate_config!
+
+        @context_properties[:encryptor] = Encryptor.new(compress: @compress) unless @compress
+        @context_properties[:encryptor] = Encryptor.new(compressor: compressor) if compressor
       end
 
       def ignore_case?
@@ -54,7 +59,7 @@ module ActiveRecord
       end
 
       def merge(other_scheme)
-        self.class.new(**to_h.merge(other_scheme.to_h))
+        self.class.new(**to_h, **other_scheme.to_h)
       end
 
       def to_h
@@ -78,6 +83,8 @@ module ActiveRecord
         def validate_config!
           raise Errors::Configuration, "ignore_case: can only be used with deterministic encryption" if @ignore_case && !@deterministic
           raise Errors::Configuration, "key_provider: and key: can't be used simultaneously" if @key_provider_param && @key
+          raise Errors::Configuration, "compressor: can't be used with compress: false" if !@compress && @compressor
+          raise Errors::Configuration, "compressor: can't be used with encryptor" if @compressor && @context_properties[:encryptor]
         end
 
         def key_provider_from_key

data/lib/active_record/enum.rb

@@ -214,34 +214,16 @@ module ActiveRecord
         attr_reader :name, :mapping
     end
 
-    def enum(name = nil, values = nil, **options)
-      if name
-        values, options = options, {} unless values
-        return _enum(name, values, **options)
-      end
-
-      definitions = options.slice!(:_prefix, :_suffix, :_scopes, :_default, :_instance_methods)
-      options.transform_keys! { |key| :"#{key[1..-1]}" }
-
-      definitions.each { |name, values| _enum(name, values, **options) }
-
-      ActiveRecord.deprecator.warn(<<~MSG)
-        Defining enums with keyword arguments is deprecated and will be removed
-        in Rails 8.0. Positional arguments should be used instead:
-
-        #{definitions.map { |name, values| "enum :#{name}, #{values}" }.join("\n")}
-      MSG
+    def enum(name, values = nil, **options)
+      values, options = options, {} unless values
+      _enum(name, values, **options)
     end
 
     private
-      def inherited(base)
-        base.defined_enums = defined_enums.deep_dup
-        super
-      end
-
       def _enum(name, values, prefix: nil, suffix: nil, scopes: true, instance_methods: true, validate: false, **options)
-        assert_valid_enum_definition_values(values)
+        values = assert_valid_enum_definition_values(values)
         assert_valid_enum_options(options)
+
         # statuses = { }
         enum_values = ActiveSupport::HashWithIndifferentAccess.new
         name = name.to_s
@@ -305,6 +287,11 @@ module ActiveRecord
         enum_values.freeze
       end
 
+      def inherited(base)
+        base.defined_enums = defined_enums.deep_dup
+        super
+      end
+
       class EnumMethods < Module # :nodoc:
         def initialize(klass)
           @klass = klass
@@ -355,6 +342,20 @@ module ActiveRecord
           if values.keys.any?(&:blank?)
             raise ArgumentError, "Enum values #{values} must not contain a blank name."
           end
+
+          values = values.transform_values do |value|
+            value.is_a?(Symbol) ? value.name : value
+          end
+
+          values.each_value do |value|
+            case value
+            when String, Integer, Float, true, false, nil
+              # noop
+            else
+              raise ArgumentError, "Enum values #{values} must be only booleans, integers, floats, symbols or strings, got: #{value.class}"
+            end
+          end
+
         when Array
           if values.empty?
             raise ArgumentError, "Enum values #{values} must not be empty."
@@ -370,6 +371,8 @@ module ActiveRecord
         else
           raise ArgumentError, "Enum values #{values} must be either a non-empty hash or an array."
         end
+
+        values
       end
 
       def assert_valid_enum_options(options)
@@ -381,25 +384,25 @@ module ActiveRecord
 
       ENUM_CONFLICT_MESSAGE = \
         "You tried to define an enum named \"%{enum}\" on the model \"%{klass}\", but " \
-        "this will generate a %{type} method \"%{method}\", which is already defined " \
+        "this will generate %{type} method \"%{method}\", which is already defined " \
         "by %{source}."
       private_constant :ENUM_CONFLICT_MESSAGE
 
       def detect_enum_conflict!(enum_name, method_name, klass_method = false)
         if klass_method && dangerous_class_method?(method_name)
-          raise_conflict_error(enum_name, method_name, type: "class")
+          raise_conflict_error(enum_name, method_name, "a class")
         elsif klass_method && method_defined_within?(method_name, Relation)
-          raise_conflict_error(enum_name, method_name, type: "class", source: Relation.name)
+          raise_conflict_error(enum_name, method_name, "a class", source: Relation.name)
         elsif klass_method && method_name.to_sym == :id
-          raise_conflict_error(enum_name, method_name)
+          raise_conflict_error(enum_name, method_name, "an instance")
         elsif !klass_method && dangerous_attribute_method?(method_name)
-          raise_conflict_error(enum_name, method_name)
+          raise_conflict_error(enum_name, method_name, "an instance")
         elsif !klass_method && method_defined_within?(method_name, _enum_methods_module, Module)
-          raise_conflict_error(enum_name, method_name, source: "another enum")
+          raise_conflict_error(enum_name, method_name, "an instance", source: "another enum")
         end
       end
 
-      def raise_conflict_error(enum_name, method_name, type: "instance", source: "Active Record")
+      def raise_conflict_error(enum_name, method_name, type, source: "Active Record")
         raise ArgumentError, ENUM_CONFLICT_MESSAGE % {
           enum: enum_name,
           klass: name,

data/lib/active_record/errors.rb

@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-require "active_support/deprecation"
 
 module ActiveRecord
   include ActiveSupport::Deprecation::DeprecatedConstantAccessor
@@ -84,6 +83,19 @@ module ActiveRecord
   class ConnectionTimeoutError < ConnectionNotEstablished
   end
 
+  # Raised when a database connection pool is requested but
+  # has not been defined.
+  class ConnectionNotDefined < ConnectionNotEstablished
+    def initialize(message = nil, connection_name: nil, role: nil, shard: nil)
+      super(message)
+      @connection_name = connection_name
+      @role = role
+      @shard = shard
+    end
+
+    attr_reader :connection_name, :role, :shard
+  end
+
   # Raised when connection to the database could not been established because it was not
   # able to connect to the host or when the authorization failed.
   class DatabaseConnectionError < ConnectionNotEstablished
@@ -280,6 +292,14 @@ module ActiveRecord
   class NotNullViolation < StatementInvalid
   end
 
+  # Raised when a record cannot be inserted or updated because it would violate a check constraint.
+  class CheckViolation < StatementInvalid
+  end
+
+  # Raised when a record cannot be inserted or updated because it would violate an exclusion constraint.
+  class ExclusionViolation < StatementInvalid
+  end
+
   # Raised when a record cannot be inserted or updated because a value too long for a column type.
   class ValueTooLong < StatementInvalid
   end
@@ -326,15 +346,15 @@ module ActiveRecord
     class << self
       def db_error(db_name)
         NoDatabaseError.new(<<~MSG)
-          We could not find your database: #{db_name}. Available database configurations can be found in config/database.yml.
+          Database not found: #{db_name}. Available database configurations can be found in config/database.yml.
 
           To resolve this error:
 
-          - Did you not create the database, or did you delete it? To create the database, run:
+          - Create the database by running:
 
               bin/rails db:create
 
-          - Has the database name changed? Verify that config/database.yml contains the correct database name.
+          - Verify that config/database.yml contains the correct database name.
         MSG
       end
     end
@@ -477,6 +497,7 @@ module ActiveRecord
   #   end
   #
   #   relation = Task.all
+  #   relation.load
   #   relation.loaded? # => true
   #
   #   # Methods which try to mutate a loaded relation fail.
@@ -484,11 +505,6 @@ module ActiveRecord
   #   relation.limit!(5)              # => ActiveRecord::UnmodifiableRelation
   class UnmodifiableRelation < ActiveRecordError
   end
-  deprecate_constant(
-    :ImmutableRelation,
-    "ActiveRecord::UnmodifiableRelation",
-    deprecator: ActiveRecord.deprecator
-  )
 
   # TransactionIsolationError will be raised under the following conditions:
   #
@@ -544,6 +560,11 @@ module ActiveRecord
   class Deadlocked < TransactionRollbackError
   end
 
+  # MissingRequiredOrderError is raised when a relation requires ordering but
+  # lacks any +order+ values in scope or any model order columns to use.
+  class MissingRequiredOrderError < ActiveRecordError
+  end
+
   # IrreversibleOrderError is raised when a relation's order is too complex for
   # +reverse_order+ to automatically reverse.
   class IrreversibleOrderError < ActiveRecordError
@@ -601,6 +622,9 @@ module ActiveRecord
   # the database version cannot be determined.
   class DatabaseVersionError < ActiveRecordError
   end
+
+  class DeprecatedAssociationError < ActiveRecordError
+  end
 end
 
 require "active_record/associations/errors"

data/lib/active_record/explain.rb

@@ -7,7 +7,7 @@ module ActiveRecord
     # Executes the block with the collect flag enabled. Queries are collected
     # asynchronously by the subscriber and returned.
     def collecting_queries_for_explain # :nodoc:
-      ExplainRegistry.collect = true
+      ExplainRegistry.start
       yield
       ExplainRegistry.queries
     ensure

data/lib/active_record/explain_registry.rb

@@ -1,6 +1,5 @@
 # frozen_string_literal: true
 
-require "active_support/core_ext/module/delegation"
 
 module ActiveRecord
   # This is a thread locals registry for EXPLAIN. For example
@@ -9,8 +8,53 @@ module ActiveRecord
   #
   # returns the collected queries local to the current thread.
   class ExplainRegistry # :nodoc:
+    class Subscriber
+      MUTEX = Mutex.new
+      @subscribed = false
+
+      class << self
+        def ensure_subscribed
+          return if @subscribed
+          MUTEX.synchronize do
+            return if @subscribed
+
+            ActiveSupport::Notifications.subscribe("sql.active_record", new)
+            @subscribed = true
+          end
+        end
+      end
+
+      def start(name, id, payload)
+        # unused
+      end
+
+      def finish(name, id, payload)
+        if ExplainRegistry.collect? && !ignore_payload?(payload)
+          ExplainRegistry.queries << payload.values_at(:sql, :binds)
+        end
+      end
+
+      def silenced?(_name)
+        !ExplainRegistry.collect?
+      end
+
+      # SCHEMA queries cannot be EXPLAINed, also we do not want to run EXPLAIN on
+      # our own EXPLAINs no matter how loopingly beautiful that would be.
+      #
+      # On the other hand, we want to monitor the performance of our real database
+      # queries, not the performance of the access to the query cache.
+      IGNORED_PAYLOADS = %w(SCHEMA EXPLAIN)
+      EXPLAINED_SQLS = /\A\s*(\/\*.*\*\/)?\s*(with|select|update|delete|insert)\b/i
+      def ignore_payload?(payload)
+        payload[:exception] ||
+          payload[:cached] ||
+          IGNORED_PAYLOADS.include?(payload[:name]) ||
+          !payload[:sql].match?(EXPLAINED_SQLS)
+      end
+    end
+
     class << self
-      delegate :reset, :collect, :collect=, :collect?, :queries, to: :instance
+      delegate :start, :reset, :collect, :collect=, :collect?, :queries, to: :instance
 
       private
         def instance
@@ -25,6 +69,11 @@ module ActiveRecord
       reset
     end
 
+    def start
+      Subscriber.ensure_subscribed
+      @collect = true
+    end
+
     def collect?
       @collect
     end

data/lib/active_record/filter_attribute_handler.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  class FilterAttributeHandler # :nodoc:
+    class << self
+      def on_sensitive_attribute_declared(&block)
+        @sensitive_attribute_declaration_listeners ||= Concurrent::Array.new
+        @sensitive_attribute_declaration_listeners << block
+      end
+
+      def sensitive_attribute_was_declared(klass, list)
+        @sensitive_attribute_declaration_listeners&.each do |block|
+          block.call(klass, list)
+        end
+      end
+    end
+
+    def initialize(app)
+      @app = app
+      @attributes_by_class = Concurrent::Map.new
+      @collecting = true
+    end
+
+    def enable
+      install_collecting_hook
+
+      apply_collected_attributes
+      @collecting = false
+    end
+
+    private
+      attr_reader :app
+
+      def install_collecting_hook
+        self.class.on_sensitive_attribute_declared do |klass, list|
+          attribute_was_declared(klass, list)
+        end
+      end
+
+      def attribute_was_declared(klass, list)
+        if collecting?
+          collect_for_later(klass, list)
+        else
+          apply_filter(klass, list)
+        end
+      end
+
+      def apply_collected_attributes
+        @attributes_by_class.each do |klass, list|
+          apply_filter(klass, list)
+        end
+      end
+
+      def collecting?
+        @collecting
+      end
+
+      def collect_for_later(klass, list)
+        @attributes_by_class[klass] ||= Concurrent::Array.new
+        @attributes_by_class[klass] += list
+      end
+
+      def apply_filter(klass, list)
+        list.each do |attribute|
+          next if klass.abstract_class? || klass == Base
+
+          klass_name = klass.name ? klass.model_name.element : nil
+          filter = [klass_name, attribute.to_s].compact.join(".")
+          app.config.filter_parameters << filter unless app.config.filter_parameters.include?(filter)
+        end
+      end
+  end
+end

data/lib/active_record/fixtures.rb

@@ -2,8 +2,6 @@
 
 require "erb"
 require "yaml"
-require "zlib"
-require "set"
 require "active_support/dependencies"
 require "active_support/core_ext/digest/uuid"
 require "active_record/test_fixtures"
@@ -244,10 +242,10 @@ module ActiveRecord
   # and one for the humans. Why don't we generate the primary key instead?
   # Hashing each fixture's label yields a consistent ID:
   #
-  #   george: # generated id: 503576764
+  #   george: # generated id: 380982691
   #     name: George the Monkey
   #
-  #   reginald: # generated id: 324201669
+  #   reginald: # generated id: 41001176
   #     name: Reginald the Pirate
   #
   # Active Record looks at the fixture's model class, discovers the correct

data/lib/active_record/future_result.rb

@@ -100,17 +100,23 @@ module ActiveRecord
     def execute_or_skip
       return unless pending?
 
-      @pool.with_connection do |connection|
-        return unless @mutex.try_lock
-        begin
-          if pending?
-            @event_buffer = EventBuffer.new(self, @instrumenter)
-            connection.with_instrumenter(@event_buffer) do
+      @session.synchronize do
+        return unless pending?
+
+        @pool.with_connection do |connection|
+          return unless @mutex.try_lock
+          previous_instrumenter = ActiveSupport::IsolatedExecutionState[:active_record_instrumenter]
+          begin
+            if pending?
+              @event_buffer = EventBuffer.new(self, @instrumenter)
+              ActiveSupport::IsolatedExecutionState[:active_record_instrumenter] = @event_buffer
+
               execute_query(connection, async: true)
             end
+          ensure
+            ActiveSupport::IsolatedExecutionState[:active_record_instrumenter] = previous_instrumenter
+            @mutex.unlock
           end
-        ensure
-          @mutex.unlock
         end
       end
     end
@@ -163,7 +169,7 @@ module ActiveRecord
       end
 
       def exec_query(connection, *args, **kwargs)
-        connection.internal_exec_query(*args, **kwargs)
+        connection.raw_exec_query(*args, **kwargs)
       end
 
       class SelectAll < FutureResult # :nodoc: