Merge branch 'multiple_db_schema_migratons' into 'master'

Support multiple DBs for loading/dumping of schema_migrations/

See merge request gitlab-org/gitlab!64884

db/ci_schema_migrations/20210617101848

+1b74312f59f6f8937cd0dd754d22dc72e9bdc7302e6254a2fda5762afebe303c
\ No newline at end of file

db/ci_structure.sql

-SET statement_timeout = 0;
-SET lock_timeout = 0;
-SET idle_in_transaction_session_timeout = 0;
-SET client_encoding = 'UTF8';
-SET standard_conforming_strings = on;
-SELECT pg_catalog.set_config('search_path', '', false);
-SET check_function_bodies = false;
-SET xmloption = content;
-SET client_min_messages = warning;
-SET row_security = off;
-
-SET default_tablespace = '';
-
-SET default_table_access_method = heap;
-
---
--- Name: ar_internal_metadata; Type: TABLE; Schema: public; Owner: -
---
-
-CREATE TABLE public.ar_internal_metadata (
+CREATE TABLE ar_internal_metadata (
     key character varying NOT NULL,
     value character varying,
     created_at timestamp(6) without time zone NOT NULL,
     updated_at timestamp(6) without time zone NOT NULL
 );
 
-
---
--- Name: ci_instance_variables; Type: TABLE; Schema: public; Owner: -
---
-
-CREATE TABLE public.ci_instance_variables (
+CREATE TABLE ci_instance_variables (
     id bigint NOT NULL,
     variable_type smallint DEFAULT 1 NOT NULL,
     masked boolean DEFAULT false,
@@ -42,80 +18,28 @@ CREATE TABLE public.ci_instance_variables (
     CONSTRAINT check_956afd70f1 CHECK ((char_length(encrypted_value) <= 13579))
 );
 
-
---
--- Name: ci_instance_variables_id_seq; Type: SEQUENCE; Schema: public; Owner: -
---
-
-CREATE SEQUENCE public.ci_instance_variables_id_seq
+CREATE SEQUENCE ci_instance_variables_id_seq
     START WITH 1
     INCREMENT BY 1
     NO MINVALUE
     NO MAXVALUE
     CACHE 1;
 
+ALTER SEQUENCE ci_instance_variables_id_seq OWNED BY ci_instance_variables.id;
 
---
--- Name: ci_instance_variables_id_seq; Type: SEQUENCE OWNED BY; Schema: public; Owner: -
---
-
-ALTER SEQUENCE public.ci_instance_variables_id_seq OWNED BY public.ci_instance_variables.id;
-
-
---
--- Name: schema_migrations; Type: TABLE; Schema: public; Owner: -
---
-
-CREATE TABLE public.schema_migrations (
+CREATE TABLE schema_migrations (
     version character varying NOT NULL
 );
 
+ALTER TABLE ONLY ci_instance_variables ALTER COLUMN id SET DEFAULT nextval('ci_instance_variables_id_seq'::regclass);
 
---
--- Name: ci_instance_variables id; Type: DEFAULT; Schema: public; Owner: -
---
-
-ALTER TABLE ONLY public.ci_instance_variables ALTER COLUMN id SET DEFAULT nextval('public.ci_instance_variables_id_seq'::regclass);
-
-
---
--- Name: ar_internal_metadata ar_internal_metadata_pkey; Type: CONSTRAINT; Schema: public; Owner: -
---
-
-ALTER TABLE ONLY public.ar_internal_metadata
+ALTER TABLE ONLY ar_internal_metadata
     ADD CONSTRAINT ar_internal_metadata_pkey PRIMARY KEY (key);
 
-
---
--- Name: ci_instance_variables ci_instance_variables_pkey; Type: CONSTRAINT; Schema: public; Owner: -
---
-
-ALTER TABLE ONLY public.ci_instance_variables
+ALTER TABLE ONLY ci_instance_variables
     ADD CONSTRAINT ci_instance_variables_pkey PRIMARY KEY (id);
 
-
---
--- Name: schema_migrations schema_migrations_pkey; Type: CONSTRAINT; Schema: public; Owner: -
---
-
-ALTER TABLE ONLY public.schema_migrations
+ALTER TABLE ONLY schema_migrations
     ADD CONSTRAINT schema_migrations_pkey PRIMARY KEY (version);
 
-
---
--- Name: index_ci_instance_variables_on_key; Type: INDEX; Schema: public; Owner: -
---
-
-CREATE UNIQUE INDEX index_ci_instance_variables_on_key ON public.ci_instance_variables USING btree (key);
-
-
---
--- PostgreSQL database dump complete
---
-
-SET search_path TO "$user", public;
-
-INSERT INTO "schema_migrations" (version) VALUES
-('20210617101848');
-
-
+CREATE UNIQUE INDEX index_ci_instance_variables_on_key ON ci_instance_variables USING btree (key);

db/structure.sql

@@ -27974,6 +27974,4 @@ ALTER TABLE ONLY user_follow_users
     ADD CONSTRAINT user_follow_users_followee_id_fkey FOREIGN KEY (followee_id) REFERENCES users(id) ON DELETE CASCADE;
 
 ALTER TABLE ONLY user_follow_users
-    ADD CONSTRAINT user_follow_users_follower_id_fkey FOREIGN KEY (follower_id) REFERENCES users(id) ON DELETE CASCADE;-- schema_migrations.version information is no longer stored in this file,
--- but instead tracked in the db/schema_migrations directory
--- see https://gitlab.com/gitlab-org/gitlab/-/issues/218590 for details
+    ADD CONSTRAINT user_follow_users_follower_id_fkey FOREIGN KEY (follower_id) REFERENCES users(id) ON DELETE CASCADE;

lib/gitlab/database/postgresql_adapter/dump_schema_versions_mixin.rb

@@ -7,10 +7,7 @@ module Gitlab
         extend ActiveSupport::Concern
 
         def dump_schema_information # :nodoc:
-          return super unless ActiveRecord::Base.configurations.primary?(pool.db_config.name)
-
-          versions = schema_migration.all_versions
-          Gitlab::Database::SchemaVersionFiles.touch_all(versions) if versions.any?
+          Gitlab::Database::SchemaMigrations.touch_all(self)
 
           nil
         end

lib/gitlab/database/postgresql_database_tasks/load_schema_versions_mixin.rb

@@ -7,13 +7,9 @@ module Gitlab
         extend ActiveSupport::Concern
 
         def structure_load(...)
-          result = super(...)
+          super(...)
 
-          if ActiveRecord::Base.configurations.primary?(connection.pool.db_config.name)
-            Gitlab::Database::SchemaVersionFiles.load_all
-          else
-            result
-          end
+          Gitlab::Database::SchemaMigrations.load_all(connection)
         end
       end
     end

lib/gitlab/database/schema_cleaner.rb

@@ -30,11 +30,7 @@ module Gitlab
         structure.gsub!(/\n{3,}/, "\n\n")
 
         io << structure.strip
-        io << <<~MSG
-          -- schema_migrations.version information is no longer stored in this file,
-          -- but instead tracked in the db/schema_migrations directory
-          -- see https://gitlab.com/gitlab-org/gitlab/-/issues/218590 for details
-        MSG
+        io << "\n"
 
         nil
       end

lib/gitlab/database/schema_migrations.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module SchemaMigrations
+      def self.touch_all(connection)
+        context = Gitlab::Database::SchemaMigrations::Context.new(connection)
+
+        Gitlab::Database::SchemaMigrations::Migrations.new(context).touch_all
+      end
+
+      def self.load_all(connection)
+        context = Gitlab::Database::SchemaMigrations::Context.new(connection)
+
+        Gitlab::Database::SchemaMigrations::Migrations.new(context).load_all
+      end
+    end
+  end
+end

lib/gitlab/database/schema_migrations/context.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module SchemaMigrations
+      class Context
+        attr_reader :connection
+
+        def initialize(connection)
+          @connection = connection
+        end
+
+        def schema_directory
+          @schema_directory ||=
+            if ActiveRecord::Base.configurations.primary?(database_name)
+              File.join(db_dir, 'schema_migrations')
+            else
+              File.join(db_dir, "#{database_name}_schema_migrations")
+            end
+        end
+
+        def versions_to_create
+          versions_from_database = @connection.schema_migration.all_versions
+          versions_from_migration_files = @connection.migration_context.migrations.map { |m| m.version.to_s }
+
+          versions_from_database & versions_from_migration_files
+        end
+
+        private
+
+        def database_name
+          @database_name ||= @connection.pool.db_config.name
+        end
+
+        def db_dir
+          @db_dir ||= Rails.application.config.paths["db"].first
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/schema_migrations/migrations.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module SchemaMigrations
+      class Migrations
+        MIGRATION_VERSION_GLOB = '20[0-9][0-9]*'
+
+        def initialize(context)
+          @context = context
+        end
+
+        def touch_all
+          return unless @context.versions_to_create.any?
+
+          version_filepaths = version_filenames.map { |f| File.join(schema_directory, f) }
+          FileUtils.rm(version_filepaths)
+
+          @context.versions_to_create.each do |version|
+            version_filepath = File.join(schema_directory, version)
+
+            File.open(version_filepath, 'w') do |file|
+              file << Digest::SHA256.hexdigest(version)
+            end
+          end
+        end
+
+        def load_all
+          return if version_filenames.empty?
+
+          values = version_filenames.map { |vf| "('#{@context.connection.quote_string(vf)}')" }
+
+          @context.connection.execute(<<~SQL)
+          INSERT INTO schema_migrations (version)
+          VALUES #{values.join(',')}
+          ON CONFLICT DO NOTHING
+          SQL
+        end
+
+        private
+
+        def schema_directory
+          @context.schema_directory
+        end
+
+        def version_filenames
+          @version_filenames ||= Dir.glob(MIGRATION_VERSION_GLOB, base: schema_directory)
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/schema_version_files.rb

-# frozen_string_literal: true
-
-module Gitlab
-  module Database
-    class SchemaVersionFiles
-      SCHEMA_DIRECTORY = 'db/schema_migrations'
-      MIGRATION_DIRECTORIES = %w[db/migrate db/post_migrate].freeze
-      MIGRATION_VERSION_GLOB = '20[0-9][0-9]*'
-
-      def self.touch_all(versions_from_database)
-        versions_from_migration_files = find_versions_from_migration_files
-
-        version_filepaths = find_version_filenames.map { |f| schema_directory.join(f) }
-        FileUtils.rm(version_filepaths)
-
-        versions_to_create = versions_from_database & versions_from_migration_files
-        versions_to_create.each do |version|
-          version_filepath = schema_directory.join(version)
-
-          File.open(version_filepath, 'w') do |file|
-            file << Digest::SHA256.hexdigest(version)
-          end
-        end
-      end
-
-      def self.load_all
-        version_filenames = find_version_filenames
-        return if version_filenames.empty?
-
-        values = version_filenames.map { |vf| "('#{connection.quote_string(vf)}')" }
-        connection.execute(<<~SQL)
-          INSERT INTO schema_migrations (version)
-          VALUES #{values.join(',')}
-          ON CONFLICT DO NOTHING
-        SQL
-      end
-
-      def self.schema_directory
-        @schema_directory ||= Rails.root.join(SCHEMA_DIRECTORY)
-      end
-
-      def self.migration_directories
-        @migration_directories ||= MIGRATION_DIRECTORIES.map { |dir| Rails.root.join(dir) }
-      end
-
-      def self.find_version_filenames
-        Dir.glob(MIGRATION_VERSION_GLOB, base: schema_directory)
-      end
-
-      def self.find_versions_from_migration_files
-        migration_directories.each_with_object([]) do |directory, migration_versions|
-          directory_migrations = Dir.glob(MIGRATION_VERSION_GLOB, base: directory)
-          directory_versions = directory_migrations.map! { |m| m.split('_').first }
-
-          migration_versions.concat(directory_versions)
-        end
-      end
-
-      def self.connection
-        ActiveRecord::Base.connection
-      end
-    end
-  end
-end

lib/tasks/gitlab/db.rake

@@ -90,12 +90,15 @@ namespace :gitlab do
 
     desc 'This adjusts and cleans db/structure.sql - it runs after db:structure:dump'
     task :clean_structure_sql do |task_name|
-      structure_file = 'db/structure.sql'
+      ActiveRecord::Base.configurations.configs_for(env_name: ActiveRecord::Tasks::DatabaseTasks.env).each do |db_config|
+        structure_file = ActiveRecord::Tasks::DatabaseTasks.dump_filename(db_config.name)
+
         schema = File.read(structure_file)
 
         File.open(structure_file, 'wb+') do |io|
           Gitlab::Database::SchemaCleaner.new(schema).clean(io)
         end
+      end
 
       # Allow this task to be called multiple times, as happens when running db:migrate:redo
       Rake::Task[task_name].reenable

spec/fixtures/gitlab/database/structure_example_cleaned.sql

@@ -23,6 +23,4 @@ ALTER TABLE ONLY abuse_reports ALTER COLUMN id SET DEFAULT nextval('abuse_report
 ALTER TABLE ONLY abuse_reports
     ADD CONSTRAINT abuse_reports_pkey PRIMARY KEY (id);
 
-CREATE INDEX index_abuse_reports_on_user_id ON abuse_reports USING btree (user_id);-- schema_migrations.version information is no longer stored in this file,
--- but instead tracked in the db/schema_migrations directory
--- see https://gitlab.com/gitlab-org/gitlab/-/issues/218590 for details
+CREATE INDEX index_abuse_reports_on_user_id ON abuse_reports USING btree (user_id);

spec/lib/gitlab/database/postgresql_adapter/dump_schema_versions_mixin_spec.rb

@@ -3,10 +3,6 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::PostgresqlAdapter::DumpSchemaVersionsMixin do
-  let(:schema_migration) { double('schema_migration', all_versions: versions) }
-  let(:db_name) { 'primary' }
-  let(:versions) { %w(5 2 1000 200 4 93 2) }
-
   let(:instance_class) do
     klass = Class.new do
       def dump_schema_information
@@ -24,43 +20,10 @@ RSpec.describe Gitlab::Database::PostgresqlAdapter::DumpSchemaVersionsMixin do
 
   let(:instance) { instance_class.new }
 
-  before do
-    allow(instance).to receive(:schema_migration).and_return(schema_migration)
-
-    # pool is from ActiveRecord::ConnectionAdapters::PostgreSQLAdapter
-    allow(instance).to receive_message_chain(:pool, :db_config, :name).and_return(db_name)
-  end
-
-  context 'when database name is primary' do
-    context 'when version files exist' do
-      it 'touches version files' do
-        expect(Gitlab::Database::SchemaVersionFiles).to receive(:touch_all).with(versions)
+  it 'calls SchemaMigrations touch_all and skips original implementation' do
+    expect(Gitlab::Database::SchemaMigrations).to receive(:touch_all).with(instance)
     expect(instance).not_to receive(:original_dump_schema_information)
 
     instance.dump_schema_information
   end
-    end
-
-    context 'when version files do not exist' do
-      let(:versions) { [] }
-
-      it 'does not touch version files' do
-        expect(Gitlab::Database::SchemaVersionFiles).not_to receive(:touch_all)
-        expect(instance).not_to receive(:original_dump_schema_information)
-
-        instance.dump_schema_information
-      end
-    end
-  end
-
-  context 'when database name is ci' do
-    let(:db_name) { 'ci' }
-
-    it 'does not touch version files' do
-      expect(Gitlab::Database::SchemaVersionFiles).not_to receive(:touch_all)
-      expect(instance).to receive(:original_dump_schema_information)
-
-      instance.dump_schema_information
-    end
-  end
 end

spec/lib/gitlab/database/postgresql_database_tasks/load_schema_versions_mixin_spec.rb

@@ -3,8 +3,6 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::PostgresqlDatabaseTasks::LoadSchemaVersionsMixin do
-  let(:db_name) { 'primary' }
-
   let(:instance_class) do
     klass = Class.new do
       def structure_load
@@ -22,28 +20,13 @@ RSpec.describe Gitlab::Database::PostgresqlDatabaseTasks::LoadSchemaVersionsMixi
 
   let(:instance) { instance_class.new }
 
-  before do
-    # connection is available in ActiveRecord::Tasks::PostgreSQLDatabaseTasks
-    allow(instance).to receive_message_chain(:connection, :pool, :db_config, :name).and_return(db_name)
-  end
+  it 'calls SchemaMigrations load_all' do
+    connection = double('connection')
+    allow(instance).to receive(:connection).and_return(connection)
 
-  context 'when database is primary' do
-    it 'loads version files' do
-      expect(Gitlab::Database::SchemaVersionFiles).to receive(:load_all)
-      expect(instance).to receive(:original_structure_load)
+    expect(instance).to receive(:original_structure_load).ordered
+    expect(Gitlab::Database::SchemaMigrations).to receive(:load_all).with(connection).ordered
 
     instance.structure_load
   end
-  end
-
-  context 'when the database is ci' do
-    let(:db_name) { 'ci' }
-
-    it 'does not load version files' do
-      expect(Gitlab::Database::SchemaVersionFiles).not_to receive(:load_all)
-      expect(instance).to receive(:original_structure_load)
-
-      instance.structure_load
-    end
-  end
 end

spec/lib/gitlab/database/schema_migrations/context_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::SchemaMigrations::Context do
+  let(:connection) { ActiveRecord::Base.connection }
+
+  let(:context) { described_class.new(connection) }
+
+  describe '#schema_directory' do
+    it 'returns db/schema_migrations' do
+      expect(context.schema_directory).to eq(File.join(Rails.root, 'db/schema_migrations'))
+    end
+
+    context 'multiple databases' do
+      let(:connection) { Ci::BaseModel.connection }
+
+      it 'returns a directory path that is database specific' do
+        skip_if_multiple_databases_not_setup
+
+        expect(context.schema_directory).to eq(File.join(Rails.root, 'db/ci_schema_migrations'))
+      end
+    end
+  end
+
+  describe '#versions_to_create' do
+    before do
+      allow(connection).to receive_message_chain(:schema_migration, :all_versions).and_return(migrated_versions)
+
+      migrations_struct = Struct.new(:version)
+      migrations = file_versions.map { |version| migrations_struct.new(version) }
+      allow(connection).to receive_message_chain(:migration_context, :migrations).and_return(migrations)
+    end
+
+    let(:version1) { '20200123' }
+    let(:version2) { '20200410' }
+    let(:version3) { '20200602' }
+    let(:version4) { '20200809' }
+
+    let(:migrated_versions) { file_versions }
+    let(:file_versions) { [version1, version2, version3, version4] }
+
+    context 'migrated versions is the same as migration file versions' do
+      it 'returns migrated versions' do
+        expect(context.versions_to_create).to eq(migrated_versions)
+      end
+    end
+
+    context 'migrated versions is subset of migration file versions' do
+      let(:migrated_versions) { [version1, version2] }
+
+      it 'returns migrated versions' do
+        expect(context.versions_to_create).to eq(migrated_versions)
+      end
+    end
+
+    context 'migrated versions is superset of migration file versions' do
+      let(:migrated_versions) { file_versions + ['20210809'] }
+
+      it 'returns file versions' do
+        expect(context.versions_to_create).to eq(file_versions)
+      end
+    end
+
+    context 'migrated versions has slightly different versions to migration file versions' do
+      let(:migrated_versions) { [version1, version2, version3, version4, '20210101'] }
+      let(:file_versions) { [version1, version2, version3, version4, '20210102'] }
+
+      it 'returns the common set' do
+        expect(context.versions_to_create).to eq([version1, version2, version3, version4])
+      end
+    end
+  end
+
+  def skip_if_multiple_databases_not_setup
+    skip 'Skipping because multiple databases not set up' unless Gitlab::Database.has_config?(:ci)
+  end
+end

spec/lib/gitlab/database/schema_version_files_spec.rb → spec/lib/gitlab/database/schema_migrations/migrations_spec.rb

@@ -2,43 +2,37 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Database::SchemaVersionFiles do
-  describe '.touch_all' do
+RSpec.describe Gitlab::Database::SchemaMigrations::Migrations do
+  let(:connection) { ApplicationRecord.connection }
+  let(:context) { Gitlab::Database::SchemaMigrations::Context.new(connection) }
+
+  let(:migrations) { described_class.new(context) }
+
+  describe '#touch_all' do
     let(:version1) { '20200123' }
     let(:version2) { '20200410' }
     let(:version3) { '20200602' }
     let(:version4) { '20200809' }
+
     let(:relative_schema_directory) { 'db/schema_migrations' }
-    let(:relative_migrate_directory) { 'db/migrate' }
-    let(:relative_post_migrate_directory) { 'db/post_migrate' }
 
     it 'creates a file containing a checksum for each version with a matching migration' do
       Dir.mktmpdir do |tmpdir|
         schema_directory = Pathname.new(tmpdir).join(relative_schema_directory)
-        migrate_directory = Pathname.new(tmpdir).join(relative_migrate_directory)
-        post_migrate_directory = Pathname.new(tmpdir).join(relative_post_migrate_directory)
-
-        FileUtils.mkdir_p(migrate_directory)
-        FileUtils.mkdir_p(post_migrate_directory)
         FileUtils.mkdir_p(schema_directory)
 
-        migration1_filepath = migrate_directory.join("#{version1}_migration.rb")
-        FileUtils.touch(migration1_filepath)
-
-        migration2_filepath = post_migrate_directory.join("#{version2}_post_migration.rb")
-        FileUtils.touch(migration2_filepath)
-
         old_version_filepath = schema_directory.join('20200101')
         FileUtils.touch(old_version_filepath)
 
         expect(File.exist?(old_version_filepath)).to be(true)
 
-        allow(described_class).to receive(:schema_directory).and_return(schema_directory)
-        allow(described_class).to receive(:migration_directories).and_return([migrate_directory, post_migrate_directory])
+        allow(context).to receive(:schema_directory).and_return(schema_directory)
+        allow(context).to receive(:versions_to_create).and_return([version1, version2])
 
-        described_class.touch_all([version1, version2, version3, version4])
+        migrations.touch_all
 
         expect(File.exist?(old_version_filepath)).to be(false)
+
         [version1, version2].each do |version|
           version_filepath = schema_directory.join(version)
           expect(File.exist?(version_filepath)).to be(true)
@@ -55,12 +49,9 @@ RSpec.describe Gitlab::Database::SchemaVersionFiles do
     end
   end
 
-  describe '.load_all' do
-    let(:connection) { double('connection') }
-
+  describe '#load_all' do
     before do
-      allow(described_class).to receive(:connection).and_return(connection)
-      allow(described_class).to receive(:find_version_filenames).and_return(filenames)
+      allow(migrations).to receive(:version_filenames).and_return(filenames)
     end
 
     context 'when there are no version files' do
@@ -70,7 +61,7 @@ RSpec.describe Gitlab::Database::SchemaVersionFiles do
         expect(connection).not_to receive(:quote_string)
         expect(connection).not_to receive(:execute)
 
-        described_class.load_all
+        migrations.load_all
       end
     end
 
@@ -88,7 +79,7 @@ RSpec.describe Gitlab::Database::SchemaVersionFiles do
           ON CONFLICT DO NOTHING
         SQL
 
-        described_class.load_all
+        migrations.load_all
       end
     end
   end

spec/tasks/gitlab/db_rake_spec.rb

@@ -124,14 +124,19 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
   describe 'clean_structure_sql' do
     let_it_be(:clean_rake_task) { 'gitlab:db:clean_structure_sql' }
     let_it_be(:test_task_name) { 'gitlab:db:_test_multiple_structure_cleans' }
-    let_it_be(:structure_file) { 'db/structure.sql' }
     let_it_be(:input) { 'this is structure data' }
 
     let(:output) { StringIO.new }
 
     before do
+      structure_files = %w[db/structure.sql db/ci_structure.sql]
+
+      allow(File).to receive(:open).and_call_original
+
+      structure_files.each do |structure_file|
         stub_file_read(structure_file, content: input)
-      allow(File).to receive(:open).with(structure_file, any_args).and_yield(output)
+        allow(File).to receive(:open).with(Rails.root.join(structure_file).to_s, any_args).and_yield(output)
+      end
     end
 
     after do
@@ -139,8 +144,10 @@ RSpec.describe 'gitlab:db namespace rake task', :silence_stdout do
     end
 
     it 'can be executed multiple times within another rake task' do
-      expect_multiple_executions_of_task(test_task_name, clean_rake_task) do
-        expect_next_instance_of(Gitlab::Database::SchemaCleaner) do |cleaner|
+      expect_multiple_executions_of_task(test_task_name, clean_rake_task, count: 2) do
+        database_count = ActiveRecord::Base.configurations.configs_for(env_name: Rails.env).size
+
+        expect_next_instances_of(Gitlab::Database::SchemaCleaner, database_count) do |cleaner|
           expect(cleaner).to receive(:clean).with(output)
         end
       end