graphql: make loaders shard-aware

closes RECNVS-376

Test plan:
  * have accounts/courses on multiple shards
  * make sure you can retrieve objects from multiple shards in one query

Change-Id: I8428345e741ed5198c29d32fb70d57248ec2aca4
Reviewed-on: https://gerrit.instructure.com/145423
Reviewed-by: Michael Jasper <mjasper@instructure.com>
Tested-by: Jenkins
QA-Review: Collin Parrish <cparrish@instructure.com>
Product-Review: Cameron Matheson <cameron@instructure.com>

app/graphql/loaders/foreign_key_loader.rb

@@ -22,14 +22,20 @@ class Loaders::ForeignKeyLoader < GraphQL::Batch::Loader
     @column = fk
   end
 
-  def load(key)
-    super(key.to_s)
+  def load(id)
+    super(Shard.global_id_for(id))
   end
 
   def perform(ids)
-    records = @scope.where(@column => ids).group_by { |o| o.send(@column).to_s }
+    Shard.partition_by_shard(ids) { |sharded_ids|
+      @scope.where(@column => sharded_ids).
+        group_by { |o| o.send(@column).to_s }.
+        each { |id, os|
+          fulfill(Shard.global_id_for(id), os)
+        }
+    }
     ids.each { |id|
-      fulfill(id, records[id])
+      fulfill(id, nil) unless fulfilled?(id)
     }
   end
 end

app/graphql/loaders/id_loader.rb

@@ -21,14 +21,16 @@ class Loaders::IDLoader < GraphQL::Batch::Loader
     @scope = scope
   end
 
-  def load(key)
-    # since we might load an id that is a number or a string, we need to coerce
-    # here to keep things consistent
-    super(key.to_s)
+  def load(id)
+    super(Shard.global_id_for(id))
   end
 
   def perform(ids)
-    @scope.where(id: ids).each { |o| fulfill(o.id.to_s, o) }
+    Shard.partition_by_shard(ids) { |sharded_ids|
+      @scope.where(id: sharded_ids).each { |o|
+        fulfill(o.global_id, o)
+      }
+    }
     ids.each { |id| fulfill(id, nil) unless fulfilled?(id) }
   end
 end

spec/graphql/loaders/foreign_key_loader_spec.rb

@@ -0,0 +1,80 @@
+#
+# Copyright (C) 2018 - 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 File.expand_path(File.dirname(__FILE__) + '/../../spec_helper')
+
+describe Loaders::ForeignKeyLoader do
+  it "works" do
+    course_with_student(active_all: true)
+    GraphQL::Batch.batch do
+      enrollments_loader = Loaders::ForeignKeyLoader.for(Enrollment, :course_id)
+      enrollments_loader.load(@course.id).then { |enrollments|
+        expect(enrollments).to match_array [@teacher.enrollments.first, @student.enrollments.first]
+      }
+      enrollments_loader.load(-1).then { |enrollments|
+        expect(enrollments).to be_nil
+      }
+    end
+  end
+
+  context "multiple shards" do
+    specs_require_sharding
+
+    before(:once) do
+      course_with_student(active_all: true)
+      @shard_a_course = @course
+      @shard_a_student = @student
+      @shard_a_teacher = @teacher
+
+      @shard1.activate {
+        shard_b_account = Account.create! name: "shard b account"
+        course_with_student(active_all: true, account: shard_b_account)
+        @shard_b_course = @course
+        @shard_b_student = @student
+        @shard_b_teacher = @teacher
+      }
+    end
+
+    it "works across multiple shards" do
+      GraphQL::Batch.batch do
+        enrollments_loader = Loaders::ForeignKeyLoader.for(Enrollment, :course_id)
+        enrollments_loader.load(@shard_a_course.id).then { |enrollments|
+          expect(enrollments).to match_array [@shard_a_teacher.enrollments.first, @shard_a_student.enrollments.first]
+        }
+        enrollments_loader.load(@shard_a_course.global_id).then { |enrollments|
+          expect(enrollments).to match_array [@shard_a_teacher.enrollments.first, @shard_a_student.enrollments.first]
+        }
+        enrollments_loader.load(@shard_b_course.global_id).then { |enrollments|
+          expect(enrollments).to match_array [@shard_b_teacher.enrollments.first, @shard_b_student.enrollments.first]
+        }
+      end
+    end
+
+    it "doesn't get cross-shard data when scoped" do
+      GraphQL::Batch.batch do
+        enrollments_loader = Loaders::ForeignKeyLoader.for(@shard_a_course.enrollments, :user_id)
+        enrollments_loader.load(@shard_a_student.id).then { |students|
+          expect(students).to match_array [@shard_a_student.enrollments.first]
+        }
+        enrollments_loader.load(@shard_b_student.global_id).then { |students|
+          expect(students).to be_nil
+        }
+      end
+    end
+  end
+end

spec/graphql/loaders/id_loader_spec.rb

@@ -0,0 +1,78 @@
+#
+# Copyright (C) 2018 - 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 File.expand_path(File.dirname(__FILE__) + '/../../spec_helper')
+
+describe Loaders::IDLoader do
+  it "works" do
+    course_with_student(active_all: true)
+    GraphQL::Batch.batch do
+      course_loader = Loaders::IDLoader.for(Course)
+      course_loader.load(@course.id).then { |course|
+        expect(course).to eq @course
+      }
+      course_loader.load(-1).then { |course|
+        expect(course).to be_nil
+      }
+    end
+  end
+
+  context "multiple shards" do
+    specs_require_sharding
+
+    before(:once) do
+      course_with_student(active_all: true)
+      @shard_a_course = @course
+      @shard_a_student = @student
+
+      @shard1.activate {
+        shard_b_account = Account.create! name: "shard b  account"
+        course_with_student(active_all: true, account: shard_b_account)
+        @shard_b_course = @course
+        @shard_b_student = @student
+      }
+    end
+
+    it "works across multiple shards" do
+      GraphQL::Batch.batch do
+        course_loader = Loaders::IDLoader.for(Course)
+        course_loader.load(@shard_a_course.id).then { |course|
+          expect(course).to eq @shard_a_course
+        }
+        course_loader.load(@shard_a_course.global_id).then { |course|
+          expect(course).to eq @shard_a_course
+        }
+        course_loader.load(@shard_b_course.global_id).then { |course|
+          expect(course).to eq @shard_b_course
+        }
+      end
+    end
+
+    it "doesn't get cross-shard data when scoped" do
+      GraphQL::Batch.batch do
+        student_loader = Loaders::IDLoader.for(@shard_a_course.students)
+        student_loader.load(@shard_a_student.id).then { |student|
+          expect(student).to eq @shard_a_student
+        }
+        student_loader.load(@shard_b_student.global_id).then { |student|
+          expect(student).to be_nil
+        }
+      end
+    end
+  end
+end