Improve db migrate with_data Rake task

lib/data_migrate.rb

@@ -12,6 +12,7 @@
 require File.join(File.dirname(__FILE__), "data_migrate", "tasks/data_migrate_tasks")
 require File.join(File.dirname(__FILE__), "data_migrate", "config")
 require File.join(File.dirname(__FILE__), "data_migrate", "schema_migration")
+require File.join(File.dirname(__FILE__), "data_migrate", "database_configurations_wrapper")
 
 module DataMigrate
   def self.root

lib/data_migrate/database_configurations_wrapper.rb

@@ -0,0 +1,11 @@
+module DataMigrate
+  # This wrapper is used to differentiate between
+  # a data and schema db config when running migrations
+  class DatabaseConfigurationWrapper
+    attr_reader :db_config
+
+    def initialize(db_config)
+      @db_config = db_config
+    end
+  end
+end

lib/data_migrate/database_tasks.rb

@@ -5,87 +5,142 @@
 module DataMigrate
   ##
   # This class extends DatabaseTasks to add a schema_file method.
-  class DatabaseTasks
+  module DatabaseTasks
     extend ActiveRecord::Tasks::DatabaseTasks
-
-    class << self
-      def schema_file(_format = nil)
-        File.join(db_dir, "data_schema.rb")
-      end
-
-      def schema_file_type(_format = nil)
-        "data_schema.rb"
-      end
-
-      # This method is removed in Rails 7.0
-      def dump_filename(spec_name, format = ActiveRecord::Base.schema_format)
-        filename = if spec_name == "primary"
-          schema_file_type(format)
+    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:
+        if name
+          db_config = ActiveRecord::Base.configurations.configs_for(env_name: env, name: name)
+          with_temporary_connection(db_config, &block)
         else
-          "#{spec_name}_#{schema_file_type(format)}"
+          ActiveRecord::Base.configurations.configs_for(env_name: env, name: name).each do |db_config|
+            with_temporary_connection(db_config, &block)
+          end
         end
-
-        ENV["DATA_SCHEMA"] || File.join(db_dir, filename)
       end
 
