proper multi DB support

lib/parallel_tests/tasks.rb

@@ -1,7 +1,10 @@
 # frozen_string_literal: true
+require 'active_record'
 require 'rake'
 require 'shellwords'
 
+databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml
+
 module ParallelTests
   module Tasks
     class << self
@@ -149,14 +152,25 @@ def rails_61_or_greater?
     ParallelTests::Tasks.run_in_parallel(ParallelTests::Tasks.suppress_schema_load_output(command), args)
   end
 
-  desc "Create test databases via db:create --> parallel:create[num_cpus,database]"
-  task :create, [:count, :database] do |_, args|
+  desc "Create test databases via db:create --> parallel:create[num_cpus]"
+  task :create, :count, do |_, args|
     ParallelTests::Tasks.run_in_parallel(
       [$0, "db:create#{":#{args[:database]}" if args[:database]}", "RAILS_ENV=#{ParallelTests::Tasks.rails_env}"],
       args
     )
   end
 
+  # create individual databases in multi DB setups
+  ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
+    desc "Create test #{name} database via db:create:#{name} --> parallel:create:#{name}[num_cpus]"
+    task "create:#{name}":to_sym, :count do |_, args|
+      ParallelTests::Tasks.run_in_parallel(
+        [$0, "db:create:#{name}", "RAILS_ENV=#{ParallelTests::Tasks.rails_env}"],
+        args
+      )
+    end
+  end
+
   desc "Drop test databases via db:drop --> parallel:drop[num_cpus]"
   task :drop, :count do |_, args|
     ParallelTests::Tasks.run_in_parallel(
@@ -170,6 +184,17 @@ def rails_61_or_greater?
     )
   end
 
+  # drop individual databases in multi DB setups
+  ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
+    desc "Drop test #{name} database via db:drop:#{name} --> parallel:drop:#{name}[num_cpus]"
+    task "drop:#{name}":to_sym, :count do |_, args|
+      ParallelTests::Tasks.run_in_parallel(
+        [$0, "db:drop:#{name}", "RAILS_ENV=#{ParallelTests::Tasks.rails_env}"],
+        args
+      )
+    end
+  end
+
   desc "Update test databases by dumping and loading --> parallel:prepare[num_cpus]"
   task(:prepare, [:count]) do |_, args|
     ParallelTests::Tasks.check_for_pending_migrations
@@ -194,14 +219,25 @@ def rails_61_or_greater?
   end
 
   # when dumping/resetting takes too long
-  desc "Update test databases via db:migrate --> parallel:migrate[num_cpus,database]"
-  task :migrate, [:count, :database] do |_, args|
+  desc "Update test databases via db:migrate --> parallel:migrate[num_cpus]"
+  task :migrate, :count do |_, args|
     ParallelTests::Tasks.run_in_parallel(
-      [$0, "db:migrate#{":#{args[:database]}" if args[:database]}", "RAILS_ENV=#{ParallelTests::Tasks.rails_env}"],
+      [$0, "db:migrate", "RAILS_ENV=#{ParallelTests::Tasks.rails_env}"],
       args
     )
   end
 
+  # migrate individual databases in multi DB setups
+  ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
+    desc "Update test #{name} database via db:migrate:#{name} --> parallel:migrate:#{name}[num_cpus]"
+    task "migrate:#{name}".to_sym, :count do |_, args|
+      ParallelTests::Tasks.run_in_parallel(
+        [$0, "db:migrate:#{name}", "RAILS_ENV=#{ParallelTests::Tasks.rails_env}"],
+        args
+      )
+    end
+  end
+
   desc "Rollback test databases via db:rollback --> parallel:rollback[num_cpus]"
   task :rollback, :count do |_, args|
     ParallelTests::Tasks.run_in_parallel(
@@ -211,18 +247,33 @@ def rails_61_or_greater?
   end
 
   # just load the schema (good for integration server <-> no development db)
-  desc "Load dumped schema for test databases via db:schema:load --> parallel:load_schema[num_cpus,database]"
-  task :load_schema, [:count, :database] do |_, args|
+  desc "Load dumped schema for test databases via db:schema:load --> parallel:load_schema[num_cpus]"
+  task :load_schema, :count do |_, args|
     command = [
       $0,
       ParallelTests::Tasks.purge_before_load,
-      "db:schema:load#{":#{args[:database]}" if args[:database]}",
+      "db:schema:load",
       "RAILS_ENV=#{ParallelTests::Tasks.rails_env}",
       "DISABLE_DATABASE_ENVIRONMENT_CHECK=1"
     ]
     ParallelTests::Tasks.run_in_parallel(ParallelTests::Tasks.suppress_schema_load_output(command), args)
   end
 
+  # load the schema for individual databases in multi DB setups
+  ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
+    desc "Load dumped schema for test #{name} database via db:schema:load:#{name} --> parallel:load_schema:#{name}[num_cpus]"
+    task "load_schema:#{name}".to_sym, :count do |_, args|
+      command = [
+        $0,
+        ParallelTests::Tasks.purge_before_load,
+        "db:schema:load:#{name}",
+        "RAILS_ENV=#{ParallelTests::Tasks.rails_env}",
+        "DISABLE_DATABASE_ENVIRONMENT_CHECK=1"
+      ]
+      ParallelTests::Tasks.run_in_parallel(ParallelTests::Tasks.suppress_schema_load_output(command), args)
+    end
+  end
+
   # load the structure from the structure.sql file
   # (faster for rails < 6.1, deprecated after and only configured by `ActiveRecord::Base.schema_format`)
   desc "Load structure for test databases via db:schema:load --> parallel:load_structure[num_cpus]"