Add support for multiple DBs in Rails >= 6.1 (#930)

* Add support for multiple DBs in Rails >= 6.1 * changelog entry * rubocop * Return nil to represent 'all' databases * No core extensions * ActiveRecord check in for_each_database * don't return too early now
grosser · Feb 6, 2024 · f776f62 · f776f62
1 parent 8129c4c
commit f776f62
Show file tree

Hide file tree

Showing 3 changed files with 108 additions and 33 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@
 ### Breaking Changes
 
 ### Added
+- Support for running tasks against individual databases in a multi-database setup with Rails >= 6.1 ([#930](https://github.com/grosser/parallel_tests/pull/930))
 
 ## 4.4.0 - 2023-12-24
 

diff --git a/Readme.md b/Readme.md
@@ -33,18 +33,28 @@ test:
 ### Create additional database(s)
     rake parallel:create
 
+### (Multi-DB) Create individual database
+    rake parallel:create:<database>
+    rake parallel:create:secondary
+
 ### Copy development schema (repeat after migrations)
     rake parallel:prepare
 
 ### Run migrations in additional database(s) (repeat after migrations)
     rake parallel:migrate
 
+### (Multi-DB) Run migrations in individual database
+    rake parallel:migrate:<database>
+
 ### Setup environment from scratch (create db and loads schema, useful for CI)
     rake parallel:setup
 
 ### Drop all test databases
     rake parallel:drop
 
+### (Multi-DB) Drop individual test database
+    rake parallel:drop:<database>
+
 ### Run!
     rake parallel:test          # Minitest
     rake parallel:spec          # RSpec

diff --git a/lib/parallel_tests/tasks.rb b/lib/parallel_tests/tasks.rb
@@ -48,6 +48,9 @@ def suppress_output(command, ignore_regex)
         activate_pipefail = "set -o pipefail"
         remove_ignored_lines = %{(grep -v #{Shellwords.escape(ignore_regex)} || true)}
 
+        # remove nil values (ex: #purge_before_load returns nil)
+        command.compact!
+
         if system('/bin/bash', '-c', "#{activate_pipefail} 2>/dev/null")
           shell_command = "#{activate_pipefail} && (#{Shellwords.shelljoin(command)}) | #{remove_ignored_lines}"
           ['/bin/bash', '-c', shell_command]
@@ -125,6 +128,23 @@ def build_run_command(type, args)
         command
       end
 
+      def configured_databases
+        return [] unless defined?(ActiveRecord) && rails_61_or_greater?
+
+        @@configured_databases ||= ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml
+      end
+
+      def for_each_database(&block)
+        # Use nil to represent all databases
+        block&.call(nil)
+
+        return unless defined?(ActiveRecord)
+
+        ActiveRecord::Tasks::DatabaseTasks.for_each(configured_databases) do |name|
+          block&.call(name)
+        end
+      end
+
       private
 
       def rails_7_or_greater?
@@ -145,25 +165,45 @@ 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]"
-  task :create, :count do |_, args|
-    ParallelTests::Tasks.run_in_parallel(
-      [$0, "db:create", "RAILS_ENV=#{ParallelTests::Tasks.rails_env}"],
-      args
-    )
+  ParallelTests::Tasks.for_each_database do |name|
+    task_name = 'create'
+    task_name += ":#{name}" if name
+    description = if name
+                    "Create test #{name} database via db:#{task_name} --> parallel:#{task_name}[num_cpus]"
+                  else
+                    "Create test databases via db:#{task_name} --> parallel:#{task_name}[num_cpus]"
+                  end
+
+    desc description
+    task task_name.to_sym, :count do |_, args|
+      ParallelTests::Tasks.run_in_parallel(
+        [$0, "db:#{task_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(
-      [
-        $0,
-        "db:drop",
-        "RAILS_ENV=#{ParallelTests::Tasks.rails_env}",
-        "DISABLE_DATABASE_ENVIRONMENT_CHECK=1"
-      ],
-      args
-    )
+  ParallelTests::Tasks.for_each_database do |name|
+    task_name = 'drop'
+    task_name += ":#{name}" if name
+    description = if name
+                    "Drop test #{name} database via db:#{task_name} --> parallel:#{task_name}[num_cpus]"
+                  else
+                    "Drop test databases via db:#{task_name} --> parallel:#{task_name}[num_cpus]"
+                  end
+
+    desc description
+    task task_name.to_sym, :count do |_, args|
+      ParallelTests::Tasks.run_in_parallel(
+        [
+          $0,
+          "db:#{task_name}",
+          "RAILS_ENV=#{ParallelTests::Tasks.rails_env}",
+          "DISABLE_DATABASE_ENVIRONMENT_CHECK=1"
+        ],
+        args
+      )
+    end
   end
 
   desc "Update test databases by dumping and loading --> parallel:prepare[num_cpus]"
@@ -190,12 +230,22 @@ def rails_61_or_greater?
   end
 
   # when dumping/resetting takes too long
-  desc "Update test databases via db:migrate --> parallel:migrate[num_cpus]"
-  task :migrate, :count do |_, args|
-    ParallelTests::Tasks.run_in_parallel(
-      [$0, "db:migrate", "RAILS_ENV=#{ParallelTests::Tasks.rails_env}"],
-      args
-    )
+  ParallelTests::Tasks.for_each_database do |name|
+    task_name = 'migrate'
+    task_name += ":#{name}" if name
+    description = if name
+                    "Update test #{name} database via db:#{task_name} --> parallel:#{task_name}[num_cpus]"
+                  else
+                    "Update test databases via db:#{task_name} --> parallel:#{task_name}[num_cpus]"
+                  end
+
+    desc description
+    task task_name.to_sym, :count do |_, args|
+      ParallelTests::Tasks.run_in_parallel(
+        [$0, "db:#{task_name}", "RAILS_ENV=#{ParallelTests::Tasks.rails_env}"],
+        args
+      )
+    end
   end
 
   desc "Rollback test databases via db:rollback --> parallel:rollback[num_cpus]"
@@ -207,16 +257,30 @@ 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]"
-  task :load_schema, :count do |_, args|
-    command = [
-      $0,
-      ParallelTests::Tasks.purge_before_load,
-      "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)
+  ParallelTests::Tasks.for_each_database do |name|
+    rails_task = 'db:schema:load'
+    rails_task += ":#{name}" if name
+
+    task_name = 'load_schema'
+    task_name += ":#{name}" if name
+
+    description = if name
+                    "Load dumped schema for test #{name} database via #{rails_task} --> parallel:#{task_name}[num_cpus]"
+                  else
+                    "Load dumped schema for test databases via #{rails_task} --> parallel:#{task_name}[num_cpus]"
+                  end
+
+    desc description
+    task task_name.to_sym, :count do |_, args|
+      command = [
+        $0,
+        ParallelTests::Tasks.purge_before_load,
+        rails_task,
+        "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