activerecord 7.0.6 → 7.1.3.4

data/lib/active_record/query_logs.rb

@@ -1,13 +1,25 @@
 # frozen_string_literal: true
 
 require "active_support/core_ext/module/attribute_accessors_per_thread"
+require "active_record/query_logs_formatter"
 
 module ActiveRecord
   # = Active Record Query Logs
   #
-  # Automatically tag SQL queries with runtime information.
+  # Automatically append comments to SQL queries with runtime information tags. This can be used to trace troublesome
+  # SQL statements back to the application code that generated these statements.
   #
-  # Default tags available for use:
+  # Query logs can be enabled via \Rails configuration in <tt>config/application.rb</tt> or an initializer:
+  #
+  #     config.active_record.query_log_tags_enabled = true
+  #
+  # By default the name of the application, the name and action of the controller, or the name of the job are logged.
+  # The default format is {SQLCommenter}[https://open-telemetry.github.io/opentelemetry-sqlcommenter/].
+  # The tags shown in a query comment can be configured via \Rails configuration:
+  #
+  #     config.active_record.query_log_tags = [ :application, :controller, :action, :job ]
+  #
+  # Active Record defines default tags available for use:
   #
   # * +application+
   # * +pid+
@@ -15,48 +27,40 @@ module ActiveRecord
   # * +db_host+
   # * +database+
   #
-  # _Action Controller and Active Job tags are also defined when used in Rails:_
+  # Action Controller adds default tags when loaded:
   #
   # * +controller+
   # * +action+
-  # * +job+
-  #
-  # The tags used in a query can be configured directly:
+  # * +namespaced_controller+
   #
-  #     ActiveRecord::QueryLogs.tags = [ :application, :controller, :action, :job ]
+  # Active Job adds default tags when loaded:
   #
-  # or via Rails configuration:
-  #
-  #     config.active_record.query_log_tags = [ :application, :controller, :action, :job ]
+  # * +job+
   #
-  # To add new comment tags, add a hash to the tags array containing the keys and values you
-  # want to add to the comment. Dynamic content can be created by setting a proc or lambda value in a hash,
-  # and can reference any value stored in the +context+ object.
+  # New comment tags can be defined by adding them in a +Hash+ to the tags +Array+. Tags can have dynamic content by
+  # setting a +Proc+ or lambda value in the +Hash+, and can reference any value stored by \Rails in the +context+ object.
+  # ActiveSupport::CurrentAttributes can be used to store application values. Tags with +nil+ values are
+  # omitted from the query comment.
   #
   # Escaping is performed on the string returned, however untrusted user input should not be used.
   #
   # Example:
   #
-  #    tags = [
-  #      :application,
-  #      {
-  #        custom_tag: ->(context) { context[:controller]&.controller_name },
-  #        custom_value: -> { Custom.value },
-  #      }
-  #    ]
-  #    ActiveRecord::QueryLogs.tags = tags
-  #
-  # The QueryLogs +context+ can be manipulated via the +ActiveSupport::ExecutionContext.set+ method.
-  #
-  # Temporary updates limited to the execution of a block:
-  #
-  #    ActiveSupport::ExecutionContext.set(foo: Bar.new) do
-  #      posts = Post.all
-  #    end
-  #
-  # Direct updates to a context value:
-  #
-  #    ActiveSupport::ExecutionContext[:foo] = Bar.new
+  #     config.active_record.query_log_tags = [
+  #       :namespaced_controller,
+  #       :action,
+  #       :job,
+  #       {
+  #         request_id: ->(context) { context[:controller]&.request&.request_id },
+  #         job_id: ->(context) { context[:job]&.job_id },
+  #         tenant_id: -> { Current.tenant&.id },
+  #         static: "value",
+  #       },
+  #     ]
+  #
+  # By default the name of the application, the name and action of the controller, or the name of the job are logged
+  # using the {SQLCommenter}[https://open-telemetry.github.io/opentelemetry-sqlcommenter/] format. This can be changed
+  # via {config.active_record.query_log_tags_format}[https://guides.rubyonrails.org/configuring.html#config-active-record-query-log-tags-format]
   #
   # Tag comments can be prepended to the query:
   #
