Merge pull request #40615 from eileencodes/fix-preventing-writes

Fix preventing_writes for granular swapping

activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb

@@ -23,6 +23,10 @@ module ActiveRecord
       include ConnectionAdapters::AbstractPool
 
       attr_accessor :schema_cache
+
+      def owner_name
+        nil
+      end
     end
 
     # Connection pool base class for managing Active Record database
@@ -356,7 +360,7 @@ module ActiveRecord
       include ConnectionAdapters::AbstractPool
 
       attr_accessor :automatic_reconnect, :checkout_timeout
-      attr_reader :db_config, :size, :reaper, :pool_config
+      attr_reader :db_config, :size, :reaper, :pool_config, :owner_name
 
       delegate :schema_cache, :schema_cache=, to: :pool_config
 
@@ -371,6 +375,7 @@ module ActiveRecord
 
         @pool_config = pool_config
         @db_config = pool_config.db_config
+        @owner_name = pool_config.connection_specification_name
 
         @checkout_timeout = db_config.checkout_timeout
         @idle_timeout = db_config.idle_timeout

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -113,14 +113,20 @@ module ActiveRecord
 
       # Determines whether writes are currently being prevents.
       #
-      # Returns true if the connection is a replica, or if +prevent_writes+
-      # is set to true.
+      # Returns true if the connection is a replica.
+      #
+      # If the application is using legacy handling, returns
+      # true if `connection_handler.prevent_writes` is set.
+      #
+      # If the application is using the new connection handling
+      # will return true based on `current_preventing_writes`.
       def preventing_writes?
-        if ActiveRecord::Base.legacy_connection_handling
-          replica? || ActiveRecord::Base.connection_handler.prevent_writes
-        else
-          replica? || ActiveRecord::Base.current_preventing_writes
-        end
+        return true if replica?
+        return ActiveRecord::Base.connection_handler.prevent_writes if ActiveRecord::Base.legacy_connection_handling
+        return false if owner_name.nil?
+
+        klass = self.owner_name.safe_constantize
+        klass&.current_preventing_writes
       end
 
       def migrations_paths # :nodoc:
@@ -196,6 +202,10 @@ module ActiveRecord
         @owner = Thread.current
       end
 
+      def owner_name # :nodoc:
+        @pool.owner_name
+      end
+
       def schema_cache
         @pool.get_schema_cache(self)
       end

activerecord/test/cases/adapters/sqlite3/sqlite3_adapter_prevent_writes_test.rb

@@ -11,9 +11,7 @@ module ActiveRecord
       self.use_transactional_tests = false
 
       def setup
-        @conn = Base.sqlite3_connection database: ":memory:",
-                                        adapter: "sqlite3",
-                                        timeout: 100
+        @conn = ActiveRecord::Base.connection
       end
 
       def test_errors_when_an_insert_query_is_called_while_preventing_writes
@@ -100,9 +98,7 @@ module ActiveRecord
         @old_value = ActiveRecord::Base.legacy_connection_handling
         ActiveRecord::Base.legacy_connection_handling = true
 
-        @conn = Base.sqlite3_connection database: ":memory:",
-                                        adapter: "sqlite3",
-                                        timeout: 100
+        @conn = ActiveRecord::Base.connection
 
         @connection_handler = ActiveRecord::Base.connection_handler
       end

activerecord/test/cases/connection_adapters/connection_swapping_nested_test.rb

@@ -349,6 +349,71 @@ module ActiveRecord
           ActiveRecord::Base.establish_connection(:arunit)
           ENV["RAILS_ENV"] = previous_env
         end
+
+        def test_prevent_writes_can_be_changed_granularly
+          previous_env, ENV["RAILS_ENV"] = ENV["RAILS_ENV"], "default_env"
+
+          # replica: true is purposefully left out so we can test the pools behavior
+          config = {
+            "default_env" => {
+              "primary" => { "adapter" => "sqlite3", "database" => "test/db/primary.sqlite3" },
+              "primary_replica" => { "adapter" => "sqlite3", "database" => "test/db/primary.sqlite3" },
+              "secondary" => { "adapter" => "sqlite3", "database" => "test/db/secondary.sqlite3" },
+              "secondary_replica" => { "adapter" => "sqlite3", "database" => "test/db/secondary_replica.sqlite3" }
+            }
+          }
+
+          @prev_configs, ActiveRecord::Base.configurations = ActiveRecord::Base.configurations, config
+
+          PrimaryBase.connects_to database: { writing: :primary, reading: :primary_replica }
+          SecondaryBase.connects_to database: { writing: :secondary, reading: :secondary_replica }
+
+          # 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?
+
+            # Switch only primary to reading
+            PrimaryBase.connected_to(role: :reading) do
+              assert_predicate PrimaryBase.connection, :preventing_writes?
+              assert_not_predicate SecondaryBase.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?
+
+                # Switch only secondary to writing
+                SecondaryBase.connected_to(role: :writing) do
+                  assert_predicate PrimaryBase.connection, :preventing_writes?
+                  assert_not_predicate SecondaryBase.connection, :preventing_writes?
+                end
+
+                # Ensure restored to global reading
+                assert_predicate PrimaryBase.connection, :preventing_writes?
+                assert_predicate SecondaryBase.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?
+              end
+
+              assert_predicate PrimaryBase.connection, :preventing_writes?
+              assert_not_predicate SecondaryBase.connection, :preventing_writes?
+            end
+
+            # Ensure restored to global writing
+            assert_not_predicate PrimaryBase.connection, :preventing_writes?
+            assert_not_predicate SecondaryBase.connection, :preventing_writes?
+          end
+        ensure
+          ActiveRecord::Base.configurations = @prev_configs
+          ActiveRecord::Base.establish_connection(:arunit)
+          ENV["RAILS_ENV"] = previous_env
+        end
       end
     end
   end