Update gitlab:db:configure for multiple databases

Update the configure rake task to work when multiple databases are
configured. The behavior should remain the same if only one database is
configured (or one database with database_tasks are enabled).

lib/tasks/gitlab/db.rake

@@ -84,16 +84,38 @@ namespace :gitlab do
 
     desc 'GitLab | DB | Configures the database by running migrate, or by loading the schema and seeding if needed'
     task configure: :environment do
-      # Check if we have existing db tables
-      # The schema_migrations table will still exist if drop_tables was called
-      if ActiveRecord::Base.connection.tables.count > 1
-        Rake::Task['db:migrate'].invoke
+      databases_with_tasks = ActiveRecord::Base.configurations.configs_for(env_name: Rails.env)
+
+      databases_loaded = []
+
+      if databases_with_tasks.size == 1
+        next unless databases_with_tasks.first.name == 'main'
+
+        connection = Gitlab::Database.database_base_models['main'].connection
+        databases_loaded << configure_database(connection)
       else
-        # Add post-migrate paths to ensure we mark all migrations as up
+        Gitlab::Database.database_base_models.each do |name, model|
+          next unless databases_with_tasks.any? { |db_with_tasks| db_with_tasks.name == name }
+
+          databases_loaded << configure_database(model.connection, database_name: name)
+        end
+      end
+
+      Rake::Task['db:seed_fu'].invoke if databases_loaded.present? && databases_loaded.all?
+    end
+
+    def configure_database(connection, database_name: nil)
+      database_name = ":#{database_name}" if database_name
+      load_database = connection.tables.count <= 1
+
+      if load_database
         Gitlab::Database.add_post_migrate_path_to_rails(force: true)
-        Rake::Task['db:structure:load'].invoke
-        Rake::Task['db:seed_fu'].invoke
+        Rake::Task["db:schema:load#{database_name}"].invoke
+      else
+        Rake::Task["db:migrate#{database_name}"].invoke
       end
+
+      load_database
     end
 
     desc 'GitLab | DB | Run database migrations and print `unattended_migrations_completed` if action taken'

spec/tasks/gitlab/db_rake_spec.rb

@@ -51,7 +51,7 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
       let(:base_models) { { 'main' => main_model, 'ci' => ci_model } }
 
       before do
-        skip_if_multiple_databases_not_setup
+        skip_unless_ci_uses_database_tasks
 
         allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
       end
@@ -133,82 +133,231 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
   end
 
   describe 'configure' do
-    it 'invokes db:migrate when schema has already been loaded' do
-      allow(ActiveRecord::Base.connection).to receive(:tables).and_return(%w[table1 table2])
+    context 'with a single database' do
+      let(:connection) { Gitlab::Database.database_base_models[:main].connection }
+      let(:main_config) { double(:config, name: 'main') }
+
+      before do
+        skip_if_multiple_databases_are_setup
+      end
+
+      context 'when geo is not configured' do
+        before do
+          allow(ActiveRecord::Base).to receive_message_chain('configurations.configs_for').and_return([main_config])
+        end
+
+        context 'when the schema is already loaded' do
+          it 'migrates the database' do
+            allow(connection).to receive(:tables).and_return(%w[table1 table2])
+
             expect(Rake::Task['db:migrate']).to receive(:invoke)
-      expect(Rake::Task['db:structure:load']).not_to receive(:invoke)
+            expect(Rake::Task['db:schema:load']).not_to receive(:invoke)
             expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
-      expect { run_rake_task('gitlab:db:configure') }.not_to raise_error
+
+            run_rake_task('gitlab:db:configure')
           end
+        end
+
+        context 'when the schema is not loaded' do
+          it 'loads the schema and seeds the database' do
+            allow(connection).to receive(:tables).and_return([])
 
