Merge branch '118674-composite-queries' into 'master'

118674 - Add composite ID query support

See merge request gitlab-org/gitlab!22055

app/models/concerns/where_composite.rb

+# frozen_string_literal: true
+
+module WhereComposite
+  extend ActiveSupport::Concern
+
+  class TooManyIds < ArgumentError
+    LIMIT = 100
+
+    def initialize(no_of_ids)
+      super(<<~MSG)
+      At most #{LIMIT} identifier sets at a time please! Got #{no_of_ids}.
+      Have you considered splitting your request into batches?
+      MSG
+    end
+
+    def self.guard(collection)
+      n = collection.size
+      return collection if n <= LIMIT
+
+      raise self, n
+    end
+  end
+
+  class_methods do
+    # Apply a set of constraints that function as composite IDs.
+    #
+    # This is the plural form of the standard ActiveRecord idiom:
+    # `where(foo: x, bar: y)`, except it allows multiple pairs of `x` and
+    # `y` to be specified, with the semantics that translate to:
+    #
+    # ```sql
+    # WHERE
+    #     (foo = x_0 AND bar = y_0)
+    #  OR (foo = x_1 AND bar = y_1)
+    #  OR ...
+    # ```
+    #
+    # or the equivalent:
+    #
+    # ```sql
+    # WHERE
+    #   (foo, bar) IN ((x_0, y_0), (x_1, y_1), ...)
+    # ```
+    #
+    # @param permitted_keys [Array<Symbol>] The keys each hash must have. There
+    #                                       must be at least one key (but really,
+    #                                       it ought to be at least two)
+    # @param hashes [Array<#to_h>|#to_h] The constraints. Each parameter must have a
+    #                                    value for the keys named in `permitted_keys`
+    #
+    # e.g.:
+    # ```
+    #   where_composite(%i[foo bar], [{foo: 1, bar: 2}, {foo: 1, bar: 3}])
+    # ```
+    #
+    def where_composite(permitted_keys, hashes)
+      raise ArgumentError, 'no permitted_keys' unless permitted_keys.present?
+
+      # accept any hash-like thing, such as Structs
+      hashes = TooManyIds.guard(Array.wrap(hashes)).map(&:to_h)
+
+      return none if hashes.empty?
+
+      case permitted_keys.size
+      when 1
+        key = permitted_keys.first
+        where(key => hashes.map { |hash| hash.fetch(key) })
+      else
+        clauses = hashes.map do |hash|
+          permitted_keys.map do |key|
+            arel_table[key].eq(hash.fetch(key))
+          end.reduce(:and)
+        end
+
+        where(clauses.reduce(:or))
+      end
+    rescue KeyError
+      raise ArgumentError, "all arguments must contain #{permitted_keys}"
+    end
+  end
+end

app/models/issue.rb

@@ -16,6 +16,7 @@ class Issue < ApplicationRecord
   include LabelEventable
   include IgnorableColumns
   include MilestoneEventable
+  include WhereComposite
 
   DueDateStruct                   = Struct.new(:title, :name).freeze
   NoDueDate                       = DueDateStruct.new('No Due Date', '0').freeze
@@ -78,6 +79,26 @@ class Issue < ApplicationRecord
 
   scope :counts_by_state, -> { reorder(nil).group(:state_id).count }
 
+  # An issue can be uniquely identified by project_id and iid
+  # Takes one or more sets of composite IDs, expressed as hash-like records of
+  # `{project_id: x, iid: y}`.
+  #
+  # @see WhereComposite::where_composite
+  #
+  # e.g:
+  #
+  #   .by_project_id_and_iid({project_id: 1, iid: 2})
+  #   .by_project_id_and_iid([]) # returns ActiveRecord::NullRelation
+  #   .by_project_id_and_iid([
+  #     {project_id: 1, iid: 1},
+  #     {project_id: 2, iid: 1},
+  #     {project_id: 1, iid: 2}
+  #   ])
+  #
+  scope :by_project_id_and_iid, ->(composites) do
+    where_composite(%i[project_id iid], composites)
+  end
+
   after_commit :expire_etag_cache, unless: :importing?
   after_save :ensure_metrics, unless: :importing?
 

