Merge branch 'pb-multi-db-rake-tasks-2' into 'master'

Update db:drop_tables to have multi-db versions

See merge request gitlab-org/gitlab!80327

lib/gitlab/database/each_database.rb

@@ -4,8 +4,11 @@ module Gitlab
   module Database
     module EachDatabase
       class << self
-        def each_database_connection
-          Gitlab::Database.database_base_models.each_pair do |connection_name, model|
+        def each_database_connection(only: nil)
+          selected_names = Array.wrap(only)
+          base_models = select_base_models(selected_names)
+
+          base_models.each_pair do |connection_name, model|
             connection = model.connection
 
             with_shared_connection(connection, connection_name) do
@@ -28,6 +31,18 @@ module Gitlab
 
         private
 
+        def select_base_models(names)
+          base_models = Gitlab::Database.database_base_models
+
+          return base_models if names.empty?
+
+          names.each_with_object(HashWithIndifferentAccess.new) do |name, hash|
+            raise ArgumentError, "#{name} is not a valid database name" unless base_models.key?(name)
+
+            hash[name] = base_models[name]
+          end
+        end
+
         def with_shared_model_connections(shared_model, &blk)
           Gitlab::Database.database_base_models.each_pair do |connection_name, connection_model|
             if shared_model.limit_connection_names

lib/tasks/gitlab/db.rake

@@ -4,30 +4,28 @@ databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml
 
 namespace :gitlab do
   namespace :db do
-    desc 'GitLab | DB | Manually insert schema migration version'
+    desc 'GitLab | DB | Manually insert schema migration version on all configured databases'
     task :mark_migration_complete, [:version] => :environment do |_, args|
       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)
+      ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |database|
+        desc "Gitlab | DB | Manually insert schema migration version on #{database} database"
+        task database, [:version] => :environment do |_, args|
+          mark_migration_complete(args[:version], only_on: database)
         end
       end
     end
 
-    def mark_migration_complete(version, database: nil)
+    def mark_migration_complete(version, only_on: nil)
       if version.to_i == 0
         puts 'Must give a version argument that is a non-zero integer'.color(:red)
         exit 1
       end
 
-      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)})")
+      Gitlab::Database::EachDatabase.each_database_connection(only: only_on) do |connection, name|
+        connection.execute("INSERT INTO schema_migrations (version) VALUES (#{connection.quote(version)})")
 
         puts "Successfully marked '#{version}' as complete on database #{name}".color(:green)
       rescue ActiveRecord::RecordNotUnique
@@ -35,32 +33,44 @@ namespace :gitlab do
       end
     end
 
-    desc 'GitLab | DB | Drop all tables'
+    desc 'GitLab | DB | Drop all tables on all configured databases'
     task drop_tables: :environment do
-      connection = ActiveRecord::Base.connection
+      drop_tables
+    end
+
+    namespace :drop_tables do
+      ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |database|
+        desc "GitLab | DB | Drop all tables on the #{database} database"
+        task database => :environment do
+          drop_tables(only_on: database)
+        end
+      end
+    end
 
-      # In PostgreSQLAdapter, data_sources returns both views and tables, so use
-      # #tables instead
-      tables = connection.tables
+    def drop_tables(only_on: nil)
+      Gitlab::Database::EachDatabase.each_database_connection(only: only_on) do |connection, name|
+        # In PostgreSQLAdapter, data_sources returns both views and tables, so use tables instead
+        tables = connection.tables
 
-      # Removes the entry from the array
-      tables.delete 'schema_migrations'
-      # Truncate schema_migrations to ensure migrations re-run
-      connection.execute('TRUNCATE schema_migrations') if connection.table_exists? 'schema_migrations'
+        # Removes the entry from the array
+        tables.delete 'schema_migrations'
+        # Truncate schema_migrations to ensure migrations re-run
+        connection.execute('TRUNCATE schema_migrations') if connection.table_exists? 'schema_migrations'
 
