Extract ReltuplesCountStrategy.

lib/gitlab/database/count.rb

@@ -54,27 +54,43 @@ module Gitlab
       # @param [Array] table names
       # @returns [Hash] Table name to count mapping (e.g. { 'projects' => 5, 'users' => 100 })
       def self.reltuples_from_recently_updated(table_names)
+        ReltuplesCountStrategy.new(table_names).count
+      end
+
+      class ReltuplesCountStrategy
+        attr_reader :table_names
+
+        # @param [Array] table names
+        def initialize(table_names)
+          @table_names = table_names
+        end
+
+        # Returns a hash of the table names that have recently updated tuples.
+        #
+        # @returns [Hash] Table name to count mapping (e.g. { 'projects' => 5, 'users' => 100 })
+        def count
           query = postgresql_estimate_query(table_names)
           rows = []
 
           # Querying tuple stats only works on the primary. Due to load
-        # balancing, we need to ensure this query hits the load balancer.  The
           # easiest way to do this is to start a transaction.
           ActiveRecord::Base.transaction do
             rows = ActiveRecord::Base.connection.select_all(query)
           end
 
           rows.each_with_object({}) { |row, data| data[row['table_name']] = row['estimate'].to_i }
-      rescue *CONNECTION_ERRORS
+        rescue *CONNECTION_ERRORS => e
           {}
         end
 
+        private
+
         # Generates the PostgreSQL query to return the tuples for tables
         # that have been vacuumed or analyzed in the last hour.
         #
         # @param [Array] table names
         # @returns [Hash] Table name to count mapping (e.g. { 'projects' => 5, 'users' => 100 })
-      def self.postgresql_estimate_query(table_names)
+        def postgresql_estimate_query(table_names)
           time = "to_timestamp(#{1.hour.ago.to_i})"
           <<~SQL
           SELECT pg_class.relname AS table_name, reltuples::bigint AS estimate FROM pg_class
@@ -85,4 +101,5 @@ module Gitlab
         end
       end
     end
+  end
 end

spec/lib/gitlab/database/count_spec.rb

@@ -35,7 +35,7 @@ describe Gitlab::Database::Count do
 
       describe 'no permission' do
         it 'falls back to standard query' do
-          allow(described_class).to receive(:postgresql_estimate_query).and_raise(PG::InsufficientPrivilege)
+          allow(ActiveRecord::Base).to receive(:transaction).and_raise(PG::InsufficientPrivilege)
 
           expect(Project).to receive(:count).and_call_original
           expect(Identity).to receive(:count).and_call_original