1 files changed, 123 insertions, 54 deletions

diff --git a/spec/lib/gitlab/database/batch_count_spec.rb b/spec/lib/gitlab/database/batch_count_spec.rb
index 811d4fad95c..a87b0c1a3a8 100644
--- a/spec/lib/gitlab/database/batch_count_spec.rb
+++ b/spec/lib/gitlab/database/batch_count_spec.rb
@@ -86,48 +86,48 @@ RSpec.describe Gitlab::Database::BatchCount do
           query: batch_count_query,
           message: 'Query has been canceled with message: query timed out'
         )
-        expect(subject.call(model, column, batch_size: batch_size, start: 0)).to eq(-1)
+        expect(subject.call(model, column, batch_size: batch_size, start: 0)).to eq(fallback)
       end
     end
   end
 
   describe '#batch_count' do
     it 'counts table' do
-      expect(described_class.batch_count(model)).to eq(5)
+      expect(described_class.batch_count(model)).to eq(model.count)
     end
 
     it 'counts with :id field' do
-      expect(described_class.batch_count(model, :id)).to eq(5)
+      expect(described_class.batch_count(model, :id)).to eq(model.count)
     end
 
     it 'counts with "id" field' do
-      expect(described_class.batch_count(model, 'id')).to eq(5)
+      expect(described_class.batch_count(model, 'id')).to eq(model.count)
     end
 
     it 'counts with table.id field' do
-      expect(described_class.batch_count(model, "#{model.table_name}.id")).to eq(5)
+      expect(described_class.batch_count(model, "#{model.table_name}.id")).to eq(model.count)
     end
 
     it 'counts with Arel column' do
-      expect(described_class.batch_count(model, model.arel_table[:id])).to eq(5)
+      expect(described_class.batch_count(model, model.arel_table[:id])).to eq(model.count)
     end
 
     it 'counts table with batch_size 50K' do
-      expect(described_class.batch_count(model, batch_size: 50_000)).to eq(5)
+      expect(described_class.batch_count(model, batch_size: 50_000)).to eq(model.count)
     end
 
     it 'will not count table with a batch size less than allowed' do
       expect(described_class.batch_count(model, batch_size: small_batch_size)).to eq(fallback)
     end
 
-    it 'counts with a small edge case batch_sizes than result' do
+    it 'produces the same result with different batch sizes' do
       stub_const('Gitlab::Database::BatchCounter::MIN_REQUIRED_BATCH_SIZE', 0)
 
-      [1, 2, 4, 5, 6].each { |i| expect(described_class.batch_count(model, batch_size: i)).to eq(5) }
+      [1, 2, 4, 5, 6].each { |i| expect(described_class.batch_count(model, batch_size: i)).to eq(model.count) }
     end
 
     it 'counts with a start and finish' do
-      expect(described_class.batch_count(model, start: model.minimum(:id), finish: model.maximum(:id))).to eq(5)
+      expect(described_class.batch_count(model, start: model.minimum(:id), finish: model.maximum(:id))).to eq(model.count)
     end
 
     it 'stops counting when finish value is reached' do
@@ -217,6 +217,113 @@ RSpec.describe Gitlab::Database::BatchCount do
     end
   end
 
