Resolve "Set traversal_ids for every namespace"

changelogs/unreleased/325508-set-traversal_ids-for-every-namespace.yml

+---
+title: Set traversal_ids for every namespace
+merge_request: 57318
+author:
+type: performance

db/post_migrate/20210506065000_schedule_backfill_traversal_ids.rb

+# frozen_string_literal: true
+
+class ScheduleBackfillTraversalIds < ActiveRecord::Migration[6.0]
+  include Gitlab::Database::MigrationHelpers
+
+  ROOTS_MIGRATION = 'BackfillNamespaceTraversalIdsRoots'
+  CHILDREN_MIGRATION = 'BackfillNamespaceTraversalIdsChildren'
+  DOWNTIME = false
+  BATCH_SIZE = 1_000
+  SUB_BATCH_SIZE = 100
+  DELAY_INTERVAL = 2.minutes
+
+  disable_ddl_transaction!
+
+  def up
+    # Personal namespaces and top-level groups
+    final_delay = queue_background_migration_jobs_by_range_at_intervals(
+      ::Gitlab::BackgroundMigration::BackfillNamespaceTraversalIdsRoots::Namespace.base_query,
+      ROOTS_MIGRATION,
+      DELAY_INTERVAL,
+      batch_size: BATCH_SIZE,
+      other_job_arguments: [SUB_BATCH_SIZE],
+      track_jobs: true
+    )
+    final_delay += DELAY_INTERVAL
+
+    # Subgroups
+    queue_background_migration_jobs_by_range_at_intervals(
+      ::Gitlab::BackgroundMigration::BackfillNamespaceTraversalIdsChildren::Namespace.base_query,
+      CHILDREN_MIGRATION,
+      DELAY_INTERVAL,
+      batch_size: BATCH_SIZE,
+      initial_delay: final_delay,
+      other_job_arguments: [SUB_BATCH_SIZE],
+      track_jobs: true
+    )
+  end
+end

db/schema_migrations/20210506065000

+d286628cce50c469afe899d5ac40f20df8dceb6ee10c6cf49c64fbaeea7e4a2e
\ No newline at end of file

