diff --git a/Rakefile b/Rakefile index 995d4f1..4778582 100644 --- a/Rakefile +++ b/Rakefile @@ -43,7 +43,11 @@ namespace :db do require './spec/models/application_record' ActiveRecord::Base.configurations.configs_for(env_name: "test").each do |config| ActiveRecord::Base.establish_connection(config) - schema_migration = ActiveRecord::Base.connection.schema_migration + schema_migration = if ActiveRecord.gem_version >= Gem::Version.new(7.2) + ActiveRecord::Base.connection.pool.schema_migration + else + ActiveRecord::Base.connection.schema_migration + end ActiveRecord::MigrationContext.new("spec/migration", schema_migration) .migrate(config.database == 'octoball_shard_5' ? 2 : 1) end diff --git a/lib/octoball/connection_adapters.rb b/lib/octoball/connection_adapters.rb index 4893129..49c147b 100644 --- a/lib/octoball/connection_adapters.rb +++ b/lib/octoball/connection_adapters.rb @@ -13,6 +13,33 @@ module ConnectionHasCurrentShard attr_accessor :current_shard end + if ActiveRecord.gem_version >= Gem::Version.new('7.2.0') + module ConnectionPoolSetCurrentShard + def with_connection(prevent_permanent_checkout: false) + lease = connection_lease + if lease.connection + lease.connection.current_shard = lease.connection.shard + end + + super + end + + def active_connection? + conn = connection_lease.connection + conn.current_shard = conn.shard + conn + end + + def active_connection + conn = connection_lease.connection + conn.current_shard = conn.shard + conn + end + end + + ::ActiveRecord::ConnectionAdapters::ConnectionPool.prepend(ConnectionPoolSetCurrentShard) + end + ::ActiveRecord::ConnectionAdapters::ConnectionHandler.prepend(ConnectionHandlerSetCurrentShard) ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ConnectionHasCurrentShard) end