Merge branch '62214-follow-up-schedule-aggregation-in-other-contexts' into 'master'

Update group statistics when some columns on ProjectStatistics are updated See merge request gitlab-org/gitlab-ce!30329

Merge branch '62214-follow-up-schedule-aggregation-in-other-contexts' into 'master'
Update group statistics when some columns on ProjectStatistics are updated See merge request gitlab-org/gitlab-ce!30329
a0c3254d · Mayra Cabrera · ededb334 · ec39a1d6 · a0c3254d · a0c3254d
--- a/app/models/project_statistics.rb
+++ b/app/models/project_statistics.rb
 # frozen_string_literal: true

 class ProjectStatistics < ApplicationRecord
+  include AfterCommitQueue
+
  belongs_to :project
  belongs_to :namespace

@@ -15,6 +17,7 @@ def wiki_size

  COLUMNS_TO_REFRESH = [:repository_size, :wiki_size, :lfs_objects_size, :commit_count].freeze
  INCREMENTABLE_COLUMNS = { build_artifacts_size: %i[storage_size], packages_size: %i[storage_size] }.freeze
+  NAMESPACE_RELATABLE_COLUMNS = [:repository_size, :wiki_size, :lfs_objects_size].freeze

  scope :for_project_ids, ->(project_ids) { where(project_id: project_ids) }

@@ -22,13 +25,17 @@ def total_repository_size
    repository_size + lfs_objects_size
  end

-  def refresh!(only: nil)
+  def refresh!(only: [])
    COLUMNS_TO_REFRESH.each do |column, generator|
-      if only.blank? || only.include?(column)
+      if only.empty? || only.include?(column)
        public_send("update_#{column}") # rubocop:disable GitlabSecurity/PublicSend
      end
    end

+    if only.empty? || only.any? { |column| NAMESPACE_RELATABLE_COLUMNS.include?(column) }
+      schedule_namespace_aggregation_worker
+    end
+
    save!
  end

@@ -81,4 +88,18 @@ def self.columns_to_increment(key, amount)

    update_all(updates.join(', '))
  end
+
+  private
+
+  def schedule_namespace_aggregation_worker
+    run_after_commit do
+      next unless schedule_aggregation_worker?
+
+      Namespaces::ScheduleAggregationWorker.perform_async(project.namespace_id)
+    end
+  end
+
+  def schedule_aggregation_worker?
+    Feature.enabled?(:update_statistics_namespace, project&.root_ancestor)
+  end
 end
--- a/spec/models/project_statistics_spec.rb
+++ b/spec/models/project_statistics_spec.rb
@@ -135,6 +135,49 @@
        expect(statistics.wiki_size).to eq(0)
      end
    end
+
+    context 'when the column is namespace relatable' do
+      let(:namespace) { create(:group) }
+      let(:project) { create(:project, namespace: namespace) }
+
+      context 'when the feature flag is off' do
+        it 'does not schedule the aggregation worker' do
+          stub_feature_flags(update_statistics_namespace: false, namespace: namespace)
+
+          expect(Namespaces::ScheduleAggregationWorker)
+            .not_to receive(:perform_async)
+
+          statistics.refresh!(only: [:lfs_objects_size])
+        end
+      end
+
+      context 'when the feature flag is on' do
+        it 'schedules the aggregation worker' do
+          expect(Namespaces::ScheduleAggregationWorker)
+            .to receive(:perform_async)
+
+          statistics.refresh!(only: [:lfs_objects_size])
+        end
+      end
+
+      context 'when no argument is passed' do
+        it 'schedules the aggregation worker' do
+          expect(Namespaces::ScheduleAggregationWorker)
+            .to receive(:perform_async)
+
+          statistics.refresh!
+        end
+      end
+    end
+
+    context 'when the column is not namespace relatable' do
+      it 'does not schedules an aggregation worker' do
+        expect(Namespaces::ScheduleAggregationWorker)
+          .not_to receive(:perform_async)
+
+        statistics.refresh!(only: [:commit_count])
+      end
+    end
  end

  describe '#update_commit_count' do

--- a/spec/workers/project_cache_worker_spec.rb
+++ b/spec/workers/project_cache_worker_spec.rb
@@ -36,6 +36,11 @@
    end

    context 'with an existing project' do
+      before do
+        lease_key = "namespace:namespaces_root_statistics:#{project.namespace_id}"
+        stub_exclusive_lease_taken(lease_key, timeout: Namespace::AggregationSchedule::DEFAULT_LEASE_TIMEOUT)
+      end
+
      it 'refreshes the method caches' do
        expect_any_instance_of(Repository).to receive(:refresh_method_caches)
          .with(%i(readme))
@@ -81,6 +86,10 @@

        expect(UpdateProjectStatisticsWorker).not_to receive(:perform_in)

+        expect(Namespaces::ScheduleAggregationWorker)
+          .not_to receive(:perform_async)
+          .with(project.namespace_id)
+
        worker.update_statistics(project, statistics)
      end
    end
@@ -98,6 +107,11 @@
          .with(lease_timeout, project.id, statistics)
          .and_call_original

+        expect(Namespaces::ScheduleAggregationWorker)
+          .to receive(:perform_async)
+          .with(project.namespace_id)
+          .twice
+
        worker.update_statistics(project, statistics)
      end
    end