Reduce N+1s by caching submitted project count (Part 1 of 2)

app/models/lesson.rb

@@ -8,6 +8,8 @@ class Lesson < ApplicationRecord
   belongs_to :parent, optional: true, class_name: :Lesson, foreign_key: :copied_from_id, inverse_of: :copies
   has_many :copies, dependent: :nullify, class_name: :Lesson, foreign_key: :copied_from_id, inverse_of: :parent
   has_one :project, dependent: :destroy
+  has_many :remixes, through: :project
+  has_many :school_projects, through: :remixes
   accepts_nested_attributes_for :project
 
   before_validation :assign_school_from_school_class
@@ -38,6 +40,13 @@ def submitted_count
     project.remixes.count { |remix| remix.school_project&.submitted? }
   end
 
+  def recalculate_submitted_projects_count!
+    with_lock do
+      count = school_projects.in_state(:submitted).count
+      update!(submitted_projects_count: count)
+    end
+  end
+
   private
 
   def assign_school_from_school_class

app/models/school_project.rb

@@ -11,6 +11,10 @@ class SchoolProject < ApplicationRecord
     initial_state: :unsubmitted
   ]
 
+  def lesson
+    project.lesson || project.parent&.lesson
+  end
+
   def status
     state_machine.current_state
   end
@@ -23,6 +27,10 @@ def unread_feedback?
     feedback.exists?(read_at: nil)
   end
 
+  def recalculate_lesson_submitted_projects_count!(_transition = nil)
+    lesson&.recalculate_submitted_projects_count!
+  end
+
   # Add convenience methods for each state
   def unsubmitted?
     state_machine.in_state?(:unsubmitted)

app/state_machines/school_project_state_machine.rb

@@ -14,4 +14,7 @@ class SchoolProjectStateMachine
   transition from: :submitted, to: %i[unsubmitted returned complete]
   transition from: :returned, to: %i[submitted complete]
   transition from: :complete, to: [:unsubmitted]
+
+  after_transition(to: :submitted, &:recalculate_lesson_submitted_projects_count!)
+  after_transition(from: :submitted, &:recalculate_lesson_submitted_projects_count!)
 end

db/migrate/20260429120000_add_submitted_projects_count_to_lessons.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+class AddSubmittedProjectsCountToLessons < ActiveRecord::Migration[7.2]
+  def change
+    add_column :lessons, :submitted_projects_count, :integer, null: false, default: 0
+  end
+end

lib/tasks/lessons.rake

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+namespace :lessons do
+  desc 'Backfill cached submitted project counts for lessons'
+  task backfill_submitted_projects_count: :environment do
+    Lesson.connection.execute <<~SQL.squish
+      WITH submitted_counts AS (
+        SELECT lesson_projects.lesson_id, COUNT(*) AS submitted_projects_count
+        FROM projects lesson_projects
+        INNER JOIN projects remixes ON remixes.remixed_from_id = lesson_projects.id
+        INNER JOIN school_projects ON school_projects.project_id = remixes.id
+        INNER JOIN school_project_transitions ON school_project_transitions.school_project_id = school_projects.id
+        WHERE school_project_transitions.most_recent = TRUE
+          AND school_project_transitions.to_state = 'submitted'
+          AND lesson_projects.lesson_id IS NOT NULL
+        GROUP BY lesson_projects.lesson_id
+      )
+      UPDATE lessons
+      SET submitted_projects_count = COALESCE(submitted_counts.submitted_projects_count, 0)
+      FROM lessons target_lessons
+      LEFT JOIN submitted_counts ON submitted_counts.lesson_id = target_lessons.id
+      WHERE lessons.id = target_lessons.id
+    SQL
+  end
+end

spec/lib/tasks/lessons_spec.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+require 'rails_helper'
+require 'rake'
+
+RSpec.describe 'lessons', type: :task do
+  describe ':backfill_submitted_projects_count' do
+    let(:task) { Rake::Task['lessons:backfill_submitted_projects_count'] }
+    let(:school) { create(:school) }
+    let(:teacher) { create(:teacher, school:) }
+    let(:student) { create(:student, school:) }
+    let(:lesson) { create(:lesson, school:, user_id: teacher.id) }
+
+    before do
+      task.reenable
+    end
+
+    it 'sets cached submitted project counts for all lessons' do
+      submitted_remix = create(:project, school:, remixed_from_id: lesson.project.id, user_id: student.id)
+      submitted_remix.school_project.transition_status_to!(:submitted, student.id)
+
+      returned_remix = create(:project, school:, remixed_from_id: lesson.project.id, user_id: student.id)
+      returned_remix.school_project.transition_status_to!(:submitted, student.id)
+      returned_remix.school_project.transition_status_to!(:returned, teacher.id)
+
+      other_lesson = create(:lesson, school:, user_id: teacher.id, submitted_projects_count: 7)
+
+      lesson.update!(submitted_projects_count: 0)
+
+      task.invoke
+
+      expect(lesson.reload.submitted_projects_count).to eq(1)
+      expect(other_lesson.reload.submitted_projects_count).to eq(0)
+    end
+  end
+end

spec/models/lesson_spec.rb

@@ -227,4 +227,32 @@
       expect(lesson.submitted_count).to eq(2)
     end
   end
+
+  describe '#recalculate_submitted_projects_count!' do
+    it 'sets the submitted projects count to 0 if there is no project' do
+      lesson = create(:lesson, project: nil, submitted_projects_count: 3)
+
+      lesson.recalculate_submitted_projects_count!
+
+      expect(lesson.reload.submitted_projects_count).to eq(0)
+    end
+
+    it 'returns the count of submitted remixes of the lesson project' do
+      student = create(:student, school:)
+      lesson = create(:lesson, school:, user_id: teacher.id)
+
+      remix_1 = create(:project, school:, remixed_from_id: lesson.project.id, user_id: student.id)
+      remix_1.school_project.transition_status_to!(:submitted, remix_1.user_id)
+
+      remix_2 = create(:project, school:, remixed_from_id: lesson.project.id, user_id: student.id)
+      remix_2.school_project.transition_status_to!(:submitted, remix_2.user_id)
+
+      create(:project, school:, remixed_from_id: lesson.project.id, user_id: student.id) # Not submitted
+
+      lesson.update!(submitted_projects_count: 0)
+      lesson.recalculate_submitted_projects_count!
+
+      expect(lesson.reload.submitted_projects_count).to eq(2)
+    end
+  end
 end

spec/state_machines/school_project_state_machine_spec.rb

@@ -16,6 +16,21 @@
   end
 
   describe 'transitions' do
+    it 'recalculates the parent lesson submitted projects count' do
+      teacher = create(:teacher, school:)
+      lesson = create(:lesson, school:, user_id: teacher.id)
+      remix = create(:project, school:, user_id: student.id, remixed_from_id: lesson.project.id)
+      state_machine = described_class.new(remix.school_project, transition_class: SchoolProjectTransition)
+
+      expect do
+        state_machine.transition_to!(:submitted)
+      end.to change { lesson.reload.submitted_projects_count }.from(0).to(1)
+
+      expect do
+        state_machine.transition_to!(:returned)
+      end.to change { lesson.reload.submitted_projects_count }.from(1).to(0)
+    end
+
     context 'when in unsubmitted state' do
       it 'can transition to submitted' do
         expect(state_machine.can_transition_to?(:submitted)).to be true