Update namespace statistics after personal snippet update/removal

After each update/removal of any project or personal snippet,
we need to update the project statistics (in case of project
snippets) and also the namespace root storage statistics.

app/models/concerns/update_project_statistics.rb

@@ -35,8 +35,8 @@ module UpdateProjectStatistics
       @project_statistics_name = project_statistics_name
       @statistic_attribute = statistic_attribute
 
-      after_save(:update_project_statistics_after_save, if: :update_project_statistics_attribute_changed?)
-      after_destroy(:update_project_statistics_after_destroy, unless: :project_destroyed?)
+      after_save(:update_project_statistics_after_save, if: :update_project_statistics_after_save?)
+      after_destroy(:update_project_statistics_after_destroy, if: :update_project_statistics_after_destroy?)
     end
 
     private :update_project_statistics
@@ -45,6 +45,14 @@ module UpdateProjectStatistics
   included do
     private
 
+    def update_project_statistics_after_save?
+      update_project_statistics_attribute_changed?
+    end
+
+    def update_project_statistics_after_destroy?
+      !project_destroyed?
+    end
+
     def update_project_statistics_after_save
       attr = self.class.statistic_attribute
       delta = read_attribute(attr).to_i - attribute_before_last_save(attr).to_i

app/models/project_statistics.rb

@@ -12,7 +12,7 @@ class ProjectStatistics < ApplicationRecord
   before_save :update_storage_size
 
   COLUMNS_TO_REFRESH = [:repository_size, :wiki_size, :lfs_objects_size, :commit_count, :snippets_size].freeze
-  INCREMENTABLE_COLUMNS = { build_artifacts_size: %i[storage_size], packages_size: %i[storage_size] }.freeze
+  INCREMENTABLE_COLUMNS = { build_artifacts_size: %i[storage_size], packages_size: %i[storage_size], snippets_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) }

app/models/snippet.rb

@@ -44,7 +44,9 @@ class Snippet < ApplicationRecord
   has_many :notes, as: :noteable, dependent: :destroy # rubocop:disable Cop/ActiveRecordDependent
   has_many :user_mentions, class_name: "SnippetUserMention", dependent: :delete_all # rubocop:disable Cop/ActiveRecordDependent
   has_one :snippet_repository, inverse_of: :snippet
-  has_one :statistics, class_name: 'SnippetStatistics'
+
+  # We need to add the `dependent` in order to call the after_destroy callback
+  has_one :statistics, class_name: 'SnippetStatistics', dependent: :destroy # rubocop:disable Cop/ActiveRecordDependent
 
   delegate :name, :email, to: :author, prefix: true, allow_nil: true
 

app/models/snippet_statistics.rb

 # frozen_string_literal: true
 
 class SnippetStatistics < ApplicationRecord
+  include AfterCommitQueue
+  include UpdateProjectStatistics
+
   belongs_to :snippet
 
   validates :snippet, presence: true
 
-  delegate :repository, to: :snippet
+  update_project_statistics project_statistics_name: :snippets_size, statistic_attribute: :repository_size
+
+  delegate :repository, :project, :project_id, to: :snippet
+
+  after_save :update_author_root_storage_statistics, if: :update_author_root_storage_statistics?
+  after_destroy :update_author_root_storage_statistics, unless: :project_snippet?
 
   def update_commit_count
     self.commit_count = repository.commit_count
@@ -32,4 +40,30 @@ class SnippetStatistics < ApplicationRecord
 
     save!
   end
+
+  private
+
+  alias_method :original_update_project_statistics_after_save?, :update_project_statistics_after_save?
+  def update_project_statistics_after_save?
+    project_snippet? && original_update_project_statistics_after_save?
+  end
+
+  alias_method :original_update_project_statistics_after_destroy?, :update_project_statistics_after_destroy?
+  def update_project_statistics_after_destroy?
+    project_snippet? && original_update_project_statistics_after_destroy?
+  end
+
+  def update_author_root_storage_statistics?
+    !project_snippet? && saved_change_to_repository_size?
+  end
+
+  def update_author_root_storage_statistics
+    run_after_commit do
+      Namespaces::ScheduleAggregationWorker.perform_async(snippet.author.namespace_id)
+    end
+  end
+
+  def project_snippet?
+    snippet.is_a?(ProjectSnippet)
+  end
 end