+  describe '#batch_count_with_timeout' do
+    it 'counts table' do
+      expect(described_class.batch_count_with_timeout(model)).to eq({ status: :completed, count: model.count })
+    end
+
+    it 'counts with :id field' do
+      expect(described_class.batch_count_with_timeout(model, :id)).to eq({ status: :completed, count: model.count })
+    end
+
+    it 'counts with "id" field' do
+      expect(described_class.batch_count_with_timeout(model, 'id')).to eq({ status: :completed, count: model.count })
+    end
+
+    it 'counts with table.id field' do
+      expect(described_class.batch_count_with_timeout(model, "#{model.table_name}.id")).to eq({ status: :completed, count: model.count })
+    end
+
+    it 'counts with Arel column' do
+      expect(described_class.batch_count_with_timeout(model, model.arel_table[:id])).to eq({ status: :completed, count: model.count })
+    end
+
+    it 'counts table with batch_size 50K' do
+      expect(described_class.batch_count_with_timeout(model, batch_size: 50_000)).to eq({ status: :completed, count: model.count })
+    end
+
+    it 'will not count table with a batch size less than allowed' do
+      expect(described_class.batch_count_with_timeout(model, batch_size: small_batch_size)).to eq({ status: :bad_config })
+    end
+
+    it 'produces the same result with different batch sizes' do
+      stub_const('Gitlab::Database::BatchCounter::MIN_REQUIRED_BATCH_SIZE', 0)
+
+      [1, 2, 4, 5, 6].each { |i| expect(described_class.batch_count_with_timeout(model, batch_size: i)).to eq({ status: :completed, count: model.count }) }
+    end
+
+    it 'counts with a start and finish' do
+      expect(described_class.batch_count_with_timeout(model, start: model.minimum(:id), finish: model.maximum(:id))).to eq({ status: :completed, count: model.count })
+    end
+
+    it 'stops counting when finish value is reached' do
+      stub_const('Gitlab::Database::BatchCounter::MIN_REQUIRED_BATCH_SIZE', 0)
+
+      expect(described_class.batch_count_with_timeout(model,
+        start: model.minimum(:id),
+        finish: model.maximum(:id) - 1, # Do not count the last record
+        batch_size: model.count - 2 # Ensure there are multiple batches
+      )).to eq({ status: :completed, count: model.count - 1 })
+    end
+
+    it 'returns a partial count when timeout elapses' do
+      stub_const('Gitlab::Database::BatchCounter::MIN_REQUIRED_BATCH_SIZE', 0)
+
+      expect(::Gitlab::Metrics::System).to receive(:monotonic_time).and_return(1, 10, 300)
+
+      expect(
+        described_class.batch_count_with_timeout(model, batch_size: 1, timeout: 250.seconds)
+      ).to eq({ status: :timeout, partial_results: 1, continue_from: model.minimum(:id) + 1 })
+    end
+
+    it 'starts counting from a given partial result' do
+      expect(described_class.batch_count_with_timeout(model, partial_results: 3)).to eq({ status: :completed, count: 3 + model.count })
+    end
+
+    it_behaves_like 'when a transaction is open' do
+      subject { described_class.batch_count_with_timeout(model) }
+    end
+
+    it_behaves_like 'when batch fetch query is canceled' do
+      let(:mode) { :itself }
+      let(:operation) { :count }
+      let(:operation_args) { nil }
+      let(:column) { nil }
+      let(:fallback) { { status: :cancelled } }
+
+      subject { described_class.method(:batch_count_with_timeout) }
+    end
+
+    context 'disallowed_configurations' do
+      include_examples 'disallowed configurations', :batch_count do
+        let(:args) { [Issue] }
+        let(:default_batch_size) { Gitlab::Database::BatchCounter::DEFAULT_BATCH_SIZE }
+      end
+
+      it 'raises an error if distinct count is requested' do
+        expect { described_class.batch_count_with_timeout(model.distinct(column)) }.to raise_error 'Use distinct count for optimized distinct counting'
+      end
+    end
+
+    context 'when a relation is grouped' do
+      let!(:one_more_issue) { create(:issue, author: user, project: model.first.project) }
+
+      before do
+        stub_const('Gitlab::Database::BatchCounter::MIN_REQUIRED_BATCH_SIZE', 1)
+      end
+
+      context 'count by default column' do
+        let(:count) do
+          described_class.batch_count_with_timeout(model.group(column), batch_size: 2)
+        end
+
+        it 'counts grouped records' do
+          expect(count).to eq({ status: :completed, count: { user.id => 4, another_user.id => 2 } })
+        end
+      end
+    end
+  end
+
   describe '#batch_distinct_count' do
     it 'counts with column field' do
       expect(described_class.batch_distinct_count(model, column)).to eq(2)
@@ -242,7 +349,7 @@ RSpec.describe Gitlab::Database::BatchCount do
       expect(described_class.batch_distinct_count(model, column, batch_size: small_batch_size)).to eq(fallback)
     end
 
-    it 'counts with a small edge case batch_sizes than result' do
+    it 'produces the same result with different batch sizes' do
       stub_const('Gitlab::Database::BatchCounter::MIN_REQUIRED_BATCH_SIZE', 0)
 
       [1, 2, 4, 5, 6].each { |i| expect(described_class.batch_distinct_count(model, column, batch_size: i)).to eq(2) }
@@ -386,56 +493,18 @@ RSpec.describe Gitlab::Database::BatchCount do
   end
 
   describe '#batch_average' do
-    let(:model) { Issue }
     let(:column) { :weight }
 
     before do
-      Issue.update_all(weight: 2)
-    end
-
-    it 'returns the average of values in the given column' do
-      expect(described_class.batch_average(model, column)).to eq(2)
-    end
-
-    it 'works when given an Arel column' do
-      expect(described_class.batch_average(model, model.arel_table[column])).to eq(2)
-    end
-
-    it 'works with a batch size of 50K' do
-      expect(described_class.batch_average(model, column, batch_size: 50_000)).to eq(2)
-    end
-
-    it 'works with start and finish provided' do
-      expect(described_class.batch_average(model, column, start: model.minimum(:id), finish: model.maximum(:id))).to eq(2)
+      allow_next_instance_of(Gitlab::Database::BatchAverageCounter) do |instance|
+        allow(instance).to receive(:count).and_return
+      end
     end
 
-    it "defaults the batch size to #{Gitlab::Database::BatchCounter::DEFAULT_AVERAGE_BATCH_SIZE}" do
-      min_id = model.minimum(:id)
-      relation = instance_double(ActiveRecord::Relation)
-      allow(model).to receive_message_chain(:select, public_send: relation)
-      batch_end_id = min_id + calculate_batch_size(Gitlab::Database::BatchCounter::DEFAULT_AVERAGE_BATCH_SIZE)
-
-      expect(relation).to receive(:where).with("id" => min_id..batch_end_id).and_return(double(send: 1))
+    it 'calls BatchAverageCounter' do
+      expect(Gitlab::Database::BatchAverageCounter).to receive(:new).with(model, column).and_call_original
 
       described_class.batch_average(model, column)
     end
-
-    it_behaves_like 'when a transaction is open' do
-      subject { described_class.batch_average(model, column) }
-    end
-
-    it_behaves_like 'disallowed configurations', :batch_average do
-      let(:args) { [model, column] }
-      let(:default_batch_size) { Gitlab::Database::BatchCounter::DEFAULT_AVERAGE_BATCH_SIZE }
-      let(:small_batch_size)   { Gitlab::Database::BatchCounter::DEFAULT_AVERAGE_BATCH_SIZE - 1 }
-    end
-
-    it_behaves_like 'when batch fetch query is canceled' do
-      let(:mode) { :itself }
-      let(:operation) { :average }
-      let(:operation_args) { [column] }
-
-      subject { described_class.method(:batch_average) }
-    end
   end
 end