Properly enhance multi-db partitioning rake tasks

We enhance some rake tasks to create dynamic partitions after they
execute, such as `db:migrate` and `db:schema:load`. We should also
enhance the multi-database versions of these, such as `db:migrate:main`
or `db:schema:load:main`.

lib/tasks/gitlab/db.rake

@@ -152,6 +152,12 @@ namespace :gitlab do
       Rake::Task['gitlab:db:create_dynamic_partitions'].invoke
     end
 
+    ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
+      Rake::Task["db:migrate:#{name}"].enhance do
+        Rake::Task['gitlab:db:create_dynamic_partitions'].invoke
+      end
+    end
+
     # When we load the database schema from db/structure.sql
     # we don't have any dynamic partitions created. We don't really need to
     # because application initializers/sidekiq take care of that, too.
@@ -160,14 +166,14 @@ namespace :gitlab do
     #
     # Other than that it's helpful to create partitions early when bootstrapping
     # a new installation.
-    #
-    # Rails 6.1 deprecates db:structure:load in favor of db:schema:load
-    Rake::Task['db:structure:load'].enhance do
+    Rake::Task['db:schema:load'].enhance do
       Rake::Task['gitlab:db:create_dynamic_partitions'].invoke
     end
 
-    Rake::Task['db:schema:load'].enhance do
-      Rake::Task['gitlab:db:create_dynamic_partitions'].invoke
+    ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
+      Rake::Task["db:schema:load:#{name}"].enhance do
+        Rake::Task['gitlab:db:create_dynamic_partitions'].invoke
+      end
     end
 
     # During testing, db:test:load restores the database schema from scratch

spec/tasks/gitlab/db_rake_spec.rb

@@ -422,13 +422,11 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
 
   context 'with multiple databases', :reestablished_active_record_base do
     before do
-      allow(ActiveRecord::Tasks::DatabaseTasks).to receive(:setup_initial_database_yaml).and_return([:main, :geo])
+      skip_if_multiple_databases_not_setup
     end
 
     describe 'db:structure:dump' do
       it 'invokes gitlab:db:clean_structure_sql' do
-        skip unless Gitlab.ee?
-
         expect(Rake::Task['gitlab:db:clean_structure_sql']).to receive(:invoke).twice.and_return(true)
 
         expect { run_rake_task('db:structure:dump:main') }.not_to raise_error
@@ -437,13 +435,19 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
 
     describe 'db:schema:dump' do
       it 'invokes gitlab:db:clean_structure_sql' do
-        skip unless Gitlab.ee?
-
         expect(Rake::Task['gitlab:db:clean_structure_sql']).to receive(:invoke).once.and_return(true)
 
         expect { run_rake_task('db:schema:dump:main') }.not_to raise_error
       end
     end
+
+    describe 'db:migrate' do
+      it 'invokes gitlab:db:create_dynamic_partitions' do
+        expect(Rake::Task['gitlab:db:create_dynamic_partitions']).to receive(:invoke).once.and_return(true)
+
+        expect { run_rake_task('db:migrate:main') }.not_to raise_error
+      end
+    end
   end
 
   describe 'gitlab:db:reset_as_non_superuser' do