@@ -65,46 +69,65 @@ module ActiveRecord
   # For applications where the content will not change during the lifetime of
   # the request or job execution, the tags can be cached for reuse in every query:
   #
-  #    ActiveRecord::QueryLogs.cache_query_log_tags = true
-  #
-  # This option can be set during application configuration or in a Rails initializer:
-  #
   #    config.active_record.cache_query_log_tags = true
   module QueryLogs
     mattr_accessor :taggings, instance_accessor: false, default: {}
     mattr_accessor :tags, instance_accessor: false, default: [ :application ]
     mattr_accessor :prepend_comment, instance_accessor: false, default: false
     mattr_accessor :cache_query_log_tags, instance_accessor: false, default: false
+    mattr_accessor :tags_formatter, instance_accessor: false
     thread_mattr_accessor :cached_comment, instance_accessor: false
 
     class << self
-      def call(sql) # :nodoc:
-        if prepend_comment
-          "#{self.comment} #{sql}"
+      def call(sql, connection) # :nodoc:
+        comment = self.comment(connection)
+
+        if comment.blank?
+          sql
+        elsif prepend_comment
+          "#{comment} #{sql}"
         else
-          "#{sql} #{self.comment}"
-        end.strip
+          "#{sql} #{comment}"
+        end
       end
 
       def clear_cache # :nodoc:
         self.cached_comment = nil
       end
 
+      # Updates the formatter to be what the passed in format is.
+      def update_formatter(format)
+        self.tags_formatter =
+          case format
+          when :legacy
+            LegacyFormatter.new
+          when :sqlcommenter
+            SQLCommenter.new
+          else
+            raise ArgumentError, "Formatter is unsupported: #{formatter}"
+          end
+      end
+
       ActiveSupport::ExecutionContext.after_change { ActiveRecord::QueryLogs.clear_cache }
 
       private
         # Returns an SQL comment +String+ containing the query log tags.
         # Sets and returns a cached comment if <tt>cache_query_log_tags</tt> is +true+.
-        def comment
+        def comment(connection)
           if cache_query_log_tags
-            self.cached_comment ||= uncached_comment
+            self.cached_comment ||= uncached_comment(connection)
           else
-            uncached_comment
+            uncached_comment(connection)
           end
         end
 
-        def uncached_comment
-          content = tag_content
+        def formatter
+          self.tags_formatter || self.update_formatter(:legacy)
+        end
+
+        def uncached_comment(connection)
+          content = tag_content(connection)
+
           if content.present?
             "/*#{escape_sql_comment(content)}*/"
           end
@@ -113,7 +136,7 @@ module ActiveRecord
         def escape_sql_comment(content)
           # Sanitize a string to appear within a SQL comment
           # For compatibility, this also surrounding "/*+", "/*", and "*/"
-          # charcacters, possibly with single surrounding space.
+          # characters, possibly with single surrounding space.
           # Then follows that by replacing any internal "*/" or "/ *" with
           # "* /" or "/ *"
           comment = content.to_s.dup
@@ -123,10 +146,11 @@ module ActiveRecord
           comment
         end
 
-        def tag_content
+        def tag_content(connection)
           context = ActiveSupport::ExecutionContext.to_h
+          context[:connection] ||= connection
 
-          tags.flat_map { |i| [*i] }.filter_map do |tag|
+          pairs = tags.flat_map { |i| [*i] }.filter_map do |tag|
             key, handler = tag
             handler ||= taggings[key]
 
@@ -141,8 +165,9 @@ module ActiveRecord
             else
               handler
             end