ee/app/models/design_management/design.rb

@@ -8,6 +8,7 @@ module DesignManagement
     include Gitlab::Utils::StrongMemoize
     include Referable
     include Mentionable
+    include WhereComposite
 
     belongs_to :project, inverse_of: :designs
     belongs_to :issue
@@ -30,6 +31,26 @@ module DesignManagement
     # reference using `to_reference`.
     scope :for_reference, -> { includes(issue: [{ project: [:route, :namespace] }]) }
 
+    # A design can be uniquely identified by issue_id and filename
+    # Takes one or more sets of composite IDs of the form:
+    # `{issue_id: Integer, filename: String}`.
+    #
+    # @see WhereComposite::where_composite
+    #
+    # e.g:
+    #
+    #   by_issue_id_and_filename(issue_id: 1, filename: 'homescreen.jpg')
+    #   by_issue_id_and_filename([]) # returns ActiveRecord::NullRelation
+    #   by_issue_id_and_filename([
+    #     { issue_id: 1, filename: 'homescreen.jpg' },
+    #     { issue_id: 2, filename: 'homescreen.jpg' },
+    #     { issue_id: 1, filename: 'menu.png' }
+    #   ])
+    #
+    scope :by_issue_id_and_filename, ->(composites) do
+      where_composite(%i[issue_id filename], composites)
+    end
+
     # Find designs visible at the given version
     #
     # @param version [nil, DesignManagement::Version]:

ee/spec/models/design_management/design_spec.rb

@@ -553,4 +553,23 @@ describe DesignManagement::Design do
       end
     end
   end
+
+  describe '.by_issue_id_and_filename' do
+    let_it_be(:issue_a) { create(:issue) }
+    let_it_be(:issue_b) { create(:issue) }
+
+    let_it_be(:design_a) { create(:design, issue: issue_a) }
+    let_it_be(:design_b) { create(:design, issue: issue_a) }
+    let_it_be(:design_c) { create(:design, issue: issue_b, filename: design_a.filename) }
+    let_it_be(:design_d) { create(:design, issue: issue_b, filename: design_b.filename) }
+
+    it_behaves_like 'a where_composite scope', :by_issue_id_and_filename do
+      let(:all_results) { [design_a, design_b, design_c, design_d] }
+      let(:first_result) { design_a }
+
+      let(:composite_ids) do
+        all_results.map { |design| { issue_id: design.issue_id, filename: design.filename } }
+      end
+    end
+  end
 end