app/services/snippets/destroy_service.rb

@@ -27,11 +27,6 @@ module Snippets
 
       attempt_destroy!
 
-      # Update project statistics if the snippet is a Project one
-      if snippet.project_id
-        ProjectCacheWorker.perform_async(snippet.project_id, [], [:snippets_size])
-      end
-
       ServiceResponse.success(message: 'Snippet was deleted.')
     rescue DestroyError
       service_response_error('Failed to remove snippet repository.', 400)

app/services/snippets/update_statistics_service.rb

@@ -16,11 +16,6 @@ module Snippets
       snippet.repository.expire_statistics_caches
       statistics.refresh!
 
-      # Update project statistics if the snippet is a Project one
-      if snippet.project_id
-        ProjectCacheWorker.perform_async(snippet.project_id, [], [:snippets_size])
-      end
-
       ServiceResponse.success(message: 'Snippet statistics successfully updated.')
     end
 

changelogs/unreleased/fj-rethink-snippet-storage-callbacks.yml

+---
+title: Update namespace statistics after personal snippet update/removal
+merge_request: 36031
+author:
+type: changed

spec/factories/snippet_statistics.rb

+# frozen_string_literal: true
+
+FactoryBot.define do
+  factory :snippet_statistics do
+    snippet
+
+    initialize_with do
+      # statistics are automatically created when a snippet is created
+      snippet&.statistics || new
+    end
+
+    transient do
+      with_data { false }
+      size_multiplier { 1 }
+    end
+
+    after(:build) do |snippet_statistics, evaluator|
+      if evaluator.with_data
+        snippet_statistics.repository_size = evaluator.size_multiplier
+        snippet_statistics.commit_count = evaluator.size_multiplier * 2
+        snippet_statistics.file_count = evaluator.size_multiplier * 3
+      end
+    end
+  end
+end

spec/models/snippet_spec.rb

@@ -20,7 +20,7 @@ RSpec.describe Snippet do
     it { is_expected.to have_many(:award_emoji).dependent(:destroy) }
     it { is_expected.to have_many(:user_mentions).class_name("SnippetUserMention") }
     it { is_expected.to have_one(:snippet_repository) }
-    it { is_expected.to have_one(:statistics).class_name('SnippetStatistics') }
+    it { is_expected.to have_one(:statistics).class_name('SnippetStatistics').dependent(:destroy) }
   end
 
   describe 'validation' do

spec/models/snippet_statistics_spec.rb

@@ -86,4 +86,64 @@ RSpec.describe SnippetStatistics do
       subject
     end
   end
+
+  context 'with a PersonalSnippet' do
+    let!(:snippet) { create(:personal_snippet, :repository) }
+
+    shared_examples 'personal snippet statistics updates' do
+      it 'schedules a namespace statistics worker' do
+        expect(Namespaces::ScheduleAggregationWorker)
+          .to receive(:perform_async).once
+
+        statistics.save!
+      end
+
+      it 'does not try to update project stats' do
+        expect(statistics).not_to receive(:schedule_update_project_statistic)
+
+        statistics.save!
+      end
+    end
+
+    context 'when creating' do
+      let(:statistics) { build(:snippet_statistics, snippet_id: snippet.id, with_data: true) }
+
+      before do
+        snippet.statistics.delete
+      end
+
+      it_behaves_like 'personal snippet statistics updates'
+    end
+
+    context 'when updating' do
+      let(:statistics) { snippet.statistics }
+
+      before do
+        snippet.statistics.repository_size = 123
+      end
+
+      it_behaves_like 'personal snippet statistics updates'
+    end
+  end
+
+  context 'with a ProjectSnippet' do
+    let!(:snippet) { create(:project_snippet) }
+
+    it_behaves_like 'UpdateProjectStatistics' do
+      subject { build(:snippet_statistics, snippet: snippet, id: snippet.id, with_data: true) }
+
+      before do
+        # The shared examples requires the snippet statistics not to be present
+        snippet.statistics.delete
+        snippet.reload
+      end
+    end
+
+    it 'does not call personal snippet callbacks' do
+      expect(snippet.statistics).not_to receive(:update_author_root_storage_statistics)
+      expect(snippet.statistics).to receive(:schedule_update_project_statistic)
+
+      snippet.statistics.update!(repository_size: 123)
+    end
+  end
 end

