Set up replica identity index for EnrollmentState

also moves the RootAccountResolver specs out of enrollment_state_spec.rb
and into a more generic root_account_resolver_spec.rb

refs FOO-1171
flag=none

test plan:
 - migration works up and down
 - tests pass

Change-Id: Ia3ee89d63974e4c787ca8d1c6ab80eba5f54beb4
Reviewed-on: https://gerrit.instructure.com/c/canvas-lms/+/260526
Tested-by: Service Cloud Jenkins <svc.cloudjenkins@instructure.com>
Reviewed-by: Simon Williams <simon@instructure.com>
QA-Review: Michael Ziwisky <mziwisky@instructure.com>
Product-Review: Michael Ziwisky <mziwisky@instructure.com>

db/migrate/20210312061721_add_replica_identity_for_enrollment_states.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+#
+# Copyright (C) 2021 - present Instructure, Inc.
+#
+# This file is part of Canvas.
+#
+# Canvas is free software: you can redistribute it and/or modify it under
+# the terms of the GNU Affero General Public License as published by the Free
+# Software Foundation, version 3 of the License.
+#
+# Canvas is distributed in the hope that it will be useful, but WITHOUT ANY
+# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
+# A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
+# details.
+#
+# You should have received a copy of the GNU Affero General Public License along
+# with this program. If not, see <http://www.gnu.org/licenses/>.
+
+class AddReplicaIdentityForEnrollmentStates < ActiveRecord::Migration[6.0]
+  tag :postdeploy
+  disable_ddl_transaction!
+
+  def up
+    add_replica_identity 'EnrollmentState', :root_account_id, 0
+    remove_index :enrollment_states, column: :root_account_id, if_exists: true
+  end
+
+  def down
+    add_index :enrollment_states, :root_account_id, algorithm: :concurrently, if_not_exists: true
+    remove_replica_identity 'EnrollmentState'
+    change_column_null :enrollment_states, :root_account_id, true
+  end
+end

lib/data_fixup/backfill_nulls.rb

@@ -51,7 +51,7 @@ module DataFixup
       end
 
       klass.find_ids_in_ranges(batch_size: batch_size) do |start_id, end_id|
-        update_count = scope.where(id: start_id..end_id).update_all(updates)
+        update_count = scope.where(klass.primary_key => start_id..end_id).update_all(updates)
         sleep_interval_per_batch = Setting.get("sleep_interval_per_backfill_nulls_batch", nil).presence&.to_f
         sleep(sleep_interval_per_batch) if update_count > 0 && sleep_interval_per_batch
       end

spec/factories/course_factory.rb

@@ -162,8 +162,7 @@ module Factories
       type = options[:enrollment_type] || "TeacherEnrollment"
       role_id = Role.get_built_in_role(type, root_account_id: account.resolved_root_account_id).id
       result = create_records(Enrollment, course_ids.each_with_index.map{ |id, i| {course_id: id, user_id: user.id, type: type, course_section_id: section_ids[i], root_account_id: account.id, workflow_state: 'active', :role_id => role_id}})
-      create_enrollment_states(result, {state: "active"})
-      result
+      create_enrollment_states(result, {state: "active", root_account_id: account.id})
     end
     course_data
   end

spec/factories/enrollment_factory.rb

@@ -78,7 +78,7 @@ module Factories
         limit_privileges_to_course_section: limit_privileges_to_course_section
       }
     }, options[:return_type])
-    create_enrollment_states(result, {state: enrollment_state})
+    create_enrollment_states(result, {state: enrollment_state, root_account_id: course.root_account_id})
     result
   end
 end

spec/lib/data_fixup/populate_root_account_id_on_models_spec.rb

@@ -412,13 +412,6 @@ describe DataFixup::PopulateRootAccountIdOnModels do
       end
     end
 
-    context 'with EnrollmentState' do
-      it_behaves_like 'a datafixup that populates root_account_id' do
-        let(:record) { reference_record.enrollment_state }
-        let(:reference_record) { enrollment_model }
-      end
-    end
-
     context 'with GradingPeriod' do
       it_behaves_like 'a datafixup that populates root_account_id' do
         let(:record) { grading_periods(count: 1).first }

