spec/lib/gitlab/database/count/tablesample_count_strategy_spec.rb


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77

require 'spec_helper'

describe Gitlab::Database::Count::TablesampleCountStrategy do
  before do
    create_list(:project, 3)
    create(:identity)
    create(:group)
  end

  let(:models) { [Project, Identity, Group, Namespace] }
  let(:strategy) { described_class.new(models) }

  subject { strategy.count }

  describe '#count', :postgresql do
    let(:estimates) do
      {
        Project => threshold + 1,
        Identity => threshold - 1,
        Group => threshold + 1,
        Namespace => threshold + 1
      }
    end
    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)
        expect(Group).not_to receive(:count)
        expect(Namespace).not_to receive(:count)

        result = subject
        expect(result[Project]).to eq(3)
        expect(result[Group]).to eq(1)
        expect(result[Namespace]).to eq(4)
      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
    before do
      stub_feature_flags(tablesample_counts: true)
    end

    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