Add more storage methods to project and namespace

Add methods to namespace and project related to excess storage. These
methods are a building block to displaying necessary data in
getStorageCounter graphql query.

ee/app/models/ee/namespace.rb

@@ -197,23 +197,36 @@ module EE
     def total_repository_size_excess
       strong_memoize(:total_repository_size_excess) do
         namespace_size_limit = actual_size_limit
-
         namespace_limit_arel = Arel::Nodes::SqlLiteral.new(namespace_size_limit.to_s.presence || 'NULL')
-        condition = 'projects.repository_size_limit != 0 AND project_statistics.repository_size > projects.repository_size_limit'
 
-        total_excess = total_repository_size_excess_calculation(condition, ::Project.arel_table[:repository_size_limit])
+        total_excess = total_repository_size_excess_calculation(::Project.arel_table[:repository_size_limit])
+        total_excess += total_repository_size_excess_calculation(namespace_limit_arel, project_level: false) if namespace_size_limit.to_i > 0
+        total_excess
+      end
+    end
 
-        if namespace_size_limit.to_i > 0
-          condition = 'projects.repository_size_limit IS NULL AND project_statistics.repository_size > :namespace_size_limit'
-          sanitized_condition = self.class.sanitize_sql_array([condition, namespace_size_limit: namespace_size_limit])
+    def repository_size_excess_project_count
+      strong_memoize(:repository_size_excess_project_count) do
+        namespace_size_limit = actual_size_limit
 
-          total_excess += total_repository_size_excess_calculation(sanitized_condition, namespace_limit_arel)
-        end
+        count = projects_for_repository_size_excess.count
+        count += projects_for_repository_size_excess(namespace_size_limit).count if namespace_size_limit.to_i > 0
+        count
+      end
+    end
 
-        total_excess
+    def total_repository_size
+      strong_memoize(:total_repository_size) do
+        all_projects
+          .joins(:statistics)
+          .pluck(total_repository_size_arel.sum).first || 0 # rubocop:disable Rails/Pick
       end
     end
 
+    def contains_locked_projects?
+      total_repository_size_excess > additional_purchased_storage_size.megabytes
+    end
+
     def actual_size_limit
       ::Gitlab::CurrentSettings.repository_size_limit
     end
@@ -432,14 +445,27 @@ module EE
       [actual_shared_runners_minutes_limit.to_f - shared_runners_minutes.to_f, 0].max
     end
 
-    def total_repository_size_excess_calculation(condition, limit)
-      select_sql = Arel::Nodes::NamedFunction.new('SUM', [::ProjectStatistics.arel_table[:repository_size] - limit]).to_sql
+    def total_repository_size_excess_calculation(repository_size_limit, project_level: true)
+      total_excess = (total_repository_size_arel - repository_size_limit).sum
+      relation = projects_for_repository_size_excess((repository_size_limit unless project_level))
+      relation.pluck(total_excess).first || 0 # rubocop:disable Rails/Pick
+    end
+
+    def total_repository_size_arel
+      arel_table = ::ProjectStatistics.arel_table
+      arel_table[:repository_size] + arel_table[:lfs_objects_size]
+    end
 
-      all_projects
-      .joins(:statistics)
-      .where(condition)
-      .pluck(Arel.sql(select_sql)) # rubocop:disable Rails/Pick
-      .first || 0
+    def projects_for_repository_size_excess(limit = nil)
+      if limit
+        all_projects
+          .with_total_repository_size_greater_than(limit)
+          .without_repository_size_limit
+      else
+        all_projects
+          .with_total_repository_size_greater_than(::Project.arel_table[:repository_size_limit])
+          .without_unlimited_repository_size_limit
+      end
     end
   end
 end

ee/app/models/ee/project.rb

@@ -155,6 +155,15 @@ module EE
 
       scope :with_group_saml_provider, -> { preload(group: :saml_provider) }
 
