Include personal snippets in RootStorageStatistics

At the moment, in RootStorageStatistics we're only counting
the project snippets size. Nevertheless, when the namespace
belongs to a user and not a group, we must include the user's
snippets size as well in the `snippets_size` stat.

app/models/namespace/root_storage_statistics.rb

 # frozen_string_literal: true
 
 class Namespace::RootStorageStatistics < ApplicationRecord
-  STATISTICS_ATTRIBUTES = %w(storage_size repository_size wiki_size lfs_objects_size build_artifacts_size packages_size snippets_size).freeze
+  SNIPPETS_SIZE_STAT_NAME = 'snippets_size'.freeze
+  STATISTICS_ATTRIBUTES = %W(storage_size repository_size wiki_size lfs_objects_size build_artifacts_size packages_size #{SNIPPETS_SIZE_STAT_NAME}).freeze
 
   self.primary_key = :namespace_id
 
@@ -13,11 +14,15 @@ class Namespace::RootStorageStatistics < ApplicationRecord
   delegate :all_projects, to: :namespace
 
   def recalculate!
-    update!(attributes_from_project_statistics)
+    update!(merged_attributes)
   end
 
   private
 
+  def merged_attributes
+    attributes_from_project_statistics.merge!(attributes_from_personal_snippets) { |key, v1, v2| v1 + v2 }
+  end
+
   def attributes_from_project_statistics
     from_project_statistics
       .take
@@ -35,7 +40,21 @@ class Namespace::RootStorageStatistics < ApplicationRecord
         'COALESCE(SUM(ps.lfs_objects_size), 0) AS lfs_objects_size',
         'COALESCE(SUM(ps.build_artifacts_size), 0) AS build_artifacts_size',
         'COALESCE(SUM(ps.packages_size), 0) AS packages_size',
-        'COALESCE(SUM(ps.snippets_size), 0) AS snippets_size'
+        "COALESCE(SUM(ps.snippets_size), 0) AS #{SNIPPETS_SIZE_STAT_NAME}"
       )
   end
+
+  def attributes_from_personal_snippets
+    # Return if the type of namespace does not belong to a user
+    return {} unless namespace.type.nil?
+
+    from_personal_snippets.take.slice(SNIPPETS_SIZE_STAT_NAME)
+  end
+
+  def from_personal_snippets
+    PersonalSnippet
+      .joins('INNER JOIN snippet_statistics s ON s.snippet_id = snippets.id')
+      .where(author: namespace.owner_id)
+      .select("COALESCE(SUM(s.repository_size), 0) AS #{SNIPPETS_SIZE_STAT_NAME}")
+  end
 end

changelogs/unreleased/fj-225964-include-personal-snippets-in-snippets-size.yml

+---
+title: Include personal snippets size in RootStorageStatistics
+merge_request: 35984
+author:
+type: changed

spec/models/namespace/root_storage_statistics_spec.rb

@@ -70,7 +70,16 @@ RSpec.describe Namespace::RootStorageStatistics, type: :model do
       end
     end
 
+    shared_examples 'does not include personal snippets' do
+      specify do
+        expect(root_storage_statistics).not_to receive(:from_personal_snippets)
+
+        root_storage_statistics.recalculate!
+      end
+    end
+
     it_behaves_like 'data refresh'
+    it_behaves_like 'does not include personal snippets'
 
     context 'with subgroups' do
       let(:subgroup1) { create(:group, parent: namespace)}
@@ -80,12 +89,45 @@ RSpec.describe Namespace::RootStorageStatistics, type: :model do
       let(:project2) { create(:project, namespace: subgroup2) }
 
       it_behaves_like 'data refresh'
+      it_behaves_like 'does not include personal snippets'
     end
 
     context 'with a personal namespace' do
-      let(:namespace) { create(:user).namespace }
+      let_it_be(:user) { create(:user) }
+      let(:namespace) { user.namespace }
 
       it_behaves_like 'data refresh'
+
+      context 'when user has personal snippets' do
+        let(:total_project_snippets_size) { stat1.snippets_size + stat2.snippets_size }
+
+        it 'aggregates personal and project snippets size' do
+          # This is just a a snippet authored by other user
+          # to ensure we only pick snippets from the namespace
+          # user
+          create(:personal_snippet, :repository).statistics.refresh!
+
+          snippets = create_list(:personal_snippet, 3, :repository, author: user)
+          snippets.each { |s| s.statistics.refresh! }
+
+          total_personal_snippets_size = snippets.map { |s| s.statistics.repository_size }.sum
+
+          root_storage_statistics.recalculate!
+
+          expect(root_storage_statistics.snippets_size).to eq(total_personal_snippets_size + total_project_snippets_size)
+        end
+
+        context 'when personal snippets do not have statistics' do
+          it 'does not raise any error' do
+            snippets = create_list(:personal_snippet, 2, :repository, author: user)
+            snippets.last.statistics.refresh!
+
+            root_storage_statistics.recalculate!
+
+            expect(root_storage_statistics.snippets_size).to eq(total_project_snippets_size + snippets.last.statistics.repository_size)
+          end
+        end
+      end
     end
   end
 end