diff options
Diffstat (limited to 'lib/gitlab/pagination/keyset')
| -rw-r--r-- | lib/gitlab/pagination/keyset/iterator.rb | 40 | ||||
| -rw-r--r-- | lib/gitlab/pagination/keyset/order.rb | 33 | ||||
| -rw-r--r-- | lib/gitlab/pagination/keyset/simple_order_builder.rb | 137 |
3 files changed, 205 insertions, 5 deletions
diff --git a/lib/gitlab/pagination/keyset/iterator.rb b/lib/gitlab/pagination/keyset/iterator.rb new file mode 100644 index 00000000000..3bc8c0bf616 --- /dev/null +++ b/lib/gitlab/pagination/keyset/iterator.rb @@ -0,0 +1,40 @@ +# frozen_string_literal: true + +module Gitlab + module Pagination + module Keyset + class Iterator + def initialize(scope:, use_union_optimization: false) + @scope = scope + @order = Gitlab::Pagination::Keyset::Order.extract_keyset_order_object(scope) + @use_union_optimization = use_union_optimization + end + + # rubocop: disable CodeReuse/ActiveRecord + def each_batch(of: 1000) + cursor_attributes = {} + + loop do + current_scope = scope.dup.limit(of) + relation = order + .apply_cursor_conditions(current_scope, cursor_attributes, { use_union_optimization: @use_union_optimization }) + .reorder(order) + .limit(of) + + yield relation + + last_record = relation.last + break unless last_record + + cursor_attributes = order.cursor_attributes_for_node(last_record) + end + end + # rubocop: enable CodeReuse/ActiveRecord + + private + + attr_reader :scope, :order + end + end + end +end diff --git a/lib/gitlab/pagination/keyset/order.rb b/lib/gitlab/pagination/keyset/order.rb index e596e1bac9d..cef3a7b291a 100644 --- a/lib/gitlab/pagination/keyset/order.rb +++ b/lib/gitlab/pagination/keyset/order.rb @@ -135,7 +135,7 @@ module Gitlab # # (id < 3 AND created_at IS NULL) OR (created_at IS NOT NULL) def build_where_values(values) - return if values.blank? + return [] if values.blank? verify_incoming_values!(values) @@ -156,13 +156,26 @@ module Gitlab end end - build_or_query(where_values) + where_values + end + + def where_values_with_or_query(values) + build_or_query(build_where_values(values.with_indifferent_access)) end # rubocop: disable CodeReuse/ActiveRecord - def apply_cursor_conditions(scope, values = {}) + def apply_cursor_conditions(scope, values = {}, options = { use_union_optimization: false }) + values ||= {} + transformed_values = values.with_indifferent_access scope = apply_custom_projections(scope) - scope.where(build_where_values(values.with_indifferent_access)) + + where_values = build_where_values(transformed_values) + + if options[:use_union_optimization] && where_values.size > 1 + build_union_query(scope, where_values).reorder(self) + else + scope.where(build_or_query(where_values)) # rubocop: disable CodeReuse/ActiveRecord + end end # rubocop: enable CodeReuse/ActiveRecord @@ -170,6 +183,8 @@ module Gitlab self.class.build(column_definitions.map(&:reverse)) end + alias_method :to_sql, :to_s + private # Adds extra columns to the SELECT clause @@ -210,11 +225,19 @@ module Gitlab end def build_or_query(expressions) - or_expression = expressions.reduce { |or_expression, expression| Arel::Nodes::Or.new(or_expression, expression) } + return [] if expressions.blank? + or_expression = expressions.reduce { |or_expression, expression| Arel::Nodes::Or.new(or_expression, expression) } Arel::Nodes::Grouping.new(or_expression) end + def build_union_query(scope, where_values) + scopes = where_values.map do |where_value| + scope.dup.where(where_value).reorder(self) # rubocop: disable CodeReuse/ActiveRecord + end + scope.model.from_union(scopes, remove_duplicates: false, remove_order: false) + end + def to_sql_literal(column_definitions) column_definitions.map do |column_definition| if column_definition.order_expression.respond_to?(:to_sql) diff --git a/lib/gitlab/pagination/keyset/simple_order_builder.rb b/lib/gitlab/pagination/keyset/simple_order_builder.rb new file mode 100644 index 00000000000..5ac5737c3be --- /dev/null +++ b/lib/gitlab/pagination/keyset/simple_order_builder.rb @@ -0,0 +1,137 @@ +# frozen_string_literal: true + +module Gitlab + module Pagination + module Keyset + # This class transforms the `order()` values from an Activerecord scope into a + # Gitlab::Pagination::Keyset::Order instance so the query later can be used in + # keyset pagination. + # + # Return values: + # [transformed_scope, true] # true indicates that the new scope was successfully built + # [orginal_scope, false] # false indicates that the order values are not supported in this class + class SimpleOrderBuilder + def self.build(scope) + new(scope: scope).build + end + + def initialize(scope:) + @scope = scope + @order_values = scope.order_values + @model_class = scope.model + @arel_table = @model_class.arel_table + @primary_key = @model_class.primary_key + end + + def build + order = if order_values.empty? + primary_key_descending_order + elsif ordered_by_primary_key? + primary_key_order + elsif ordered_by_other_column? + column_with_tie_breaker_order + elsif ordered_by_other_column_with_tie_breaker? + tie_breaker_attribute = order_values.second + + tie_breaker_column_order = Gitlab::Pagination::Keyset::ColumnOrderDefinition.new( + attribute_name: model_class.primary_key, + order_expression: tie_breaker_attribute + ) + + column_with_tie_breaker_order(tie_breaker_column_order) + end + + order ? [scope.reorder!(order), true] : [scope, false] # [scope, success] + end + + private + + attr_reader :scope, :order_values, :model_class, :arel_table, :primary_key + + def primary_key_descending_order + Gitlab::Pagination::Keyset::Order.build([ + Gitlab::Pagination::Keyset::ColumnOrderDefinition.new( + attribute_name: model_class.primary_key, + order_expression: arel_table[primary_key].desc + ) + ]) + end + + def primary_key_order + Gitlab::Pagination::Keyset::Order.build([ + Gitlab::Pagination::Keyset::ColumnOrderDefinition.new( + attribute_name: model_class.primary_key, + order_expression: order_values.first + ) + ]) + end + + def column_with_tie_breaker_order(tie_breaker_column_order = default_tie_breaker_column_order) + order_expression = order_values.first + attribute_name = order_expression.expr.name + + column_nullable = model_class.columns.find { |column| column.name == attribute_name }.null + + nullable = if column_nullable && order_expression.is_a?(Arel::Nodes::Ascending) + :nulls_last + elsif column_nullable && order_expression.is_a?(Arel::Nodes::Descending) + :nulls_first + else + :not_nullable + end + + Gitlab::Pagination::Keyset::Order.build([ + Gitlab::Pagination::Keyset::ColumnOrderDefinition.new( + attribute_name: attribute_name, + order_expression: order_expression, + nullable: nullable, + distinct: false + ), + tie_breaker_column_order + ]) + end + + def ordered_by_primary_key? + return unless order_values.one? + + attribute = order_values.first.try(:expr) + + return unless attribute + + arel_table[primary_key].to_s == attribute.to_s + end + + def ordered_by_other_column? + return unless order_values.one? + + attribute = order_values.first.try(:expr) + + return unless attribute + return unless attribute.try(:name) + + model_class.column_names.include?(attribute.name.to_s) + end + + def ordered_by_other_column_with_tie_breaker? + return unless order_values.size == 2 + + attribute = order_values.first.try(:expr) + tie_breaker_attribute = order_values.second.try(:expr) + + return unless attribute + return unless tie_breaker_attribute + + model_class.column_names.include?(attribute.name.to_s) && + arel_table[primary_key].to_s == tie_breaker_attribute.to_s + end + + def default_tie_breaker_column_order + Gitlab::Pagination::Keyset::ColumnOrderDefinition.new( + attribute_name: model_class.primary_key, + order_expression: arel_table[primary_key].desc + ) + end + end + end + end +end |