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