ibm_db 5.6.1 → 5.7.0

data/test/cases/connection_adapters/connection_handlers_sharding_db_test.rb

@@ -10,14 +10,6 @@ module ActiveRecord
 
       fixtures :people
 
-      def setup
-        @handler = ConnectionHandler.new
-        @owner_name = "ActiveRecord::Base"
-        db_config = ActiveRecord::Base.configurations.configs_for(env_name: "arunit", name: "primary")
-        @rw_pool = @handler.establish_connection(db_config)
-        @ro_pool = @handler.establish_connection(db_config, role: :reading)
-      end
-
       def teardown
         clean_up_connection_handler
       end
@@ -29,7 +21,7 @@ module ActiveRecord
             ActiveRecord::Base.establish_connection(db_config)
             assert_nothing_raised { Person.first }
 
-            assert_equal [:default, :shard_one], ActiveRecord::Base.connection_handler.send(:owner_to_pool_manager).fetch("ActiveRecord::Base").instance_variable_get(:@name_to_role_mapping).values.flat_map(&:keys).uniq
+            assert_equal [:default, :shard_one].sort, ActiveRecord::Base.connection_handler.send(:get_pool_manager, "ActiveRecord::Base").shard_names.sort
           end
         end
 
@@ -46,13 +38,14 @@ module ActiveRecord
           @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config
 
           ActiveRecord::Base.connects_to(shards: {
-            default: { writing: :primary },
-            shard_one: { writing: :primary_shard_one }
+            default: { writing: :primary, reading: :primary },
+            shard_one: { writing: :primary_shard_one, reading: :primary_shard_one }
           })
 
           base_pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool("ActiveRecord::Base")
           default_pool = ActiveRecord::Base.connection_handler.retrieve_connection_pool("ActiveRecord::Base", shard: :default)
 
+          assert_equal [:default, :shard_one], ActiveRecord::Base.connection_handler.send(:get_pool_manager, "ActiveRecord::Base").shard_names
           assert_equal base_pool, default_pool
           assert_equal "test/db/primary.sqlite3", default_pool.db_config.database
           assert_equal "primary", default_pool.db_config.name
@@ -135,7 +128,7 @@ module ActiveRecord
             assert_not ActiveRecord::Base.connected_to?(role: :writing, shard: :default)
             assert_not ActiveRecord::Base.connected_to?(role: :writing, shard: :shard_one)
             assert_not ActiveRecord::Base.connected_to?(role: :reading, shard: :shard_one)
-            assert_predicate ActiveRecord::Base.connection, :preventing_writes?
+            assert_predicate ActiveRecord::Base.lease_connection, :preventing_writes?
           end
 
           ActiveRecord::Base.connected_to(role: :writing, shard: :default) do
@@ -144,7 +137,7 @@ module ActiveRecord
             assert_not ActiveRecord::Base.connected_to?(role: :reading, shard: :default)
             assert_not ActiveRecord::Base.connected_to?(role: :reading, shard: :shard_one)
             assert_not ActiveRecord::Base.connected_to?(role: :writing, shard: :shard_one)
-            assert_not_predicate ActiveRecord::Base.connection, :preventing_writes?
+            assert_not_predicate ActiveRecord::Base.lease_connection, :preventing_writes?
           end
 
           ActiveRecord::Base.connected_to(role: :reading, shard: :shard_one) do
@@ -153,7 +146,7 @@ module ActiveRecord
             assert_not ActiveRecord::Base.connected_to?(role: :writing, shard: :shard_one)
             assert_not ActiveRecord::Base.connected_to?(role: :writing, shard: :default)
             assert_not ActiveRecord::Base.connected_to?(role: :reading, shard: :default)
-            assert_predicate ActiveRecord::Base.connection, :preventing_writes?
+            assert_predicate ActiveRecord::Base.lease_connection, :preventing_writes?
           end
 
           ActiveRecord::Base.connected_to(role: :writing, shard: :shard_one) do
@@ -162,7 +155,7 @@ module ActiveRecord
             assert_not ActiveRecord::Base.connected_to?(role: :reading, shard: :shard_one)
             assert_not ActiveRecord::Base.connected_to?(role: :reading, shard: :default)
             assert_not ActiveRecord::Base.connected_to?(role: :writing, shard: :default)