-      # Drop any views
-      connection.views.each do |view|
-        connection.execute("DROP VIEW IF EXISTS #{connection.quote_table_name(view)} CASCADE")
-      end
+        # Drop any views
+        connection.views.each do |view|
+          connection.execute("DROP VIEW IF EXISTS #{connection.quote_table_name(view)} CASCADE")
+        end
 
-      # Drop tables with cascade to avoid dependent table errors
-      # PG: http://www.postgresql.org/docs/current/static/ddl-depend.html
-      # Add `IF EXISTS` because cascade could have already deleted a table.
-      tables.each { |t| connection.execute("DROP TABLE IF EXISTS #{connection.quote_table_name(t)} CASCADE") }
+        # Drop tables with cascade to avoid dependent table errors
+        # PG: http://www.postgresql.org/docs/current/static/ddl-depend.html
+        # Add `IF EXISTS` because cascade could have already deleted a table.
+        tables.each { |t| connection.execute("DROP TABLE IF EXISTS #{connection.quote_table_name(t)} CASCADE") }
 
-      # Drop all extra schema objects GitLab owns
-      Gitlab::Database::EXTRA_SCHEMAS.each do |schema|
-        connection.execute("DROP SCHEMA IF EXISTS #{connection.quote_table_name(schema)} CASCADE")
+        # Drop all extra schema objects GitLab owns
+        Gitlab::Database::EXTRA_SCHEMAS.each do |schema|
+          connection.execute("DROP SCHEMA IF EXISTS #{connection.quote_table_name(schema)} CASCADE")
+        end
       end
     end
 
@@ -208,8 +218,6 @@ namespace :gitlab do
     end
 
     namespace :reindex do
-      databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml
-
       ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |database_name|
         desc "Reindex #{database_name} database without downtime to eliminate bloat"
         task database_name => :environment do

spec/lib/gitlab/database/each_database_spec.rb

@@ -3,13 +3,13 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::EachDatabase do
-  describe '.each_database_connection' do
+  describe '.each_database_connection', :add_ci_connection do
     before do
       allow(Gitlab::Database).to receive(:database_base_models)
         .and_return({ main: ActiveRecord::Base, ci: Ci::ApplicationRecord }.with_indifferent_access)
     end
 
-    it 'yields each connection after connecting SharedModel', :add_ci_connection do
+    it 'yields each connection after connecting SharedModel' do
       expect(Gitlab::Database::SharedModel).to receive(:using_connection)
         .with(ActiveRecord::Base.connection).ordered.and_yield
 
@@ -22,6 +22,42 @@ RSpec.describe Gitlab::Database::EachDatabase do
           [Ci::ApplicationRecord.connection, 'ci']
         )
     end
+
+    context 'when only certain databases are selected' do
+      it 'yields the selected connections after connecting SharedModel' do
+        expect(Gitlab::Database::SharedModel).to receive(:using_connection)
+          .with(Ci::ApplicationRecord.connection).ordered.and_yield
+
+        expect { |b| described_class.each_database_connection(only: 'ci', &b) }
+          .to yield_successive_args([Ci::ApplicationRecord.connection, 'ci'])
+      end
+
+      context 'when the selected names are passed as symbols' do
+        it 'yields the selected connections after connecting SharedModel' do
+          expect(Gitlab::Database::SharedModel).to receive(:using_connection)
+            .with(Ci::ApplicationRecord.connection).ordered.and_yield
+
+          expect { |b| described_class.each_database_connection(only: :ci, &b) }
+            .to yield_successive_args([Ci::ApplicationRecord.connection, 'ci'])
+        end
+      end
+
+      context 'when the selected names are invalid' do
+        it 'does not yield any connections' do
+          expect do |b|
+            described_class.each_database_connection(only: :notvalid, &b)
+          rescue ArgumentError => e
+            expect(e.message).to match(/notvalid is not a valid database name/)
+          end.not_to yield_control
+        end
+
+        it 'raises an error' do
+          expect do
+            described_class.each_database_connection(only: :notvalid) {}
+          end.to raise_error(ArgumentError, /notvalid is not a valid database name/)
+        end
+      end
+    end
   end
 
   describe '.each_model_connection' do

