Remove EE-specific ApplicationRecord code

This moves the database load balancing functionality into CE.

app/models/application_record.rb

@@ -53,10 +53,12 @@ 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(timeout_ms = 5000)
-    transaction(requires_new: true) do
-      connection.exec_query("SET LOCAL statement_timeout = #{timeout_ms}")
+    ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
+      transaction(requires_new: true) do
+        connection.exec_query("SET LOCAL statement_timeout = #{timeout_ms}")
 
-      yield
+        yield
+      end
     end
   end
 
@@ -85,5 +87,3 @@ class ApplicationRecord < ActiveRecord::Base
     enum(enum_mod.key => values)
   end
 end
-
-ApplicationRecord.prepend_mod_with('ApplicationRecordHelpers')

ee/app/models/ee/application_record_helpers.rb

-# frozen_string_literal: true
-
-module EE
-  # Intentionally pick a different name, to prevent naming conflicts
-  module ApplicationRecordHelpers
-    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.fallback_to_replicas_for_ambiguous_queries do
-          super
-        end
-      end
-    end
-  end
-end

ee/spec/models/ee/application_record_helpers_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(:fallback_to_replicas_for_ambiguous_queries).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

@@ -132,5 +132,47 @@ RSpec.describe ApplicationRecord do
         end.to raise_error(ActiveRecord::QueryCanceled)
       end
     end
+
+    context 'with database load balancing' do
+      let(:session) { double(:session) }
+
+      before do
+        allow(::Gitlab::Database::LoadBalancing::Session).to receive(:current).and_return(session)
+        allow(session).to receive(:fallback_to_replicas_for_ambiguous_queries).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
 end