-            assert_not_predicate ActiveRecord::Base.connection, :preventing_writes?
+            assert_not_predicate ActiveRecord::Base.lease_connection, :preventing_writes?
           end
         ensure
           ActiveRecord::Base.configurations = @prev_configs
@@ -231,11 +224,8 @@ module ActiveRecord
         end
 
         def test_retrieve_connection_pool_with_invalid_shard
-          assert_not_nil @handler.retrieve_connection_pool("ActiveRecord::Base")
-          assert_nil @handler.retrieve_connection_pool("ActiveRecord::Base", shard: :foo)
-
-          assert_not_nil @handler.retrieve_connection_pool("ActiveRecord::Base", role: :reading)
-          assert_nil @handler.retrieve_connection_pool("ActiveRecord::Base", role: :reading, shard: :foo)
+          assert_not_nil ActiveRecord::Base.connection_handler.retrieve_connection_pool("ActiveRecord::Base")
+          assert_nil ActiveRecord::Base.connection_handler.retrieve_connection_pool("ActiveRecord::Base", shard: :foo)
         end
 
         def test_calling_connected_to_on_a_non_existent_shard_raises
@@ -243,13 +233,106 @@ module ActiveRecord
             default: { writing: :arunit, reading: :arunit }
           })
 
-          error = assert_raises ActiveRecord::ConnectionNotEstablished do
+          error = assert_raises ActiveRecord::ConnectionNotDefined do
             ActiveRecord::Base.connected_to(role: :reading, shard: :foo) do
               Person.first
             end
           end
 
-          assert_equal "No connection pool for 'ActiveRecord::Base' found for the 'foo' shard.", error.message
+          assert_equal "No database connection defined for 'foo' shard and 'reading' role.", error.message
+          assert_equal "ActiveRecord::Base", error.connection_name
+          assert_equal :foo, error.shard
+          assert_equal :reading, error.role
+        end
+
+        def test_calling_connected_to_on_a_non_existent_role_for_shard_raises
+          ActiveRecord::Base.connects_to(shards: {
+            default: { writing: :arunit, reading: :arunit },
+            shard_one: { writing: :arunit, reading: :arunit }
+          })
+
+          error = assert_raises ActiveRecord::ConnectionNotDefined do
+            ActiveRecord::Base.connected_to(role: :non_existent, shard: :shard_one) do
+              Person.first
+            end
+          end
+
+          assert_equal "No database connection defined for 'shard_one' shard and 'non_existent' role.", error.message
+          assert_equal "ActiveRecord::Base", error.connection_name
+          assert_equal :shard_one, error.shard
+          assert_equal :non_existent, error.role
+        end
+
+        def test_calling_connected_to_on_a_default_role_for_non_existent_shard_raises
+          ActiveRecord::Base.connects_to(shards: {
+            default: { writing: :arunit, reading: :arunit }
+          })
+
+          error = assert_raises ActiveRecord::ConnectionNotDefined do
+            ActiveRecord::Base.connected_to(shard: :foo) do
+              Person.first
+            end
+          end
+
+          assert_equal "No database connection defined for 'foo' shard.", error.message
+          assert_equal "ActiveRecord::Base", error.connection_name
+          assert_equal :foo, error.shard
+          assert_equal :writing, error.role
+        end
+
+        def test_cannot_swap_shards_while_prohibited
+          previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
+
+          config = {
+            "default_env" => {
+              "primary" => { "adapter" => "sqlite3", "database" => "test/db/primary.sqlite3" },
+              "primary_shard_one" => { "adapter" => "sqlite3", "database" => "test/db/primary_shard_one.sqlite3" }
+            }
+          }
+
+          @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config
+
+          ActiveRecord::Base.connects_to(shards: {
+            default: { writing: :primary },
+            shard_one: { writing: :primary_shard_one }
+          })
+
+          assert_raises(ArgumentError) do
+            ActiveRecord::Base.prohibit_shard_swapping do
+              ActiveRecord::Base.connected_to(role: :reading, shard: :default) do
+              end
+            end
+          end
+        ensure
+          ActiveRecord::Base.configurations = @prev_configs
+          ActiveRecord::Base.establish_connection(:arunit)
+          ENV["RAILS_ENV"] = previous_env
+        end
+
+        def test_can_swap_roles_while_shard_swapping_is_prohibited
+          previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
+
+          config = {
+            "default_env" => {
+              "primary" => { "adapter" => "sqlite3", "database" => "test/db/primary.sqlite3" },
+              "primary_replica" => { "adapter" => "sqlite3", "database" => "test/db/primary.sqlite3", "replica" => true }
+            }
+          }
+
+          @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config
+
+          ActiveRecord::Base.connects_to(shards: { default: { writing: :primary, reading: :primary_replica } })
+
+          assert_nothing_raised do
+            ActiveRecord::Base.prohibit_shard_swapping do # no exception
+              ActiveRecord::Base.connected_to(role: :reading) do
+              end
+            end
+          end
+        ensure
+          ActiveRecord::Base.configurations = @prev_configs
+          ActiveRecord::Base.establish_connection(:arunit)
+          ENV["RAILS_ENV"] = previous_env
         end
       end
 
