Move load-balancing logic in ApplicationRecord to EE

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

app/models/application_record.rb

@@ -52,13 +52,11 @@ class ApplicationRecord < ActiveRecord::Base
 
   # Start a new transaction with a shorter-than-usual statement timeout. This is
   # currently one third of the default 15-second timeout
-  def self.with_fast_read_statement_timeout
-    ::Gitlab::Database::LoadBalancing::Session.current.use_replica_if_possible do
-      transaction(requires_new: true) do
-        connection.exec_query("SET LOCAL statement_timeout = 5000")
+  def self.with_fast_read_statement_timeout(timeout_ms = 5000)
+    transaction(requires_new: true) do
+      connection.exec_query("SET LOCAL statement_timeout = #{timeout_ms}")
 
-        yield
-      end
+      yield
     end
   end
 
@@ -81,3 +79,5 @@ class ApplicationRecord < ActiveRecord::Base
     enum(enum_mod.key => values)
   end
 end
+
+ApplicationRecord.prepend_if_ee('EE::ApplicationRecord')

ee/app/models/ee/application_record.rb

+# frozen_string_literal: true
+
+module EE
+  module ApplicationRecord
+    extend ActiveSupport::Concern
+
+    class_methods do
+      extend ::Gitlab::Utils::Override
+
+      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
+        end
+      end
+    end
+  end
+end

ee/spec/models/ee/application_record_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe ApplicationRecord do
+  describe '.with_fast_read_statement_timeout' do
+    let(:session) { double(:session) }
+
+    before do
+      allow(::Gitlab::Database::LoadBalancing::Session).to receive(:current).and_return(session)
+      allow(session).to receive(:use_replica_if_possible).and_yield
+    end
+
+    it 'yields control' do
+      expect do |blk|
+        described_class.with_fast_read_statement_timeout(&blk)
+      end.to yield_control.once
+    end
+
+    context 'when the query runs faster than configured timeout' do
+      it 'executes the query without error' do
+        result = nil
+
+        expect do
+          described_class.with_fast_read_statement_timeout(100) do
+            result = described_class.connection.exec_query('SELECT 1')
+          end
+        end.not_to raise_error
+
+        expect(result).not_to be_nil
+      end
+    end
+
+    # This query hangs for 10ms and then gets cancelled.  As there is no
+    # other way to test the timeout for sure, 10ms of waiting seems to be
+    # reasonable!
+    context 'when the query runs longer than configured timeout' do
+      it 'cancels the query and raiss an exception' do
+        expect do
+          described_class.with_fast_read_statement_timeout(10) do
+            described_class.connection.exec_query('SELECT pg_sleep(0.1)')
+          end
+        end.to raise_error(ActiveRecord::QueryCanceled)
+      end
+    end
+  end
+end

spec/models/application_record_spec.rb

@@ -100,4 +100,33 @@ RSpec.describe ApplicationRecord do
       expect(User.where_exists(User.limit(1))).to eq([user])
     end
   end
+
+  describe '.with_fast_read_statement_timeout' do
+    context 'when the query runs faster than configured timeout' do
+      it 'executes the query without error' do
+        result = nil
+
+        expect do
+          described_class.with_fast_read_statement_timeout(100) do
+            result = described_class.connection.exec_query('SELECT 1')
+          end
+        end.not_to raise_error
+
+        expect(result).not_to be_nil
+      end
+    end
+
+    # This query hangs for 10ms and then gets cancelled.  As there is no
+    # other way to test the timeout for sure, 10ms of waiting seems to be
+    # reasonable!
+    context 'when the query runs longer than configured timeout' do
+      it 'cancels the query and raiss an exception' do
+        expect do
+          described_class.with_fast_read_statement_timeout(10) do
+            described_class.connection.exec_query('SELECT pg_sleep(0.1)')
+          end
+        end.to raise_error(ActiveRecord::QueryCanceled)
+      end
+    end
+  end
 end