+      scope :with_total_repository_size_greater_than, -> (value) do
+        statistics = ::ProjectStatistics.arel_table
+
+        joins(:statistics)
+          .where((statistics[:repository_size] + statistics[:lfs_objects_size]).gt(value))
+      end
+      scope :without_unlimited_repository_size_limit, -> { where.not(repository_size_limit: 0) }
+      scope :without_repository_size_limit, -> { where(repository_size_limit: nil) }
+
       delegate :shared_runners_minutes, :shared_runners_seconds, :shared_runners_seconds_last_reset,
         to: :statistics, allow_nil: true
 
@@ -507,16 +516,28 @@ module EE
       ::Gitlab::UrlSanitizer.new(bare_url, credentials: { user: import_data&.user }).full_url
     end
 
+    def actual_size_limit
+      strong_memoize(:actual_size_limit) do
+        repository_size_limit || namespace.actual_size_limit
+      end
+    end
+
     def repository_size_checker
       strong_memoize(:repository_size_checker) do
         ::Gitlab::RepositorySizeChecker.new(
           current_size_proc: -> { statistics.total_repository_size },
-          limit: (repository_size_limit || namespace.actual_size_limit),
+          limit: actual_size_limit,
           enabled: License.feature_available?(:repository_size_limit)
         )
       end
     end
 
+    def repository_size_excess
+      return 0 unless actual_size_limit.to_i > 0
+
+      [statistics.total_repository_size - actual_size_limit, 0].max
+    end
+
     def username_only_import_url=(value)
       unless ::Gitlab::UrlSanitizer.valid?(value)
         self.import_url = value

ee/spec/models/ee/namespace_spec.rb

@@ -1434,36 +1434,20 @@ RSpec.describe Namespace do
   describe '#total_repository_size_excess' do
     let_it_be(:namespace) { create(:namespace) }
 
-    def create_project(repository_size:, repository_size_limit:)
-      create(:project, namespace: namespace, repository_size_limit: repository_size_limit).tap do |project|
-        create(:project_statistics, project: project, repository_size: repository_size)
-      end
-    end
-
     before do
       namespace.clear_memoization(:total_repository_size_excess)
     end
 
     context 'projects with a variety of repository sizes and limits' do
       before_all do
-        create_project(repository_size: 100, repository_size_limit: nil)
-        create_project(repository_size: 150, repository_size_limit: nil)
-        create_project(repository_size: 200, repository_size_limit: nil)
-
-        create_project(repository_size: 100, repository_size_limit: 0)
-        create_project(repository_size: 150, repository_size_limit: 0)
-        create_project(repository_size: 200, repository_size_limit: 0)
-
-        create_project(repository_size: 300, repository_size_limit: 400)
-        create_project(repository_size: 400, repository_size_limit: 400)
-        create_project(repository_size: 500, repository_size_limit: 300)
+        create_storage_excess_example_projects
       end
 
       context 'when namespace-level repository_size_limit is not set' do
         it 'returns the total excess size of projects with repositories that exceed the size limit' do
           allow(namespace).to receive(:actual_size_limit).and_return(nil)
 
-          expect(namespace.total_repository_size_excess).to eq(200)
+          expect(namespace.total_repository_size_excess).to eq(400)
         end
       end
 
@@ -1471,7 +1455,7 @@ RSpec.describe Namespace do
         it 'returns the total excess size of projects with repositories that exceed the size limit' do
           allow(namespace).to receive(:actual_size_limit).and_return(0)
 
-          expect(namespace.total_repository_size_excess).to eq(200)
+          expect(namespace.total_repository_size_excess).to eq(400)
         end
       end
 
@@ -1479,16 +1463,16 @@ RSpec.describe Namespace do
         it 'returns the total excess size of projects with repositories that exceed the size limit' do
           allow(namespace).to receive(:actual_size_limit).and_return(150)
 
