Update Gitlab::Database::Count to work with decomposed database

Split models that need to be counted per database, count, and combine
the results.

https://gitlab.com/gitlab-org/gitlab/-/issues/345279

lib/gitlab/database/count/reltuples_count_strategy.rb

@@ -32,12 +32,12 @@ module Gitlab
         # Models using single-type inheritance (STI) don't work with
         # reltuple count estimates. We just have to ignore them and
         # use another strategy to compute them.
-        def non_sti_models
+        def non_sti_models(models)
           models.reject { |model| sti_model?(model) }
         end
 
-        def non_sti_table_names
-          non_sti_models.map(&:table_name)
+        def non_sti_table_names(models)
+          non_sti_models(models).map(&:table_name)
         end
 
         def sti_model?(model)
@@ -45,21 +45,34 @@ module Gitlab
             model.base_class != model
         end
 
-        def table_names
-          models.map(&:table_name)
+        def table_to_model_mapping
+          @table_to_model_mapping ||= models.each_with_object({}) { |model, h| h[model.table_name] = model }
+        end
+
+        def table_to_model(table_name)
+          table_to_model_mapping[table_name]
         end
 
         def size_estimates(check_statistics: true)
-          table_to_model = models.each_with_object({}) { |model, h| h[model.table_name] = model }
-
-          # Querying tuple stats only works on the primary. Due to load balancing, the
-          # easiest way to do this is to start a transaction.
-          ActiveRecord::Base.transaction do # rubocop: disable Database/MultipleDatabases
-            get_statistics(non_sti_table_names, check_statistics: check_statistics).each_with_object({}) do |row, data|
-              model = table_to_model[row.table_name]
-              data[model] = row.estimate
+          results = {}
+
+          models.group_by { |model| model.connection_db_config.name }.map do |db_name, models_for_db|
+            base_model = Gitlab::Database.database_base_models[db_name]
+            tables = non_sti_table_names(models_for_db)
+
+            # Querying tuple stats only works on the primary. Due to load balancing, the
+            # easiest way to do this is to start a transaction.
+            base_model.transaction do
+              Gitlab::Database::SharedModel.using_connection(base_model.connection) do
+                get_statistics(tables, check_statistics: check_statistics).each do |row|
+                  model = table_to_model(row.table_name)
+                  results[model] = row.estimate
+                end
+              end
             end
           end
+
+          results
         end
 
         # Generates the PostgreSQL query to return the tuples for tables

lib/gitlab/database/count/tablesample_count_strategy.rb

@@ -61,7 +61,7 @@ module Gitlab
             #{where_clause(model)}
           SQL
 
-          rows = ActiveRecord::Base.connection.select_all(query) # rubocop: disable Database/MultipleDatabases
+          rows = model.connection.select_all(query)
 
           Integer(rows.first['count'])
         end

lib/gitlab/database/migrations/background_migration_helpers.rb

@@ -258,7 +258,9 @@ module Gitlab
 
           # We keep track of the estimated number of tuples to reason later
           # about the overall progress of a migration.
-          migration.total_tuple_count = Gitlab::Database::PgClass.for_table(batch_table_name)&.cardinality_estimate
+          migration.total_tuple_count = Gitlab::Database::SharedModel.using_connection(connection) do
+            Gitlab::Database::PgClass.for_table(batch_table_name)&.cardinality_estimate
+          end
           migration.save!
 
           migration

lib/gitlab/database/pg_class.rb

@@ -2,7 +2,7 @@
 
 module Gitlab
   module Database
-    class PgClass < ActiveRecord::Base
+    class PgClass < SharedModel
       self.table_name = 'pg_class'
 
       def self.for_table(relname)

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

@@ -5,24 +5,24 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::Count::ReltuplesCountStrategy do
   before do
     create_list(:project, 3)
-    create(:identity)
+    create_list(:ci_instance_variable, 2)
   end
 
   subject { described_class.new(models).count }
 
   describe '#count' do
-    let(:models) { [Project, Identity] }
+    let(:models) { [Project, Ci::InstanceVariable] }
 
     context 'when reltuples is up to date' do
       before do
-        ActiveRecord::Base.connection.execute('ANALYZE projects')
-        ActiveRecord::Base.connection.execute('ANALYZE identities')
+        Project.connection.execute('ANALYZE projects')
+        Ci::InstanceVariable.connection.execute('ANALYZE ci_instance_variables')
       end
 
       it 'uses statistics to do the count' do
         models.each { |model| expect(model).not_to receive(:count) }
 
-        expect(subject).to eq({ Project => 3, Identity => 1 })
+        expect(subject).to eq({ Project => 3, Ci::InstanceVariable => 2 })
       end
     end
 
@@ -31,7 +31,7 @@ RSpec.describe Gitlab::Database::Count::ReltuplesCountStrategy do
 
       before do
         models.each do |model|
-          ActiveRecord::Base.connection.execute("ANALYZE #{model.table_name}")
+          model.connection.execute("ANALYZE #{model.table_name}")
         end
       end
 
@@ -45,7 +45,9 @@ RSpec.describe Gitlab::Database::Count::ReltuplesCountStrategy do
 
     context 'insufficient permissions' do
       it 'returns an empty hash' do
-        allow(ActiveRecord::Base).to receive(:transaction).and_raise(PG::InsufficientPrivilege)
+        Gitlab::Database.database_base_models.each_value do |base_model|
+          allow(base_model).to receive(:transaction).and_raise(PG::InsufficientPrivilege)
+        end
 
         expect(subject).to eq({})
       end

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

@@ -5,11 +5,12 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::Count::TablesampleCountStrategy do
   before do
     create_list(:project, 3)
+    create_list(:ci_instance_variable, 2)
     create(:identity)
     create(:group)
   end
 
-  let(:models) { [Project, Identity, Group, Namespace] }
+  let(:models) { [Project, Ci::InstanceVariable, Identity, Group, Namespace] }
   let(:strategy) { described_class.new(models) }
 
   subject { strategy.count }
@@ -20,7 +21,8 @@ RSpec.describe Gitlab::Database::Count::TablesampleCountStrategy do
         Project => threshold + 1,
         Identity => threshold - 1,
         Group => threshold + 1,
-        Namespace => threshold + 1
+        Namespace => threshold + 1,
+        Ci::InstanceVariable => threshold + 1
       }
     end
 
@@ -43,12 +45,14 @@ RSpec.describe Gitlab::Database::Count::TablesampleCountStrategy do
         expect(Project).not_to receive(:count)
         expect(Group).not_to receive(:count)
         expect(Namespace).not_to receive(:count)
+        expect(Ci::InstanceVariable).not_to receive(:count)
 
         result = subject
         expect(result[Project]).to eq(3)
         expect(result[Group]).to eq(1)
         # 1-Group, 3 namespaces for each project and 3 project namespaces for each project
         expect(result[Namespace]).to eq(7)
+        expect(result[Ci::InstanceVariable]).to eq(2)
       end
     end