Include wiki_size and storage_size in NamespaceStatistics

In this commit, we include wiki_size and storage_size in
NamespaceStatistics, and also several callbacks to update
these columns when certain actions happen.

ee/app/models/namespace_statistics.rb

 # frozen_string_literal: true
 
 class NamespaceStatistics < ApplicationRecord
+  include AfterCommitQueue
+
   belongs_to :namespace
 
   validates :namespace, presence: true
 
   scope :for_namespaces, -> (namespaces) { where(namespace: namespaces) }
   scope :with_any_ci_minutes_used, -> { where.not(shared_runners_seconds: 0) }
+
+  before_save :update_storage_size
+  after_save :update_root_storage_statistics, if: :saved_change_to_storage_size?
+  after_destroy :update_root_storage_statistics
+
+  delegate :group?, to: :namespace
+
+  COLUMNS_TO_REFRESH = [:wiki_size].freeze
+
+  def refresh!(only: [])
+    return if Gitlab::Database.read_only?
+    return unless group?
+
+    COLUMNS_TO_REFRESH.each do |column|
+      if only.empty? || only.include?(column)
+        public_send("update_#{column}") # rubocop:disable GitlabSecurity/PublicSend
+      end
+    end
+
+    save!
+  end
+
+  def update_storage_size
+    self.storage_size = wiki_size
+  end
+
+  def update_wiki_size
+    return unless group_wiki_available?
+
+    self.wiki_size = namespace.wiki.repository.size.megabytes
+  end
+
+  private
+
+  def group_wiki_available?
+    group? && namespace.feature_available?(:group_wikis)
+  end
+
+  def update_root_storage_statistics
+    return unless group?
+
+    run_after_commit do
+      Namespaces::ScheduleAggregationWorker.perform_async(namespace.id)
+    end
+  end
 end

ee/changelogs/unreleased/230465-fj-add-new-logic-to-namespace-statistics.yml

+---
+title: Include wiki_size and storage_size in NamespaceStatistics
+merge_request: 54967
+author:
+type: added

ee/spec/models/namespace_statistics_spec.rb

@@ -3,7 +3,227 @@
 require 'spec_helper'
 
 RSpec.describe NamespaceStatistics do
+  include WikiHelpers
+
+  let_it_be(:user) { create(:user) }
+  let_it_be(:group) { create(:group) }
+  let_it_be(:group_wiki) do
+    create(:group_wiki, group: group).tap do |group_wiki|
+      group_wiki.create_page('home', 'test content')
+    end
+  end
+
   it { is_expected.to belong_to(:namespace) }
 
   it { is_expected.to validate_presence_of(:namespace) }
