Fix sub-optimal query batchingajk-batchmodelloader

2 files changed, 51 insertions, 0 deletions

diff --git a/lib/gitlab/batch_model_loader.rb b/lib/gitlab/batch_model_loader.rb
new file mode 100644
index 00000000000..72941962850
--- /dev/null
+++ b/lib/gitlab/batch_model_loader.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Gitlab
+  class BatchModelLoader
+    attr_reader :model_class
+
+    def self.for(model_class)
+      Gitlab::SafeRequestStore.store.fetch("#{self}/for/#{model_class}") do
+        new(model_class)
+      end
+    end
+
+    def initialize(model_class)
+      @model_class = model_class
+    end
+
+    # rubocop: disable CodeReuse/ActiveRecord
+    def find(model_id)
+      BatchLoader.for(model_id.to_i).batch(key: model_class) do |ids, found|
+        model_class.where(id: ids).each { |model| found[model.id, model] }
+      end
+    end
+    # rubocop: enable CodeReuse/ActiveRecord
+  end
+end
diff --git a/spec/lib/gitlab/batch_model_loader_spec.rb b/spec/lib/gitlab/batch_model_loader_spec.rb
new file mode 100644
index 00000000000..443d0975903
--- /dev/null
+++ b/spec/lib/gitlab/batch_model_loader_spec.rb
@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Gitlab::BatchModelLoader do
+  it 'optimises query access by batching', :request_cache do
+    project_a_id = create(:project, title: 'a').id
+    project_b_id = create(:project, title: 'b').id
+
+    user_a_id = create(:user, username: 'user-a').id
+
+    expect do
+      project_a = described_class.for(Project).find(project_a_id)
+      user_a = described_class.for(User).find(user_a_id)
+
+      # force user-a
+      expect(user_a.username).to eq('user-a')
+
+      project_b = described_class.for(Project).find(project_b_id)
+
+      # Force projects
+      expect(project_a.title).to eq('a')
+      expect(project_b.title).to eq('b')
+    end.not_to exceed_query_limit(2)
+  end
+end