-          expect(namespace.total_repository_size_excess).to eq(250)
+          expect(namespace.total_repository_size_excess).to eq(560)
         end
       end
     end
 
     context 'when all projects have repository_size_limit of 0 (unlimited)' do
       before do
-        create_project(repository_size: 100, repository_size_limit: 0)
-        create_project(repository_size: 150, repository_size_limit: 0)
-        create_project(repository_size: 200, repository_size_limit: 0)
+        create_project(repository_size: 100, lfs_objects_size: 0, repository_size_limit: 0)
+        create_project(repository_size: 150, lfs_objects_size: 0, repository_size_limit: 0)
+        create_project(repository_size: 200, lfs_objects_size: 100, repository_size_limit: 0)
 
         allow(namespace).to receive(:actual_size_limit).and_return(150)
       end
@@ -1499,6 +1483,104 @@ RSpec.describe Namespace do
     end
   end
 
+  describe '#repository_size_excess_project_count' do
+    let_it_be(:namespace) { create(:namespace) }
+
+    before do
+      namespace.clear_memoization(:repository_size_excess_project_count)
+    end
+
+    context 'projects with a variety of repository sizes and limits' do
+      before_all do
+        create_storage_excess_example_projects
+      end
+
+      context 'when namespace-level repository_size_limit is not set' do
+        before do
+          allow(namespace).to receive(:actual_size_limit).and_return(nil)
+        end
+
+        it 'returns the count of projects with repositories that exceed the size limit' do
+          expect(namespace.repository_size_excess_project_count).to eq(2)
+        end
+      end
+
+      context 'when namespace-level repository_size_limit is 0 (unlimited)' do
+        before do
+          allow(namespace).to receive(:actual_size_limit).and_return(0)
+        end
+
+        it 'returns the count of projects with repositories that exceed the size limit' do
+          expect(namespace.repository_size_excess_project_count).to eq(2)
+        end
+      end
+
+      context 'when namespace-level repository_size_limit is a positive number' do
+        before do
+          allow(namespace).to receive(:actual_size_limit).and_return(150)
+        end
+
+        it 'returns the count of projects with repositories that exceed the size limit' do
+          expect(namespace.repository_size_excess_project_count).to eq(4)
+        end
+      end
+    end
+
+    context 'when all projects have repository_size_limit of 0 (unlimited)' do
+      before do
+        create_project(repository_size: 100, lfs_objects_size: 0, repository_size_limit: 0)
+        create_project(repository_size: 150, lfs_objects_size: 0, repository_size_limit: 0)
+        create_project(repository_size: 200, lfs_objects_size: 100, repository_size_limit: 0)
+
+        allow(namespace).to receive(:actual_size_limit).and_return(150)
+      end
+
+      it 'returns zero regardless of the namespace or instance-level repository_size_limit' do
+        expect(namespace.repository_size_excess_project_count).to eq(0)
+      end
+    end
+  end
+
+  describe '#total_repository_size' do
+    let(:namespace) { create(:namespace) }
+
+    before do
+      create_project(repository_size: 100, lfs_objects_size: 0, repository_size_limit: nil)
+      create_project(repository_size: 150, lfs_objects_size: 100, repository_size_limit: 0)
+      create_project(repository_size: 325, lfs_objects_size: 200, repository_size_limit: 400)
+    end
+
+    it 'returns the total size of all project repositories' do
+      expect(namespace.total_repository_size).to eq(875)
+    end
+  end
+
+  describe '#contains_locked_projects?' do
+    using RSpec::Parameterized::TableSyntax
+
+    let_it_be(:namespace) { create(:namespace) }
+
+    before_all do
+      create(:namespace_limit, namespace: namespace, additional_purchased_storage_size: 10)
+    end
+
+    where(:total_excess, :result) do
+      5.megabytes  | false
+      10.megabytes | false
+      15.megabytes | true
+    end
+
+    with_them do
+      before do
+        allow(namespace).to receive(:total_repository_size_excess).and_return(total_excess)
+      end
+
+      it 'returns a boolean indicating whether the root namespace contains locked projects' do
+        expect(namespace.contains_locked_projects?).to be result
+      end
+    end
+  end
+
   describe '#actual_size_limit' do
     let(:namespace) { build(:namespace) }
 
