Add strict_loading mode to prevent lazy loading

Add `#strict_loading` to any record to prevent lazy loading of associations.
`strict_loading` will cascade down from the parent record to all the
associations to help you catch any places where you may want to use
`preload` instead of lazy loading. This is useful for preventing N+1's.

Co-authored-by: Aaron Patterson <aaron.patterson@gmail.com>

activerecord/CHANGELOG.md

@@ -1,3 +1,17 @@
+*   Add support for `strict_loading` mode to prevent lazy loading of records.
+
+    Raise an error if a parent record is marked as `strict_loading` and attempts to lazily load its associations. This is useful for finding places you may want to preload an association and avoid additional queries.
+
+    Usage:
+
+    ```
+    >> Developer.strict_loading.first
+    >> dev.audit_logs.to_a
+    => ActiveRecord::StrictLoadingViolationError: Developer is marked as strict_loading and AuditLog cannot be lazily loaded.
+    ```
+
+    *Eileen M. Uchitelle*, *Aaron Patterson*
+
 *   Add support for PostgreSQL 11+ partitioned indexes when using `upsert_all`.
 
     *Sebastián Palma*

activerecord/lib/active_record/associations/association.rb

@@ -207,6 +207,10 @@ module ActiveRecord
 
       private
         def find_target
+          if owner.strict_loading?
+            raise StrictLoadingViolationError, "#{owner.class} is marked as strict_loading and #{klass} cannot be lazily loaded."
+          end
+
           scope = self.scope
           return scope.to_a if skip_statement_cache?(scope)
 

activerecord/lib/active_record/associations/collection_association.rb

@@ -308,6 +308,7 @@ module ActiveRecord
 
       def find_from_target?
         loaded? ||
+          owner.strict_loading? ||
           owner.new_record? ||
           target.any? { |record| record.new_record? || record.changed? }
       end

activerecord/lib/active_record/associations/join_dependency.rb

@@ -94,7 +94,7 @@ module ActiveRecord
         }
       end
 
