From 231bcc7cd2916c6247550e4a7dce21f56e6a0d52 Mon Sep 17 00:00:00 2001
From: eileencodes <eileencodes@gmail.com>
Date: Thu, 29 Oct 2020 16:45:50 -0400
Subject: [PATCH] Implement `connecting_to` method

Sometimes you need to have a different default connection but aren't
calling the connection with a block. An example is booting a console in
`reading` mode. This adds the ability for a script to set a specific
connection on boot while preserving the behavior of `connected_to` for
application code.
---
 .../lib/active_record/connection_handling.rb  | 17 +++++++++
 activerecord/test/cases/base_test.rb          | 37 +++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/activerecord/lib/active_record/connection_handling.rb b/activerecord/lib/active_record/connection_handling.rb
index 7a5ba005f0..0b9dc05690 100644
--- a/activerecord/lib/active_record/connection_handling.rb
+++ b/activerecord/lib/active_record/connection_handling.rb
@@ -172,6 +172,23 @@ module ActiveRecord
       end
     end
 
+    # Use a specified connection.
+    #
+    # This method is useful for ensuring that a specific connection is
+    # being used. For example, when booting a console in readonly mode.
+    #
+    # It is not recommended to use this method in a request since it
+    # does not yield to a block like `connected_to`.
+    def connecting_to(role: default_role, shard: default_shard, prevent_writes: false)
+      if legacy_connection_handling
+        raise NotImplementedError, "`connecting_to` is not available with `legacy_connection_handling`."
+      end
+
+      prevent_writes = true if role == reading_role
+
+      self.connected_to_stack << { role: role, shard: shard, prevent_writes: prevent_writes, klass: self }
+    end
+
     def while_preventing_writes(enabled = true, &block)
       connected_to(role: current_role, prevent_writes: enabled, &block)
     end
diff --git a/activerecord/test/cases/base_test.rb b/activerecord/test/cases/base_test.rb
index 603802bc9d..9f0f65498a 100644
--- a/activerecord/test/cases/base_test.rb
+++ b/activerecord/test/cases/base_test.rb
@@ -1651,4 +1651,41 @@ class BasicsTest < ActiveRecord::TestCase
       assert AbstractCompany.connected_to?(role: :reading, shard: :default)
     end
   end
+
+  test "#connecting_to with role" do
+    AbstractCompany.connecting_to(role: :reading)
+
+    assert AbstractCompany.connected_to?(role: :reading)
+    assert AbstractCompany.current_preventing_writes
+  ensure
+    ActiveRecord::Base.connected_to_stack.pop
+  end
+
+  test "#connecting_to with role and shard" do
+    AbstractCompany.connecting_to(role: :reading, shard: :default)
+
+    assert AbstractCompany.connected_to?(role: :reading, shard: :default)
+  ensure
+    ActiveRecord::Base.connected_to_stack.pop
+  end
+
+  test "#connecting_to with prevent_writes" do
+    AbstractCompany.connecting_to(role: :writing, prevent_writes: true)
+
+    assert AbstractCompany.connected_to?(role: :writing)
+    assert AbstractCompany.current_preventing_writes
+  ensure
+    ActiveRecord::Base.connected_to_stack.pop
+  end
+
+  test "#connecting_to doesn't work with legacy connection handling" do
+    old_value = ActiveRecord::Base.legacy_connection_handling
+    ActiveRecord::Base.legacy_connection_handling = true
+
+    assert_raises NotImplementedError do
+      AbstractCompany.connecting_to(role: :writing, prevent_writes: true)
+    end
+  ensure
+    ActiveRecord::Base.legacy_connection_handling = old_value
+  end
 end