Include namespace storage stats in root statistics

Once we have storage sizes in the namespace's statistics,
we can now sum the storage values in root statistics and
the ones in namespace statistics.

app/models/namespace/root_storage_statistics.rb

@@ -57,8 +57,7 @@ class Namespace::RootStorageStatistics < ApplicationRecord
   end
 
   def attributes_from_personal_snippets
-    # Return if the type of namespace does not belong to a user
-    return {} unless namespace.type.nil?
+    return {} unless namespace.user?
 
     from_personal_snippets.take.slice(SNIPPETS_SIZE_STAT_NAME)
   end
@@ -70,3 +69,5 @@ class Namespace::RootStorageStatistics < ApplicationRecord
       .select("COALESCE(SUM(s.repository_size), 0) AS #{SNIPPETS_SIZE_STAT_NAME}")
   end
 end
+
+Namespace::RootStorageStatistics.prepend_if_ee('EE::Namespace::RootStorageStatistics')

ee/app/models/ee/namespace/root_storage_statistics.rb

+# frozen_string_literal: true
+
+module EE
+  module Namespace
+    module RootStorageStatistics
+      extend ActiveSupport::Concern
+      extend ::Gitlab::Utils::Override
+
+      prepended do
+        NAMESPACE_STATISTICS_ATTRIBUTES = %w[storage_size wiki_size].freeze
+      end
+
+      override :merged_attributes
+      def merged_attributes
+        super.merge!(attributes_from_namespace_statistics) { |key, v1, v2| v1 + v2 }
+      end
+
+      private
+
+      def attributes_from_namespace_statistics
+        # At the moment, only groups can have some storage data because of group wikis.
+        # Therefore, if the namespace is not a group one, there is no need to perform
+        # the query. If this changes in the future and we add some sort of resource to
+        # users that it's store in NamespaceStatistics, we will need to remove this
+        # guard clause.
+        return {} unless namespace.group?
+
+        from_namespace_statistics.take.slice(*NAMESPACE_STATISTICS_ATTRIBUTES)
+      end
+
+      def from_namespace_statistics
+        namespace
+          .self_and_descendants
+          .joins("INNER JOIN namespace_statistics ns ON ns.namespace_id  = namespaces.id")
+          .select(
+            'COALESCE(SUM(ns.storage_size), 0) AS storage_size',
+            'COALESCE(SUM(ns.wiki_size), 0) AS wiki_size'
+          )
+      end
+    end
+  end
+end

ee/changelogs/unreleased/230465-fj-update-root-storage-stats-to-include-namespace-storage.yml

+---
+title: Include namespace storage stats in root statistics
+merge_request: 54788
+author:
+type: added