-      def instantiate(result_set, &block)
+      def instantiate(result_set, strict_loading_value, &block)
         primary_key = aliases.column_alias(join_root, join_root.primary_key)
 
         seen = Hash.new { |i, object_id|
@@ -120,7 +120,7 @@ module ActiveRecord
           result_set.each { |row_hash|
             parent_key = primary_key ? row_hash[primary_key] : row_hash
             parent = parents[parent_key] ||= join_root.instantiate(row_hash, column_aliases, &block)
-            construct(parent, join_root, row_hash, seen, model_cache)
+            construct(parent, join_root, row_hash, seen, model_cache, strict_loading_value)
           }
         end
 
@@ -215,7 +215,7 @@ module ActiveRecord
           end
         end
 
-        def construct(ar_parent, parent, row, seen, model_cache)
+        def construct(ar_parent, parent, row, seen, model_cache, strict_loading_value)
           return if ar_parent.nil?
 
           parent.children.each do |node|
@@ -224,7 +224,7 @@ module ActiveRecord
               other.loaded!
             elsif ar_parent.association_cached?(node.reflection.name)
               model = ar_parent.association(node.reflection.name).target
-              construct(model, node, row, seen, model_cache)
+              construct(model, node, row, seen, model_cache, strict_loading_value)
               next
             end
 
@@ -239,21 +239,22 @@ module ActiveRecord
             model = seen[ar_parent.object_id][node][id]
 
             if model
-              construct(model, node, row, seen, model_cache)
+              construct(model, node, row, seen, model_cache, strict_loading_value)
             else
-              model = construct_model(ar_parent, node, row, model_cache, id)
+              model = construct_model(ar_parent, node, row, model_cache, id, strict_loading_value)
 
               seen[ar_parent.object_id][node][id] = model
-              construct(model, node, row, seen, model_cache)
+              construct(model, node, row, seen, model_cache, strict_loading_value)
             end
           end
         end
 
-        def construct_model(record, node, row, model_cache, id)
+        def construct_model(record, node, row, model_cache, id, strict_loading_value)
           other = record.association(node.reflection.name)
 
           model = model_cache[node][id] ||=
             node.instantiate(row, aliases.column_aliases(node)) do |m|
+              m.strict_loading! if strict_loading_value
               other.set_inverse_instance(m)
             end
 
@@ -264,6 +265,7 @@ module ActiveRecord
           end
 
           model.readonly! if node.readonly?
+          model.strict_loading! if node.strict_loading?
           model
         end
     end

activerecord/lib/active_record/associations/join_dependency/join_association.rb

@@ -65,6 +65,12 @@ module ActiveRecord
           @readonly = reflection.scope && reflection.scope_for(base_klass.unscoped).readonly_value
         end
 
+        def strict_loading?
+          return @strict_loading if defined?(@strict_loading)
+
+          @strict_loading = reflection.scope && reflection.scope_for(base_klass.unscoped).strict_loading_value
+        end
+
         private
           def append_constraints(join, constraints)
             if join.is_a?(Arel::Nodes::StringJoin)

activerecord/lib/active_record/associations/preloader/association.rb

@@ -143,8 +143,16 @@ module ActiveRecord
             end
 
             scope.merge!(reflection_scope) if reflection.scope
-            scope.merge!(preload_scope) if preload_scope
-            scope
+
+            if preload_scope && !preload_scope.empty_scope?
+              scope.merge!(preload_scope)
+            end
+
+            if preload_scope && preload_scope.strict_loading_value
+              scope.strict_loading
+            else
+              scope
+            end
           end
       end
     end

activerecord/lib/active_record/core.rb

@@ -491,6 +491,14 @@ module ActiveRecord
       @readonly
     end
 
+    def strict_loading?
+      @strict_loading
+    end
+
+    def strict_loading!
+      @strict_loading = true
+    end
+
     # Marks this record as read only.
     def readonly!
       @readonly = true
@@ -575,6 +583,7 @@ module ActiveRecord
         @destroyed_by_association = nil
         @_start_transaction_state = nil
         @transaction_state        = nil
+        @strict_loading           = false
 
         self.class.define_attribute_methods
       end

activerecord/lib/active_record/errors.rb

@@ -228,6 +228,10 @@ module ActiveRecord
   class ReadOnlyRecord < ActiveRecordError
   end
 
+  # Raised on attempt to lazily load records that are marked as strict loading.
+  class StrictLoadingViolationError < ActiveRecordError
+  end
+
   # {ActiveRecord::Base.transaction}[rdoc-ref:Transactions::ClassMethods#transaction]
   # uses this exception to distinguish a deliberate rollback from other exceptional situations.
   # Normally, raising an exception will cause the

activerecord/lib/active_record/querying.rb

@@ -16,7 +16,7 @@ module ActiveRecord
       :where, :rewhere, :preload, :extract_associated, :eager_load, :includes, :from, :lock, :readonly, :extending, :or,
       :having, :create_with, :distinct, :references, :none, :unscope, :optimizer_hints, :merge, :except, :only,
       :count, :average, :minimum, :maximum, :sum, :calculate, :annotate,
-      :pluck, :pick, :ids
+      :pluck, :pick, :ids, :strict_loading
     ].freeze # :nodoc:
     delegate(*QUERYING_METHODS, to: :all)
 

activerecord/lib/active_record/relation.rb

@@ -7,7 +7,7 @@ module ActiveRecord
                             :order, :joins, :left_outer_joins, :references,
                             :extending, :unscope, :optimizer_hints, :annotate]
 