@@ -267,15 +350,23 @@ module ActiveRecord
       class ShardConnectionTestModelB < SomeOtherBase
       end
 
+      def test_default_shard_is_chosen_by_first_key_or_default
+        SecondaryBase.connects_to shards: { not_default: { writing: { database: ":memory:", adapter: "sqlite3" } } }
+        SomeOtherBase.connects_to database: { writing: { database: ":memory:", adapter: "sqlite3" } }
+
+        assert_equal :not_default, SecondaryBase.default_shard
+        assert_equal :default, SomeOtherBase.default_shard
+      end
+
       def test_same_shards_across_clusters
         SecondaryBase.connects_to shards: { one: { writing: { database: ":memory:", adapter: "sqlite3" } } }
         SomeOtherBase.connects_to shards: { one: { writing: { database: ":memory:", adapter: "sqlite3" } } }
 
         ActiveRecord::Base.connected_to(role: :writing, shard: :one) do
-          ShardConnectionTestModel.connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
+          ShardConnectionTestModel.lease_connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
           ShardConnectionTestModel.create!(shard_key: "test_model_default")
 
-          ShardConnectionTestModelB.connection.execute("CREATE TABLE `shard_connection_test_model_bs` (shard_key VARCHAR (255))")
+          ShardConnectionTestModelB.lease_connection.execute("CREATE TABLE `shard_connection_test_model_bs` (shard_key VARCHAR (255))")
           ShardConnectionTestModelB.create!(shard_key: "test_model_b_default")
 
           assert_equal "test_model_default", ShardConnectionTestModel.where(shard_key: "test_model_default").first.shard_key
@@ -291,7 +382,7 @@ module ActiveRecord
 
         [:default, :one].each do |shard_name|
           ActiveRecord::Base.connected_to(role: :writing, shard: shard_name) do
-            ShardConnectionTestModel.connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
+            ShardConnectionTestModel.lease_connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
           end
         end
 
@@ -327,27 +418,27 @@ module ActiveRecord
 
         [:default, :one].each do |shard_name|
           ActiveRecord::Base.connected_to(role: :writing, shard: shard_name) do
-            ShardConnectionTestModel.connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
-            ShardConnectionTestModel.connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}')")
+            ShardConnectionTestModel.lease_connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
+            ShardConnectionTestModel.lease_connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}')")
           end
         end
 
         shard_one_latch = Concurrent::CountDownLatch.new
         shard_default_latch = Concurrent::CountDownLatch.new
 
-        ShardConnectionTestModel.connection
+        ShardConnectionTestModel.lease_connection
 
         thread = Thread.new do
-          ShardConnectionTestModel.connection
+          ShardConnectionTestModel.lease_connection
 
           shard_default_latch.wait
-          assert_equal "shard_key_default", ShardConnectionTestModel.connection.select_value("SELECT shard_key from shard_connection_test_models")
+          assert_equal "shard_key_default", ShardConnectionTestModel.lease_connection.select_value("SELECT shard_key from shard_connection_test_models")
           shard_one_latch.count_down
         end
 
         ActiveRecord::Base.connected_to(role: :writing, shard: :one) do
           shard_default_latch.count_down
-          assert_equal "shard_key_one", ShardConnectionTestModel.connection.select_value("SELECT shard_key from shard_connection_test_models")
+          assert_equal "shard_key_one", ShardConnectionTestModel.lease_connection.select_value("SELECT shard_key from shard_connection_test_models")
           shard_one_latch.wait
         end
 
@@ -372,32 +463,32 @@ module ActiveRecord
 
         [:default, :one].each do |shard_name|
           ActiveRecord::Base.connected_to(role: :writing, shard: shard_name) do
