From 4d6fbe2934e94384e722ff6ca16e97c8978d4665 Mon Sep 17 00:00:00 2001
From: Jason Nochlin <hundredwatt@gmail.com>
Date: Sat, 7 Feb 2015 15:33:13 -0500
Subject: 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

The warning is outputed a module which is prepended in an initializer,
so there will be no performance impact if
`config.active_record.warn_on_records_fetched_greater_than` is not set.
---
 activerecord/lib/active_record/core.rb             |  9 ++++
 activerecord/lib/active_record/railtie.rb          |  8 ++++
 .../active_record/relation/record_fetch_warning.rb | 50 ++++++++++++++++++++++
 3 files changed, 67 insertions(+)
 create mode 100644 activerecord/lib/active_record/relation/record_fetch_warning.rb

(limited to 'activerecord/lib/active_record')

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
-- 
cgit v1.2.3