-            "#{key}:#{val}" unless val.nil?
-          end.join(",")
+            [key, val] unless val.nil?
+          end
+          self.formatter.format(pairs)
         end
     end
   end

data/lib/active_record/query_logs_formatter.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module QueryLogs
+    class LegacyFormatter # :nodoc:
+      def initialize
+        @key_value_separator = ":"
+      end
+
+      # Formats the key value pairs into a string.
+      def format(pairs)
+        pairs.map! do |key, value|
+          "#{key}#{key_value_separator}#{format_value(value)}"
+        end.join(",")
+      end
+
+      private
+        attr_reader :key_value_separator
+
+        def format_value(value)
+          value
+        end
+    end
+
+    class SQLCommenter < LegacyFormatter # :nodoc:
+      def initialize
+        @key_value_separator = "="
+      end
+
+      def format(pairs)
+        pairs.sort_by!(&:first)
+        super
+      end
+
+      private
+        def format_value(value)
+          "'#{ERB::Util.url_encode(value)}'"
+        end
+    end
+  end
+end

data/lib/active_record/querying.rb

@@ -12,12 +12,13 @@ module ActiveRecord
       :create_or_find_by, :create_or_find_by!,
       :destroy_all, :delete_all, :update_all, :touch_all, :destroy_by, :delete_by,
       :find_each, :find_in_batches, :in_batches,
-      :select, :reselect, :order, :in_order_of, :reorder, :group, :limit, :offset, :joins, :left_joins, :left_outer_joins,
+      :select, :reselect, :order, :regroup, :in_order_of, :reorder, :group, :limit, :offset, :joins, :left_joins, :left_outer_joins,
       :where, :rewhere, :invert_where, :preload, :extract_associated, :eager_load, :includes, :from, :lock, :readonly,
       :and, :or, :annotate, :optimizer_hints, :extending,
       :having, :create_with, :distinct, :references, :none, :unscope, :merge, :except, :only,
       :count, :average, :minimum, :maximum, :sum, :calculate,
-      :pluck, :pick, :ids, :strict_loading, :excluding, :without
+      :pluck, :pick, :ids, :async_ids, :strict_loading, :excluding, :without, :with,
+      :async_count, :async_average, :async_minimum, :async_maximum, :async_sum, :async_pluck, :async_pick,
     ].freeze # :nodoc:
     delegate(*QUERYING_METHODS, to: :all)
 
@@ -50,6 +51,13 @@ module ActiveRecord
       _load_from_sql(_query_by_sql(sql, binds, preparable: preparable), &block)
     end
 
+    # Same as <tt>#find_by_sql</tt> but perform the query asynchronously and returns an ActiveRecord::Promise.
+    def async_find_by_sql(sql, binds = [], preparable: nil, &block)
+      _query_by_sql(sql, binds, preparable: preparable, async: true).then do |result|
+        _load_from_sql(result, &block)
+      end
+    end
+
     def _query_by_sql(sql, binds = [], preparable: nil, async: false) # :nodoc:
       connection.select_all(sanitize_sql(sql), "#{name} Load", binds, preparable: preparable, async: async)
     end
@@ -93,5 +101,10 @@ module ActiveRecord
     def count_by_sql(sql)
       connection.select_value(sanitize_sql(sql), "#{name} Count").to_i
     end
+
+    # Same as <tt>#count_by_sql</tt> but perform the query asynchronously and returns an ActiveRecord::Promise.
+    def async_count_by_sql(sql)
+      connection.select_value(sanitize_sql(sql), "#{name} Count", async: true).then(&:to_i)
+    end
   end
 end

data/lib/active_record/railtie.rb

@@ -34,7 +34,11 @@ module ActiveRecord
     config.active_record.sqlite3_production_warning = true
     config.active_record.query_log_tags_enabled = false
     config.active_record.query_log_tags = [ :application ]
