Commit 0f73c047 authored by Alex Kalderimis's avatar Alex Kalderimis Committed by Mayra Cabrera

Add composite ID query support

This allows us to search for sets of items by multiple values in a
single query, e.g. a query semantically equivalent to:

```sql
select from things
  where (a,b,c) in ((?,?,?), ...)
```

or the expanded form:

```
select from things
  where ((a = ? and b = ? and c = ?)
         or
         (a = ? and b = ? and c = ?)
         ...
         )
```

This is provided by a concern `WhereComposite` which allows us to search
in this way.

We support structs as well as hashes as input.
parent f470b52b
# 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
......@@ -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?
......
......@@ -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]:
......
......@@ -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
# 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
......@@ -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
......
......@@ -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
......
# 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
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment