Commit 90a345b7 authored by Alex Kalderimis's avatar Alex Kalderimis

Ensure we never run more than MAX_UNION_SIZE queries per union

As a sanity check, we limit the number of queries in any individual
union.
parent d0762262
......@@ -44,7 +44,10 @@
#
# Classes may implement:
# - #item_found(A, R) (return value is ignored)
# - max_union_size Integer (the maximum number of queries to run in any one union)
module CachingArrayResolver
MAX_UNION_SIZE = 50
def resolve(**args)
key = query_input(**args)
......@@ -56,8 +59,9 @@ module CachingArrayResolver
else
queries = keys.map { |key| query_for(key) }
queries.in_groups_of(max_union_size, false).each do |group|
by_id = model_class
.from_union(tag(queries), remove_duplicates: false)
.from_union(tag(group), remove_duplicates: false)
.group_by { |r| r[primary_key] }
by_id.values.each do |item_group|
......@@ -69,11 +73,16 @@ module CachingArrayResolver
end
end
end
end
# Override this to intercept the items once they are found
def item_found(query_input, item)
end
def max_union_size
MAX_UNION_SIZE
end
private
def primary_key
......
......@@ -8,6 +8,7 @@ RSpec.describe ::CachingArrayResolver do
let_it_be(:non_admins) { create_list(:user, 4, admin: false) }
let(:query_context) { {} }
let(:max_page_size) { 10 }
let(:field) { double('Field', max_page_size: max_page_size) }
let(:schema) { double('Schema', default_max_page_size: 3) }
let_it_be(:caching_resolver) do
......@@ -35,6 +36,49 @@ RSpec.describe ::CachingArrayResolver do
end
describe '#resolve' do
context 'there are more than MAX_UNION_SIZE queries' do
let_it_be(:max_union) { 3 }
let_it_be(:resolver) do
mod = described_class
max = max_union
Class.new(::Resolvers::BaseResolver) do
include mod
def query_input(username:)
username
end
def query_for(username)
if username.nil?
model_class.all
else
model_class.where(username: username)
end
end
def model_class
User # Happens to include FromUnion, and is cheap-ish to create
end
define_method :max_union_size do
max
end
end
end
it 'executes the queries in multiple batches' do
users = create_list(:user, (max_union * 2) + 1)
expect(User).to receive(:from_union).twice.and_call_original
results = users.in_groups_of(2, false).map do |users|
resolve(resolver, args: { username: users.map(&:username) }, field: field, schema: schema)
end
expect(results.flat_map(&method(:force))).to match_array(users)
end
end
context 'all queries return results' do
let_it_be(:admins) { create_list(:admin, 3) }
......@@ -154,7 +198,6 @@ RSpec.describe ::CachingArrayResolver do
end
def resolve_users(is_admin, resolver = caching_resolver)
field = double('Field', max_page_size: max_page_size)
args = { is_admin: is_admin }
resolve(resolver, args: args, field: field, ctx: query_context, schema: schema)
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