+
+  describe '#refresh!' do
+    let(:namespace) { group }
+    let(:statistics) { create(:namespace_statistics, namespace: namespace) }
+    let(:columns) { [] }
+
+    subject { statistics.refresh!(only: columns) }
+
+    context 'when database is read_only' do
+      it 'does not save the object' do
+        allow(Gitlab::Database).to receive(:read_only?).and_return(true)
+
+        expect(statistics).not_to receive(:save!)
+
+        subject
+      end
+    end
+
+    context 'when namespace belong to a user' do
+      let(:namespace) { user.namespace }
+
+      it 'does not save the object' do
+        expect(statistics).not_to receive(:save!)
+
+        subject
+      end
+    end
+
+    shared_examples 'creates the namespace statistics' do
+      specify do
+        expect(statistics).to receive(:save!)
+
+        subject
+      end
+    end
+
+    context 'when invalid option is passed' do
+      let(:columns) { [:foo] }
+
+      it 'does not update any column' do
+        expect(statistics).not_to receive(:update_wiki_size)
+
+        subject
+      end
+
+      it_behaves_like 'creates the namespace statistics'
+    end
+
+    context 'when no option is passed' do
+      it 'updates the wiki size' do
+        expect(statistics).to receive(:update_wiki_size)
+
+        subject
+      end
+
+      it_behaves_like 'creates the namespace statistics'
+    end
+
+    context 'when wiki_size option is passed' do
+      let(:columns) { [:wiki_size] }
+
+      it 'updates the wiki size' do
+        expect(statistics).to receive(:update_wiki_size)
+
+        subject
+      end
+
+      it_behaves_like 'creates the namespace statistics'
+    end
+  end
+
+  describe '#update_storage_size' do
+    let_it_be(:statistics, reload: true) { create(:namespace_statistics, namespace: group) }
+
+    it 'sets storage_size to the wiki_size' do
+      statistics.wiki_size = 3
+
+      statistics.update_storage_size
+
+      expect(statistics.storage_size).to eq 3
+    end
+  end
+
+  describe '#update_wiki_size' do
+    let_it_be(:statistics, reload: true) { create(:namespace_statistics, namespace: group) }
+
+    subject { statistics.update_wiki_size }
+
+    context 'when group_wikis feature is not enabled' do
+      it 'does not update the wiki size' do
+        stub_group_wikis(false)
+
+        subject
+
+        expect(statistics.wiki_size).to be_zero
+      end
+    end
+
+    context 'when group_wikis feature is enabled enabled' do
+      before do
+        stub_group_wikis(true)
+      end
+
+      it 'updates the wiki size' do
+        subject
+
+        expect(statistics.wiki_size).to eq group.wiki.repository.size.megabytes.to_i
+      end
+
+      context 'when namespace does not belong to a group' do
+        let(:statistics) { create(:namespace_statistics, namespace: user.namespace) }
+
+        it 'does not update the wiki size' do
+          expect(statistics).not_to receive(:wiki)
+
+          subject
+
+          expect(statistics.wiki_size).to be_zero
+        end
+      end
+    end
+  end
+
+  context 'before saving statistics' do
+    let(:statistics) { create(:namespace_statistics, namespace: group, wiki_size: 10) }
+
+    it 'updates storage size' do
+      expect(statistics).to receive(:update_storage_size).and_call_original
+
+      statistics.save!
+
+      expect(statistics.storage_size).to eq 10
+    end
+  end
+
+  context 'after saving statistics', :aggregate_failures do
+    let(:statistics) { create(:namespace_statistics, namespace: namespace) }
+    let(:namespace) { group }
+
+    context 'when storage_size is not updated' do
+      it 'does not enqueue the job to update root storage statistics' do
+        expect(statistics).not_to receive(:update_root_storage_statistics)
+        expect(Namespaces::ScheduleAggregationWorker).not_to receive(:perform_async)
+
+        statistics.save!
+      end
+    end
+
+    context 'when storage_size is updated' do
+      before do
+        # we have to update this value instead of `storage_size` because the before_save
+        # hook we have. If we don't do it, storage_size will be set to the wiki_size value
+        # which is 0.
+        statistics.wiki_size = 10
+      end
+
+      it 'enqueues the job to update root storage statistics' do
+        expect(statistics).to receive(:update_root_storage_statistics).and_call_original
+        expect(Namespaces::ScheduleAggregationWorker).to receive(:perform_async).with(group.id)
+
+        statistics.save!
+      end
+
+      context 'when namespace does not belong to a group' do
+        let(:namespace) { user.namespace }
+
+        it 'does not enqueue the job to update root storage statistics' do
+          expect(statistics).to receive(:update_root_storage_statistics).and_call_original
+          expect(Namespaces::ScheduleAggregationWorker).not_to receive(:perform_async)
+
+          statistics.save!
+        end
+      end
+    end
+
+    context 'when other columns are updated' do
+      it 'does not enqueue the job to update root storage statistics' do
+        columns_to_update = NamespaceStatistics.columns_hash.except('id', 'namespace_id', 'wiki_size', 'storage_size').keys
+        columns_to_update.each { |c| statistics[c] = 10 }
+
+        expect(statistics).not_to receive(:update_root_storage_statistics)
+        expect(Namespaces::ScheduleAggregationWorker).not_to receive(:perform_async)
+
+        statistics.save!
+      end
+    end
+  end
+
+  context 'after destroy statistics', :aggregate_failures do
+    let(:statistics) { create(:namespace_statistics, namespace: namespace) }
+    let(:namespace) { group }
+
+    it 'enqueues the job to update root storage statistics' do
+      expect(statistics).to receive(:update_root_storage_statistics).and_call_original
+      expect(Namespaces::ScheduleAggregationWorker).to receive(:perform_async).with(group.id)
+
+      statistics.destroy!
+    end
+
+    context 'when namespace belongs to a group' do
+      let(:namespace) { user.namespace }
+
+      it 'does not enqueue the job to update root storage statistics' do
+        expect(statistics).to receive(:update_root_storage_statistics).and_call_original
+        expect(Namespaces::ScheduleAggregationWorker).not_to receive(:perform_async)
+
+        statistics.destroy!
+      end
+    end
+  end
 end