-            ShardConnectionTestModel.connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
-            ShardConnectionTestModel.connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}')")
+            ShardConnectionTestModel.lease_connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
+            ShardConnectionTestModel.lease_connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}')")
           end
 
           ActiveRecord::Base.connected_to(role: :secondary, shard: shard_name) do
-            ShardConnectionTestModel.connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
-            ShardConnectionTestModel.connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}_secondary')")
+            ShardConnectionTestModel.lease_connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
+            ShardConnectionTestModel.lease_connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}_secondary')")
           end
         end
 
         shard_one_latch = Concurrent::CountDownLatch.new
         shard_default_latch = Concurrent::CountDownLatch.new
 
-        ShardConnectionTestModel.connection
+        ShardConnectionTestModel.lease_connection
 
         thread = Thread.new do
-          ShardConnectionTestModel.connection
+          ShardConnectionTestModel.lease_connection
 
           shard_default_latch.wait
-          assert_equal "shard_key_default", ShardConnectionTestModel.connection.select_value("SELECT shard_key from shard_connection_test_models")
+          assert_equal "shard_key_default", ShardConnectionTestModel.lease_connection.select_value("SELECT shard_key from shard_connection_test_models")
           shard_one_latch.count_down
         end
 
         ActiveRecord::Base.connected_to(shard: :one, role: :secondary) do
           shard_default_latch.count_down
-          assert_equal "shard_key_one_secondary", ShardConnectionTestModel.connection.select_value("SELECT shard_key from shard_connection_test_models")
+          assert_equal "shard_key_one_secondary", ShardConnectionTestModel.lease_connection.select_value("SELECT shard_key from shard_connection_test_models")
           shard_one_latch.wait
         end
 
@@ -435,42 +526,42 @@ module ActiveRecord
 
         [:default, :one].each do |shard_name|
           ActiveRecord::Base.connected_to(role: :writing, shard: shard_name) do
-            ShardConnectionTestModel.connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
-            ShardConnectionTestModel.connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}')")
-            ShardConnectionTestModelB.connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
-            ShardConnectionTestModelB.connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}_b')")
+            ShardConnectionTestModel.lease_connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
+            ShardConnectionTestModel.lease_connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}')")
+            ShardConnectionTestModelB.lease_connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
+            ShardConnectionTestModelB.lease_connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}_b')")
           end
 
           ActiveRecord::Base.connected_to(role: :secondary, shard: shard_name) do
-            ShardConnectionTestModel.connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
-            ShardConnectionTestModel.connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}_secondary')")
-            ShardConnectionTestModelB.connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
-            ShardConnectionTestModelB.connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}_secondary_b')")
+            ShardConnectionTestModel.lease_connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
+            ShardConnectionTestModel.lease_connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}_secondary')")
+            ShardConnectionTestModelB.lease_connection.execute("CREATE TABLE `shard_connection_test_models` (shard_key VARCHAR (255))")
+            ShardConnectionTestModelB.lease_connection.execute("INSERT INTO `shard_connection_test_models` VALUES ('shard_key_#{shard_name}_secondary_b')")
           end
         end
 
         shard_one_latch = Concurrent::CountDownLatch.new
         shard_default_latch = Concurrent::CountDownLatch.new
 
-        ShardConnectionTestModel.connection
+        ShardConnectionTestModel.lease_connection
 
         thread = Thread.new do
-          ShardConnectionTestModel.connection
+          ShardConnectionTestModel.lease_connection
 
           shard_default_latch.wait
-          assert_equal "shard_key_default", ShardConnectionTestModel.connection.select_value("SELECT shard_key from shard_connection_test_models")
-          assert_equal "shard_key_default_b", ShardConnectionTestModelB.connection.select_value("SELECT shard_key from shard_connection_test_models")
+          assert_equal "shard_key_default", ShardConnectionTestModel.lease_connection.select_value("SELECT shard_key from shard_connection_test_models")
+          assert_equal "shard_key_default_b", ShardConnectionTestModelB.lease_connection.select_value("SELECT shard_key from shard_connection_test_models")
           shard_one_latch.count_down
         end
 
         SecondaryBase.connected_to(shard: :one, role: :secondary) do
           shard_default_latch.count_down