+    config.active_record.query_log_tags_format = :legacy
     config.active_record.cache_query_log_tags = false
+    config.active_record.raise_on_assign_to_attr_readonly = false
+    config.active_record.belongs_to_required_validates_foreign_key = true
+    config.active_record.generate_secure_token_on = :create
 
     config.active_record.queues = ActiveSupport::InheritableOptions.new
 
@@ -63,7 +67,7 @@ module ActiveRecord
       unless ActiveSupport::Logger.logger_outputs_to?(Rails.logger, STDERR, STDOUT)
         console = ActiveSupport::Logger.new(STDERR)
         console.level = Rails.logger.level
-        Rails.logger.extend ActiveSupport::Logger.broadcast console
+        Rails.logger.broadcast_to(console)
       end
       ActiveRecord.verbose_query_logs = false
     end
@@ -72,6 +76,10 @@ module ActiveRecord
       require "active_record/base"
     end
 
+    initializer "active_record.deprecator", before: :load_environment_config do |app|
+      app.deprecators[:active_record] = ActiveRecord.deprecator
+    end
+
     initializer "active_record.initialize_timezone" do
       ActiveSupport.on_load(:active_record) do
         self.time_zone_aware_attributes = true
@@ -100,7 +108,7 @@ module ActiveRecord
       end
     end
 
-    initializer "Check for cache versioning support" do
+    initializer "active_record.cache_versioning_support" do
       config.after_initialize do |app|
         ActiveSupport.on_load(:active_record) do
           if app.config.active_record.cache_versioning && Rails.cache
@@ -125,9 +133,15 @@ To keep using the current cache store, you can turn off cache versioning entirel
       end
     end
 
+    initializer "active_record.use_schema_cache_dump" do
+      ActiveRecord::ConnectionAdapters::SchemaReflection.use_schema_cache_dump = config.active_record.use_schema_cache_dump
+    end
+
     initializer "active_record.check_schema_cache_dump" do
       check_schema_cache_dump_version = config.active_record.check_schema_cache_dump_version
 
+      ActiveRecord::ConnectionAdapters::SchemaReflection.check_schema_cache_dump_version = check_schema_cache_dump_version
+
       if config.active_record.use_schema_cache_dump && !config.active_record.lazily_load_schema_cache
         config.after_initialize do |app|
           ActiveSupport.on_load(:active_record) do
@@ -135,10 +149,10 @@ To keep using the current cache store, you can turn off cache versioning entirel
 
             filename = ActiveRecord::Tasks::DatabaseTasks.cache_dump_filename(
               db_config.name,
-              schema_cache_path: db_config&.schema_cache_path
+              schema_cache_path: db_config.schema_cache_path
             )
 
-            cache = ActiveRecord::ConnectionAdapters::SchemaCache.load_from(filename)
+            cache = ActiveRecord::ConnectionAdapters::SchemaCache._load_from(filename)
             next if cache.nil?
 
             if check_schema_cache_dump_version
@@ -150,34 +164,47 @@ To keep using the current cache store, you can turn off cache versioning entirel
               end
               next if current_version.nil?
 
-              if cache.version != current_version
-                warn "Ignoring #{filename} because it has expired. The current schema version is #{current_version}, but the one in the schema cache file is #{cache.version}."
+              if cache.schema_version != current_version
+                warn "Ignoring #{filename} because it has expired. The current schema version is #{current_version}, but the one in the schema cache file is #{cache.schema_version}."
                 next
               end
             end
 
             Rails.logger.info("Using schema cache file #{filename}")
-            connection_pool.set_schema_cache(cache)
+            connection_pool.schema_reflection.set_schema_cache(cache)
           end
         end
       end
     end
 
     initializer "active_record.define_attribute_methods" do |app|
