PERF: N+1 when assignments status are enabled (#614)

This properly preload assignment statuses on topic lists when they're enabled to prevent N+1.

Internal ref - t/142804

plugin.rb

@@ -157,7 +157,7 @@ after_initialize do
     Site.preloaded_category_custom_fields << "enable_unassigned_filter"
   end
 
-  BookmarkQuery.on_preload do |bookmarks, bookmark_query|
+  BookmarkQuery.on_preload do |bookmarks, _bookmark_query|
     if SiteSetting.assign_enabled?
       topics =
         Bookmark
@@ -173,8 +173,10 @@ after_initialize do
           .index_by(&:topic_id)
 
       topics.each do |topic|
-        assigned_to = assignments[topic.id]&.assigned_to
+        assignment = assignments[topic.id]
-        topic.preload_assigned_to(assigned_to)
+        # NOTE: preloading to `nil` is necessary to avoid N+1 queries
+        topic.preload_assigned_to(assignment&.assigned_to)
+        topic.preload_assignment_status(assignment&.status)
       end
     end
   end
@@ -184,49 +186,46 @@ after_initialize do
   end
 
   TopicList.on_preload do |topics, topic_list|
-    if SiteSetting.assign_enabled?
+    next unless SiteSetting.assign_enabled?
-      can_assign = topic_list.current_user && topic_list.current_user.can_assign?
+
+    can_assign = topic_list.current_user&.can_assign?
     allowed_access = SiteSetting.assigns_public || can_assign
 
-      if allowed_access && topics.length > 0
+    next if !allowed_access || topics.empty?
+
     assignments =
       Assignment.strict_loading.active.where(topic: topics).includes(:target, :assigned_to)
     assignments_map = assignments.group_by(&:topic_id)
 
-        user_ids =
+    user_ids = assignments.filter(&:assigned_to_user?).map(&:assigned_to_id)
-          assignments.filter { |assignment| assignment.assigned_to_user? }.map(&:assigned_to_id)
     users_map = User.where(id: user_ids).select(UserLookup.lookup_columns).index_by(&:id)
 
-        group_ids =
+    group_ids = assignments.filter(&:assigned_to_group?).map(&:assigned_to_id)
-          assignments.filter { |assignment| assignment.assigned_to_group? }.map(&:assigned_to_id)
     groups_map = Group.where(id: group_ids).index_by(&:id)
 
     topics.each do |topic|
-          assignments = assignments_map[topic.id]
+      if assignments = assignments_map[topic.id]
-          direct_assignment =
+        topic_assignments, post_assignments = assignments.partition { _1.target_type == "Topic" }
-            assignments&.find do |assignment|
+
-              assignment.target_type == "Topic" && assignment.target_id == topic.id
+        direct_assignment = topic_assignments.find { _1.target_id == topic.id }
-            end
+
         indirectly_assigned_to = {}
-          assignments
+
-            &.each do |assignment|
+        post_assignments.each do |assignment|
-              next if assignment.target_type == "Topic"
+          next unless assignment.target
-              next if !assignment.target
+
-              next(
+          if assignment.assigned_to_user?
             indirectly_assigned_to[assignment.target_id] = {
               assigned_to: users_map[assignment.assigned_to_id],
               post_number: assignment.target.post_number,
             }
-              ) if assignment&.assigned_to_user?
+          elsif assignment.assigned_to_group?
-              next(
             indirectly_assigned_to[assignment.target_id] = {
               assigned_to: groups_map[assignment.assigned_to_id],
               post_number: assignment.target.post_number,
             }
-              ) if assignment&.assigned_to_group?
           end
-            &.compact
+        end
-            &.uniq
 
         assigned_to =
           if direct_assignment&.assigned_to_user?
@@ -234,53 +233,57 @@ after_initialize do
           elsif direct_assignment&.assigned_to_group?
             groups_map[direct_assignment.assigned_to_id]
           end
+      end
+
+      # NOTE: preloading to `nil` is necessary to avoid N+1 queries
       topic.preload_assigned_to(assigned_to)
+      topic.preload_assignment_status(direct_assignment&.status)
       topic.preload_indirectly_assigned_to(indirectly_assigned_to)
     end
   end
-    end
-  end
 
   Search.on_preload do |results, search|
-    if SiteSetting.assign_enabled?
+    next unless SiteSetting.assign_enabled?
+
     can_assign = search.guardian&.can_assign?
     allowed_access = SiteSetting.assigns_public || can_assign
 
-      if allowed_access && results.posts.length > 0
+    next if !allowed_access || results.posts.empty?
+
     topics = results.posts.map(&:topic)
+
     assignments =
       Assignment
         .strict_loading
-            .where(topic: topics, active: true)
+        .active
+        .where(topic: topics)
         .includes(:assigned_to, :target)
         .group_by(&:topic_id)
 
     results.posts.each do |post|
-          topic_assignments = assignments[post.topic.id]
+      if topic_assignments = assignments[post.topic_id]
-          direct_assignment =
+        direct_assignment = topic_assignments.find { _1.target_type == "Topic" }
-            topic_assignments&.find { |assignment| assignment.target_type == "Topic" }
+        indirect_assignments = topic_assignments.select { _1.target_type == "Post" }
-          indirect_assignments =
+      end
-            topic_assignments&.select { |assignment| assignment.target_type == "Post" }
 
-          post.topic.preload_assigned_to(direct_assignment&.assigned_to)
-          post.topic.preload_indirectly_assigned_to(nil)
       if indirect_assignments.present?
-            indirect_assignment_map =
+        indirect_assignment_map = {}
-              indirect_assignments.reduce({}) do |acc, assignment|
+
-                if assignment.target
+        indirect_assignments.each do |assignment|
-                  acc[assignment.target_id] = {
+          next unless assignment.target
+          indirect_assignment_map[assignment.target_id] = {
             assigned_to: assignment.assigned_to,
             post_number: assignment.target.post_number,
           }
         end
-                acc
       end
+
+      # NOTE: preloading to `nil` is necessary to avoid N+1 queries
+      post.topic.preload_assigned_to(direct_assignment&.assigned_to)
+      post.topic.preload_assignment_status(direct_assignment&.status)
       post.topic.preload_indirectly_assigned_to(indirect_assignment_map)
     end
   end
-      end
-    end
-  end
 
   # TopicQuery
   TopicQuery.add_custom_filter(:assigned) do |results, topic_query|
@@ -442,6 +445,11 @@ after_initialize do
     @assigned_to = assignment.assigned_to if assignment&.active
   end
 
+  add_to_class(:topic, :assignment_status) do
+    return @assignment_status if defined?(@assignment_status)
+    @assignment_status = assignment.status if SiteSetting.enable_assign_status && assignment&.active
+  end
+
   add_to_class(:topic, :indirectly_assigned_to) do
     return @indirectly_assigned_to if defined?(@indirectly_assigned_to)
     @indirectly_assigned_to =
@@ -465,6 +473,10 @@ after_initialize do
 
   add_to_class(:topic, :preload_assigned_to) { |assigned_to| @assigned_to = assigned_to }
 
+  add_to_class(:topic, :preload_assignment_status) do |assignment_status|
+    @assignment_status = assignment_status
+  end
+
   add_to_class(:topic, :preload_indirectly_assigned_to) do |indirectly_assigned_to|
     @indirectly_assigned_to = indirectly_assigned_to
   end
@@ -531,9 +543,9 @@ after_initialize do
     :assignment_status,
     include_condition: -> do
       SiteSetting.enable_assign_status && (SiteSetting.assigns_public || scope.can_assign?) &&
-        object.topic.assignment.present?
+        object.topic.assignment_status.present?
     end,
-  ) { object.topic.assignment.status }
+  ) { object.topic.assignment_status }
 
   # SuggestedTopic serializer
   add_to_serializer(
@@ -590,9 +602,9 @@ after_initialize do
     :assignment_status,
     include_condition: -> do
       SiteSetting.enable_assign_status && (SiteSetting.assigns_public || scope.can_assign?) &&
-        object.assignment.present?
+        object.assignment_status.present?
     end,
-  ) { object.assignment.status }
+  ) { object.assignment_status }
 
   # SearchTopicListItem serializer
   add_to_serializer(