diff --git a/lib/gitlab/markdown/reference_filter.rb b/lib/gitlab/markdown/reference_filter.rb index ede9f8865ff..adaca78ba27 100644 --- a/lib/gitlab/markdown/reference_filter.rb +++ b/lib/gitlab/markdown/reference_filter.rb @@ -14,10 +14,18 @@ module Gitlab class ReferenceFilter < HTML::Pipeline::Filter LazyReference = Struct.new(:klass, :ids) do def self.load(refs) - refs.group_by(&:klass).flat_map do |klass, refs| + lazy_references, values = refs.partition { |ref| ref.is_a?(self) } + + lazy_values = lazy_references.group_by(&:klass).flat_map do |klass, refs| ids = refs.flat_map(&:ids) klass.where(id: ids) end + + values + lazy_values + end + + def load + self.klass.where(id: self.ids) end end diff --git a/lib/gitlab/markdown/reference_gatherer_filter.rb b/lib/gitlab/markdown/reference_gatherer_filter.rb index 18df5db94d6..31fb71a98a3 100644 --- a/lib/gitlab/markdown/reference_gatherer_filter.rb +++ b/lib/gitlab/markdown/reference_gatherer_filter.rb @@ -12,8 +12,7 @@ module Gitlab def initialize(*) super - result[:lazy_references] ||= Hash.new { |hash, type| hash[type] = [] } - result[:references] ||= Hash.new { |hash, type| hash[type] = [] } + result[:references] ||= Hash.new { |hash, type| hash[type] = [] } end def call @@ -41,23 +40,16 @@ module Gitlab references.each do |type, values| Array.wrap(values).each do |value| - refs = - if value.is_a?(ReferenceFilter::LazyReference) - result[:lazy_references] - else - result[:references] - end - - refs[type] << value + result[:references][type] << value end end end # Will load all references of one type using one query. def load_lazy_references - result[:lazy_references].each do |type, refs| - values = ReferenceFilter::LazyReference.load(refs) - result[:references][type].concat(values) + refs = result[:references] + refs.each do |type, values| + refs[type] = ReferenceFilter::LazyReference.load(values) end end diff --git a/lib/gitlab/reference_extractor.rb b/lib/gitlab/reference_extractor.rb index 8100f2675a7..d6b739d7b9a 100644 --- a/lib/gitlab/reference_extractor.rb +++ b/lib/gitlab/reference_extractor.rb @@ -54,19 +54,13 @@ module Gitlab pipeline = HTML::Pipeline.new([filter, Gitlab::Markdown::ReferenceGathererFilter], context) - values = [] - lazy_references = [] - - @texts.each do |text| + values = @texts.flat_map do |text| result = pipeline.call(text) - values.concat(result[:references][filter_type]) - lazy_references.concat(result[:lazy_references][filter_type]) + result[:references][filter_type] end - lazy_values = Gitlab::Markdown::ReferenceFilter::LazyReference.load(lazy_references) - values.concat(lazy_values) - values + Gitlab::Markdown::ReferenceFilter::LazyReference.load(values) end end end