-    SINGLE_VALUE_METHODS = [:limit, :offset, :lock, :readonly, :reordering,
+    SINGLE_VALUE_METHODS = [:limit, :offset, :lock, :readonly, :reordering, :strict_loading,
                             :reverse_order, :distinct, :create_with, :skip_query_cache]
 
     CLAUSE_METHODS = [:where, :having, :from]
@@ -751,13 +751,24 @@ module ActiveRecord
       ActiveRecord::Associations::AliasTracker.create(connection, table.name, joins)
     end
 
+    class StrictLoadingScope
+      def self.empty_scope?
+        true
+      end
+
+      def self.strict_loading_value
+        true
+      end
+    end
+
     def preload_associations(records) # :nodoc:
       preload = preload_values
       preload += includes_values unless eager_loading?
       preloader = nil
+      scope = strict_loading_value ? StrictLoadingScope : nil
       preload.each do |associations|
         preloader ||= build_preloader
-        preloader.preload records, associations
+        preloader.preload records, associations, scope
       end
     end
 
@@ -831,7 +842,7 @@ module ActiveRecord
                 else
                   relation = join_dependency.apply_column_aliases(relation)
                   rows = connection.select_all(relation.arel, "SQL")
-                  join_dependency.instantiate(rows, &block)
+                  join_dependency.instantiate(rows, strict_loading_value, &block)
                 end.freeze
               end
             else
@@ -841,6 +852,7 @@ module ActiveRecord
           preload_associations(records) unless skip_preloading_value
 
           records.each(&:readonly!) if readonly_value
+          records.each(&:strict_loading!) if strict_loading_value
 
           records
         end

activerecord/lib/active_record/relation/query_methods.rb

@@ -852,6 +852,21 @@ module ActiveRecord
       self
     end
 
+    # Sets the returned relation to strict_loading mode. This will raise an error
+    # if the record tries to lazily load an association.
+    #
+    #   user = User.first.strict_loading
+    #   user.comments.to_a
+    #   => ActiveRecord::StrictLoadingViolationError
+    def strict_loading(value = true)
+      spawn.strict_loading!(value)
+    end
+
+    def strict_loading!(value = true) # :nodoc:
+      self.strict_loading_value = value
+      self
+    end
+
     # Sets attributes to be used when creating new records from a
     # relation object.
     #

activerecord/test/cases/relation/mutation_test.rb

@@ -59,7 +59,7 @@ module ActiveRecord
       assert_equal [], relation.extending_values
     end
 
-    (Relation::SINGLE_VALUE_METHODS - [:lock, :reordering, :reverse_order, :create_with, :skip_query_cache]).each do |method|
+    (Relation::SINGLE_VALUE_METHODS - [:lock, :reordering, :reverse_order, :create_with, :skip_query_cache, :strict_loading]).each do |method|
       test "##{method}!" do
         assert relation.public_send("#{method}!", :foo).equal?(relation)
         assert_equal :foo, relation.public_send("#{method}_value")

activerecord/test/cases/strict_loading_test.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+require "cases/helper"
+require "models/developer"
+require "models/computer"
+
+class StrictLoadingTest < ActiveRecord::TestCase
+  fixtures :developers
+
+  def test_strict_loading
+    Developer.all.each { |d| assert_not d.strict_loading? }
+    Developer.strict_loading.each { |d| assert d.strict_loading? }
+  end
+
+  def test_raises_if_strict_loading_and_lazy_loading
+    dev = Developer.strict_loading.first
+    assert_predicate dev, :strict_loading?
+
+    assert_raises ActiveRecord::StrictLoadingViolationError do
+      dev.audit_logs.to_a
+    end
+  end
+
+  def test_preload_audit_logs_are_strict_loading_because_parent_is_strict_loading
+    developer = Developer.first
+
+    3.times do
+      AuditLog.create(developer: developer, message: "I am message")
+    end
+
+    dev = Developer.includes(:audit_logs).strict_loading.first
+
+    assert_predicate dev, :strict_loading?
+    assert dev.audit_logs.all?(&:strict_loading?), "Expected all audit logs to be strict_loading"
+  end
+
+  def test_eager_load_audit_logs_are_strict_loading_because_parent_is_strict_loading_in_hm_relation
+    developer = Developer.first
+
+    3.times do
+      AuditLog.create(developer: developer, message: "I am message")
+    end
+
+    dev = Developer.eager_load(:strict_loading_audit_logs).first
+
+    assert dev.strict_loading_audit_logs.all?(&:strict_loading?), "Expected all audit logs to be strict_loading"
+  end
+
+  def test_eager_load_audit_logs_are_strict_loading_because_parent_is_strict_loading
+    developer = Developer.first
+
+    3.times do
+      AuditLog.create(developer: developer, message: "I am message")
+    end
+
+    dev = Developer.eager_load(:audit_logs).strict_loading.first
+
+    assert_predicate dev, :strict_loading?
+    assert dev.audit_logs.all?(&:strict_loading?), "Expected all audit logs to be strict_loading"
+  end
+
+  def test_raises_on_unloaded_relation_methods_if_strict_loading
+    dev = Developer.strict_loading.first
+    assert_predicate dev, :strict_loading?
+
+    assert_raises ActiveRecord::StrictLoadingViolationError do
+      dev.audit_logs.first
+    end
+  end
+end

activerecord/test/models/developer.rb

@@ -52,6 +52,7 @@ class Developer < ActiveRecord::Base
                           class_name: "SpecialProject"
 
   has_many :audit_logs
+  has_many :strict_loading_audit_logs, -> { strict_loading }, class_name: "AuditLog"
   has_many :contracts
   has_many :firms, through: :contracts, source: :firm
   has_many :comments, ->(developer) { where(body: "I'm #{developer.name}") }