+      # For resiliency, it is critical that a Rails application should be
+      # able to boot without depending on the database (or any other service)
+      # being responsive.
+      #
+      # Otherwise a bad deploy adding a lot of load on the database may require to
+      # entirely shutdown the application so the database can recover before a fixed
+      # version can be deployed again.
+      #
+      # This is why this initializer tries hard not to query the database, and if it
+      # does, it makes sure to rescue any possible database error.
+      check_schema_cache_dump_version = config.active_record.check_schema_cache_dump_version
       config.after_initialize do
         ActiveSupport.on_load(:active_record) do
-          if app.config.eager_load
+          # In development and test we shouldn't eagerly define attribute methods because
+          # db:test:prepare will trigger later and might change the schema.
+          #
+          # Additionally if `check_schema_cache_dump_version` is enabled (which is the default),
+          # loading the schema cache dump trigger a database connection to compare the schema
+          # versions.
+          # This means the attribute methods will be lazily defined whent the model is accessed,
+          # likely as part of the first few requests or jobs. This isn't good for performance
+          # but we unfortunately have to arbitrate between resiliency and performance, and chose
+          # resiliency.
+          if !check_schema_cache_dump_version && app.config.eager_load && !Rails.env.local?
             begin
               descendants.each do |model|
-                # If the schema cache was loaded from a dump, we can use it without connecting
-                schema_cache = model.connection_pool.schema_cache
-
-                # If there's no connection yet, we avoid connecting.
-                schema_cache ||= model.connected? && model.connection.schema_cache
-
-                # If the schema cache doesn't have the columns
-                # hash for the model cached, `define_attribute_methods` would trigger a query.
-                if schema_cache && schema_cache.columns_hash?(model.table_name)
+                if model.connection_pool.schema_reflection.cached?(model.table_name)
                   model.define_attribute_methods
                 end
               end
@@ -210,6 +237,16 @@ To keep using the current cache store, you can turn off cache versioning entirel
       end
     end
 
+    initializer "active_record.sqlite3_adapter_strict_strings_by_default" do
+      config.after_initialize do
+        if config.active_record.sqlite3_adapter_strict_strings_by_default
+          ActiveSupport.on_load(:active_record_sqlite3adapter) do
+            self.strict_strings_by_default = true
+          end
+        end
+      end
+    end
+
     initializer "active_record.set_configs" do |app|
       configs = app.config.active_record
 
@@ -234,8 +271,10 @@ To keep using the current cache store, you can turn off cache versioning entirel
           :shard_resolver,
           :query_log_tags_enabled,
           :query_log_tags,
+          :query_log_tags_format,
           :cache_query_log_tags,
           :sqlite3_production_warning,
+          :sqlite3_adapter_strict_strings_by_default,
           :check_schema_cache_dump_version,
           :use_schema_cache_dump
         )
@@ -260,21 +299,23 @@ To keep using the current cache store, you can turn off cache versioning entirel
     # and then establishes the connection.
     initializer "active_record.initialize_database" do
       ActiveSupport.on_load(:active_record) do
-        if ActiveRecord.legacy_connection_handling
-          self.connection_handlers = { ActiveRecord.writing_role => ActiveRecord::Base.default_connection_handler }
-        end
         self.configurations = Rails.application.config.database_configuration
 
         establish_connection
       end
     end
 
-    # Expose database runtime to controller for logging.
+    # Expose database runtime for logging.
     initializer "active_record.log_runtime" do
       require "active_record/railties/controller_runtime"
       ActiveSupport.on_load(:action_controller) do
         include ActiveRecord::Railties::ControllerRuntime
       end
+
+      require "active_record/railties/job_runtime"
+      ActiveSupport.on_load(:active_job) do
+        include ActiveRecord::Railties::JobRuntime
+      end
     end
 
     initializer "active_record.set_reloader_hooks" do
@@ -282,7 +323,7 @@ To keep using the current cache store, you can turn off cache versioning entirel
         ActiveSupport::Reloader.before_class_unload do
           if ActiveRecord::Base.connected?
             ActiveRecord::Base.clear_cache!