spec/models/enrollment_state_spec.rb

@@ -20,43 +20,6 @@
 require_relative "../spec_helper"
 
 describe EnrollmentState do
-  describe "root_account_id" do
-    let(:course) { course_factory && @course }
-    let(:enrollment) { student_in_course(course: course) }
-
-    it 'assigns it on save if it is not set' do
-      enrollment.enrollment_state.root_account_id = nil
-
-      expect {
-        enrollment.enrollment_state.save!
-      }.to change {
-        enrollment.enrollment_state.root_account_id
-      }.from(nil).to(enrollment.root_account_id)
-    end
-
-    it 'preserves it on save if it was already set' do
-      expect(enrollment).not_to receive(:root_account_id)
-
-      expect {
-        enrollment.enrollment_state.save!
-      }.not_to change {
-        EnrollmentState.find_by(enrollment_id: enrollment.id).root_account_id
-      }
-    end
-
-    it 'does nothing on save if it is not set and could not be resolved' do
-      enrollment.enrollment_state.update_column(:root_account_id, nil)
-
-      expect(enrollment).to receive(:root_account_id).and_return(nil)
-
-      expect {
-        enrollment.enrollment_state.save!
-      }.not_to change {
-        EnrollmentState.find(enrollment.enrollment_state.id).root_account_id
-      }
-    end
-  end
-
   describe "#enrollments_needing_calculation" do
     it "should find enrollments that need calculation" do
       course_factory

spec/models/root_account_resolver_spec.rb

@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+#
+# Copyright (C) 2021 - present Instructure, Inc.
+#
+# This file is part of Canvas.
+#
+# Canvas is free software: you can redistribute it and/or modify it under
+# the terms of the GNU Affero General Public License as published by the Free
+# Software Foundation, version 3 of the License.
+#
+# Canvas is distributed in the hope that it will be useful, but WITHOUT ANY
+# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
+# A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
+# details.
+#
+# You should have received a copy of the GNU Affero General Public License along
+# with this program. If not, see <http://www.gnu.org/licenses/>.
+
+require_relative "../spec_helper"
+
+describe RootAccountResolver do
+  let(:minimal_active_record) do
+    Class.new do
+      extend ActiveModel::Callbacks
+      define_model_callbacks :save
+
+      attr_accessor :root_account_id
+
+      def save
+        run_callbacks :save
+      end
+
+      def attributes
+        { 'root_account_id' => root_account_id }
+      end
+
+      def self.belongs_to(relation, **opts)
+      end
+    end
+  end
+  let(:test_class) do
+    Class.new(minimal_active_record) do
+      extend RootAccountResolver
+
+      def initialize(root_account_id)
+        @root_account_id = root_account_id
+      end
+
+      def knower_of_account_1337
+        Struct.new(:root_account_id).new(1337)
+      end
+    end
+  end
+
+  describe "resolution via member" do
+    before do
+      test_class.resolves_root_account through: :knower_of_account_1337
+    end
+
+    it 'assigns it on save if it is not set' do
+      subject = test_class.new(nil)
+
+      expect {
+        subject.save
+      }.to change {
+        subject.root_account_id
+      }.from(nil).to(1337)
+    end
+
+    it 'preserves it on save if it was already set' do
+      subject = test_class.new(666)
+
+      expect {
+        subject.save
+      }.not_to change {
+        subject.root_account_id
+      }
+    end
+  end
+
+  describe "resolution via proc" do
+    before do
+      test_class.resolves_root_account through: ->(instance) { 9000 }
+    end
+
+    it 'assigns it on save if it is not set' do
+      subject = test_class.new(nil)
+
+      expect {
+        subject.save
+      }.to change {
+        subject.root_account_id
+      }.from(nil).to(9000)
+    end
+
+    it 'preserves it on save if it was already set' do
+      subject = test_class.new(666)
+
+      expect {
+        subject.save
+      }.not_to change {
+        subject.root_account_id
+      }
+    end
+  end
+end