Refactor relation preloading (#37217)

app/controllers/api/v1/statuses/quotes_controller.rb

@@ -41,8 +41,8 @@ class Api::V1::Statuses::QuotesController < Api::V1::Statuses::BaseController
     if current_account&.id != @status.account_id
       domains = @statuses.filter_map(&:account_domain).uniq
       account_ids = @statuses.map(&:account_id).uniq
-      relations = current_account&.relations_map(account_ids, domains) || {}
-      @statuses.reject! { |status| StatusFilter.new(status, current_account, relations).filtered? }
+      current_account&.preload_relations!(account_ids, domains)
+      @statuses.reject! { |status| StatusFilter.new(status, current_account).filtered? }
     end
   end
 

app/lib/status_filter.rb

@@ -3,10 +3,9 @@
 class StatusFilter
   attr_reader :status, :account
 
-  def initialize(status, account, preloaded_relations = {})
+  def initialize(status, account)
     @status              = status
     @account             = account
-    @preloaded_relations = preloaded_relations
   end
 
   def filtered?
@@ -40,15 +39,15 @@ class StatusFilter
   end
 
   def blocking_account?
-    @preloaded_relations[:blocking] ? @preloaded_relations[:blocking][status.account_id] : account.blocking?(status.account_id)
+    account.blocking?(status.account_id)
   end
 
   def blocking_domain?
-    @preloaded_relations[:domain_blocking_by_domain] ? @preloaded_relations[:domain_blocking_by_domain][status.account_domain] : account.domain_blocking?(status.account_domain)
+    account.domain_blocking?(status.account_domain)
   end
 
   def muting_account?
-    @preloaded_relations[:muting] ? @preloaded_relations[:muting][status.account_id] : account.muting?(status.account_id)
+    account.muting?(status.account_id)
   end
 
   def silenced_account?
@@ -60,7 +59,7 @@ class StatusFilter
   end
 
   def account_following_status_account?
-    @preloaded_relations[:following] ? @preloaded_relations[:following][status.account_id] : account&.following?(status.account_id)
+    account&.following?(status.account_id)
   end
 
   def blocked_by_policy?
@@ -68,6 +67,6 @@ class StatusFilter
   end
 
   def policy_allows_show?
-    StatusPolicy.new(account, status, @preloaded_relations).show?
+    StatusPolicy.new(account, status).show?
   end
 end

app/models/concerns/account/interactions.rb

@@ -123,7 +123,11 @@ module Account::Interactions
   end
 
   def following?(other_account)
-    active_relationships.exists?(target_account: other_account)
+    other_id = other_account.is_a?(Account) ? other_account.id : other_account
+
+    preloaded_relation(:following, other_id) do
+      active_relationships.exists?(target_account: other_account)
+    end
   end
 
   def following_anyone?
@@ -139,15 +143,33 @@ module Account::Interactions
   end
 
   def blocking?(other_account)
-    block_relationships.exists?(target_account: other_account)
+    other_id = other_account.is_a?(Account) ? other_account.id : other_account
+
+    preloaded_relation(:blocking, other_id) do
+      block_relationships.exists?(target_account: other_account)
+    end
+  end
+
+  def blocked_by?(other_account)
+    other_id = other_account.is_a?(Account) ? other_account.id : other_account
+
+    preloaded_relation(:blocked_by, other_id) do
+      other_account.block_relationships.exists?(target_account: self)
+    end
   end
 
   def domain_blocking?(other_domain)
-    domain_blocks.exists?(domain: other_domain)
+    preloaded_relation(:domain_blocking_by_domain, other_domain) do
+      domain_blocks.exists?(domain: other_domain)
+    end
   end
 
   def muting?(other_account)
-    mute_relationships.exists?(target_account: other_account)
+    other_id = other_account.is_a?(Account) ? other_account.id : other_account
+
+    preloaded_relation(:muting, other_id) do
+      mute_relationships.exists?(target_account: other_account)
+    end
   end
 
   def muting_conversation?(conversation)
@@ -226,4 +248,10 @@ module Account::Interactions
   def normalized_domain(domain)
     TagManager.instance.normalize_domain(domain)
   end
+
+  private
+
+  def preloaded_relation(type, key)
+    @preloaded_relations && @preloaded_relations[type] ? @preloaded_relations[type][key].present? : yield
+  end
 end

app/models/concerns/account/mappings.rb

@@ -91,6 +91,12 @@ module Account::Mappings
     end
   end
 
+  def preload_relations!(...)
+    @preloaded_relations = relations_map(...)
+  end
+
+  private
+
   def relations_map(account_ids, domains = nil, **options)
     relations = {
       blocked_by: Account.blocked_by_map(account_ids, id),

app/models/concerns/status/interaction_policy_concern.rb

@@ -26,7 +26,7 @@ module Status::InteractionPolicyConcern
   end
 
   # Returns `:automatic`, `:manual`, `:unknown` or `:denied`
-  def quote_policy_for_account(other_account, preloaded_relations: {})
+  def quote_policy_for_account(other_account)
     return :denied if other_account.nil? || direct_visibility? || reblog?
 
     following_author = nil
@@ -41,7 +41,7 @@ module Status::InteractionPolicyConcern
     return :automatic if automatic_policy.anybits?(QUOTE_APPROVAL_POLICY_FLAGS[:public])
 
     if automatic_policy.anybits?(QUOTE_APPROVAL_POLICY_FLAGS[:followers])
-      following_author = preloaded_relations[:following] ? preloaded_relations[:following][account_id] : other_account.following?(account) if following_author.nil?
+      following_author = other_account.following?(account) if following_author.nil?
       return :automatic if following_author
     end
 
@@ -54,7 +54,7 @@ module Status::InteractionPolicyConcern
     return :manual if manual_policy.anybits?(QUOTE_APPROVAL_POLICY_FLAGS[:public])
 
     if manual_policy.anybits?(QUOTE_APPROVAL_POLICY_FLAGS[:followers])
-      following_author = preloaded_relations[:following] ? preloaded_relations[:following][account_id] : other_account.following?(account) if following_author.nil?
+      following_author = other_account.following?(account) if following_author.nil?
       return :manual if following_author
     end
 

app/models/concerns/status/threading_concern.rb

@@ -8,9 +8,10 @@ module Status::ThreadingConcern
       statuses    = Status.with_accounts(ids).to_a
       account_ids = statuses.map(&:account_id).uniq
       domains     = statuses.filter_map(&:account_domain).uniq
-      relations   = account&.relations_map(account_ids, domains) || {}
 
-      statuses.reject! { |status| StatusFilter.new(status, account, relations).filtered? }
+      account&.preload_relations!(account_ids, domains)
+
+      statuses.reject! { |status| StatusFilter.new(status, account).filtered? }
 
       if stable
         statuses.sort_by! { |status| ids.index(status.id) }

app/policies/admin/status_policy.rb

@@ -1,12 +1,6 @@
 # frozen_string_literal: true
 
 class Admin::StatusPolicy < ApplicationPolicy
-  def initialize(current_account, record, preloaded_relations = {})
-    super(current_account, record)
-
-    @preloaded_relations = preloaded_relations
-  end
-
   def index?
     role.can?(:manage_reports, :manage_users)
   end
@@ -34,6 +28,6 @@ class Admin::StatusPolicy < ApplicationPolicy
   end
 
   def viewable_through_normal_policy?
-    StatusPolicy.new(current_account, record, @preloaded_relations).show?
+    StatusPolicy.new(current_account, record).show?
   end
 end

app/policies/status_policy.rb

@@ -1,12 +1,6 @@
 # frozen_string_literal: true
 
 class StatusPolicy < ApplicationPolicy
-  def initialize(current_account, record, preloaded_relations = {})
-    super(current_account, record)
-
-    @preloaded_relations = preloaded_relations
-  end
-
   def show?
     return false if author.unavailable?
 
@@ -21,7 +15,7 @@ class StatusPolicy < ApplicationPolicy
 
   # This is about requesting a quote post, not validating it
   def quote?
-    show? && record.quote_policy_for_account(current_account, preloaded_relations: @preloaded_relations) != :denied
+    show? && record.quote_policy_for_account(current_account) != :denied
   end
 
   def reblog?
@@ -75,19 +69,19 @@ class StatusPolicy < ApplicationPolicy
   def blocking_author?
     return false if current_account.nil?
 
-    @preloaded_relations[:blocking] ? @preloaded_relations[:blocking][author.id] : current_account.blocking?(author)
+    current_account.blocking?(author)
   end
 
   def author_blocking?
     return false if current_account.nil?
 
-    @preloaded_relations[:blocked_by] ? @preloaded_relations[:blocked_by][author.id] : author.blocking?(current_account)
+    current_account.blocked_by?(author)
   end
 
   def following_author?
     return false if current_account.nil?
 
-    @preloaded_relations[:following] ? @preloaded_relations[:following][author.id] : current_account.following?(author)
+    current_account.following?(author)
   end
 
   def author

app/presenters/status_relationships_presenter.rb

@@ -41,14 +41,8 @@ class StatusRelationshipsPresenter
   end
 
   # This one is currently on-demand as it is only used for quote posts
-  def preloaded_account_relations
-    @preloaded_account_relations ||= begin
-      accounts = @statuses.compact.flat_map { |s| [s.account, s.proper.account, s.proper.quote&.quoted_account] }.uniq.compact
-
-      account_ids = accounts.pluck(:id)
-      account_domains = accounts.pluck(:domain).uniq
-      Account.find(@current_account_id).relations_map(account_ids, account_domains)
-    end
+  def authoring_accounts
+    @authoring_accounts ||= @statuses.compact.flat_map { |s| [s.account, s.proper.account, s.proper.quote&.quoted_account] }.uniq.compact
   end
 
   private

app/serializers/rest/base_quote_serializer.rb

@@ -19,6 +19,13 @@ class REST::BaseQuoteSerializer < ActiveModel::Serializer
   private
 
   def status_filter
-    @status_filter ||= StatusFilter.new(object.quoted_status, current_user&.account, instance_options[:relationships]&.preloaded_account_relations || {})
+    @status_filter ||= begin
+      if current_user && instance_options[:relationships]
+        account_ids = instance_options[:relationships].authoring_accounts.pluck(:id)
+        domains = instance_options[:relationships].authoring_accounts.pluck(:domain).uniq
+        current_user.account.preload_relations!(account_ids, domains)
+      end
+      StatusFilter.new(object.quoted_status, current_user&.account)
+    end
   end
 end

app/services/statuses_search_service.rb

@@ -29,9 +29,10 @@ class StatusesSearchService < BaseService
     results             = request.collapse(field: :id).order(id: { order: :desc }).limit(@limit).offset(@offset).objects.compact
     account_ids         = results.map(&:account_id)
     account_domains     = results.map(&:account_domain)
-    preloaded_relations = @account.relations_map(account_ids, account_domains)
 
-    results.reject { |status| StatusFilter.new(status, @account, preloaded_relations).filtered? }
+    @account.preload_relations!(account_ids, account_domains)
+
+    results.reject { |status| StatusFilter.new(status, @account).filtered? }
   rescue Faraday::ConnectionFailed, Parslet::ParseFailed, Errno::ENETUNREACH
     []
   end

spec/models/concerns/account/interactions_spec.rb

@@ -302,9 +302,24 @@ RSpec.describe Account::Interactions do
     subject { account.following?(target_account) }
 
     context 'when following target_account' do
-      it 'returns true' do
+      before do
         account.active_relationships.create(target_account: target_account)
-        expect(subject).to be true
+      end
+
+      it 'returns true' do
+        result = nil
+        expect { result = subject }.to execute_queries
+        expect(result).to be true
+      end
+
+      context 'when relations are preloaded' do
+        it 'does not query the database to get the result' do
+          account.preload_relations!([target_account.id])
+
+          result = nil
+          expect { result = subject }.to_not execute_queries
+          expect(result).to be true
+        end
       end
     end
 
@@ -336,9 +351,26 @@ RSpec.describe Account::Interactions do
     subject { account.blocking?(target_account) }
 
     context 'when blocking target_account' do
-      it 'returns true' do
+      before do
         account.block_relationships.create(target_account: target_account)
-        expect(subject).to be true
+      end
+
+      it 'returns true' do
+        result = nil
+        expect { result = subject }.to execute_queries
+
+        expect(result).to be true
+      end
+
+      context 'when relations are preloaded' do
+        it 'does not query the database to get the result' do
+          account.preload_relations!([target_account.id])
+
+          result = nil
+          expect { result = subject }.to_not execute_queries
+
+          expect(result).to be true
+        end
       end
     end
 
@@ -349,16 +381,65 @@ RSpec.describe Account::Interactions do
     end
   end
 
+  describe '#blocked_by?' do
+    subject { account.blocked_by?(target_account) }
+
+    context 'when blocked by target_account' do
+      before do
+        target_account.block_relationships.create(target_account: account)
+      end
+
+      it 'returns true' do
+        result = nil
+        expect { result = subject }.to execute_queries
+
+        expect(result).to be true
+      end
+
+      context 'when relations are preloaded' do
+        it 'does not query the database to get the result' do
+          account.preload_relations!([target_account.id])
+
+          result = nil
+          expect { result = subject }.to_not execute_queries
+
+          expect(result).to be true
+        end
+      end
+    end
+
+    context 'when not blocked by target_account' do
+      it 'returns false' do
+        expect(subject).to be false
+      end
+    end
+  end
+
   describe '#domain_blocking?' do
     subject { account.domain_blocking?(domain) }
 
     let(:domain) { 'example.com' }
 
     context 'when blocking the domain' do
-      it 'returns true' do
+      before do
         account_domain_block = Fabricate(:account_domain_block, domain: domain)
         account.domain_blocks << account_domain_block
-        expect(subject).to be true
+      end
+
+      it 'returns true' do
+        result = nil
+        expect { result = subject }.to execute_queries
+        expect(result).to be true
+      end
+
+      context 'when relations are preloaded' do
+        it 'does not query the database to get the result' do
+          account.preload_relations!([], [domain])
+
+          result = nil
+          expect { result = subject }.to_not execute_queries
+          expect(result).to be true
+        end
       end
     end
 
@@ -373,10 +454,25 @@ RSpec.describe Account::Interactions do
     subject { account.muting?(target_account) }
 
     context 'when muting target_account' do
-      it 'returns true' do
+      before do
         mute = Fabricate(:mute, account: account, target_account: target_account)
         account.mute_relationships << mute
-        expect(subject).to be true
+      end
+
+      it 'returns true' do
+        result = nil
+        expect { result = subject }.to execute_queries
+        expect(result).to be true
+      end
+
+      context 'when relations are preloaded' do
+        it 'does not query the database to get the result' do
+          account.preload_relations!([target_account.id])
+
+          result = nil
+          expect { result = subject }.to_not execute_queries
+          expect(result).to be true
+        end
       end
     end
 

spec/support/matchers/sql_queries.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+require 'active_record/testing/query_assertions'
+
+# Implement something similar to Rails' built-in assertion.
+# Can be removed once https://github.com/rspec/rspec-rails/pull/2818
+# has been merged and released.
+RSpec::Matchers.define :execute_queries do |expected = nil|
+  match do |actual|
+    counter = ActiveRecord::Assertions::QueryAssertions::SQLCounter.new
+
+    queries = ActiveSupport::Notifications.subscribed(counter, 'sql.active_record') do
+      actual.call
+      counter.log
+    end
+
+    if expected.nil?
+      queries.count >= 1
+    else
+      queries.count == expected
+    end
+  end
+
+  supports_block_expectations
+end