diff options
Diffstat (limited to 'lib/gitlab/import_export/json/streaming_serializer.rb')
-rw-r--r-- | lib/gitlab/import_export/json/streaming_serializer.rb | 88 |
1 files changed, 71 insertions, 17 deletions
diff --git a/lib/gitlab/import_export/json/streaming_serializer.rb b/lib/gitlab/import_export/json/streaming_serializer.rb index 05b7679e0ff..ec42c5e51c0 100644 --- a/lib/gitlab/import_export/json/streaming_serializer.rb +++ b/lib/gitlab/import_export/json/streaming_serializer.rb @@ -38,16 +38,6 @@ module Gitlab end end - private - - attr_reader :json_writer, :relations_schema, :exportable - - def serialize_root - attributes = exportable.as_json( - relations_schema.merge(include: nil, preloads: nil)) - json_writer.write_attributes(@exportable_path, attributes) - end - def serialize_relation(definition) raise ArgumentError, 'definition needs to be Hash' unless definition.is_a?(Hash) raise ArgumentError, 'definition needs to have exactly one Hash element' unless definition.one? @@ -64,17 +54,22 @@ module Gitlab end end + private + + attr_reader :json_writer, :relations_schema, :exportable + + def serialize_root + attributes = exportable.as_json( + relations_schema.merge(include: nil, preloads: nil)) + json_writer.write_attributes(@exportable_path, attributes) + end + def serialize_many_relations(key, records, options) enumerator = Enumerator.new do |items| key_preloads = preloads&.dig(key) - records = records.preload(key_preloads) if key_preloads - records.in_batches(of: batch_size) do |batch| # rubocop:disable Cop/InBatches - # order each batch by its primary key to ensure - # consistent and predictable ordering of each exported relation - # as additional `WHERE` clauses can impact the order in which data is being - # returned by database when no `ORDER` is specified - batch = batch.reorder(batch.klass.primary_key) + batch(records, key) do |batch| + batch = batch.preload(key_preloads) if key_preloads batch.each do |record| items << Raw.new(record.to_json(options)) @@ -85,6 +80,29 @@ module Gitlab json_writer.write_relation_array(@exportable_path, key, enumerator) end + def batch(relation, key) + opts = { of: batch_size } + order_by = reorders(relation, key) + + # we need to sort issues by non primary key column(relative_position) + # and `in_batches` does not support that + if order_by + scope = relation.reorder(order_by) + + Gitlab::Pagination::Keyset::Iterator.new(scope: scope, use_union_optimization: true).each_batch(**opts) do |batch| + yield batch + end + else + relation.in_batches(**opts) do |batch| # rubocop:disable Cop/InBatches + # order each batch by its primary key to ensure + # consistent and predictable ordering of each exported relation + # as additional `WHERE` clauses can impact the order in which data is being + # returned by database when no `ORDER` is specified + yield batch.reorder(batch.klass.primary_key) + end + end + end + def serialize_many_each(key, records, options) enumerator = Enumerator.new do |items| records.each do |record| @@ -112,6 +130,42 @@ module Gitlab def batch_size @batch_size ||= self.class.batch_size(@exportable) end + + def reorders(relation, key) + export_reorder = relations_schema[:export_reorder]&.dig(key) + return unless export_reorder + + custom_reorder(relation.klass, export_reorder) + end + + def custom_reorder(klass, order_by) + arel_table = klass.arel_table + column = order_by[:column] || klass.primary_key + direction = order_by[:direction] || :asc + nulls_position = order_by[:nulls_position] || :nulls_last + + arel_order_classes = ::Gitlab::Pagination::Keyset::ColumnOrderDefinition::AREL_ORDER_CLASSES.invert + reverse_direction = ::Gitlab::Pagination::Keyset::ColumnOrderDefinition::REVERSED_ORDER_DIRECTIONS[direction] + reverse_nulls_position = ::Gitlab::Pagination::Keyset::ColumnOrderDefinition::REVERSED_NULL_POSITIONS[nulls_position] + order_expression = ::Gitlab::Database.nulls_order(column, direction, nulls_position) + reverse_order_expression = ::Gitlab::Database.nulls_order(column, reverse_direction, reverse_nulls_position) + + ::Gitlab::Pagination::Keyset::Order.build([ + ::Gitlab::Pagination::Keyset::ColumnOrderDefinition.new( + attribute_name: column, + column_expression: arel_table[column], + order_expression: order_expression, + reversed_order_expression: reverse_order_expression, + order_direction: direction, + nullable: nulls_position, + distinct: false + ), + ::Gitlab::Pagination::Keyset::ColumnOrderDefinition.new( + attribute_name: klass.primary_key, + order_expression: arel_order_classes[direction].new(arel_table[klass.primary_key.to_sym]) + ) + ]) + end end end end |