Use lease_connection over deprecated connection for rails 8 (#353)

* Use `lease_connection` over deprecated `connection` for rails 8
Call `register_task` as extending DatabaseTasks doesn't hold onto that

* Use lease_connection over connection only when available

Author: vprigent

Changelog.md

@@ -1,7 +1,9 @@
 # Changelog
 
-# 11.2.0
+# Unreleased
+- Use lease_connection over deprecated connection for rails 8 https://github.com/ilyakatz/data-migrate/pull/353
 
+# 11.2.0
 - Remove committed Gemfile.lock, reduce bundled file list when running `gem install` https://github.com/ilyakatz/data-migrate/pull/351
 - [Bump actionpack from 7.1.3.4 to 7.1.4.1](https://github.com/ilyakatz/data-migrate/pull/348)
 - [Bump rexml from 3.3.6 to 3.3.9](https://github.com/ilyakatz/data-migrate/pull/349)

README.md

@@ -144,7 +144,7 @@ From now on capistrano will run `rake db:migrate:with_data` in every deploy.
 ## Rails Engines support
 
 This gem also has a initial support for adding data migrations inside Rails engines.
-Just add your engines directory pattern to data_migrations initializer, for example 
+Just add your engines directory pattern to data_migrations initializer, for example
 in the case your engines are located in `engines` folder you can set it up like this:
 
 ```ruby
@@ -166,6 +166,7 @@ bundle exec appraisal rails-6.1 rspec
 bundle exec appraisal rails-7.0 rspec
 bundle exec appraisal rails-7.1 rspec
 bundle exec appraisal rails-7.2 rspec
+bundle exec appraisal rails-8.0 rspec
 ```
 
 ## Releasing new version

lib/data_migrate/data_migrator.rb

@@ -48,7 +48,6 @@ def migrations_status
 
       # TODO: this was added to be backward compatible, need to re-evaluate
       def migrations(_migrations_paths)
-        #DataMigrate::MigrationContext.new(migrations_paths).migrations
         DataMigrate::MigrationContext.new(_migrations_paths).migrations
       end
 

lib/data_migrate/database_tasks.rb

@@ -9,6 +9,13 @@ module DatabaseTasks
     extend ActiveRecord::Tasks::DatabaseTasks
     extend self
 
+    if respond_to?(:register_task)
+      register_task(/mysql/,        "ActiveRecord::Tasks::MySQLDatabaseTasks")
+      register_task(/trilogy/,      "ActiveRecord::Tasks::MySQLDatabaseTasks")
+      register_task(/postgresql/,   "ActiveRecord::Tasks::PostgreSQLDatabaseTasks")
+      register_task(/sqlite/,       "ActiveRecord::Tasks::SQLiteDatabaseTasks")
+    end
+
     # These method are only introduced in Rails 7.1
     unless respond_to?(:with_temporary_pool_for_each)
       def with_temporary_pool_for_each(env: ActiveRecord::Tasks::DatabaseTasks.env, name: nil, &block) # :nodoc:
@@ -99,7 +106,7 @@ def pending_migrations
     end
 
     def sort_migrations(*migrations)
-      migrations.flatten.sort { |a, b|  sort_string(a) <=> sort_string(b) }
+      migrations.flatten.sort { |a, b| sort_string(a) <=> sort_string(b) }
     end
 
     def sort_string migration
@@ -159,7 +166,7 @@ def pending_data_migrations
       data_migrator = DataMigrate::RailsHelper.data_migrator(:up, data_migrations)
       sort_migrations(
         data_migrator.pending_migrations.map { |m| { version: m.version, name: m.name, kind: :data } }
-        )
+      )
     end
 
     def pending_schema_migrations
@@ -219,7 +226,8 @@ def self.prepare_all_with_data
         next unless primary?(db_config)
 
         with_temporary_pool(db_config) do |pool|
-          unless database_exists?(pool.connection)
+          connection = pool.respond_to?(:lease_connection) ? pool.lease_connection : pool.connection
+          unless database_exists?(connection)
             create(db_config)
             if File.exist?(schema_dump_path(db_config))
               load_schema(db_config, schema_format, nil)

spec/data_migrate/database_tasks_spec.rb

@@ -119,7 +119,7 @@
       before do
         allow(subject).to receive(:each_current_configuration).and_yield(db_config)
         allow(subject).to receive(:with_temporary_pool).with(db_config).and_yield(pool)
-        allow(pool).to receive(:connection).and_return(connection)
+        allow(pool).to receive(:lease_connection).and_return(connection)
         allow(subject).to receive(:schema_dump_path).and_return("db/data_schema.rb")
         allow(File).to receive(:exist?).and_return(true)
         allow(subject).to receive(:load_schema)