-          assert_equal "shard_key_one_secondary", ShardConnectionTestModel.connection.select_value("SELECT shard_key from shard_connection_test_models")
-          assert_equal "shard_key_default_b", ShardConnectionTestModelB.connection.select_value("SELECT shard_key from shard_connection_test_models")
+          assert_equal "shard_key_one_secondary", ShardConnectionTestModel.lease_connection.select_value("SELECT shard_key from shard_connection_test_models")
+          assert_equal "shard_key_default_b", ShardConnectionTestModelB.lease_connection.select_value("SELECT shard_key from shard_connection_test_models")
 
           SomeOtherBase.connected_to(shard: :one, role: :secondary) do
-            assert_equal "shard_key_one_secondary", ShardConnectionTestModel.connection.select_value("SELECT shard_key from shard_connection_test_models")
-            assert_equal "shard_key_one_secondary_b", ShardConnectionTestModelB.connection.select_value("SELECT shard_key from shard_connection_test_models")
+            assert_equal "shard_key_one_secondary", ShardConnectionTestModel.lease_connection.select_value("SELECT shard_key from shard_connection_test_models")
+            assert_equal "shard_key_one_secondary_b", ShardConnectionTestModelB.lease_connection.select_value("SELECT shard_key from shard_connection_test_models")
           end
           shard_one_latch.wait
         end

data/test/cases/connection_adapters/connection_swapping_nested_test.rb

@@ -14,6 +14,16 @@ module ActiveRecord
         clean_up_connection_handler
       end
 
+      def skip_unless_sqlite3_gem_available
+        @prev_configs ||= ActiveRecord::Base.configurations
+
+        skip "DB2 limitation: sqlite3-based nested connection swapping scenarios are not supported under ibm_db adapter runs" if current_adapter?(:IBM_DBAdapter)
+
+        return if Gem::Specification.find_all_by_name("sqlite3", ">= 2.1").any?
+
+        skip "requires sqlite3 gem (>= 2.1)"
+      end
+
       class PrimaryBase < ActiveRecord::Base
         self.abstract_class = true
       end
@@ -44,6 +54,8 @@ module ActiveRecord
 
       unless in_memory_db?
         def test_roles_can_be_swapped_granularly
