1 files changed, 175 insertions, 73 deletions

diff --git a/spec/tasks/gitlab/click_house/migration_rake_spec.rb b/spec/tasks/gitlab/click_house/migration_rake_spec.rb
index 75a1c1a1856..37ae9d7f090 100644
--- a/spec/tasks/gitlab/click_house/migration_rake_spec.rb
+++ b/spec/tasks/gitlab/click_house/migration_rake_spec.rb
@@ -8,125 +8,227 @@ RSpec.describe 'gitlab:clickhouse', click_house: :without_migrations, feature_ca
   # We don't need to delete data since we don't modify Postgres data
   self.use_transactional_tests = false
 
-  let(:migrations_base_dir) { 'click_house/migrations' }
-  let(:migrations_dirname) { '' }
-  let(:migrations_dir) { expand_fixture_path("#{migrations_base_dir}/#{migrations_dirname}") }
-  let(:verbose) { nil }
-
   before(:all) do
     Rake.application.rake_require 'tasks/gitlab/click_house/migration'
   end
 
-  before do
-    stub_env('VERBOSE', verbose) if verbose
-  end
+  it 'migrates and rolls back the database' do
+    expect { run_rake_task('gitlab:clickhouse:migrate:main') }.to change { active_schema_migrations_count }.from(0)
+      .and output.to_stdout
 
-  describe 'migrate' do
-    subject(:migration) { run_rake_task('gitlab:clickhouse:migrate') }
+    expect { run_rake_task('gitlab:clickhouse:rollback:main') }.to change { active_schema_migrations_count }.by(-1)
+      .and output.to_stdout
 
-    let(:target_version) { nil }
+    stub_env('VERSION', 0)
+    expect { run_rake_task('gitlab:clickhouse:rollback:main') }.to change { active_schema_migrations_count }.to(0)
+      .and output.to_stdout
+  end
 
-    around do |example|
-      ClickHouse::MigrationSupport::Migrator.migrations_paths = [migrations_dir]
+  context 'when clickhouse database is not configured' do
+    before do
+      allow(::ClickHouse::Client).to receive(:configuration).and_return(::ClickHouse::Client::Configuration.new)
+    end
 
-      example.run
+    it 'raises an error' do
+      expect { run_rake_task('gitlab:clickhouse:migrate:main') }.to raise_error(ClickHouse::Client::ConfigurationError)
+    end
 
-      clear_consts(expand_fixture_path(migrations_base_dir))
+    it 'prints the error message and exits successfully if skip_unless_configured is passed' do
+      expect do
+        run_rake_task('gitlab:clickhouse:migrate:main', true)
+      end.to output(/The 'main' ClickHouse database is not configured, skipping migrations/).to_stdout
     end
+  end
 
-    before do
-      stub_env('VERSION', target_version) if target_version
+  describe 'gitlab:clickhouse:migrate' do
+    it 'delegates to gitlab:clickhouse:migrate:main' do
+      task = Rake::Task['gitlab:clickhouse:migrate:main']
+      task.reenable # re-enabling task in case other tests already run it
+      expect(task).to receive(:invoke).with("true").and_call_original
+
+      expect do
+        run_rake_task('gitlab:clickhouse:migrate', true)
+      end.to change { active_schema_migrations_count }.from(0).and output.to_stdout
     end
+  end
+
+  context 'with migration fixtures', :silence_stdout do
+    let(:migrations_base_dir) { 'click_house/migrations' }
+    let(:migrations_dirname) { 'undefined' }
+    let(:migrations_dir) { expand_fixture_path("#{migrations_base_dir}/#{migrations_dirname}") }
+
+    describe 'migrate:main' do
+      subject(:migration) { run_rake_task('gitlab:clickhouse:migrate:main') }
+
+      let(:verbose) { nil }
+      let(:target_version) { nil }
+      let(:step) { nil }
 
-    describe 'when creating a table' do
-      let(:migrations_dirname) { 'plain_table_creation' }
+      before do
+        allow(ClickHouse::MigrationSupport::Migrator).to receive(:migrations_paths).with(:main)
+          .and_return(migrations_dir)
 
-      it 'creates a table' do
-        expect { migration }.to change { active_schema_migrations_count }.from(0).to(1)
-          .and output.to_stdout
+        stub_env('VERBOSE', verbose) if verbose
+        stub_env('VERSION', target_version) if target_version
+        stub_env('STEP', step.to_s) if step
+      end
 
-        expect(describe_table('some')).to match({
-          id: a_hash_including(type: 'UInt64'),
-          date: a_hash_including(type: 'Date')
-        })
+      after do
+        unload_click_house_migration_classes(expand_fixture_path(migrations_dir))
       end
 
-      context 'when VERBOSE is false' do
-        let(:verbose) { 'false' }
+      describe 'when creating a table' do
+        let(:migrations_dirname) { 'plain_table_creation' }
 
-        it 'does not write to stdout' do
-          expect { migration }.not_to output.to_stdout
+        it 'creates a table' do
+          expect { migration }.to change { active_schema_migrations_count }.from(0).to(1)
+            .and output.to_stdout
 
           expect(describe_table('some')).to match({
             id: a_hash_including(type: 'UInt64'),
             date: a_hash_including(type: 'Date')
           })
         end
+
+        context 'when VERBOSE is false' do
+          let(:verbose) { 'false' }
+
+          it 'does not write to stdout' do
+            expect { migration }.not_to output.to_stdout
+
+            expect(describe_table('some')).to match({
+              id: a_hash_including(type: 'UInt64'),
+              date: a_hash_including(type: 'Date')
+            })
+          end
+        end
       end
-    end
 
