Implement uncached behavior for multiple databases

Add a Gitlab::Database.all_uncached method which implements
ActiveRecord::Base.uncached behavior that works with multiple
databases.

The new method wraps the given block in an #uncached call for each base
model, so that the behavior will be correct on any database connection.

.rubocop_todo/database/multiple_databases.yml

@@ -19,7 +19,6 @@ Database/MultipleDatabases:
   - lib/gitlab/database/partitioning_migration_helpers/backfill_partitioned_table.rb
   - lib/gitlab/database.rb
   - lib/gitlab/health_checks/db_check.rb
-  - lib/gitlab/import_export/group/relation_tree_restorer.rb
   - lib/gitlab/seeder.rb
   - spec/db/schema_spec.rb
   - spec/initializers/database_config_spec.rb

lib/gitlab/database.rb

@@ -195,6 +195,16 @@ module Gitlab
       MAX_TIMESTAMP_VALUE > timestamp ? timestamp : MAX_TIMESTAMP_VALUE.dup
     end
 
+    def self.all_uncached(&block)
+      # Calls to #uncached only disable caching for the current connection. Since the load balancer
+      # can potentially upgrade from read to read-write mode (using a different connection), we specify
+      # up-front that we'll explicitly use the primary for the duration of the operation.
+      Gitlab::Database::LoadBalancing::Session.current.use_primary do
+        base_models = database_base_models.values
+        base_models.reduce(block) { |blk, model| -> { model.uncached(&blk) } }.call
+      end
+    end
+
     def self.allow_cross_joins_across_databases(url:)
       # this method is implemented in:
       # spec/support/database/prevent_cross_joins.rb

lib/gitlab/import_export/group/relation_tree_restorer.rb

@@ -29,7 +29,7 @@ module Gitlab
         end
 
         def restore
-          ActiveRecord::Base.uncached do
+          Gitlab::Database.all_uncached do
             ActiveRecord::Base.no_touching do
               update_params!
 

spec/lib/gitlab/database_spec.rb

@@ -279,6 +279,46 @@ RSpec.describe Gitlab::Database do
     end
   end
 
+  describe '.all_uncached' do
+    let(:base_model) do
+      Class.new do
+        def self.uncached
+          @uncached = true
+
+          yield
+        end
+      end
+    end
+
+    let(:model1) { Class.new(base_model) }
+    let(:model2) { Class.new(base_model) }
+
+    before do
+      allow(described_class).to receive(:database_base_models)
+        .and_return({ model1: model1, model2: model2 }.with_indifferent_access)
+    end
+
+    it 'wraps the given block in uncached calls for each primary connection', :aggregate_failures do
+      expect(model1.instance_variable_get(:@uncached)).to be_nil
+      expect(model2.instance_variable_get(:@uncached)).to be_nil
+
+      expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_primary).and_yield
+
+      expect(model2).to receive(:uncached).and_call_original
+      expect(model1).to receive(:uncached).and_call_original
+
+      yielded_to_block = false
+      described_class.all_uncached do
+        expect(model1.instance_variable_get(:@uncached)).to be(true)
+        expect(model2.instance_variable_get(:@uncached)).to be(true)
+
+        yielded_to_block = true
+      end
+
+      expect(yielded_to_block).to be(true)
+    end
+  end
+
   describe '.read_only?' do
     it 'returns false' do
       expect(described_class.read_only?).to eq(false)