+          skip_unless_sqlite3_gem_available
+
           previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
 
           config = {
@@ -116,6 +128,8 @@ module ActiveRecord
         end
 
         def test_shards_can_be_swapped_granularly
+          skip_unless_sqlite3_gem_available
+
           previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
 
           config = {
@@ -156,17 +170,17 @@ module ActiveRecord
             assert_equal "secondary", SecondaryBase.connection_pool.db_config.name
 
             # Switch only primary to shard_one
-            PrimaryBase.connected_to(role: global_role, shard: :shard_one) do
+            PrimaryBase.connected_to(shard: :shard_one) do
               assert_equal "primary_shard_one", PrimaryBase.connection_pool.db_config.name
               assert_equal "secondary", SecondaryBase.connection_pool.db_config.name
 
               # Switch global to shard_one
-              ActiveRecord::Base.connected_to(role: global_role, shard: :shard_one) do
+              ActiveRecord::Base.connected_to(shard: :shard_one) do
                 assert_equal "primary_shard_one", PrimaryBase.connection_pool.db_config.name
                 assert_equal "secondary_shard_one", SecondaryBase.connection_pool.db_config.name
 
                 # Switch only secondary to shard_two
-                SecondaryBase.connected_to(role: global_role, shard: :shard_two) do
+                SecondaryBase.connected_to(shard: :shard_two) do
                   assert_equal "primary_shard_one", PrimaryBase.connection_pool.db_config.name
                   assert_equal "secondary_shard_two", SecondaryBase.connection_pool.db_config.name
                 end
@@ -183,7 +197,7 @@ module ActiveRecord
               end
 
               # Switch everything to default
-              ActiveRecord::Base.connected_to(role: global_role, shard: :default) do
+              ActiveRecord::Base.connected_to(shard: :default) do
                 assert_equal "primary", PrimaryBase.connection_pool.db_config.name
                 assert_equal "secondary", SecondaryBase.connection_pool.db_config.name
               end
@@ -203,7 +217,63 @@ module ActiveRecord
           ENV["RAILS_ENV"] = previous_env
         end
 
+        def test_shard_swapping_prohibition_exception_recovery
+          skip_unless_sqlite3_gem_available
+
+          previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
+
+          config = {
+            "default_env" => {
+              "primary" => { "adapter" => "sqlite3", "database" => "test/db/primary.sqlite3" },
+              "primary_shard_one" => { "adapter" => "sqlite3", "database" => "test/db/primary_shard_one.sqlite3" },
+              "primary_shard_two" => { "adapter" => "sqlite3", "database" => "test/db/primary_shard_two.sqlite3" },
+            }
+          }
+
+          @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config
+
+          PrimaryBase.connects_to(shards: {
+            default: { writing: :primary },
+            shard_one: { writing: :primary_shard_one }
+          })
+
+          global_role = :writing
+
+          # Switch everything to default
+          ActiveRecord::Base.connected_to(role: global_role, shard: :default) do
+            assert_equal "primary", PrimaryBase.connection_pool.db_config.name
+
+            # Switch only primary to shard_one
+            PrimaryBase.connected_to(shard: :shard_one) do
+              assert_equal "primary_shard_one", PrimaryBase.connection_pool.db_config.name
+
+              PrimaryBase.prohibit_shard_swapping do
+                e = assert_raises(ArgumentError) do
+                  PrimaryBase.connected_to(shard: :shard_two) { }
+                end
+                assert_match(/cannot swap/, e.message)
+              end
+
+              assert_equal "primary_shard_one", PrimaryBase.connection_pool.db_config.name
+
+              PrimaryBase.prohibit_shard_swapping do
+                e = assert_raises(ArgumentError) do
+                  ActiveRecord::Base.connected_to_many([PrimaryBase], shard: :shard_two, role: :writing) { }
+                end
+                assert_match(/cannot swap/, e.message)
+              end
+
+              assert_equal "primary_shard_one", PrimaryBase.connection_pool.db_config.name
+            end
+          end
+        ensure
+          ActiveRecord::Base.configurations = @prev_configs
+          ActiveRecord::Base.establish_connection(:arunit)
+          ENV["RAILS_ENV"] = previous_env
+        end
         def test_roles_and_shards_can_be_swapped_granularly
+          skip_unless_sqlite3_gem_available
+
           previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
 
           config = {
@@ -290,6 +360,8 @@ module ActiveRecord
         end
 
         def test_connected_to_many
+          skip_unless_sqlite3_gem_available
+
           previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
 
           config = {
@@ -366,6 +438,8 @@ module ActiveRecord
         end
 
         def test_prevent_writes_can_be_changed_granularly
+          skip_unless_sqlite3_gem_available
+
           previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
 
           # replica: true is purposefully left out so we can test the pools behavior
@@ -385,44 +459,44 @@ module ActiveRecord
 
           # Switch everything to writing
           ActiveRecord::Base.connected_to(role: :writing) do
-            assert_not_predicate ActiveRecord::Base.connection, :preventing_writes?
-            assert_not_predicate PrimaryBase.connection, :preventing_writes?
-            assert_not_predicate SecondaryBase.connection, :preventing_writes?
+            assert_not_predicate ActiveRecord::Base.lease_connection, :preventing_writes?
+            assert_not_predicate PrimaryBase.lease_connection, :preventing_writes?
+            assert_not_predicate SecondaryBase.lease_connection, :preventing_writes?
 
             # Switch only primary to reading
             PrimaryBase.connected_to(role: :reading) do
-              assert_predicate PrimaryBase.connection, :preventing_writes?
-              assert_not_predicate SecondaryBase.connection, :preventing_writes?
+              assert_predicate PrimaryBase.lease_connection, :preventing_writes?
+              assert_not_predicate SecondaryBase.lease_connection, :preventing_writes?
 
               # Switch global to reading
               ActiveRecord::Base.connected_to(role: :reading) do
-                assert_predicate PrimaryBase.connection, :preventing_writes?
-                assert_predicate SecondaryBase.connection, :preventing_writes?
+                assert_predicate PrimaryBase.lease_connection, :preventing_writes?
+                assert_predicate SecondaryBase.lease_connection, :preventing_writes?
 
                 # Switch only secondary to writing
                 SecondaryBase.connected_to(role: :writing) do
-                  assert_predicate PrimaryBase.connection, :preventing_writes?
-                  assert_not_predicate SecondaryBase.connection, :preventing_writes?
+                  assert_predicate PrimaryBase.lease_connection, :preventing_writes?
+                  assert_not_predicate SecondaryBase.lease_connection, :preventing_writes?
                 end
 
                 # Ensure restored to global reading
-                assert_predicate PrimaryBase.connection, :preventing_writes?
-                assert_predicate SecondaryBase.connection, :preventing_writes?
+                assert_predicate PrimaryBase.lease_connection, :preventing_writes?
+                assert_predicate SecondaryBase.lease_connection, :preventing_writes?
               end
 
               # Switch everything to writing
               ActiveRecord::Base.connected_to(role: :writing) do
-                assert_not_predicate PrimaryBase.connection, :preventing_writes?
-                assert_not_predicate SecondaryBase.connection, :preventing_writes?
+                assert_not_predicate PrimaryBase.lease_connection, :preventing_writes?
+                assert_not_predicate SecondaryBase.lease_connection, :preventing_writes?
               end
 
-              assert_predicate PrimaryBase.connection, :preventing_writes?
-              assert_not_predicate SecondaryBase.connection, :preventing_writes?
+              assert_predicate PrimaryBase.lease_connection, :preventing_writes?
+              assert_not_predicate SecondaryBase.lease_connection, :preventing_writes?
             end
 
             # Ensure restored to global writing
-            assert_not_predicate PrimaryBase.connection, :preventing_writes?
-            assert_not_predicate SecondaryBase.connection, :preventing_writes?
+            assert_not_predicate PrimaryBase.lease_connection, :preventing_writes?
+            assert_not_predicate SecondaryBase.lease_connection, :preventing_writes?
           end
         ensure
           ActiveRecord::Base.configurations = @prev_configs
@@ -441,16 +515,55 @@ module ActiveRecord
 
           # Switch everything to writing
           ActiveRecord::Base.connected_to(role: :writing) do
-            assert_not_predicate ApplicationRecord.connection, :preventing_writes?
+            assert_not_predicate ActiveRecord::Base.lease_connection, :preventing_writes?
+            assert_not_predicate ApplicationRecord.lease_connection, :preventing_writes?
 
             ApplicationRecord.connected_to(role: :reading) do
-              assert_predicate ApplicationRecord.connection, :preventing_writes?
+              assert_predicate ApplicationRecord.lease_connection, :preventing_writes?
+            end
+
+            # reading is fine bc it's looking up by AppRec but writing is not fine
+            # bc its looking up by ARB in the stack
+            ApplicationRecord.connected_to(role: :writing, prevent_writes: true) do
+              assert_predicate ApplicationRecord.lease_connection, :preventing_writes?
             end
           end
         ensure
+          ApplicationRecord.remove_connection
+          ActiveRecord.application_record_class = nil
           Object.send(:remove_const, :ApplicationRecord)
           ActiveRecord::Base.establish_connection :arunit
         end
+
+        def test_prevent_writes_handles_class_reloading
+          skip "IBM_DB keeps preventing_writes enabled for this class-reloading override path" if current_adapter?(:IBM_DBAdapter)
+
+          # Regression test for https://github.com/rails/rails/issues/54343
+          Object.const_set(:ReloadedRecord, Class.new(ActiveRecord::Base) { self.abstract_class = true })
+          ReloadedRecord.connects_to(database: { writing: :arunit, reading: :arunit })
+
+          ActiveRecord::Base.connected_to(role: :reading, prevent_writes: true) do
+            ReloadedRecord.connected_to(role: :writing, prevent_writes: false) do
+              assert_not_predicate ReloadedRecord.lease_connection, :preventing_writes?
+            end
+          end
+
+          # emulate a reload in development mode
+          Object.send(:remove_const, :ReloadedRecord)
+          Object.const_set(:ReloadedRecord, Class.new(ActiveRecord::Base) { self.abstract_class = true })
+          ReloadedRecord.connects_to(database: { writing: :arunit, reading: :arunit })
+
+          ActiveRecord::Base.connected_to(role: :reading, prevent_writes: true) do
+            ReloadedRecord.connected_to(role: :writing, prevent_writes: false) do
+              assert_not_predicate ReloadedRecord.lease_connection, :preventing_writes?
+            end
+          end
+        ensure
+          ReloadedRecord.remove_connection if Object.const_defined?(:ReloadedRecord)
+          ActiveRecord.application_record_class = nil
+          Object.send(:remove_const, :ReloadedRecord) if Object.const_defined?(:ReloadedRecord)
+          ActiveRecord::Base.establish_connection :arunit
+        end
       end
     end
   end