-    describe 'when dropping a table' do
-      let(:migrations_dirname) { 'drop_table' }
-      let(:target_version) { 2 }
+      describe 'when dropping a table' do
+        let(:migrations_dirname) { 'drop_table' }
+
+        context 'with VERSION set' do
+          let(:target_version) { 2 }
+
+          it 'drops table' do
+            stub_env('VERSION', 1)
+            run_rake_task('gitlab:clickhouse:migrate:main')
+
+            expect(table_names).to include('some')
+
+            stub_env('VERSION', target_version)
+            migration
+            expect(table_names).not_to include('some')
+          end
+
+          context 'with STEP also set' do
+            let(:step) { 1 }
+
+            it 'ignores STEP and executes both migrations' do
+              migration
+
+              expect(table_names).not_to include('some')
+            end
+          end
+        end
+
+        context 'with STEP set to 1' do
+          let(:step) { 1 }
 
-      it 'drops table' do
-        stub_env('VERSION', 1)
-        run_rake_task('gitlab:clickhouse:migrate')
+          it 'executes only first step and creates table' do
+            migration
+
+            expect(table_names).to include('some')
+          end
+        end
+
+        context 'with STEP set to 0' do
+          let(:step) { 0 }
+
+          it 'executes only first step and creates table' do
+            expect { migration }.to raise_error ArgumentError, 'STEP should be a positive number'
+          end
+        end
 
-        expect(table_names).to include('some')
+        context 'with STEP set to not-a-number' do
+          let(:step) { 'NaN' }
 
-        stub_env('VERSION', target_version)
-        migration
-        expect(table_names).not_to include('some')
+          it 'raises an error' do
+            expect { migration }.to raise_error ArgumentError, 'invalid value for Integer(): "NaN"'
+          end
+        end
+
+        context 'with STEP set to empty string' do
+          let(:step) { '' }
+
+          it 'raises an error' do
+            expect { migration }.to raise_error ArgumentError, 'invalid value for Integer(): ""'
+          end
+        end
       end
-    end
 
-    describe 'with VERSION is invalid' do
-      let(:migrations_dirname) { 'plain_table_creation' }
-      let(:target_version) { 'invalid' }
+      context 'with VERSION is invalid' do
+        let(:migrations_dirname) { 'plain_table_creation' }
+        let(:target_version) { 'invalid' }
 
-      it { expect { migration }.to raise_error RuntimeError, 'Invalid format of target version: `VERSION=invalid`' }
+        it { expect { migration }.to raise_error RuntimeError, 'Invalid format of target version: `VERSION=invalid`' }
+      end
     end
-  end
 
-  describe 'rollback' do
-    subject(:migration) { run_rake_task('gitlab:clickhouse:rollback') }
+    describe 'rollback:main' do
+      subject(:migration) { run_rake_task('gitlab:clickhouse:rollback:main') }
 
-    let(:schema_migration) { ClickHouse::MigrationSupport::SchemaMigration }
+      let(:target_version) { nil }
+      let(:rollback_step) { nil }
+      let(:migrations_dirname) { 'table_creation_with_down_method' }
 
-    around do |example|
-      ClickHouse::MigrationSupport::Migrator.migrations_paths = [migrations_dir]
-      migrate(nil, ClickHouse::MigrationSupport::MigrationContext.new(migrations_dir, schema_migration))
+      before do
+        allow(ClickHouse::MigrationSupport::Migrator).to receive(:migrations_paths).with(:main)
+          .and_return(migrations_dir)
 
-      example.run
+        run_rake_task('gitlab:clickhouse:migrate:main')
 
-      clear_consts(expand_fixture_path(migrations_base_dir))
-    end
+        stub_env('VERSION', target_version) if target_version
+        stub_env('STEP', rollback_step.to_s) if rollback_step
+      end
 
-    context 'when migrating back all the way to 0' do
-      let(:target_version) { 0 }
+      after do
+        unload_click_house_migration_classes(expand_fixture_path(migrations_dir))
+      end
 
-      context 'when down method is present' do
-        let(:migrations_dirname) { 'table_creation_with_down_method' }
+      context 'with VERSION set' do
+        context 'when migrating back all the way to 0' do
+          let(:target_version) { 0 }
 
-        it 'removes migration' do
-          expect(table_names).to include('some')
+          it 'rolls back all migrations' do
+            expect(table_names).to include('some', 'another')
+
+            migration
+            expect(table_names).not_to include('some', 'another')
+          end
+
+          context 'with STEP also set' do
+            let(:rollback_step) { 1 }
 
-          migration
-          expect(table_names).not_to include('some')
+            it 'ignores STEP and rolls back all migrations' do
+              expect(table_names).to include('some', 'another')
+
+              migration
+              expect(table_names).not_to include('some', 'another')
+            end
+          end
         end
       end
-    end
-  end
 
-  %w[gitlab:clickhouse:migrate].each do |task|
-    context "when running #{task}" do
-      it "does run gitlab:clickhouse:prepare_schema_migration_table before" do
-        expect(Rake::Task['gitlab:clickhouse:prepare_schema_migration_table']).to receive(:execute).and_return(true)
-        expect(Rake::Task[task]).to receive(:execute).and_return(true)
+      context 'with STEP set to 1' do
+        let(:rollback_step) { 1 }
 
-        Rake::Task['gitlab:clickhouse:prepare_schema_migration_table'].reenable
-        run_rake_task(task)
+        it 'executes only first step and drops "another" table' do
+          run_rake_task('gitlab:clickhouse:rollback:main')
+
+          expect(table_names).to include('some')
+          expect(table_names).not_to include('another')
+        end
       end
     end
   end