-      def check_schema_file(filename)
-        unless File.exist?(filename)
-          message = +%{#{filename} doesn't exist yet. Run `rake data:migrate` to create it, then try again.}
-          Kernel.abort message
+      def with_temporary_connection(db_config) # :nodoc:
+        with_temporary_pool(db_config) do |pool|
+          yield pool.connection
         end
       end
 
-      def pending_migrations
-        sort_migrations(
-          pending_schema_migrations,
-          pending_data_migrations
-        )
+      def migration_class # :nodoc:
+        ActiveRecord::Base
       end
 
-      def sort_migrations(*migrations)
-        migrations.flatten.sort { |a, b|  sort_string(a) <=> sort_string(b) }
+      def migration_connection # :nodoc:
+        migration_class.connection
       end
 
-      def sort_string migration
-        "#{migration[:version]}_#{migration[:kind] == :data ? 1 : 0}"
-      end
+      private def with_temporary_pool(db_config)
+        original_db_config = migration_class.connection_db_config
+        pool = migration_class.connection_handler.establish_connection(db_config)
 
-      def data_migrations_path
-        ::DataMigrate.config.data_migrations_path
+        yield pool
+      ensure
+        migration_class.connection_handler.establish_connection(original_db_config)
       end
+    end
 
-      def run_migration(migration, direction)
-        if migration[:kind] == :data
-          ::ActiveRecord::Migration.write("== %s %s" % ['Data', "=" * 71])
-          ::DataMigrate::DataMigrator.run(direction, data_migrations_path, migration[:version])
-        else
-          ::ActiveRecord::Migration.write("== %s %s" % ['Schema', "=" * 69])
-          ::DataMigrate::SchemaMigration.run(
-            direction,
-            ::DataMigrate::SchemaMigration.migrations_paths,
-            migration[:version]
-          )
+    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
+        if db_config.primary?
+          versions_to_run = DataMigrate::DatabaseTasks.pending_data_migrations.map { |m| m[:version] }
+          target_version = ActiveRecord::Tasks::DatabaseTasks.target_version
+
+          versions_to_run.each do |version|
+            next if target_version && target_version != version
+            db_configs_with_versions[version] << DatabaseConfigurationWrapper.new(db_config)
+          end
         end
       end
 
-      def schema_dump_path(db_config, format = ActiveRecord.schema_format)
-        return ENV["DATA_SCHEMA"] if ENV["DATA_SCHEMA"]
+      db_configs_with_versions
+    end
 
-        # We only require a schema.rb file for the primary database
-        return unless db_config.primary?
+    def schema_file(_format = nil)
+      File.join(db_dir, "data_schema.rb")
+    end
 
-        File.join(File.dirname(ActiveRecord::Tasks::DatabaseTasks.schema_dump_path(db_config, format)), schema_file_type)
+    def schema_file_type(_format = nil)
+      "data_schema.rb"
+    end
+
+    # This method is removed in Rails 7.0
+    def dump_filename(spec_name, format = ActiveRecord::Base.schema_format)
+      filename = if spec_name == "primary"
+        schema_file_type(format)
+      else
+        "#{spec_name}_#{schema_file_type(format)}"
       end
 
-      # Override this method from `ActiveRecord::Tasks::DatabaseTasks`
-      # to ensure that the sha saved in ar_internal_metadata table
-      # is from the original schema.rb file
-      def schema_sha1(file)
-        ActiveRecord::Tasks::DatabaseTasks.schema_dump_path(ActiveRecord::Base.configurations.configs_for(env_name: ActiveRecord::Tasks::DatabaseTasks.env, name: "primary"))
+      ENV["DATA_SCHEMA"] || File.join(db_dir, filename)
+    end
+
+    def check_schema_file(filename)
+      unless File.exist?(filename)
+        message = +%{#{filename} doesn't exist yet. Run `rake data:migrate` to create it, then try again.}
+        Kernel.abort message
       end
     end
 
-    def self.forward(step = 1)
+    def pending_migrations
+      sort_migrations(
+        pending_schema_migrations,
+        pending_data_migrations
+      )
+    end
+
+    def sort_migrations(*migrations)
+      migrations.flatten.sort { |a, b|  sort_string(a) <=> sort_string(b) }
+    end
+
+    def sort_string migration
+      "#{migration[:version]}_#{migration[:kind] == :data ? 1 : 0}"
+    end
+
+    def data_migrations_path
+      ::DataMigrate.config.data_migrations_path
+    end
+
+    def run_migration(migration, direction)
+      if migration[:kind] == :data
+        ::ActiveRecord::Migration.write("== %s %s" % ['Data', "=" * 71])
+        ::DataMigrate::DataMigrator.run(direction, data_migrations_path, migration[:version])
+      else
+        ::ActiveRecord::Migration.write("== %s %s" % ['Schema', "=" * 69])
+        ::DataMigrate::SchemaMigration.run(
+          direction,
+          ::DataMigrate::SchemaMigration.migrations_paths,
+          migration[:version]
+        )
+      end
+    end
+
+    def schema_dump_path(db_config, format = ActiveRecord.schema_format)
+      return ENV["DATA_SCHEMA"] if ENV["DATA_SCHEMA"]
+
+      # We only require a schema.rb file for the primary database
+      return unless db_config.primary?
+
+      File.join(File.dirname(ActiveRecord::Tasks::DatabaseTasks.schema_dump_path(db_config, format)), schema_file_type)
+    end
+
+    # Override this method from `ActiveRecord::Tasks::DatabaseTasks`
+    # to ensure that the sha saved in ar_internal_metadata table
+    # is from the original schema.rb file
+    def schema_sha1(file)
+      ActiveRecord::Tasks::DatabaseTasks.schema_dump_path(ActiveRecord::Base.configurations.configs_for(env_name: ActiveRecord::Tasks::DatabaseTasks.env, name: "primary"))
+    end
+
+    def forward(step = 1)
       DataMigrate::DataMigrator.create_data_schema_table
       migrations = pending_migrations.reverse.pop(step).reverse
       migrations.each do | pending_migration |
@@ -99,19 +154,19 @@ def self.forward(step = 1)
       end
     end
 
-    def self.pending_data_migrations
+    def pending_data_migrations
       data_migrations = DataMigrate::DataMigrator.migrations(data_migrations_path)
       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 self.pending_schema_migrations
+    def pending_schema_migrations
       ::DataMigrate::SchemaMigration.pending_schema_migrations
     end
 
-    def self.past_migrations(sort = nil)
+    def past_migrations(sort = nil)
       data_versions = DataMigrate::RailsHelper.data_schema_migration.table_exists? ? DataMigrate::RailsHelper.data_schema_migration.normalized_versions : []
       schema_versions = DataMigrate::RailsHelper.schema_migration.normalized_versions
       migrations = data_versions.map { |v| { version: v.to_i, kind: :data } } + schema_versions.map { |v| { version: v.to_i, kind: :schema } }

lib/data_migrate/test.rb

@@ -0,0 +1,14 @@
+module Base
+  extend self
+
+  def foo
+    puts "Base#foo called"
+  end
+end
+
+module Child
+  extend Base
+  extend self
+
+  puts "foo: #{respond_to?(:foo)}"
+end

tasks/databases.rake

@@ -7,48 +7,31 @@ namespace :db do
     desc "Migrate the database data and schema (options: VERSION=x, VERBOSE=false)."
     task :with_data => :environment do
       DataMigrate::DataMigrator.create_data_schema_table
-
       ActiveRecord::Migration.verbose = ENV["VERBOSE"] ? ENV["VERBOSE"] == "true" : true
-      target_version = ENV["VERSION"] ? ENV["VERSION"].to_i : nil
-      migrations = []
 
-      if target_version.nil?
-        migrations = DataMigrate::DatabaseTasks.pending_migrations.map{ |m| m.merge(:direction =>:up) }
-      else
-        current_schema_version = ActiveRecord::Migrator.current_version
-        schema_migrations = if target_version > current_schema_version
-                              DataMigrate::DatabaseTasks.pending_schema_migrations.keep_if{ |m| m[:version] <= target_version }.map{ |m| m.merge(:direction =>:up) }
-                            elsif target_version < current_schema_version
-                              DataMigrate::DatabaseTasks.past_migrations.keep_if{ |m| m[:version] > target_version }.map{ |m| m.merge(:direction =>:down) }
-                            else # ==
-                              []
-                            end
-
-        current_data_version = DataMigrate::DataMigrator.current_version
-        data_migrations = if target_version > current_data_version
-                            DataMigrate::DatabaseTasks.pending_data_migrations.keep_if{ |m| m[:version] <= target_version }.map{ |m| m.merge(:direction =>:up) }
-                          elsif target_version < current_data_version
-                            DataMigrate::DatabaseTasks.past_migrations.keep_if{ |m| m[:version] > target_version }.map{ |m| m.merge(:direction =>:down) }
-                          else # ==
-                            []
-                          end
-        migrations = if schema_migrations.empty?
-                       data_migrations
-                     elsif data_migrations.empty?
-                       schema_migrations
-                     elsif target_version > current_data_version && target_version > current_schema_version
-                       DataMigrate::DatabaseTasks.sort_migrations data_migrations, schema_migrations
-                     elsif target_version < current_data_version && target_version < current_schema_version
-                       DataMigrate::DatabaseTasks.sort_migrations(data_migrations, schema_migrations).reverse
-                     elsif target_version > current_data_version && target_version < current_schema_version
-                       schema_migrations + data_migrations
-                     elsif target_version < current_data_version && target_version > current_schema_version
-                       schema_migrations + data_migrations
-                     end
+      db_configs = ActiveRecord::Base.configurations.configs_for(env_name: ActiveRecord::Tasks::DatabaseTasks.env)
+
+      schema_mapped_versions = ActiveRecord::Tasks::DatabaseTasks.db_configs_with_versions(db_configs)
+      data_mapped_versions = DataMigrate::DatabaseTasks.db_configs_with_versions
+
+      mapped_versions = schema_mapped_versions.merge(data_mapped_versions) do |_key, schema_db_configs, data_db_configs|
+        schema_db_configs + data_db_configs
       end
 
-      migrations.each do |migration|
-        DataMigrate::DatabaseTasks.run_migration(migration, migration[:direction])
+      mapped_versions.sort.each do |version, db_configs|
+        db_configs.each do |db_config|
+          if is_data_migration = db_config.is_a?(DataMigrate::DatabaseConfigurationWrapper)
+            db_config = db_config.db_config
+          end
+
+          DataMigrate::DatabaseTasks.with_temporary_connection(db_config) do
+            if is_data_migration
+              DataMigrate::DataMigrator.run(:up, DataMigrate::DatabaseTasks.data_migrations_path, version)
+            else
+              ActiveRecord::Tasks::DatabaseTasks.migrate(version)
+            end
+          end
+        end
       end
 
       Rake::Task["db:_dump"].invoke