Update rake tasks to have multi-db versions

Update the `gitlab:db:mark_migration_complete` rake task to work with
multiple databases. This follows the standard Rails way of writing rake
tasks to work with multiple database setups.

lib/tasks/gitlab/db.rake

@@ -6,23 +6,32 @@ namespace :gitlab do
   namespace :db do
     desc 'GitLab | DB | Manually insert schema migration version'
     task :mark_migration_complete, [:version] => :environment do |_, args|
-      unless args[:version]
-        puts "Must specify a migration version as an argument".color(:red)
-        exit 1
+      mark_migration_complete(args[:version])
+    end
+
+    namespace :mark_migration_complete do
+      ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
+        desc "Gitlab | DB | Manually insert schema migration version on #{name} database"
+        task name, [:version] => :environment do |_, args|
+          mark_migration_complete(args[:version], database: name)
+        end
       end
+    end
 
-      version = args[:version].to_i
-      if version == 0
-        puts "Version '#{args[:version]}' must be a non-zero integer".color(:red)
+    def mark_migration_complete(version, database: nil)
+      if version.to_i == 0
+        puts 'Must give a version argument that is a non-zero integer'.color(:red)
         exit 1
       end
 
-      sql = "INSERT INTO schema_migrations (version) VALUES (#{version})"
-      begin
-        ActiveRecord::Base.connection.execute(sql)
-        puts "Successfully marked '#{version}' as complete".color(:green)
+      Gitlab::Database.database_base_models.each do |name, model|
+        next if database && database.to_s != name
+
+        model.connection.execute("INSERT INTO schema_migrations (version) VALUES (#{model.connection.quote(version)})")
+
+        puts "Successfully marked '#{version}' as complete on database #{name}".color(:green)
       rescue ActiveRecord::RecordNotUnique
-        puts "Migration version '#{version}' is already marked complete".color(:yellow)
+        puts "Migration version '#{version}' is already marked complete on database #{name}".color(:yellow)
       end
     end
 

spec/tasks/gitlab/db_rake_spec.rb

@@ -20,6 +20,99 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
     allow(Rake::Task['db:seed_fu']).to receive(:invoke).and_return(true)
   end
 
+  describe 'mark_migration_complete' do
+    context 'with a single database' do
+      let(:main_model) { ActiveRecord::Base }
+
+      before do
+        skip_if_multiple_databases_are_setup
+      end
+
+      it 'marks the migration complete on the given database' do
+        expect(main_model.connection).to receive(:quote).and_call_original
+        expect(main_model.connection).to receive(:execute)
+          .with("INSERT INTO schema_migrations (version) VALUES ('123')")
+
+        run_rake_task('gitlab:db:mark_migration_complete', '[123]')
+      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 } }
+
+      before do
+        skip_if_multiple_databases_not_setup
+
+        allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
+      end
+
+      it 'marks the migration complete on each database' do
+        expect(main_model.connection).to receive(:quote).with('123').and_return("'123'")
+        expect(main_model.connection).to receive(:execute)
+          .with("INSERT INTO schema_migrations (version) VALUES ('123')")
+
+        expect(ci_model.connection).to receive(:quote).with('123').and_return("'123'")
+        expect(ci_model.connection).to receive(:execute)
+          .with("INSERT INTO schema_migrations (version) VALUES ('123')")
+
+        run_rake_task('gitlab:db:mark_migration_complete', '[123]')
+      end
+
+      context 'when the single database task is used' do
+        it 'marks the migration complete for the given database' do
+          expect(main_model.connection).to receive(:quote).with('123').and_return("'123'")
+          expect(main_model.connection).to receive(:execute)
+            .with("INSERT INTO schema_migrations (version) VALUES ('123')")
+
+          expect(ci_model.connection).not_to receive(:quote)
+          expect(ci_model.connection).not_to receive(:execute)
+
+          run_rake_task('gitlab:db:mark_migration_complete:main', '[123]')
+        end
+      end
+    end
+
+    context 'when the migration is already marked complete' do
+      let(:main_model) { double(:model, connection: double(:connection)) }
+      let(:base_models) { { 'main' => main_model } }
+
+      before do
+        allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
+      end
+
+      it 'prints a warning message' do
+        allow(main_model.connection).to receive(:quote).with('123').and_return("'123'")
+
+        expect(main_model.connection).to receive(:execute)
+          .with("INSERT INTO schema_migrations (version) VALUES ('123')")
+          .and_raise(ActiveRecord::RecordNotUnique)
+
+        expect { run_rake_task('gitlab:db:mark_migration_complete', '[123]') }
+          .to output(/Migration version '123' is already marked complete on database main/).to_stdout
+      end
+    end
+
+    context 'when an invalid version is given' do
+      let(:main_model) { double(:model, connection: double(:connection)) }
+      let(:base_models) { { 'main' => main_model } }
+
+      before do
+        allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
+      end
+
+      it 'prints an error and exits' do
+        expect(main_model).not_to receive(:quote)
+        expect(main_model.connection).not_to receive(:execute)
+
+        expect { run_rake_task('gitlab:db:mark_migration_complete', '[abc]') }
+          .to output(/Must give a version argument that is a non-zero integer/).to_stdout
+          .and raise_error(SystemExit) { |error| expect(error.status).to eq(1) }
+      end
+    end
+  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])