Add start/finish & more tests to batch counting

We add start/finish option and specs for the batch counting

lib/gitlab/database/batch_count.rb

@@ -19,8 +19,8 @@
 module Gitlab
   module Database
     module BatchCount
-      def batch_count(relation, column = nil, batch_size: nil)
-        BatchCounter.new(relation, column: column).count(batch_size: batch_size)
+      def batch_count(relation, column = nil, batch_size: nil, start: nil, finish: nil)
+        BatchCounter.new(relation, column: column).count(batch_size: batch_size, start: start, finish: finish)
       end
 
       def batch_distinct_count(relation, column = nil, batch_size: nil, start: nil, finish: nil)

lib/gitlab/usage_data.rb

@@ -240,9 +240,9 @@ module Gitlab
         {} # augmented in EE
       end
 
-      def count(relation, column = nil, fallback: -1, batch: true)
+      def count(relation, column = nil, fallback: -1, batch: true, start: nil, finish: nil)
         if batch && Feature.enabled?(:usage_ping_batch_counter, default_enabled: true)
-          Gitlab::Database::BatchCount.batch_count(relation, column)
+          Gitlab::Database::BatchCount.batch_count(relation, column, start: start, finish: finish)
         else
           relation.count
         end

spec/lib/gitlab/database/batch_count_spec.rb

@@ -3,6 +3,8 @@
 require 'spec_helper'
 
 describe Gitlab::Database::BatchCount do
+  let_it_be(:fallback) { ::Gitlab::Database::BatchCounter::FALLBACK }
+  let_it_be(:small_batch_size) { ::Gitlab::Database::BatchCounter::MIN_REQUIRED_BATCH_SIZE - 1 }
   let(:model) { Issue }
   let(:column) { :author_id }
 
@@ -37,9 +39,8 @@ describe Gitlab::Database::BatchCount do
       expect(described_class.batch_count(model, batch_size: 50_000)).to eq(5)
     end
 
-    it 'will not count table with batch_size 1K' do
-      fallback = ::Gitlab::Database::BatchCounter::FALLBACK
-      expect(described_class.batch_count(model, batch_size: fallback / 2)).to eq(fallback)
+    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
@@ -57,6 +58,25 @@ describe Gitlab::Database::BatchCount do
         end.to raise_error 'BatchCount can not be run inside a transaction'
       end
     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)
+    end
+
+    context 'disallowed configurations' do
+      it 'returns fallback if start is bigger than finish' do
+        expect(described_class.batch_count(model, start: 1, finish: 0)).to eq(fallback)
+      end
+
+      it 'returns fallback if loops more than allowed' do
+        large_finish = Gitlab::Database::BatchCounter::MAX_ALLOWED_LOOPS * Gitlab::Database::BatchCounter::DEFAULT_BATCH_SIZE + 1
+        expect(described_class.batch_count(model, start: 1, finish: large_finish)).to eq(fallback)
+      end
+
+      it 'returns fallback if batch size is less than min required' do
+        expect(described_class.batch_count(model, batch_size: small_batch_size)).to eq(fallback)
+      end
+    end
   end
 
   describe '#batch_distinct_count' do
@@ -80,9 +100,8 @@ describe Gitlab::Database::BatchCount do
       expect(described_class.batch_distinct_count(model, column, batch_size: 50_000)).to eq(2)
     end
 
-    it 'will not count table with batch_size 1K' do
-      fallback = ::Gitlab::Database::BatchCounter::FALLBACK
-      expect(described_class.batch_distinct_count(model, column, batch_size: fallback / 2)).to eq(fallback)
+    it 'will not count table with a batch size less than allowed' 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
@@ -98,5 +117,20 @@ describe Gitlab::Database::BatchCount do
     it 'counts with User min and max as start and finish' do
       expect(described_class.batch_distinct_count(model, column, start: User.minimum(:id), finish: User.maximum(:id))).to eq(2)
     end
+
+    context 'disallowed configurations' do
+      it 'returns fallback if start is bigger than finish' do
+        expect(described_class.batch_distinct_count(model, column, start: 1, finish: 0)).to eq(fallback)
+      end
+
+      it 'returns fallback if loops more than allowed' do
+        large_finish = Gitlab::Database::BatchCounter::MAX_ALLOWED_LOOPS * Gitlab::Database::BatchCounter::DEFAULT_DISTINCT_BATCH_SIZE + 1
+        expect(described_class.batch_distinct_count(model, column, start: 1, finish: large_finish)).to eq(fallback)
+      end
+
+      it 'returns fallback if batch size is less than min required' do
+        expect(described_class.batch_distinct_count(model, column, batch_size: small_batch_size)).to eq(fallback)
+      end
+    end
   end
 end