-            ActiveRecord::Base.clear_reloadable_connections!
+            ActiveRecord::Base.connection_handler.clear_reloadable_connections!(:all)
           end
         end
       end
@@ -309,8 +350,8 @@ To keep using the current cache store, you can turn off cache versioning entirel
           # this connection is trivial: the rest of the pool would need to be
           # populated anyway.
 
-          clear_active_connections!
-          flush_idle_connections!
+          connection_handler.clear_active_connections!(:all)
+          connection_handler.flush_idle_connections!(:all)
         end
       end
     end
@@ -327,34 +368,41 @@ To keep using the current cache store, you can turn off cache versioning entirel
       end
     end
 
+    initializer "active_record.generated_token_verifier" do
+      config.after_initialize do |app|
+        ActiveSupport.on_load(:active_record) do
+          self.generated_token_verifier ||= app.message_verifier("active_record/token_for")
+        end
+      end
+    end
+
     initializer "active_record_encryption.configuration" do |app|
-      ActiveRecord::Encryption.configure \
-         primary_key: app.credentials.dig(:active_record_encryption, :primary_key),
-         deterministic_key: app.credentials.dig(:active_record_encryption, :deterministic_key),
-         key_derivation_salt: app.credentials.dig(:active_record_encryption, :key_derivation_salt),
-         **config.active_record.encryption
+      auto_filtered_parameters = ActiveRecord::Encryption::AutoFilteredParameters.new(app)
 
-      ActiveSupport.on_load(:active_record) do
-        # Support extended queries for deterministic attributes and validations
-        if ActiveRecord::Encryption.config.extend_queries
-          ActiveRecord::Encryption::ExtendedDeterministicQueries.install_support
-          ActiveRecord::Encryption::ExtendedDeterministicUniquenessValidator.install_support
+      config.after_initialize do |app|
+        ActiveRecord::Encryption.configure \
+          primary_key: app.credentials.dig(:active_record_encryption, :primary_key),
+          deterministic_key: app.credentials.dig(:active_record_encryption, :deterministic_key),
+          key_derivation_salt: app.credentials.dig(:active_record_encryption, :key_derivation_salt),
+          **config.active_record.encryption
+
+        auto_filtered_parameters.enable if ActiveRecord::Encryption.config.add_to_filter_parameters
+
+        ActiveSupport.on_load(:active_record) do
+          # Support extended queries for deterministic attributes and validations
+          if ActiveRecord::Encryption.config.extend_queries
+            ActiveRecord::Encryption::ExtendedDeterministicQueries.install_support
+            ActiveRecord::Encryption::ExtendedDeterministicUniquenessValidator.install_support
+          end
         end
       end
 
       ActiveSupport.on_load(:active_record_fixture_set) do
-        # Encrypt active record fixtures
+        # Encrypt Active Record fixtures
         if ActiveRecord::Encryption.config.encrypt_fixtures
           ActiveRecord::Fixture.prepend ActiveRecord::Encryption::EncryptedFixtures
         end
       end
-
-      # Filtered params
-      ActiveSupport.on_load(:action_controller, run_once: true) do
-        if ActiveRecord::Encryption.config.add_to_filter_parameters
-          ActiveRecord::Encryption.install_auto_filtered_parameters_hook(app)
-        end
-      end
     end
 
     initializer "active_record.query_log_tags_config" do |app|
@@ -363,16 +411,21 @@ To keep using the current cache store, you can turn off cache versioning entirel
           ActiveRecord.query_transformers << ActiveRecord::QueryLogs
           ActiveRecord::QueryLogs.taggings.merge!(
             application:  Rails.application.class.name.split("::").first,
-            pid:          -> { Process.pid },
-            socket:       -> { ActiveRecord::Base.connection_db_config.socket },
-            db_host:      -> { ActiveRecord::Base.connection_db_config.host },
-            database:     -> { ActiveRecord::Base.connection_db_config.database }
+            pid:          -> { Process.pid.to_s },
+            socket:       ->(context) { context[:connection].pool.db_config.socket },
+            db_host:      ->(context) { context[:connection].pool.db_config.host },
+            database:     ->(context) { context[:connection].pool.db_config.database }
           )