spec/tasks/gitlab/db_rake_spec.rb

@@ -261,45 +261,78 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
   end
 
   describe 'drop_tables' do
-    subject { run_rake_task('gitlab:db:drop_tables') }
-
-    let(:tables) { %w(one two) }
+    let(:tables) { %w(one two schema_migrations) }
     let(:views) { %w(three four) }
-    let(:connection) { ActiveRecord::Base.connection }
+    let(:schemas) { Gitlab::Database::EXTRA_SCHEMAS }
 
-    before do
-      allow(connection).to receive(:execute).and_return(nil)
+    context 'with a single database' do
+      let(:connection) { ActiveRecord::Base.connection }
+
+      before do
+        skip_if_multiple_databases_are_setup
 
-      allow(connection).to receive(:tables).and_return(tables)
-      allow(connection).to receive(:views).and_return(views)
+        allow(connection).to receive(:execute).and_return(nil)
+
+        allow(connection).to receive(:tables).and_return(tables)
+        allow(connection).to receive(:views).and_return(views)
+      end
+
+      it 'drops all objects for the database', :aggregate_failures do
+        expect_objects_to_be_dropped(connection)
+
+        run_rake_task('gitlab:db:drop_tables')
+      end
     end
 
-    it 'drops all tables, except schema_migrations' do
-      expect(connection).to receive(:execute).with('DROP TABLE IF EXISTS "one" CASCADE')
-      expect(connection).to receive(:execute).with('DROP TABLE IF EXISTS "two" CASCADE')
+    context 'with multiple databases', :aggregate_failures do
+      let(:main_model) { double(:model, connection: double(:connection, tables: tables, views: views)) }
+      let(:ci_model) { double(:model, connection: double(:connection, tables: tables, views: views)) }
+      let(:base_models) { { 'main' => main_model, 'ci' => ci_model } }
 
-      subject
+      before do
+        skip_if_multiple_databases_not_setup
+
+        allow(Gitlab::Database).to receive(:database_base_models).and_return(base_models)
+
+        allow(main_model.connection).to receive(:table_exists?).with('schema_migrations').and_return(true)
+        allow(ci_model.connection).to receive(:table_exists?).with('schema_migrations').and_return(true)
+
+        (tables + views + schemas).each do |name|
+          allow(main_model.connection).to receive(:quote_table_name).with(name).and_return("\"#{name}\"")
+          allow(ci_model.connection).to receive(:quote_table_name).with(name).and_return("\"#{name}\"")
+        end
+      end
+
+      it 'drops all objects for all databases', :aggregate_failures do
+        expect_objects_to_be_dropped(main_model.connection)
+        expect_objects_to_be_dropped(ci_model.connection)
+
+        run_rake_task('gitlab:db:drop_tables')
+      end
+
+      context 'when the single database task is used' do
+        it 'drops all objects for the given database', :aggregate_failures do
+          expect_objects_to_be_dropped(main_model.connection)
+
+          expect(ci_model.connection).not_to receive(:execute)
+
+          run_rake_task('gitlab:db:drop_tables:main')
+        end
+      end
     end
 
-    it 'drops all views' do
+    def expect_objects_to_be_dropped(connection)
+      expect(connection).to receive(:execute).with('DROP TABLE IF EXISTS "one" CASCADE')
+      expect(connection).to receive(:execute).with('DROP TABLE IF EXISTS "two" CASCADE')
+
       expect(connection).to receive(:execute).with('DROP VIEW IF EXISTS "three" CASCADE')
       expect(connection).to receive(:execute).with('DROP VIEW IF EXISTS "four" CASCADE')
 
-      subject
-    end
-
-    it 'truncates schema_migrations table' do
       expect(connection).to receive(:execute).with('TRUNCATE schema_migrations')
 
-      subject
-    end
-
-    it 'drops extra schemas' do
       Gitlab::Database::EXTRA_SCHEMAS.each do |schema|
         expect(connection).to receive(:execute).with("DROP SCHEMA IF EXISTS \"#{schema}\" CASCADE")
       end
-
-      subject
     end
   end