lib/gitlab/background_migration/backfill_namespace_traversal_ids_children.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    # A job to set namespaces.traversal_ids in sub-batches, of all namespaces with
+    # a parent and not already set.
+    # rubocop:disable Style/Documentation
+    class BackfillNamespaceTraversalIdsChildren
+      class Namespace < ActiveRecord::Base
+        include ::EachBatch
+
+        self.table_name = 'namespaces'
+
+        scope :base_query, -> { where.not(parent_id: nil) }
+      end
+
+      PAUSE_SECONDS = 0.1
+
+      def perform(start_id, end_id, sub_batch_size)
+        batch_query = Namespace.base_query.where(id: start_id..end_id)
+        batch_query.each_batch(of: sub_batch_size) do |sub_batch|
+          first, last = sub_batch.pluck(Arel.sql('min(id), max(id)')).first
+          ranged_query = Namespace.unscoped.base_query.where(id: first..last)
+
+          update_sql = <<~SQL
+            UPDATE namespaces
+            SET traversal_ids = calculated_ids.traversal_ids
+            FROM #{calculated_traversal_ids(ranged_query)} calculated_ids
+            WHERE namespaces.id = calculated_ids.id
+              AND namespaces.traversal_ids = '{}'
+          SQL
+          ActiveRecord::Base.connection.execute(update_sql)
+
+          sleep PAUSE_SECONDS
+        end
+
+        # We have to add all arguments when marking a job as succeeded as they
+        #  are all used to track the job by `queue_background_migration_jobs_by_range_at_intervals`
+        mark_job_as_succeeded(start_id, end_id, sub_batch_size)
+      end
+
+      private
+
+      # Calculate the ancestor path for a given set of namespaces.
+      def calculated_traversal_ids(batch)
+        <<~SQL
+          (
+            WITH RECURSIVE cte(source_id, namespace_id, parent_id, height) AS (
+              (
+                SELECT batch.id, batch.id, batch.parent_id, 1
+                FROM (#{batch.to_sql}) AS batch
+              )
+              UNION ALL
+              (
+                SELECT cte.source_id, n.id, n.parent_id, cte.height+1
+                FROM namespaces n, cte
+                WHERE n.id = cte.parent_id
+              )
+            )
+            SELECT flat_hierarchy.source_id as id,
+                   array_agg(flat_hierarchy.namespace_id ORDER BY flat_hierarchy.height DESC) as traversal_ids
+            FROM (SELECT * FROM cte FOR UPDATE) flat_hierarchy
+            GROUP BY flat_hierarchy.source_id
+          )
+        SQL
+      end
+
+      def mark_job_as_succeeded(*arguments)
+        Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded(
+          'BackfillNamespaceTraversalIdsChildren',
+          arguments
+        )
+      end
+    end
+  end
+end

lib/gitlab/background_migration/backfill_namespace_traversal_ids_roots.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module BackgroundMigration
+    # A job to set namespaces.traversal_ids in sub-batches, of all namespaces
+    # without a parent and not already set.
+    # rubocop:disable Style/Documentation
+    class BackfillNamespaceTraversalIdsRoots
+      class Namespace < ActiveRecord::Base
+        include ::EachBatch
+
+        self.table_name = 'namespaces'
+
+        scope :base_query, -> { where(parent_id: nil) }
+      end
+
+      PAUSE_SECONDS = 0.1
+
+      def perform(start_id, end_id, sub_batch_size)
+        ranged_query = Namespace.base_query
+          .where(id: start_id..end_id)
+          .where("traversal_ids = '{}'")
+
+        ranged_query.each_batch(of: sub_batch_size) do |sub_batch|
+          sub_batch.update_all('traversal_ids = ARRAY[id]')
+          sleep PAUSE_SECONDS
+        end
+
+        mark_job_as_succeeded(start_id, end_id, sub_batch_size)
+      end
+
+      private
+
+      def mark_job_as_succeeded(*arguments)
+        Gitlab::Database::BackgroundMigrationJob.mark_all_as_succeeded(
+          'BackfillNamespaceTraversalIdsRoots',
+          arguments
+        )
+      end
+    end
+  end
+end

spec/lib/gitlab/background_migration/backfill_namespace_traversal_ids_children_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::BackfillNamespaceTraversalIdsChildren, :migration, schema: 20210506065000 do
+  let(:namespaces_table) { table(:namespaces) }
+
+  let!(:user_namespace) { namespaces_table.create!(id: 1, name: 'user', path: 'user', type: nil) }
+  let!(:root_group) { namespaces_table.create!(id: 2, name: 'group', path: 'group', type: 'Group', parent_id: nil) }
+  let!(:sub_group) { namespaces_table.create!(id: 3, name: 'subgroup', path: 'subgroup', type: 'Group', parent_id: 2) }
+
+  describe '#perform' do
+    it 'backfills traversal_ids for child namespaces' do
+      described_class.new.perform(1, 3, 5)
+
+      expect(user_namespace.reload.traversal_ids).to eq([])
+      expect(root_group.reload.traversal_ids).to eq([])
+      expect(sub_group.reload.traversal_ids).to eq([root_group.id, sub_group.id])
+    end
+  end
+end

spec/lib/gitlab/background_migration/backfill_namespace_traversal_ids_roots_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::BackgroundMigration::BackfillNamespaceTraversalIdsRoots, :migration, schema: 20210506065000 do
+  let(:namespaces_table) { table(:namespaces) }
+
+  let!(:user_namespace) { namespaces_table.create!(id: 1, name: 'user', path: 'user', type: nil) }
+  let!(:root_group) { namespaces_table.create!(id: 2, name: 'group', path: 'group', type: 'Group', parent_id: nil) }
+  let!(:sub_group) { namespaces_table.create!(id: 3, name: 'subgroup', path: 'subgroup', type: 'Group', parent_id: 2) }
+
+  describe '#perform' do
+    it 'backfills traversal_ids for root namespaces' do
+      described_class.new.perform(1, 3, 5)
+
+      expect(user_namespace.reload.traversal_ids).to eq([user_namespace.id])
+      expect(root_group.reload.traversal_ids).to eq([root_group.id])
+      expect(sub_group.reload.traversal_ids).to eq([])
+    end
+  end
+end