Simplify to use models instead of tables.

2 files changed, 25 insertions, 25 deletions

diff --git a/lib/gitlab/database/count.rb b/lib/gitlab/database/count.rb
index 7b7fd3148fa..55f0cbd18be 100644
--- a/lib/gitlab/database/count.rb
+++ b/lib/gitlab/database/count.rb
@@ -26,24 +26,17 @@ module Gitlab
       # @param [Array]
       # @return [Hash] of Model -> count mapping
       def self.approximate_counts(models)
-        table_to_model_map = models.each_with_object({}) do |model, hash|
-          hash[model.table_name] = model
-        end
-
-        table_names = table_to_model_map.keys
-        counts_by_table_name = Gitlab::Database.postgresql? ? reltuples_from_recently_updated(table_names) : {}
+        counts_by_model = {}
 
-        # Convert table -> count to Model -> count
-        counts_by_model = counts_by_table_name.each_with_object({}) do |pair, hash|
-          model = table_to_model_map[pair.first]
-          hash[model] = pair.second
+        if Gitlab::Database.postgresql?
+          #counts_by_model = ReltuplesCountStrategy.new(models).count
+          counts_by_model = reltuples_from_recently_updated(models)
         end
 
-        missing_tables = table_names - counts_by_table_name.keys
+        missing_models = models - counts_by_model.keys
 
-        missing_tables.each do |table|
-          model = table_to_model_map[table]
-          counts_by_model[model] = model.count
+        ExactCountStrategy.new(missing_models).count.each do |model, count|
+          counts_by_model[model] = count
         end
 
         counts_by_model
@@ -51,10 +44,10 @@ module Gitlab
 
       # Returns a hash of the table names that have recently updated tuples.
       #
-      # @param [Array] table names
+      # @param [Array] models to count
       # @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
+      def self.reltuples_from_recently_updated(models)
+        ReltuplesCountStrategy.new(models).count
       end
 
       class ExactCountStrategy
@@ -71,11 +64,9 @@ module Gitlab
       end
 
       class ReltuplesCountStrategy
-        attr_reader :table_names
-
-        # @param [Array] table names
-        def initialize(table_names)
-          @table_names = table_names
+        attr_reader :models
+        def initialize(models)
+          @models = models
         end
 
         # Returns a hash of the table names that have recently updated tuples.
@@ -91,13 +82,22 @@ module Gitlab
             rows = ActiveRecord::Base.connection.select_all(query)
           end
 
-          rows.each_with_object({}) { |row, data| data[row['table_name']] = row['estimate'].to_i }
+          table_to_model = models.each_with_object({}) { |model, h| h[model.table_name] = model }
+
+          rows.each_with_object({}) do |row, data|
+            model = table_to_model[row['table_name']]
+            data[model] = row['estimate'].to_i
+          end
         rescue *CONNECTION_ERRORS => e
           {}
         end
 
         private
 
+        def table_names
+          models.map(&:table_name)
+        end
+
         # Generates the PostgreSQL query to return the tuples for tables
         # that have been vacuumed or analyzed in the last hour.
         #
diff --git a/spec/lib/gitlab/database/count_spec.rb b/spec/lib/gitlab/database/count_spec.rb
index 3f4e047a9e5..9b4b61d7c94 100644
--- a/spec/lib/gitlab/database/count_spec.rb
+++ b/spec/lib/gitlab/database/count_spec.rb
@@ -25,7 +25,7 @@ describe Gitlab::Database::Count do
     context 'with PostgreSQL', :postgresql do
       describe 'when reltuples have not been updated' do
         it 'counts all models the normal way' do
-          expect(described_class).to receive(:reltuples_from_recently_updated).with(%w(projects identities)).and_return({})
+          expect(described_class).to receive(:reltuples_from_recently_updated).with(models).and_return({})
 
           expect(Project).to receive(:count).and_call_original
           expect(Identity).to receive(:count).and_call_original
@@ -45,7 +45,7 @@ describe Gitlab::Database::Count do
 
       describe 'when some reltuples have been updated' do
         it 'counts projects in the fast way' do
-          expect(described_class).to receive(:reltuples_from_recently_updated).with(%w(projects identities)).and_return({ 'projects' => 3 })
+          expect(described_class).to receive(:reltuples_from_recently_updated).with(models).and_return({ Project => 3 })
 
           expect(Project).not_to receive(:count).and_call_original
           expect(Identity).to receive(:count).and_call_original