spec/models/concerns/where_composite_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe WhereComposite do
+  describe '.where_composite' do
+    let_it_be(:test_table_name) { "test_table_#{SecureRandom.hex(10)}" }
+
+    let(:model) do
+      tbl_name = test_table_name
+      Class.new(ActiveRecord::Base) do
+        self.table_name = tbl_name
+
+        include WhereComposite
+      end
+    end
+
+    def connection
+      ActiveRecord::Base.connection
+    end
+
+    before_all do
+      connection.create_table(test_table_name) do |t|
+        t.integer :foo
+        t.integer :bar
+        t.string  :wibble
+      end
+    end
+
+    it 'requires at least one permitted key' do
+      expect { model.where_composite([], nil) }
+        .to raise_error(ArgumentError)
+    end
+
+    it 'requires all arguments to match the permitted_keys' do
+      expect { model.where_composite([:foo], [{ foo: 1 }, { bar: 2 }]) }
+        .to raise_error(ArgumentError)
+    end
+
+    it 'attaches a key error as cause if a key is missing' do
+      expect { model.where_composite([:foo], [{ foo: 1 }, { bar: 2 }]) }
+        .to raise_error(have_attributes(cause: KeyError))
+    end
+
+    it 'returns an empty relation if there are no arguments' do
+      expect(model.where_composite([:foo, :bar], nil))
+        .to be_empty
+
+      expect(model.where_composite([:foo, :bar], []))
+        .to be_empty
+    end
+
+    it 'permits extra arguments' do
+      a = model.where_composite([:foo, :bar], { foo: 1, bar: 2 })
+      b = model.where_composite([:foo, :bar], { foo: 1, bar: 2, baz: 3 })
+
+      expect(a.to_sql).to eq(b.to_sql)
+    end
+
+    it 'can handle multiple fields' do
+      fields = [:foo, :bar, :wibble]
+      args = { foo: 1, bar: 2, wibble: 'wobble' }
+      pattern = %r{
+        WHERE \s+
+          \(?
+             \s* "#{test_table_name}"\."foo" \s* = \s* 1
+             \s+ AND
+             \s+ "#{test_table_name}"\."bar" \s* = \s* 2
+             \s+ AND
+             \s+ "#{test_table_name}"\."wibble" \s* = \s* 'wobble'
+             \s*
+          \)?
+      }x
+
+      expect(model.where_composite(fields, args).to_sql).to match(pattern)
+    end
+
+    it 'is equivalent to ids.map { |attrs| model.find_by(attrs) }' do
+      10.times do |i|
+        10.times do |j|
+          model.create!(foo: i, bar: j, wibble: generate(:filename))
+        end
+      end
+
+      ids = [{ foo: 1, bar: 9 }, { foo: 9, bar: 1 }]
+
+      found = model.where_composite(%i[foo bar], ids)
+
+      expect(found).to match_array(ids.map { |attrs| model.find_by!(attrs) })
+    end
+
+    it 'constructs (A&B) for one argument' do
+      fields = [:foo, :bar]
+      args = [
+        { foo: 1, bar: 2 }
+      ]
+      pattern = %r{
+        WHERE \s+
+          \(?
+             \s* "#{test_table_name}"\."foo" \s* = \s* 1
+             \s+ AND
+             \s+ "#{test_table_name}"\."bar" \s* = \s* 2
+             \s*
+          \)?
+      }x
+
+      expect(model.where_composite(fields, args).to_sql).to match(pattern)
+      expect(model.where_composite(fields, args[0]).to_sql).to match(pattern)
+    end
+
+    it 'constructs (A&B) OR (C&D) for two arguments' do
+      args = [
+        { foo: 1, bar: 2 },
+        { foo: 3, bar: 4 }
+      ]
+      pattern = %r{
+        WHERE \s+
+          \( \s* "#{test_table_name}"\."foo" \s* = \s* 1
+             \s+ AND
+             \s+ "#{test_table_name}"\."bar" \s* = \s* 2
+             \s* \)?
+          \s* OR \s*
+          \(? \s* "#{test_table_name}"\."foo" \s* = \s* 3
+              \s+ AND
+              \s+ "#{test_table_name}"\."bar" \s* = \s* 4
+              \s* \)
+      }x
+
+      q = model.where_composite([:foo, :bar], args)
+
+      expect(q.to_sql).to match(pattern)
+    end
+
+    it 'constructs (A&B) OR (C&D) OR (E&F) for three arguments' do
+      args = [
+        { foo: 1, bar: 2 },
+        { foo: 3, bar: 4 },
+        { foo: 5, bar: 6 }
+      ]
+      pattern = %r{
+        WHERE \s+
+          \({2}
+             \s* "#{test_table_name}"\."foo" \s* = \s* 1
+             \s+ AND
+             \s+ "#{test_table_name}"\."bar" \s* = \s* 2
+             \s* \)?
+          \s* OR \s*
+          \(? \s* "#{test_table_name}"\."foo" \s* = \s* 3
+              \s+ AND
+              \s+ "#{test_table_name}"\."bar" \s* = \s* 4
+              \s* \)?
+          \s* OR \s*
+          \(? \s* "#{test_table_name}"\."foo" \s* = \s* 5
+              \s+ AND
+              \s+ "#{test_table_name}"\."bar" \s* = \s* 6
+              \s* \)
+      }x
+
+      q = model.where_composite([:foo, :bar], args)
+
+      expect(q.to_sql).to match(pattern)
+    end
+
+    describe 'large sets of IDs' do
+      def query(size)
+        args = (0..).lazy.take(size).map { |n| { foo: n, bar: n * n, wibble: 'x' * n } }.to_a
+        model.where_composite([:foo, :bar, :wibble], args)
+      end
+
+      it 'constructs correct trees of constraints' do
+        n = described_class::TooManyIds::LIMIT
+        q = query(n)
+        sql = q.to_sql
+
+        expect(sql.scan(/OR/).count).to eq(n - 1)
+        expect(sql.scan(/AND/).count).to eq(2 * n)
+      end
+
+      it 'raises errors if too many IDs are passed' do
+        expect { query(described_class::TooManyIds::LIMIT + 1) }.to raise_error(described_class::TooManyIds)
+      end
+    end
+  end
+end