-    it 'invokes db:shema:load and db:seed_fu when schema is not loaded' do
-      allow(ActiveRecord::Base.connection).to receive(:tables).and_return([])
-      expect(Rake::Task['db:structure:load']).to receive(:invoke)
+            expect(Rake::Task['db:schema:load']).to receive(:invoke)
             expect(Rake::Task['db:seed_fu']).to receive(:invoke)
             expect(Rake::Task['db:migrate']).not_to receive(:invoke)
-      expect { run_rake_task('gitlab:db:configure') }.not_to raise_error
+
+            run_rake_task('gitlab:db:configure')
           end
+        end
+
+        context 'when only a single table is present' do
+          it 'loads the schema and seeds the database' do
+            allow(connection).to receive(:tables).and_return(['default'])
 
-    it 'invokes db:shema:load and db:seed_fu when there is only a single table present' do
-      allow(ActiveRecord::Base.connection).to receive(:tables).and_return(['default'])
-      expect(Rake::Task['db:structure:load']).to receive(:invoke)
+            expect(Rake::Task['db:schema:load']).to receive(:invoke)
             expect(Rake::Task['db:seed_fu']).to receive(:invoke)
             expect(Rake::Task['db:migrate']).not_to receive(:invoke)
-      expect { run_rake_task('gitlab:db:configure') }.not_to raise_error
-    end
 
-    it 'does not invoke any other rake tasks during an error' do
-      allow(ActiveRecord::Base).to receive(:connection).and_raise(RuntimeError, 'error')
-      expect(Rake::Task['db:migrate']).not_to receive(:invoke)
-      expect(Rake::Task['db:structure:load']).not_to receive(:invoke)
-      expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
-      expect { run_rake_task('gitlab:db:configure') }.to raise_error(RuntimeError, 'error')
-      # unstub connection so that the database cleaner still works
-      allow(ActiveRecord::Base).to receive(:connection).and_call_original
+            run_rake_task('gitlab:db:configure')
+          end
         end
 
-    it 'does not invoke seed after a failed schema_load' do
-      allow(ActiveRecord::Base.connection).to receive(:tables).and_return([])
-      allow(Rake::Task['db:structure:load']).to receive(:invoke).and_raise(RuntimeError, 'error')
-      expect(Rake::Task['db:structure:load']).to receive(:invoke)
+        context 'when loading the schema fails' do
+          it 'does not seed the database' do
+            allow(connection).to receive(:tables).and_return([])
+
+            expect(Rake::Task['db:schema:load']).to receive(:invoke).and_raise('error')
             expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
             expect(Rake::Task['db:migrate']).not_to receive(:invoke)
+
             expect { run_rake_task('gitlab:db:configure') }.to raise_error(RuntimeError, 'error')
           end
+        end
 
         context 'SKIP_POST_DEPLOYMENT_MIGRATIONS environment variable set' do
           let(:rails_paths) { { 'db' => ['db'], 'db/migrate' => ['db/migrate'] } }
 
           before do
-        allow(ENV).to receive(:[]).and_call_original
-        allow(ENV).to receive(:[]).with('SKIP_POST_DEPLOYMENT_MIGRATIONS').and_return true
+            stub_env('SKIP_POST_DEPLOYMENT_MIGRATIONS', true)
 
             # Our environment has already been loaded, so we need to pretend like post_migrations were not
             allow(Rails.application.config).to receive(:paths).and_return(rails_paths)
             allow(ActiveRecord::Migrator).to receive(:migrations_paths).and_return(rails_paths['db/migrate'].dup)
           end
 
-      it 'adds post deployment migrations before schema load if the schema is not already loaded' do
-        allow(ActiveRecord::Base.connection).to receive(:tables).and_return([])
+          context 'when the schema is not loaded' do
+            it 'adds the post deployment migration path before schema load' do
+              allow(connection).to receive(:tables).and_return([])
+
               expect(Gitlab::Database).to receive(:add_post_migrate_path_to_rails).and_call_original
