Implement TablesampleCountStrategy.

A tablesample count executes in two phases:
* Estimate table sizes based on reltuples.
* Based on the estimate:
  * If the table is considered 'small', execute an exact relation count.
  * Otherwise, count on a sample of the table using TABLESAMPLE.

The size of the sample is chosen in a way that we always roughly scan
the same amount of rows (see TABLESAMPLE_ROW_TARGET).

1 files changed, 53 insertions, 0 deletions

diff --git a/spec/lib/gitlab/database/count_spec.rb b/spec/lib/gitlab/database/count_spec.rb
index 699422abb6b..d6668cee23e 100644
--- a/spec/lib/gitlab/database/count_spec.rb
+++ b/spec/lib/gitlab/database/count_spec.rb
@@ -123,4 +123,57 @@ describe Gitlab::Database::Count do
       end
     end
   end
+
+  describe Gitlab::Database::Count::TablesampleCountStrategy do
+    subject { strategy.count }
+    let(:strategy) { described_class.new(models) }
+
+    describe '#count' do
+      let(:estimates) { { Project => threshold + 1, Identity => threshold - 1 } }
+      let(:threshold) { Gitlab::Database::Count::TablesampleCountStrategy::EXACT_COUNT_THRESHOLD }
+
+      before do
+        allow(strategy).to receive(:size_estimates).with(check_statistics: false).and_return(estimates)
+      end
+
+      context 'for tables with an estimated small size' do
+        it 'performs an exact count' do
+          expect(Identity).to receive(:count).and_call_original
+
+          expect(subject).to include({ Identity => 1 })
+        end
+      end
+
+      context 'for tables with an estimated large size' do
+        it 'performs a tablesample count' do
+          expect(Project).not_to receive(:count)
+
+          result = subject
+          expect(result[Project]).to eq(3)
+        end
+      end
+
+      context 'insufficient permissions' do
+        it 'returns an empty hash' do
+          allow(strategy).to receive(:size_estimates).and_raise(PG::InsufficientPrivilege)
+
+          expect(subject).to eq({})
+        end
+      end
+    end
+
+    describe '.enabled?' do
+      it 'is enabled for PostgreSQL' do
+        allow(Gitlab::Database).to receive(:postgresql?).and_return(true)
+
+        expect(described_class.enabled?).to be_truthy
+      end
+
+      it 'is disabled for MySQL' do
+        allow(Gitlab::Database).to receive(:postgresql?).and_return(false)
+
+        expect(described_class.enabled?).to be_falsey
+      end
+    end
+  end
 end