Add n_plus_one_only mode to Core#strict_loading!

Add an optional mode argument to
Core#strict_loading! to support n_plus_one_only
mode. Currently, when we turn on strict_loading
for a single record, it will raise even if we are
loading an association that is relatively safe to
lazy load like a belongs_to. This can be helpful
for some use cases, but prevents us from using
strict_loading to identify only problematic
instances of lazy loading.

The n_plus_one_only argument allows us to turn
strict_loading on for a single record, and only
raise when a N+1 query is likely to be executed.
When loading associations on a single record,
this only happens when we go through a has_many
association type. Note that the has_many
association itself is not problematic as it only
requires one query. We do this by turning
strict_loading on for each record that is loaded
through the has_many. This ensures that any
subsequent lazy loads on these records will raise
a StrictLoadingViolationError.

For example, where a developer belongs_to a ship
and each ship has_many parts, we expect the
following behaviour:

  developer.strict_loading!(mode: :n_plus_one_only)

  # Do not raise when a belongs_to association
  # (:ship) loads its has_many association (:parts)
  assert_nothing_raised do
    developer.ship.parts.to_a
  end

  refute developer.ship.strict_loading?
  assert developer.ship.parts.all?(&:strict_loading?)
  assert_raises ActiveRecord::StrictLoadingViolationError do
    developer.ship.parts.first.trinkets.to_a
  end

activerecord/CHANGELOG.md

@@ -1,3 +1,22 @@
+*   Add mode argument to record level `strict_loading!`
+
+    This argument can be used when enabling strict loading for a single record
+    to specify that we only want to raise on n plus one queries.
+
+    ```ruby
+    developer.strict_loading!(mode: :n_plus_one_only)
+
+    developer.projects.to_a # Does not raise
+    developer.projects.first.client # Raises StrictLoadingViolationError
+    ```
+
+    Previously, enabling strict loading would cause any lazily loaded
+    association to raise an error. Using `n_plus_one_only` mode allows us to
+    lazily load belongs_to, has_many, and other associations that are fetched
+    through a single query.
+
+    *Dinah Shi*
+
 *   Prevent double saves in autosave of cyclic associations.
 
     Adds an internal saving state which tracks if a record is currently being saved.

activerecord/lib/active_record/associations/association.rb

@@ -213,7 +213,7 @@ def create!(attributes = nil, &block)
 
       private
         def find_target
-          if strict_loading? && owner.validation_context.nil?
+          if violates_strict_loading? && owner.validation_context.nil?
             Base.strict_loading_violation!(owner: owner.class, reflection: reflection)
           end
 
@@ -226,13 +226,20 @@ def find_target
           end
 
           binds = AssociationScope.get_bind_values(owner, reflection.chain)
-          sc.execute(binds, klass.connection) { |record| set_inverse_instance(record) }
+          sc.execute(binds, klass.connection) do |record|
+            set_inverse_instance(record)
+            if owner.strict_loading_n_plus_one_only? && reflection.macro == :has_many
+              record.strict_loading!
+            else
+              record.strict_loading_mode = owner.strict_loading_mode
+            end
+          end
         end
 
-        def strict_loading?
+        def violates_strict_loading?
           return reflection.strict_loading? if reflection.options.key?(:strict_loading)
 
-          owner.strict_loading?
+          owner.strict_loading? && !owner.strict_loading_n_plus_one_only?
         end
 
         # The scope for this association.

activerecord/lib/active_record/core.rb

@@ -140,6 +140,7 @@ def self.configurations
       mattr_accessor :action_on_strict_loading_violation, instance_accessor: false, default: :raise
 
       class_attribute :strict_loading_by_default, instance_accessor: false, default: false
+      class_attribute :strict_loading_mode, instance_accessor: true, default: []
 
       mattr_accessor :writing_role, instance_accessor: false, default: :writing
 
@@ -735,15 +736,30 @@ def strict_loading?
     #   user.comments
     #   => ActiveRecord::StrictLoadingViolationError
     #
-    # strict_loading! accepts a boolean argument to specify whether
-    # to enable or disable strict loading mode.
+    # === Parameters:
+    #
+    # * value - Boolean specifying whether to enable or disable strict loading.
+    # * mode - Symbol specifying strict loading mode. Defaults to :all. Using
+    #          :n_plus_one_only mode will only raise an error if an association
+    #          that will lead to an n plus one query is lazily loaded.
+    #
+    # === Example:
     #
     #   user = User.first
     #   user.strict_loading!(false) # => false
     #   user.comments
     #   => #<ActiveRecord::Associations::CollectionProxy>
-    def strict_loading!(value = true)
+    def strict_loading!(value = true, mode: :all)
+      unless [:all, :n_plus_one_only].include?(mode)
+        raise ArgumentError, "The :mode option must be one of [:all, :n_plus_one_only]."
+      end
+
       @strict_loading = value
+      @strict_loading_mode = mode
+    end
+
+    def strict_loading_n_plus_one_only? # :nodoc:
+      @strict_loading_mode == :n_plus_one_only
     end
 
     # Marks this record as read only.
@@ -830,6 +846,7 @@ def init_internals
         @destroyed_by_association = nil
         @_start_transaction_state = nil
         @strict_loading           = self.class.strict_loading_by_default
+        @strict_loading_mode = self.class.strict_loading_mode
 
         self.class.define_attribute_methods
       end

activerecord/test/cases/strict_loading_test.rb

@@ -2,17 +2,18 @@
 
 require "cases/helper"
 require "models/developer"
+require "models/company"
 require "models/computer"
 require "models/mentor"
 require "models/project"
 require "models/ship"
+require "models/ship_part"
 require "models/strict_zine"
 require "models/interest"
+require "models/treasure"
 
 class StrictLoadingTest < ActiveRecord::TestCase
-  fixtures :developers
-  fixtures :projects
-  fixtures :ships
+  fixtures :developers, :developers_projects, :projects, :ships
 
   def test_strict_loading!
     developer = Developer.first
@@ -33,6 +34,45 @@ def test_strict_loading!
     end
   end
 
+  def test_strict_loading_n_plus_one_only_mode
+    developer = Developer.first
+    ship = Ship.first
+    ShipPart.create!(name: "Stern", ship: ship)
+    firm = Firm.create!("name" => "NASA")
+    project = Project.create!(name: "Apollo", firm: firm)
+
+    ship.update_column(:developer_id, developer.id)
+    developer.projects << project
+    developer.reload
+
+    developer.strict_loading!(mode: :n_plus_one_only)
+    assert_predicate developer, :strict_loading?
+
+    # Does not raise when loading a has_many association (:projects)
+    assert_nothing_raised do
+      developer.projects.to_a
+    end
+
+    # strict_loading is enabled for has_many associations
+    assert developer.projects.all?(&:strict_loading?)
+    assert_raises ActiveRecord::StrictLoadingViolationError do
+      developer.projects.last.firm
+    end
+
+    # Does not raise when a belongs_to association (:ship) loads its
+    # has_many association (:parts)
+    assert_nothing_raised do
+      developer.ship.parts.to_a
+    end
+
+    # strict_loading is enabled for has_many through a belongs_to
+    assert_not developer.ship.strict_loading?
+    assert developer.ship.parts.all?(&:strict_loading?)
+    assert_raises ActiveRecord::StrictLoadingViolationError do
+      developer.ship.parts.first.trinkets.to_a
+    end
+  end
+
   def test_strict_loading
     Developer.all.each { |d| assert_not d.strict_loading? }
     Developer.strict_loading.each { |d| assert d.strict_loading? }