From 4ed2affa4ca4efb4bdc9e4e8408f15e553d055b1 Mon Sep 17 00:00:00 2001
From: Jacob Burroughs <jburroughs@instructure.com>
Date: Tue, 7 Jun 2022 17:47:05 -0500
Subject: [PATCH] Fix cross-shard message delivery

fixes VICE-2848

Change-Id: I14e4468d00201e7b4409e9752fbf3ae11e9ee2dd
Reviewed-on: https://gerrit.instructure.com/c/canvas-lms/+/293449
Tested-by: Service Cloud Jenkins <svc.cloudjenkins@instructure.com>
Reviewed-by: Drake Harper <drake.harper@instructure.com>
QA-Review: Drake Harper <drake.harper@instructure.com>
Product-Review: Drake Harper <drake.harper@instructure.com>
---
 gems/canvas_partman/canvas_partman.gemspec    |  2 +-
 .../canvas_partman/concerns/partitioned.rb    | 28 ++++++-------------
 spec/lib/message_dispatcher_spec.rb           | 24 ++++++++++++++++
 3 files changed, 34 insertions(+), 20 deletions(-)

diff --git a/gems/canvas_partman/canvas_partman.gemspec b/gems/canvas_partman/canvas_partman.gemspec
index 3412c47a0b1..2757e6e6c35 100644
--- a/gems/canvas_partman/canvas_partman.gemspec
+++ b/gems/canvas_partman/canvas_partman.gemspec
@@ -13,7 +13,7 @@ Gem::Specification.new do |spec|
   spec.require_paths = ["lib"]
   spec.license       = "AGPL"
 
-  spec.add_dependency "activerecord", ">= 3.2", "< 7.1"
+  spec.add_dependency "activerecord", ">= 6.1", "< 7.1"
   spec.add_dependency "activerecord-pg-extensions", "~> 0.4"
   spec.add_dependency "pg", ">= 0.17", "< 2.0"
 
diff --git a/gems/canvas_partman/lib/canvas_partman/concerns/partitioned.rb b/gems/canvas_partman/lib/canvas_partman/concerns/partitioned.rb
index 4a840ff305b..5cd48e0325d 100644
--- a/gems/canvas_partman/lib/canvas_partman/concerns/partitioned.rb
+++ b/gems/canvas_partman/lib/canvas_partman/concerns/partitioned.rb
@@ -126,20 +126,14 @@ module CanvasPartman::Concerns
       end
 
       def _insert_record(values)
-        if ::ActiveRecord.version >= Gem::Version.new("5.2")
-          begin
-            prev_table = @arel_table
-            prev_builder = @predicate_builder
-            @arel_table = arel_table_from_key_values(values)
-            @predicate_builder = nil
-            super
-          ensure
-            @arel_table = prev_table
-            @predicate_builder = prev_builder
-          end
-        else
-          super
-        end
+        prev_table = @arel_table
+        prev_builder = @predicate_builder
+        @arel_table = arel_table_from_key_values(values)
+        @predicate_builder = nil
+        super
+      ensure
+        @arel_table = prev_table
+        @predicate_builder = prev_builder
       end
 
       # :nodoc:
@@ -147,11 +141,7 @@ module CanvasPartman::Concerns
         partition_table_name = infer_partition_table_name(attributes)
 
         @arel_tables ||= {}
-        @arel_tables[partition_table_name] ||= if ::ActiveRecord.version < Gem::Version.new("5")
-                                                 Arel::Table.new(partition_table_name, { engine: arel_engine })
-                                               else
-                                                 Arel::Table.new(partition_table_name, type_caster: type_caster)
-                                               end
+        @arel_tables[partition_table_name] ||= Arel::Table.new(partition_table_name, klass: self)
       end
 
       # @internal
diff --git a/spec/lib/message_dispatcher_spec.rb b/spec/lib/message_dispatcher_spec.rb
index 7a4d9009459..f7a9914af83 100644
--- a/spec/lib/message_dispatcher_spec.rb
+++ b/spec/lib/message_dispatcher_spec.rb
@@ -84,5 +84,29 @@ describe "MessageDispatcher" do
       expect(job2.payload_object.message).to eq @messages[1]
       expect(job2.run_at.to_i).to eq @messages[1].dispatch_at.to_i
     end
+
+    describe "cross_shard" do
+      specs_require_sharding
+
+      before do
+        @shard1.activate do
+          @messages += (0...3).map { message_model(dispatch_at: Time.now, workflow_state: "staged", to: "somebody", updated_at: Time.now.utc - 11.minutes, user: user_factory, path_type: "email") }
+        end
+      end
+
+      it "loads all the matches" do
+        track_jobs { MessageDispatcher.batch_dispatch(@messages) }
+        expect(created_jobs.size).to eq 1
+        job = created_jobs.first
+
+        am_message = double
+        allow(am_message).to receive(:deliver_now).and_return(true)
+        allow(Mailer).to receive(:create_message).and_return(am_message)
+
+        track_jobs { Delayed::Worker.new.perform(job) }
+        @messages.each(&:reload)
+        expect(@messages.map(&:state)).to eq %i[sent sent sent sent sent sent]
+      end
+    end
   end
 end