@@ -1687,4 +1769,28 @@ RSpec.describe Namespace do
       expect(namespace.errors[:"namespace_limit.temporary_storage_increase_ends_on"]).to be_present
     end
   end
+
+  def create_project(repository_size:, lfs_objects_size:, repository_size_limit:)
+    create(:project, namespace: namespace, repository_size_limit: repository_size_limit).tap do |project|
+      create(:project_statistics, project: project, repository_size: repository_size, lfs_objects_size: lfs_objects_size)
+    end
+  end
+
+  def create_storage_excess_example_projects
+    [
+      { repository_size: 100, lfs_objects_size: 0, repository_size_limit: nil },
+      { repository_size: 150, lfs_objects_size: 0, repository_size_limit: nil },
+      { repository_size: 140, lfs_objects_size: 10, repository_size_limit: nil },
+      { repository_size: 150, lfs_objects_size: 10, repository_size_limit: nil },
+      { repository_size: 200, lfs_objects_size: 100, repository_size_limit: nil },
+      { repository_size: 100, lfs_objects_size: 0, repository_size_limit: 0 },
+      { repository_size: 150, lfs_objects_size: 10, repository_size_limit: 0 },
+      { repository_size: 200, lfs_objects_size: 100, repository_size_limit: 0 },
+      { repository_size: 300, lfs_objects_size: 0, repository_size_limit: 400 },
+      { repository_size: 400, lfs_objects_size: 0, repository_size_limit: 400 },
+      { repository_size: 300, lfs_objects_size: 100, repository_size_limit: 400 },
+      { repository_size: 400, lfs_objects_size: 100, repository_size_limit: 400 },
+      { repository_size: 500, lfs_objects_size: 100, repository_size_limit: 300 }
+    ].map { |attrs| create_project(**attrs) }
+  end
 end

ee/spec/models/project_spec.rb

@@ -2188,6 +2188,25 @@ RSpec.describe Project do
     end
   end
 
+  describe '#actual_size_limit' do
+    context 'when repository_size_limit is set on the project' do
+      it 'returns the repository_size_limit' do
+        project = build(:project, repository_size_limit: 10)
+
+        expect(project.actual_size_limit).to eq(10)
+      end
+    end
+
+    context 'when repository_size_limit is not set on the project' do
+      it 'returns the actual_size_limit of the namespace' do
+        group = build(:group, repository_size_limit: 20)
+        project = build(:project, namespace: group, repository_size_limit: nil)
+
+        expect(project.actual_size_limit).to eq(20)
+      end
+    end
+  end
+
   describe '#repository_size_checker' do
     let(:project) { build(:project) }
     let(:checker) { project.repository_size_checker }
@@ -2257,6 +2276,32 @@ RSpec.describe Project do
     end
   end
 
+  describe '#repository_size_excess' do
+    using RSpec::Parameterized::TableSyntax
+
+    subject { project.repository_size_excess }
+
+    let_it_be(:statistics) { create(:project_statistics) }
+    let_it_be(:project) { statistics.project }
+
+    where(:total_repository_size, :size_limit, :result) do
+      50 | nil | 0
+      50 | 0   | 0
+      50 | 60  | 0
+      50 | 50  | 0
+      50 | 10  | 40
+    end
+
+    with_them do
+      before do
+        allow(project).to receive(:actual_size_limit).and_return(size_limit)
+        allow(statistics).to receive(:total_repository_size).and_return(total_repository_size)
+      end
+
+      it { is_expected.to eq(result) }
+    end
+  end
+
   describe '#repository_size_limit column' do
     it 'support values up to 8 exabytes' do
       project = create(:project)