Add latest changes from gitlab-org/gitlab@14-4-stable-eev14.4.0-rc42

2 files changed, 74 insertions, 1 deletions

diff --git a/spec/support_specs/database/multiple_databases_spec.rb b/spec/support_specs/database/multiple_databases_spec.rb
new file mode 100644
index 00000000000..6ad15fd6594
--- /dev/null
+++ b/spec/support_specs/database/multiple_databases_spec.rb
@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe 'Database::MultipleDatabases' do
+  describe '.with_reestablished_active_record_base' do
+    context 'when doing establish_connection' do
+      context 'on ActiveRecord::Base' do
+        it 'raises exception' do
+          expect { ActiveRecord::Base.establish_connection(:main) }.to raise_error /Cannot re-establish/
+        end
+
+        context 'when using with_reestablished_active_record_base' do
+          it 'does not raise exception' do
+            with_reestablished_active_record_base do
+              expect { ActiveRecord::Base.establish_connection(:main) }.not_to raise_error
+            end
+          end
+        end
+      end
+
+      context 'on Ci::CiDatabaseRecord' do
+        before do
+          skip_if_multiple_databases_not_setup
+        end
+
+        it 'raises exception' do
+          expect { Ci::CiDatabaseRecord.establish_connection(:ci) }.to raise_error /Cannot re-establish/
+        end
+
+        context 'when using with_reestablished_active_record_base' do
+          it 'does not raise exception' do
+            with_reestablished_active_record_base do
+              expect { Ci::CiDatabaseRecord.establish_connection(:main) }.not_to raise_error
+            end
+          end
+        end
+      end
+    end
+
+    context 'when trying to access connection' do
+      context 'when reconnect is true' do
+        it 'does not raise exception' do
+          with_reestablished_active_record_base(reconnect: true) do
+            expect { ActiveRecord::Base.connection.execute("SELECT 1") }.not_to raise_error # rubocop:disable Database/MultipleDatabases
+          end
+        end
+      end
+
+      context 'when reconnect is false' do
+        it 'does raise exception' do
+          with_reestablished_active_record_base(reconnect: false) do
+            expect { ActiveRecord::Base.connection.execute("SELECT 1") }.to raise_error(ActiveRecord::ConnectionNotEstablished) # rubocop:disable Database/MultipleDatabases
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/spec/support_specs/database/prevent_cross_joins_spec.rb b/spec/support_specs/database/prevent_cross_joins_spec.rb
index e9a95fe77a5..0fbcd190c2c 100644
--- a/spec/support_specs/database/prevent_cross_joins_spec.rb
+++ b/spec/support_specs/database/prevent_cross_joins_spec.rb
@@ -22,6 +22,12 @@ RSpec.describe Database::PreventCrossJoins do
           described_class::CrossJoinAcrossUnsupportedTablesError)
       end
 
+      context 'when annotation is used' do
+        it 'does not raise exception' do
+          expect { main_and_ci_allowed_via_annotate }.not_to raise_error
+        end
+      end
+
       context 'when allow_cross_joins_across_databases is used' do
         it 'does not raise exception' do
           expect { main_and_ci_query_allowlisted }.not_to raise_error
@@ -52,6 +58,12 @@ RSpec.describe Database::PreventCrossJoins do
     end
   end
 
+  def main_and_ci_allowed_via_annotate
+    main_and_ci_query do |relation|
+      relation.allow_cross_joins_across_databases(url: 'http://issue-url')
+    end
+  end
+
   def main_only_query
     Issue.joins(:project).last
   end
@@ -61,6 +73,8 @@ RSpec.describe Database::PreventCrossJoins do
   end
 
   def main_and_ci_query
-    Ci::Build.joins(:project).last
+    relation = Ci::Build.joins(:project)
+    relation = yield(relation) if block_given?
+    relation.last
   end
 end