spec/models/issue_spec.rb

@@ -900,6 +900,22 @@ describe Issue do
     end
   end
 
+  describe '.by_project_id_and_iid' do
+    let_it_be(:issue_a) { create(:issue) }
+    let_it_be(:issue_b) { create(:issue, iid: issue_a.iid) }
+    let_it_be(:issue_c) { create(:issue, project: issue_a.project) }
+    let_it_be(:issue_d) { create(:issue, project: issue_a.project) }
+
+    it_behaves_like 'a where_composite scope', :by_project_id_and_iid do
+      let(:all_results) { [issue_a, issue_b, issue_c, issue_d] }
+      let(:first_result) { issue_a }
+
+      let(:composite_ids) do
+        all_results.map { |issue| { project_id: issue.project_id, iid: issue.iid } }
+      end
+    end
+  end
+
   it_behaves_like 'throttled touch' do
     subject { create(:issue, updated_at: 1.hour.ago) }
   end

spec/support/matchers/exceed_query_limit.rb

@@ -73,6 +73,42 @@ module ExceedQueryLimitHelpers
   end
 end
 
+RSpec::Matchers.define :issue_same_number_of_queries_as do
+  supports_block_expectations
+
+  include ExceedQueryLimitHelpers
+
+  def control
+    block_arg
+  end
+
+  def control_recorder
+    @control_recorder ||= ActiveRecord::QueryRecorder.new(&control)
+  end
+
+  def expected_count
+    @expected_count ||= control_recorder.count
+  end
+
+  def verify_count(&block)
+    @subject_block = block
+
+    (expected_count - actual_count).abs <= threshold
+  end
+
+  match do |block|
+    verify_count(&block)
+  end
+
+  failure_message_when_negated do |actual|
+    failure_message
+  end
+
+  def skip_cached
+    false
+  end
+end
+
 RSpec::Matchers.define :exceed_all_query_limit do |expected|
   supports_block_expectations
 

spec/support/shared_examples/models/concerns/composite_id_shared_examples.rb

+# frozen_string_literal: true
+
+RSpec.shared_examples 'a where_composite scope' do |scope_name|
+  let(:result) { described_class.public_send(scope_name, ids) }
+
+  context 'we pass an empty array' do
+    let(:ids) { [] }
+
+    it 'returns a null relation' do
+      expect(result).to be_empty
+    end
+  end
+
+  context 'we pass nil' do
+    let(:ids) { nil }
+
+    it 'returns a null relation' do
+      expect(result).to be_empty
+    end
+  end
+
+  context 'we pass a singleton composite id' do
+    let(:ids) { composite_ids.first }
+
+    it 'finds the first result' do
+      expect(result).to contain_exactly(first_result)
+    end
+  end
+
+  context 'we pass group of ids' do
+    let(:ids) { composite_ids }
+
+    it 'finds all the results' do
+      expect(result).to contain_exactly(*all_results)
+    end
+  end
+
+  describe 'performance' do
+    it 'is not O(N)' do
+      all_ids = composite_ids
+      one_id = composite_ids.first
+
+      expect { described_class.public_send(scope_name, all_ids) }
+        .to issue_same_number_of_queries_as { described_class.public_send(scope_name, one_id) }
+    end
+  end
+end