diff --git a/lib/data_migrate/database_tasks.rb b/lib/data_migrate/database_tasks.rb
index 90f9a27..96c6e35 100644
--- a/lib/data_migrate/database_tasks.rb
+++ b/lib/data_migrate/database_tasks.rb
@@ -10,21 +10,21 @@ module DatabaseTasks
     extend self
 
     # These method are only introduced in Rails 7.1
-    unless respond_to?(:with_temporary_connection_for_each)
-      def with_temporary_connection_for_each(env: ActiveRecord::Tasks::DatabaseTasks.env, name: nil, &block) # :nodoc:
+    unless respond_to?(:with_temporary_pool_for_each)
+      def with_temporary_pool_for_each(env: ActiveRecord::Tasks::DatabaseTasks.env, name: nil, &block) # :nodoc:
         if name
           db_config = ActiveRecord::Base.configurations.configs_for(env_name: env, name: name)
-          with_temporary_connection(db_config, &block)
+           with_temporary_pool(db_config, &block)
         else
           ActiveRecord::Base.configurations.configs_for(env_name: env, name: name).each do |db_config|
-            with_temporary_connection(db_config, &block)
+             with_temporary_pool(db_config, &block)
           end
         end
       end
 
-      def with_temporary_connection(db_config) # :nodoc:
+      def with_temporary_connection(db_config, &block) # :nodoc:
         with_temporary_pool(db_config) do |pool|
-          yield pool.connection
+          pool.with_connection(&block)
         end
       end
 
@@ -49,8 +49,8 @@ def migration_connection # :nodoc:
     def db_configs_with_versions
       db_configs_with_versions = Hash.new { |h, k| h[k] = [] }
 
-      with_temporary_connection_for_each do |conn|
-        db_config = conn.pool.db_config
+      with_temporary_pool_for_each do |pool|
+        db_config = pool.db_config
         if db_config.primary?
           versions_to_run = DataMigrate::DatabaseTasks.pending_data_migrations.map { |m| m[:version] }
           target_version = ActiveRecord::Tasks::DatabaseTasks.target_version