-        expect(Rake::Task['db:structure:load']).to receive(:invoke)
+              expect(Rake::Task['db:schema:load']).to receive(:invoke)
               expect(Rake::Task['db:seed_fu']).to receive(:invoke)
               expect(Rake::Task['db:migrate']).not_to receive(:invoke)
-        expect { run_rake_task('gitlab:db:configure') }.not_to raise_error
+
+              run_rake_task('gitlab:db:configure')
+
               expect(rails_paths['db/migrate'].include?(File.join(Rails.root, 'db', 'post_migrate'))).to be(true)
             end
+          end
+
+          context 'when the schema is loaded' do
+            it 'ignores post deployment migrations' do
+              allow(connection).to receive(:tables).and_return(%w[table1 table2])
 
-      it 'ignores post deployment migrations  when schema has already been loaded' do
-        allow(ActiveRecord::Base.connection).to receive(:tables).and_return(%w[table1 table2])
               expect(Rake::Task['db:migrate']).to receive(:invoke)
               expect(Gitlab::Database).not_to receive(:add_post_migrate_path_to_rails)
-        expect(Rake::Task['db:structure:load']).not_to receive(:invoke)
+              expect(Rake::Task['db:schema:load']).not_to receive(:invoke)
               expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
-        expect { run_rake_task('gitlab:db:configure') }.not_to raise_error
+
+              run_rake_task('gitlab:db:configure')
+
               expect(rails_paths['db/migrate'].include?(File.join(Rails.root, 'db', 'post_migrate'))).to be(false)
             end
           end
         end
+      end
+
+      context 'when geo is configured' do
+        context 'when the main database is also configured' do
+          before do
+            skip_unless_geo_configured
+          end
+
+          it 'only configures the main database' do
+            allow(connection).to receive(:tables).and_return(%w[table1 table2])
+
+            expect(Rake::Task['db:migrate:main']).to receive(:invoke)
+
+            expect(Rake::Task['db:migrate:geo']).not_to receive(:invoke)
+            expect(Rake::Task['db:schema:load:geo']).not_to receive(:invoke)
+
+            run_rake_task('gitlab:db:configure')
+          end
+        end
+      end
+    end
+
+    context 'with multiple databases' do
+      let(:main_model) { double(:model, connection: double(:connection)) }
+      let(:ci_model) { double(:model, connection: double(:connection)) }
+      let(:base_models) { { 'main' => main_model, 'ci' => ci_model }.with_indifferent_access }
+
+      let(:main_config) { double(:config, name: 'main') }
+      let(:ci_config) { double(:config, name: 'ci') }
+
+      before do
+        skip_unless_ci_uses_database_tasks
+
+        allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
+      end
+
+      context 'when geo is not configured' do
+        before do
+          allow(ActiveRecord::Base).to receive_message_chain('configurations.configs_for')
+            .and_return([main_config, ci_config])
+        end
+
+        context 'when no database has the schema loaded' do
+          before do
+            allow(main_model.connection).to receive(:tables).and_return(%w[schema_migrations])
+            allow(ci_model.connection).to receive(:tables).and_return([])
+          end
+
+          it 'loads the schema and seeds all the databases' do
+            expect(Rake::Task['db:schema:load:main']).to receive(:invoke)
+            expect(Rake::Task['db:schema:load:ci']).to receive(:invoke)
+
+            expect(Rake::Task['db:migrate:main']).not_to receive(:invoke)
+            expect(Rake::Task['db:migrate:ci']).not_to receive(:invoke)
+
+            expect(Rake::Task['db:seed_fu']).to receive(:invoke)
+
+            run_rake_task('gitlab:db:configure')
+          end
+        end
+
+        context 'when both databases have the schema loaded' do
+          before do
+            allow(main_model.connection).to receive(:tables).and_return(%w[table1 table2])
+            allow(ci_model.connection).to receive(:tables).and_return(%w[table1 table2])
+          end
+
+          it 'migrates the databases without seeding them' do
+            expect(Rake::Task['db:migrate:main']).to receive(:invoke)
+            expect(Rake::Task['db:migrate:ci']).to receive(:invoke)
+
+            expect(Rake::Task['db:schema:load:main']).not_to receive(:invoke)
+            expect(Rake::Task['db:schema:load:ci']).not_to receive(:invoke)
+
+            expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
+
+            run_rake_task('gitlab:db:configure')
+          end
+        end
+
+        context 'when only one database has the schema loaded' do
+          before do
+            allow(main_model.connection).to receive(:tables).and_return(%w[table1 table2])
+            allow(ci_model.connection).to receive(:tables).and_return([])
+          end
+
+          it 'migrates and loads the schema correctly, without seeding the databases' do
+            expect(Rake::Task['db:migrate:main']).to receive(:invoke)
+            expect(Rake::Task['db:schema:load:main']).not_to receive(:invoke)
+
+            expect(Rake::Task['db:schema:load:ci']).to receive(:invoke)
+            expect(Rake::Task['db:migrate:ci']).not_to receive(:invoke)
+
+            expect(Rake::Task['db:seed_fu']).not_to receive(:invoke)
+
+            run_rake_task('gitlab:db:configure')
+          end
+        end
+      end
+
+      context 'when geo is configured' do
+        let(:geo_config) { double(:config, name: 'geo') }
+
+        before do
+          skip_unless_geo_configured
+
+          allow(main_model.connection).to receive(:tables).and_return(%w[schema_migrations])
+          allow(ci_model.connection).to receive(:tables).and_return(%w[schema_migrations])
+        end
+
+        it 'does not run tasks against geo' do
+          expect(Rake::Task['db:schema:load:main']).to receive(:invoke)
+          expect(Rake::Task['db:schema:load:ci']).to receive(:invoke)
+          expect(Rake::Task['db:seed_fu']).to receive(:invoke)
+
+          expect(Rake::Task['db:migrate:geo']).not_to receive(:invoke)
+          expect(Rake::Task['db:schema:load:geo']).not_to receive(:invoke)
+
+          run_rake_task('gitlab:db:configure')
+        end
+      end
+    end
+  end
 
   describe 'unattended' do
     using RSpec::Parameterized::TableSyntax