ee/spec/models/ee/namespace/root_storage_statistics_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe EE::Namespace::RootStorageStatistics do
+  describe '#recalculate!' do
+    let(:root_storage_statistics) { create(:namespace_root_storage_statistics, namespace: namespace) }
+
+    context 'when namespace belongs to a group' do
+      let_it_be(:root_group) { create(:group) }
+      let_it_be(:group1) { create(:group, parent: root_group) }
+      let_it_be(:subgroup1) { create(:group, parent: group1) }
+      let_it_be(:group2) { create(:group, parent: root_group) }
+      let_it_be(:project1) { create(:project, namespace: group1) }
+      let_it_be(:project2) { create(:project, namespace: group2) }
+      let_it_be(:project_stat1) { create(:project_statistics, project: project1, with_data: true, size_multiplier: 100) }
+      let_it_be(:project_stat2) { create(:project_statistics, project: project2, with_data: true, size_multiplier: 100) }
+      let_it_be(:root_namespace_stat) { create(:namespace_statistics, namespace: root_group, storage_size: 100, wiki_size: 100) }
+      let_it_be(:group1_namespace_stat) { create(:namespace_statistics, namespace: group1, storage_size: 200, wiki_size: 200) }
+      let_it_be(:group2_namespace_stat) { create(:namespace_statistics, namespace: group2, storage_size: 300, wiki_size: 300) }
+      let_it_be(:subgroup1_namespace_stat) { create(:namespace_statistics, namespace: subgroup1, storage_size: 300, wiki_size: 100) }
+
+      let(:namespace) { root_group }
+
+      it 'aggregates namespace statistics' do
+        # This group is not a descendant of the root_group so it shouldn't be included in the final stats.
+        other_group = create(:group)
+        create(:namespace_statistics, namespace: other_group, storage_size: 500, wiki_size: 500)
+
+        root_storage_statistics.recalculate!
+
+        root_storage_statistics.reload
+
+        total_repository_size = project_stat1.repository_size + project_stat2.repository_size
+        total_lfs_objects_size = project_stat1.lfs_objects_size + project_stat2.lfs_objects_size
+        total_build_artifacts_size = project_stat1.build_artifacts_size + project_stat2.build_artifacts_size
+        total_packages_size = project_stat1.packages_size + project_stat2.packages_size
+        total_snippets_size = project_stat1.snippets_size + project_stat2.snippets_size
+        total_pipeline_artifacts_size = project_stat1.pipeline_artifacts_size + project_stat2.pipeline_artifacts_size
+        total_uploads_size = project_stat1.uploads_size + project_stat2.uploads_size
+        total_wiki_size = project_stat1.wiki_size + project_stat2.wiki_size + root_namespace_stat.wiki_size + group1_namespace_stat.wiki_size + group2_namespace_stat.wiki_size + subgroup1_namespace_stat.wiki_size
+        total_storage_size = project_stat1.storage_size + project_stat2.storage_size + root_namespace_stat.storage_size + group1_namespace_stat.storage_size + group2_namespace_stat.storage_size + subgroup1_namespace_stat.storage_size
+
+        expect(root_storage_statistics.repository_size).to eq(total_repository_size)
+        expect(root_storage_statistics.lfs_objects_size).to eq(total_lfs_objects_size)
+        expect(root_storage_statistics.build_artifacts_size).to eq(total_build_artifacts_size)
+        expect(root_storage_statistics.packages_size).to eq(total_packages_size)
+        expect(root_storage_statistics.snippets_size).to eq(total_snippets_size)
+        expect(root_storage_statistics.pipeline_artifacts_size).to eq(total_pipeline_artifacts_size)
+        expect(root_storage_statistics.uploads_size).to eq(total_uploads_size)
+        expect(root_storage_statistics.storage_size).to eq(total_storage_size)
+        expect(root_storage_statistics.wiki_size).to eq(total_wiki_size)
+      end
+
+      it 'works when there are no namespace statistics' do
+        NamespaceStatistics.delete_all
+
+        root_storage_statistics.recalculate!
+
+        root_storage_statistics.reload
+
+        total_wiki_size = project_stat1.wiki_size + project_stat2.wiki_size
+        total_storage_size = project_stat1.storage_size + project_stat2.storage_size
+
+        expect(root_storage_statistics.storage_size).to eq(total_storage_size)
+        expect(root_storage_statistics.wiki_size).to eq(total_wiki_size)
+      end
+    end
+
+    context 'when namespace belong to a user' do
+      let_it_be(:user) { create(:user) }
+      let_it_be(:project) { create(:project, namespace: user.namespace) }
+      let_it_be(:project_stat) { create(:project_statistics, project: project, with_data: true, size_multiplier: 100) }
+
+      let(:namespace) { user.namespace }
+
+      it 'does not aggregate namespace statistics' do
+        create(:namespace_statistics, namespace: user.namespace, storage_size: 200, wiki_size: 200)
+
+        root_storage_statistics.recalculate!
+
+        root_storage_statistics.reload
+
+        expect(root_storage_statistics.storage_size).to eq(project_stat.storage_size)
+        expect(root_storage_statistics.wiki_size).to eq(project_stat.wiki_size)
+      end
+    end
+  end
+end