spec/services/snippets/destroy_service_spec.rb

@@ -106,11 +106,24 @@ RSpec.describe Snippets::DestroyService do
         it_behaves_like 'a successful destroy'
         it_behaves_like 'deletes the snippet repository'
 
-        it 'schedules a project cache update for snippet_size' do
-          expect(ProjectCacheWorker).to receive(:perform_async)
-                                          .with(snippet.project_id, [], [:snippets_size])
+        context 'project statistics' do
+          before do
+            snippet.statistics.refresh!
+          end
 
-          subject
+          it 'updates stats after deletion' do
+            expect(project.reload.statistics.snippets_size).not_to be_zero
+
+            subject
+
+            expect(project.reload.statistics.snippets_size).to be_zero
+          end
+
+          it 'schedules a namespace statistics update' do
+            expect(Namespaces::ScheduleAggregationWorker).to receive(:perform_async).with(project.namespace_id).once
+
+            subject
+          end
         end
       end
 
@@ -130,8 +143,8 @@ RSpec.describe Snippets::DestroyService do
         it_behaves_like 'a successful destroy'
         it_behaves_like 'deletes the snippet repository'
 
-        it 'does not schedule a project cache update' do
-          expect(ProjectCacheWorker).not_to receive(:perform_async)
+        it 'schedules a namespace statistics update' do
+          expect(Namespaces::ScheduleAggregationWorker).to receive(:perform_async).with(author.namespace_id)
 
           subject
         end

spec/services/snippets/update_statistics_service_spec.rb

@@ -17,17 +17,6 @@ RSpec.describe Snippets::UpdateStatisticsService do
         subject
       end
 
-      it 'schedules project cache worker based on type' do
-        if snippet.project_id
-          expect(ProjectCacheWorker).to receive(:perform_async)
-                                          .with(snippet.project_id, [], [:snippets_size])
-        else
-          expect(ProjectCacheWorker).not_to receive(:perform_async)
-        end
-
-        subject
-      end
-
       context 'when snippet statistics does not exist' do
         it 'creates snippet statistics' do
           snippet.statistics.delete
@@ -64,6 +53,13 @@ RSpec.describe Snippets::UpdateStatisticsService do
           expect(subject).to be_error
         end
       end
+
+      it 'schedules a namespace storage statistics update' do
+        expect(Namespaces::ScheduleAggregationWorker)
+            .to receive(:perform_async).once
+
+        subject
+      end
     end
 
     context 'with PersonalSnippet' do
@@ -74,8 +70,17 @@ RSpec.describe Snippets::UpdateStatisticsService do
 
     context 'with ProjectSnippet' do
       let!(:snippet) { create(:project_snippet, :repository) }
+      let(:project_statistics) { snippet.project.statistics }
 
       it_behaves_like 'updates statistics'
+
+      it 'updates projects statistics "snippets_size"' do
+        expect(project_statistics.snippets_size).to be_zero
+
+        subject
+
+        expect(snippet.reload.statistics.repository_size).to eq project_statistics.reload.snippets_size
+      end
     end
   end
 end