+          ActiveRecord.disable_prepared_statements = true
 
           if app.config.active_record.query_log_tags.present?
             ActiveRecord::QueryLogs.tags = app.config.active_record.query_log_tags
           end
 
+          if app.config.active_record.query_log_tags_format
+            ActiveRecord::QueryLogs.update_formatter(app.config.active_record.query_log_tags_format)
+          end
+
           if app.config.active_record.cache_query_log_tags
             ActiveRecord::QueryLogs.cache_query_log_tags = true
           end
@@ -382,7 +435,7 @@ To keep using the current cache store, you can turn off cache versioning entirel
 
     initializer "active_record.unregister_current_scopes_on_unload" do |app|
       config.after_initialize do
-        unless app.config.cache_classes
+        if app.config.reloading_enabled?
           Rails.autoloaders.main.on_unload do |_cpath, value, _abspath|
             # Conditions are written this way to be robust against custom
             # implementations of value#is_a? or value#<.
@@ -393,5 +446,14 @@ To keep using the current cache store, you can turn off cache versioning entirel
         end
       end
     end
+
+    initializer "active_record.message_pack" do
+      ActiveSupport.on_load(:message_pack) do
+        ActiveSupport.on_load(:active_record) do
+          require "active_record/message_pack"
+          ActiveRecord::MessagePack::Extensions.install(ActiveSupport::MessagePack::CacheSerializer)
+        end
+      end
+    end
   end
 end

data/lib/active_record/railties/controller_runtime.rb

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 require "active_support/core_ext/module/attr_internal"
-require "active_record/log_subscriber"
+require "active_record/runtime_registry"
 
 module ActiveRecord
   module Railties # :nodoc:
@@ -16,6 +16,11 @@ module ActiveRecord
         end
       end
 
+      def initialize(...) # :nodoc:
+        super
+        self.db_runtime = nil
+      end
+
       private
         attr_internal :db_runtime
 
@@ -23,18 +28,19 @@ module ActiveRecord
           # We also need to reset the runtime before each action
           # because of queries in middleware or in cases we are streaming
           # and it won't be cleaned up by the method below.
-          ActiveRecord::LogSubscriber.reset_runtime
+          ActiveRecord::RuntimeRegistry.reset
           super
         end
 
         def cleanup_view_runtime
-          if logger && logger.info? && ActiveRecord::Base.connected?
-            db_rt_before_render = ActiveRecord::LogSubscriber.reset_runtime
+          if logger && logger.info?
+            db_rt_before_render = ActiveRecord::RuntimeRegistry.reset
             self.db_runtime = (db_runtime || 0) + db_rt_before_render
             runtime = super
-            db_rt_after_render = ActiveRecord::LogSubscriber.reset_runtime
+            queries_rt = ActiveRecord::RuntimeRegistry.sql_runtime - ActiveRecord::RuntimeRegistry.async_sql_runtime
+            db_rt_after_render = ActiveRecord::RuntimeRegistry.reset
             self.db_runtime += db_rt_after_render
-            runtime - db_rt_after_render
+            runtime - queries_rt
           else
             super
           end
@@ -42,9 +48,8 @@ module ActiveRecord
 
         def append_info_to_payload(payload)
           super
-          if ActiveRecord::Base.connected?
-            payload[:db_runtime] = (db_runtime || 0) + ActiveRecord::LogSubscriber.reset_runtime
-          end
+
+          payload[:db_runtime] = (db_runtime || 0) + ActiveRecord::RuntimeRegistry.reset
         end
     end
   end