@@ -301,7 +450,7 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
       let(:base_models) { { 'main' => main_model, 'ci' => ci_model } }
 
       before do
-        skip_if_multiple_databases_not_setup
+        skip_unless_ci_uses_database_tasks
 
         allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
 
@@ -373,7 +522,7 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
 
     context 'with multiple databases' do
       before do
-        skip_if_multiple_databases_not_setup
+        skip_unless_ci_uses_database_tasks
       end
 
       context 'when running the multi-database variant' do
@@ -443,6 +592,10 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
       end
 
       context 'when the single database task is used' do
+        before do
+          skip_unless_ci_uses_database_tasks
+        end
+
         it 'delegates to Gitlab::Database::Reindexing with a specific database' do
           expect(Gitlab::Database::Reindexing).to receive(:invoke).with('ci')
 
@@ -576,7 +729,7 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
 
   context 'with multiple databases', :reestablished_active_record_base do
     before do
-      skip_if_multiple_databases_not_setup
+      skip_unless_ci_uses_database_tasks
     end
 
     describe 'db:structure:dump against a single database' do
@@ -671,6 +824,14 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
     run_rake_task(test_task_name)
   end
 
+  def skip_unless_ci_uses_database_tasks
+    skip "Skipping because database tasks won't run against the ci database" unless ci_database_tasks?
+  end
+
+  def ci_database_tasks?
+    !!ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, name: 'ci')&.database_tasks?
+  end
+
   def skip_unless_geo_configured
     skip 'Skipping because the geo database is not configured' unless geo_configured?
   end