Address feedback in the reviews

Issue https://gitlab.com/gitlab-org/gitlab/-/issues/322133

ee/app/models/ee/application_record_helpers.rb

@@ -11,11 +11,7 @@ module EE
       override :with_fast_read_statement_timeout
       def with_fast_read_statement_timeout(timeout_ms = 5000)
         ::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible do
-          transaction(requires_new: true) do
-            connection.exec_query("SET LOCAL statement_timeout = #{timeout_ms}")
-
-            yield
-          end
+          super
         end
       end
     end

ee/lib/gitlab/database/load_balancing/session.rb

@@ -62,15 +62,11 @@ module Gitlab
         # - The current session already performed writes
         # - It prefers to use primary, aka, use_primary or use_primary! were called
         def use_replica_if_possible(&blk)
-          return yield if use_primary? || performed_write?
-
-          begin
-            used_replica = @use_replica
-            @use_replica = true
-            yield
-          ensure
-            @use_replica = used_replica && !use_primary? && !performed_write?
-          end
+          used_replica = @use_replica
+          @use_replica = true
+          yield
+        ensure
+          @use_replica = used_replica
         end
 
         def use_replica?

ee/spec/lib/gitlab/database/load_balancing_spec.rb

@@ -409,14 +409,27 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   # of the load balancer.
   # - A real model with a table backed behind is defined
   # - The load balancing module is set up for this module only, as to prevent
-  # breaking other tests. The replica configuraiton is cloned from the test
+  # breaking other tests. The replica configuration is cloned from the test
   # configuraiton.
   # - In each test, we listen to the SQL queries (via sql.active_record
-  # instrumentaiton) while triggering real queries from the defined model.
+  # instrumentation) while triggering real queries from the defined model.
   # - We assert the desinations (replica/primary) of the queries in order.
   describe 'LoadBalancing integration tests', :delete do
+    before(:all) do
+      ActiveRecord::Schema.define do
+        create_table :load_balancing_test, force: true do |t|
+          t.string :name, null: true
+        end
+      end
+    end
+
+    after(:all) do
+      ActiveRecord::Schema.define do
+        drop_table :load_balancing_test, force: true
+      end
+    end
+
     shared_context 'LoadBalancing setup' do
-      let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
       let(:hosts) { [ActiveRecord::Base.configurations["development"]['host']] }
       let(:model) do
         Class.new(ApplicationRecord) do
@@ -425,11 +438,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       end
 
       before do
-        ActiveRecord::Schema.define do
-          create_table :load_balancing_test, force: true do |t|
-            t.string :name, null: true
-          end
-        end
+        stub_licensed_features(db_load_balancing: true)
         # Preloading testing class
         model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
 
@@ -446,9 +455,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
 
       after do
         subject.clear_configuration
-        ActiveRecord::Schema.define do
-          drop_table :load_balancing_test, force: true
-        end
       end
     end
 
@@ -531,17 +537,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
           true, [:primary, :primary, :primary, :primary]
         ],
 
-        # Read-only transaction
-        [
-          -> {
-            model.transaction do
-              model.first
-              model.where(name: 'test1').to_a
-            end
-          },
-          true, [:primary, :primary, :primary, :primary]
-        ],
-
         # use_primary
         [
           -> {