From 45add344781b3e5305f85ec10cd4056934cb99aa Mon Sep 17 00:00:00 2001 From: Guo Xiang Tan Date: Tue, 30 Jun 2020 16:26:47 +0800 Subject: [PATCH] Move advisory locks to own connection handler. Removes the use of `ActiveRecord::AdvisoryLockBase` since it inherits from `ActiveRecord::Base` and hence share module attributes that are defined in `ActiveRecord::Base`. This is problematic because establishing connections through `ActiveRecord::AdvisoryLockBase` can end up changing state of the default connection handler of `ActiveRecord::Base` leading to unexpected behaviors in a Rails application. In the case of https://github.com/rails/rails/issues/39157, Running migrations with `rails db:migrate:primary_shard_one` was not working as the application itself defined the following ``` class ApplicationRecord < ActiveRecord::Base self.abstract_class = true connects_to shards: { default: { writing: :primary }, shard_one: { writing: :primary_shard_one } } end ``` In the database migrate rake task, the default connection was established with the database config of `primary_shard_one`. However, the default connection was altered to that of `primary` because `ActiveRecord::AdvisoryLockBase.establish_connection` ended up loading `ApplicationRecord` which calls `connects_to shards:`. Since all we really need here is just a normal database connection, we can avoid accidentally altering the default connection handler state during the migration by creating a custom connection handler used for retrieving a connection. --- activerecord/lib/active_record.rb | 1 - .../lib/active_record/advisory_lock_base.rb | 18 ----------- activerecord/lib/active_record/migration.rb | 31 ++++++++++++------- activerecord/test/cases/migration_test.rb | 6 ++-- 4 files changed, 24 insertions(+), 32 deletions(-) delete mode 100644 activerecord/lib/active_record/advisory_lock_base.rb diff --git a/activerecord/lib/active_record.rb b/activerecord/lib/active_record.rb index e9a7623f4055e..4d524623a3863 100644 --- a/activerecord/lib/active_record.rb +++ b/activerecord/lib/active_record.rb @@ -36,7 +36,6 @@ module ActiveRecord extend ActiveSupport::Autoload - autoload :AdvisoryLockBase autoload :Base autoload :Callbacks autoload :Core diff --git a/activerecord/lib/active_record/advisory_lock_base.rb b/activerecord/lib/active_record/advisory_lock_base.rb deleted file mode 100644 index 7110a1f132337..0000000000000 --- a/activerecord/lib/active_record/advisory_lock_base.rb +++ /dev/null @@ -1,18 +0,0 @@ -# frozen_string_literal: true - -module ActiveRecord - # This class is used to create a connection that we can use for advisory - # locks. This will take out a "global" lock that can't be accidentally - # removed if a new connection is established during a migration. - class AdvisoryLockBase < ActiveRecord::Base # :nodoc: - self.abstract_class = true - - self.connection_specification_name = "AdvisoryLockBase" - - class << self - def _internal? - true - end - end - end -end diff --git a/activerecord/lib/active_record/migration.rb b/activerecord/lib/active_record/migration.rb index 88e8c583bf0d1..56c7495226dc1 100644 --- a/activerecord/lib/active_record/migration.rb +++ b/activerecord/lib/active_record/migration.rb @@ -1376,20 +1376,29 @@ def use_advisory_lock? def with_advisory_lock lock_id = generate_migrator_advisory_lock_id - AdvisoryLockBase.establish_connection(ActiveRecord::Base.connection_db_config) unless AdvisoryLockBase.connected? - connection = AdvisoryLockBase.connection - got_lock = connection.get_advisory_lock(lock_id) - raise ConcurrentMigrationError unless got_lock - load_migrated # reload schema_migrations to be sure it wasn't changed by another process before we got the lock - yield - ensure - if got_lock && !connection.release_advisory_lock(lock_id) - raise ConcurrentMigrationError.new( - ConcurrentMigrationError::RELEASE_LOCK_FAILED_MESSAGE - ) + + with_advisory_lock_connection do |connection| + got_lock = connection.get_advisory_lock(lock_id) + raise ConcurrentMigrationError unless got_lock + load_migrated # reload schema_migrations to be sure it wasn't changed by another process before we got the lock + yield + ensure + if got_lock && !connection.release_advisory_lock(lock_id) + raise ConcurrentMigrationError.new( + ConcurrentMigrationError::RELEASE_LOCK_FAILED_MESSAGE + ) + end end end + def with_advisory_lock_connection + pool = ActiveRecord::ConnectionAdapters::ConnectionHandler.new.establish_connection( + ActiveRecord::Base.connection_db_config + ) + + pool.with_connection { |connection| yield(connection) } + end + MIGRATOR_SALT = 2053462845 def generate_migrator_advisory_lock_id db_name_hash = Zlib.crc32(Base.connection.current_database) diff --git a/activerecord/test/cases/migration_test.rb b/activerecord/test/cases/migration_test.rb index 0d65dd7f9c16a..8946e0afb93e1 100644 --- a/activerecord/test/cases/migration_test.rb +++ b/activerecord/test/cases/migration_test.rb @@ -940,8 +940,10 @@ def test_with_advisory_lock_raises_the_right_error_when_it_fails_to_release_lock e = assert_raises(ActiveRecord::ConcurrentMigrationError) do silence_stream($stderr) do - migrator.send(:with_advisory_lock) do - ActiveRecord::AdvisoryLockBase.connection.release_advisory_lock(lock_id) + migrator.stub(:with_advisory_lock_connection, ->(&block) { block.call(ActiveRecord::Base.connection) }) do + migrator.send(:with_advisory_lock) do + ActiveRecord::Base.connection.release_advisory_lock(lock_id) + end end end end