Refactor db:migrate:with_data

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

@@ -24,7 +24,7 @@ def with_temporary_connection_for_each(env: ActiveRecord::Tasks::DatabaseTasks.e
 
       def with_temporary_connection(db_config) # :nodoc:
         with_temporary_pool(db_config) do |pool|
-          yield pool.connection
+          yield pool.connection, db_config
         end
       end
 
@@ -49,14 +49,15 @@ 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
-        versions_to_run = conn.migration_context.pending_migration_versions
-        target_version = ActiveRecord::Tasks::DatabaseTasks.target_version
+      with_temporary_connection_for_each do |conn, 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] << db_config
+          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
 

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