Merge pull request #18846 from hundredwatt/feat/warn-on-result-set-size

Add `config.active_record.warn_on_result_set_size` option

5 files changed, 102 insertions, 0 deletions

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index 6c129a1c79..be3a461eda 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,3 +1,10 @@
+*   Add `config.active_record.warn_on_records_fetched_greater_than` option
+
+    When set to an integer, a warning will be logged whenever a result set
+    larger than the specified size is returned by a query. Fixes #16463
+
+    *Jason Nochlin*
+
 *   Ignore psqlrc when loading database structure.
 
     *Jason Weathered*
diff --git a/activerecord/lib/active_record/core.rb b/activerecord/lib/active_record/core.rb
index 9a39a0e919..a17eb226f4 100644
--- a/activerecord/lib/active_record/core.rb
+++ b/activerecord/lib/active_record/core.rb
@@ -94,6 +94,15 @@ module ActiveRecord
       mattr_accessor :dump_schemas, instance_writer: false
       self.dump_schemas = :schema_search_path
 
+      ##
+      # :singleton-method:
+      # Specify a threshold for the size of query result sets. If the
+      # number of records in the set exceeds threshold, a warning is
+      # logged. This should be used to identify queries which pull
+      # thousands of records, which may cause memory bloat.
+      mattr_accessor :warn_on_records_fetched_greater_than, instance_writer: false
+      self.warn_on_records_fetched_greater_than = nil
+
       mattr_accessor :maintain_test_schema, instance_accessor: false
 
       mattr_accessor :belongs_to_required_by_default, instance_accessor: false
diff --git a/activerecord/lib/active_record/railtie.rb b/activerecord/lib/active_record/railtie.rb
index f1bdbc845c..7e907beec0 100644
--- a/activerecord/lib/active_record/railtie.rb
+++ b/activerecord/lib/active_record/railtie.rb
@@ -102,6 +102,14 @@ module ActiveRecord
       end
     end
 
+    initializer "active_record.warn_on_records_fetched_greater_than" do
+      if config.active_record.warn_on_records_fetched_greater_than
+        ActiveSupport.on_load(:active_record) do
+          require 'active_record/relation/record_fetch_warning'
+        end
+      end
+    end
+
     initializer "active_record.set_configs" do |app|
       ActiveSupport.on_load(:active_record) do
         app.config.active_record.each do |k,v|
diff --git a/activerecord/lib/active_record/relation/record_fetch_warning.rb b/activerecord/lib/active_record/relation/record_fetch_warning.rb
new file mode 100644
index 0000000000..0d31f73ddd
--- /dev/null
+++ b/activerecord/lib/active_record/relation/record_fetch_warning.rb
@@ -0,0 +1,50 @@
+module ActiveRecord
+  class Relation
+    module RecordFetchWarning
+      # When this module is prepended to ActiveRecord::Relation and
+      # `config.active_record.warn_on_records_fetched_greater_than` is
+      # set to an integer, if the number of records a query returns is
+      # greater than the value of `warn_on_records_fetched_greater_than`,
+      # a warning is logged. This allows for the dection of queries that
+      # return a large number of records, which could cause memory
+      # bloat.
+      #
+      # In most cases, fetching large number of records can be performed
+      # efficiently using the ActiveRecord::Batches methods.
+      # See active_record/lib/relation/batches.rb for more information.
+      def exec_queries
+        QueryRegistry.reset
+
+        super.tap do
+          if logger && warn_on_records_fetched_greater_than
+            if @records.length > warn_on_records_fetched_greater_than
+              logger.warn "Query fetched #{@records.size} #{@klass} records: #{QueryRegistry.queries.join(";")}"
+            end
+          end
+        end
+      end
+
+      ActiveSupport::Notifications.subscribe("sql.active_record") do |*args|
+        payload = args.last
+
+        QueryRegistry.queries << payload[:sql]
+      end
+
+      class QueryRegistry # :nodoc:
+        extend ActiveSupport::PerThreadRegistry
+
+        attr_accessor :queries
+
+        def initialize
+          reset
+        end
+
+        def reset
+          @queries = []
+        end
+      end
+    end
+  end
+end
+
+ActiveRecord::Relation.prepend ActiveRecord::Relation::RecordFetchWarning
diff --git a/activerecord/test/cases/relation/record_fetch_warning_test.rb b/activerecord/test/cases/relation/record_fetch_warning_test.rb
new file mode 100644
index 0000000000..62f0a7cc49
--- /dev/null
+++ b/activerecord/test/cases/relation/record_fetch_warning_test.rb
@@ -0,0 +1,28 @@
+require 'cases/helper'
+require 'models/post'
+
+module ActiveRecord
+  class RecordFetchWarningTest < ActiveRecord::TestCase
+    fixtures :posts
+
+    def test_warn_on_records_fetched_greater_than
+      original_logger = ActiveRecord::Base.logger
+      orginal_warn_on_records_fetched_greater_than = ActiveRecord::Base.warn_on_records_fetched_greater_than
+
+      log = StringIO.new
+      ActiveRecord::Base.logger = ActiveSupport::Logger.new(log)
+      ActiveRecord::Base.logger.level = Logger::WARN
+
+      require 'active_record/relation/record_fetch_warning'
+
+      ActiveRecord::Base.warn_on_records_fetched_greater_than = 1
+
+      Post.all.to_a
+
+      assert_match(/Query fetched/, log.string)
+    ensure
+      ActiveRecord::Base.logger = original_logger
+      ActiveRecord::Base.warn_on_records_fetched_greater_than = orginal_warn_on_records_fetched_greater_than
+    end
+  end
+end