5 files changed, 271 insertions, 97 deletions

diff --git a/spec/tasks/gitlab/db/truncate_legacy_tables_rake_spec.rb b/spec/tasks/gitlab/db/truncate_legacy_tables_rake_spec.rb
new file mode 100644
index 00000000000..f9ebb985255
--- /dev/null
+++ b/spec/tasks/gitlab/db/truncate_legacy_tables_rake_spec.rb
@@ -0,0 +1,157 @@
+# frozen_string_literal: true
+
+require 'rake_helper'
+
+RSpec.describe 'gitlab:db:truncate_legacy_tables', :silence_stdout, :reestablished_active_record_base,
+               :suppress_gitlab_schemas_validate_connection do
+  let(:main_connection) { ApplicationRecord.connection }
+  let(:ci_connection) { Ci::ApplicationRecord.connection }
+  let(:test_gitlab_main_table) { '_test_gitlab_main_table' }
+  let(:test_gitlab_ci_table) { '_test_gitlab_ci_table' }
+
+  before :all do
+    Rake.application.rake_require 'active_record/railties/databases'
+    Rake.application.rake_require 'tasks/seed_fu'
+    Rake.application.rake_require 'tasks/gitlab/db/validate_config'
+    Rake.application.rake_require 'tasks/gitlab/db/truncate_legacy_tables'
+
+    # empty task as env is already loaded
+    Rake::Task.define_task :environment
+  end
+
+  before do
+    skip_if_multiple_databases_not_setup
+
+    # Filling the table on both databases main and ci
+    Gitlab::Database.database_base_models.each_value do |base_model|
+      base_model.connection.execute(<<~SQL)
+         CREATE TABLE #{test_gitlab_main_table} (id integer NOT NULL);
+         INSERT INTO #{test_gitlab_main_table} VALUES(generate_series(1, 50));
+      SQL
+      base_model.connection.execute(<<~SQL)
+         CREATE TABLE #{test_gitlab_ci_table} (id integer NOT NULL);
+         INSERT INTO #{test_gitlab_ci_table} VALUES(generate_series(1, 50));
+      SQL
+    end
+
+    allow(Gitlab::Database::GitlabSchema).to receive(:tables_to_schema).and_return(
+      {
+        test_gitlab_main_table => :gitlab_main,
+        test_gitlab_ci_table => :gitlab_ci
+      }
+    )
+  end
+
+  shared_examples 'truncating legacy tables' do
+    before do
+      allow(ENV).to receive(:[]).and_return(nil)
+    end
+
+    context 'when tables are not locked for writes' do
+      it 'raises an error when trying to truncate the tables' do
+        error_message = /is not locked for writes. Run the rake task gitlab:db:lock_writes first/
+        expect { truncate_legacy_tables }.to raise_error(error_message)
+      end
+    end
+
+    context 'when tables are locked for writes' do
+      before do
+        # Locking ci table on the main database
+        Gitlab::Database::LockWritesManager.new(
+          table_name: test_gitlab_ci_table,
+          connection: main_connection,
+          database_name: "main"
+        ).lock_writes
+
+        # Locking main table on the ci database
+        Gitlab::Database::LockWritesManager.new(
+          table_name: test_gitlab_main_table,
+          connection: ci_connection,
+          database_name: "ci"
+        ).lock_writes
+      end
+
+      it 'calls TablesTruncate with the correct parameters and default minimum batch size' do
+        expect(Gitlab::Database::TablesTruncate).to receive(:new).with(
+          database_name: database_name,
+          min_batch_size: 5,
+          logger: anything,
+          dry_run: false,
+          until_table: nil
+        ).and_call_original
+
+        truncate_legacy_tables
+      end
+
+      it 'truncates the legacy table' do
+        expect do
+          truncate_legacy_tables
+        end.to change { connection.select_value("SELECT count(*) from #{legacy_table}") }.from(50).to(0)
+      end
+
+      it 'does not truncate the table that belongs to the connection schema' do
+        expect do
+          truncate_legacy_tables
+        end.not_to change { connection.select_value("SELECT count(*) from #{active_table}") }
+      end
+
+      context 'when running in dry_run mode' do
+        before do
+          allow(ENV).to receive(:[]).with("DRY_RUN").and_return("true")
+        end
+
+        it 'does not truncate any tables' do
+          expect do
+            truncate_legacy_tables
+          end.not_to change { connection.select_value("SELECT count(*) from #{legacy_table}") }
+        end
+
+        it 'prints the truncation sql statement to the output' do
+          expect do
+            truncate_legacy_tables
+          end.to output(/TRUNCATE TABLE #{legacy_table} RESTRICT/).to_stdout
+        end
+      end
+
+      context 'when passing until_table parameter via environment variable' do
+        before do
+          allow(ENV).to receive(:[]).with("UNTIL_TABLE").and_return(legacy_table)
+        end
+
+        it 'sends the table name to TablesTruncate' do
+          expect(Gitlab::Database::TablesTruncate).to receive(:new).with(
+            database_name: database_name,
+            min_batch_size: 5,
+            logger: anything,
+            dry_run: false,
+            until_table: legacy_table
+          ).and_call_original
+
+          truncate_legacy_tables
+        end
+      end
+    end
+  end
+
+  context 'when truncating ci tables on the main database' do
+    subject(:truncate_legacy_tables) { run_rake_task('gitlab:db:truncate_legacy_tables:main') }
+
+    let(:connection) { ApplicationRecord.connection }
+    let(:database_name) { 'main' }
+    let(:active_table) { test_gitlab_main_table }
+    let(:legacy_table) { test_gitlab_ci_table }
+
+    it_behaves_like 'truncating legacy tables'
+  end
+
+  context 'when truncating main tables on the ci database' do
+    subject(:truncate_legacy_tables) { run_rake_task('gitlab:db:truncate_legacy_tables:ci') }
+
+    let(:connection) { Ci::ApplicationRecord.connection }
+    let(:database_name) { 'ci' }
+    let(:active_table) { test_gitlab_ci_table }
+    let(:legacy_table) { test_gitlab_main_table }
+
+    it_behaves_like 'truncating legacy tables'
+  end
+end
diff --git a/spec/tasks/gitlab/db/validate_config_rake_spec.rb b/spec/tasks/gitlab/db/validate_config_rake_spec.rb
index ad15c7f0d1c..1d47c94aa77 100644
--- a/spec/tasks/gitlab/db/validate_config_rake_spec.rb
+++ b/spec/tasks/gitlab/db/validate_config_rake_spec.rb
@@ -216,7 +216,7 @@ RSpec.describe 'gitlab:db:validate_config', :silence_stdout, :suppress_gitlab_sc
       let(:exception) { ActiveRecord::StatementInvalid.new("READONLY") }
 
       before do
-        allow(exception).to receive(:cause).and_return(PG::ReadOnlySqlTransaction.new("cannot execute INSERT in a read-only transaction"))
+        allow(exception).to receive(:cause).and_return(PG::ReadOnlySqlTransaction.new("cannot execute UPSERT in a read-only transaction"))
         allow(ActiveRecord::InternalMetadata).to receive(:upsert).at_least(:once).and_raise(exception)
       end
 
diff --git a/spec/tasks/gitlab/snippets_rake_spec.rb b/spec/tasks/gitlab/snippets_rake_spec.rb
index c55bded1d5a..c50b04b4600 100644
--- a/spec/tasks/gitlab/snippets_rake_spec.rb
+++ b/spec/tasks/gitlab/snippets_rake_spec.rb
@@ -3,7 +3,7 @@
 require 'rake_helper'
 
 RSpec.describe 'gitlab:snippets namespace rake task', :silence_stdout do
-  let_it_be(:user) { create(:user)}
+  let_it_be(:user) { create(:user) }
   let_it_be(:migrated) { create(:personal_snippet, :repository, author: user) }
 
   let(:non_migrated) { create_list(:personal_snippet, 3, author: user) }
diff --git a/spec/tasks/gitlab/uploads/migrate_rake_spec.rb b/spec/tasks/gitlab/uploads/migrate_rake_spec.rb
index e293271ca67..3a368a5011b 100644
--- a/spec/tasks/gitlab/uploads/migrate_rake_spec.rb
+++ b/spec/tasks/gitlab/uploads/migrate_rake_spec.rb
@@ -2,133 +2,93 @@
 
 require 'rake_helper'
 
-RSpec.describe 'gitlab:uploads:migrate and migrate_to_local rake tasks', :silence_stdout do
-  let(:model_class) { nil }
-  let(:uploader_class) { nil }
-  let(:mounted_as) { nil }
-  let(:batch_size) { 3 }
-
+RSpec.describe 'gitlab:uploads:migrate and migrate_to_local rake tasks', :sidekiq_inline, :silence_stdout do
   before do
-    stub_env('MIGRATION_BATCH_SIZE', batch_size.to_s)
-    stub_uploads_object_storage(uploader_class)
+    stub_env('MIGRATION_BATCH_SIZE', 3.to_s)
+    stub_uploads_object_storage(AvatarUploader)
+    stub_uploads_object_storage(FileUploader)
     Rake.application.rake_require 'tasks/gitlab/uploads/migrate'
 
-    allow(ObjectStorage::MigrateUploadsWorker).to receive(:perform_async)
+    create_list(:project, 2, :with_avatar)
+    create_list(:group, 2, :with_avatar)
+    create_list(:project, 2) do |model|
+      FileUploader.new(model).store!(fixture_file_upload('spec/fixtures/doc_sample.txt'))
+    end
   end
 
-  context "for AvatarUploader" do
-    let(:uploader_class) { AvatarUploader }
-    let(:mounted_as) { :avatar }
+  let(:total_uploads_count) { 6 }
 
-    context "for Project" do
-      let(:model_class) { Project }
-      let!(:projects) { create_list(:project, 10, :with_avatar) }
+  it 'migrates all uploads to object storage in batches' do
+    expect(ObjectStorage::MigrateUploadsWorker)
+      .to receive(:perform_async).twice.and_call_original
 
-      it_behaves_like 'enqueue upload migration jobs in batch', batch: 4
-    end
+    run_rake_task('gitlab:uploads:migrate:all')
 
-    context "for Group" do
-      let(:model_class) { Group }
+    expect(Upload.with_files_stored_locally.count).to eq(0)
+    expect(Upload.with_files_stored_remotely.count).to eq(total_uploads_count)
+  end
 
-      before do
-        create_list(:group, 10, :with_avatar)
-      end
+  it 'migrates all uploads to local storage in batches' do
+    run_rake_task('gitlab:uploads:migrate')
+    expect(Upload.with_files_stored_remotely.count).to eq(total_uploads_count)
 
-      it_behaves_like 'enqueue upload migration jobs in batch', batch: 4
-    end
+    expect(ObjectStorage::MigrateUploadsWorker)
+      .to receive(:perform_async).twice.and_call_original
 
-    context "for User" do
-      let(:model_class) { User }
+    run_rake_task('gitlab:uploads:migrate_to_local:all')
 
-      before do
-        create_list(:user, 10, :with_avatar)
-      end
-
-      it_behaves_like 'enqueue upload migration jobs in batch', batch: 4
-    end
+    expect(Upload.with_files_stored_remotely.count).to eq(0)
+    expect(Upload.with_files_stored_locally.count).to eq(total_uploads_count)
   end
 
-  context "for AttachmentUploader" do
-    let(:uploader_class) { AttachmentUploader }
+  shared_examples 'migrate task with filters' do
+    it 'migrates matching uploads to object storage' do
+      run_rake_task('gitlab:uploads:migrate', task_arguments)
 
-    context "for Note" do
-      let(:model_class) { Note }
-      let(:mounted_as) { :attachment }
+      migrated_count = matching_uploads.with_files_stored_remotely.count
 
-      before do
-        create_list(:note, 10, :with_attachment)
-      end
-
-      it_behaves_like 'enqueue upload migration jobs in batch', batch: 4
+      expect(migrated_count).to eq(matching_uploads.count)
+      expect(Upload.with_files_stored_locally.count).to eq(total_uploads_count - migrated_count)
     end
 
-    context "for Appearance" do
-      let(:model_class) { Appearance }
-      let(:mounted_as) { :logo }
+    it 'migrates matching uploads to local storage' do
+      run_rake_task('gitlab:uploads:migrate')
+      expect(Upload.with_files_stored_remotely.count).to eq(total_uploads_count)
+
+      run_rake_task('gitlab:uploads:migrate_to_local', task_arguments)
 
-      before do
-        create(:appearance, :with_logos)
-      end
+      migrated_count = matching_uploads.with_files_stored_locally.count
 
-      %i(logo header_logo).each do |mount|
-        it_behaves_like 'enqueue upload migration jobs in batch', batch: 1 do
-          let(:mounted_as) { mount }
-        end
-      end
+      expect(migrated_count).to eq(matching_uploads.count)
+      expect(Upload.with_files_stored_remotely.count).to eq(total_uploads_count - migrated_count)
     end
   end
 
-  context "for FileUploader" do
-    let(:uploader_class) { FileUploader }
-    let(:model_class) { Project }
+  context 'when uploader_class is given' do
+    let(:task_arguments) { ['FileUploader'] }
+    let(:matching_uploads) { Upload.where(uploader: 'FileUploader') }
 
-    before do
-      create_list(:project, 10) do |model|
-        uploader_class.new(model)
-          .store!(fixture_file_upload('spec/fixtures/doc_sample.txt'))
-      end
-    end
-
-    it_behaves_like 'enqueue upload migration jobs in batch', batch: 4
+    it_behaves_like 'migrate task with filters'
   end
 
-  context "for PersonalFileUploader" do
-    let(:uploader_class) { PersonalFileUploader }
-    let(:model_class) { PersonalSnippet }
-
-    before do
-      create_list(:personal_snippet, 10) do |model|
-        uploader_class.new(model)
-          .store!(fixture_file_upload('spec/fixtures/doc_sample.txt'))
-      end
-    end
+  context 'when model_class is given' do
+    let(:task_arguments) { [nil, 'Project'] }
+    let(:matching_uploads) { Upload.where(model_type: 'Project') }
 
-    it_behaves_like 'enqueue upload migration jobs in batch', batch: 4
+    it_behaves_like 'migrate task with filters'
   end
 
-  context "for NamespaceFileUploader" do
-    let(:uploader_class) { NamespaceFileUploader }
-    let(:model_class) { Snippet }
+  context 'when mounted_as is given' do
+    let(:task_arguments) { [nil, nil, :avatar] }
+    let(:matching_uploads) { Upload.where(mount_point: :avatar) }
 
-    before do
-      create_list(:snippet, 10) do |model|
-        uploader_class.new(model)
-          .store!(fixture_file_upload('spec/fixtures/doc_sample.txt'))
-      end
-    end
-
-    it_behaves_like 'enqueue upload migration jobs in batch', batch: 4
+    it_behaves_like 'migrate task with filters'
   end
 
-  context 'for DesignManagement::DesignV432x230Uploader' do
-    let(:uploader_class) { DesignManagement::DesignV432x230Uploader }
-    let(:model_class) {  DesignManagement::Action }
-    let(:mounted_as) { :image_v432x230 }
-
-    before do
-      create_list(:design_action, 10, :with_image_v432x230)
-    end
+  context 'when multiple filters are given' do
+    let(:task_arguments) { %w[AvatarUploader Project] }
+    let(:matching_uploads) { Upload.where(uploader: 'AvatarUploader', model_type: 'Project') }
 
-    it_behaves_like 'enqueue upload migration jobs in batch', batch: 4
+    it_behaves_like 'migrate task with filters'
   end
 end
diff --git a/spec/tasks/gitlab/usage_data_rake_spec.rb b/spec/tasks/gitlab/usage_data_rake_spec.rb
index 442b884b313..f54d06f406f 100644
--- a/spec/tasks/gitlab/usage_data_rake_spec.rb
+++ b/spec/tasks/gitlab/usage_data_rake_spec.rb
@@ -3,15 +3,23 @@
 require 'rake_helper'
 
 RSpec.describe 'gitlab:usage data take tasks', :silence_stdout do
+  include StubRequests
   include UsageDataHelpers
 
+  let(:metrics_file) { Rails.root.join('tmp', 'test', 'sql_metrics_queries.json') }
+
   before do
     Rake.application.rake_require 'tasks/gitlab/usage_data'
+
     # stub prometheus external http calls https://gitlab.com/gitlab-org/gitlab/-/issues/245277
     stub_prometheus_queries
     stub_database_flavor_check
   end
 
+  after do
+    FileUtils.rm_rf(metrics_file)
+  end
+
   describe 'dump_sql_in_yaml' do
     it 'dumps SQL queries in yaml format' do
       expect { run_rake_task('gitlab:usage_data:dump_sql_in_yaml') }.to output(/.*recorded_at:.*/).to_stdout
@@ -23,4 +31,53 @@ RSpec.describe 'gitlab:usage data take tasks', :silence_stdout do
       expect { run_rake_task('gitlab:usage_data:dump_sql_in_json') }.to output(/.*"recorded_at":.*/).to_stdout
     end
   end
+
+  describe 'dump_non_sql_in_json' do
+    it 'dumps non SQL data in json format' do
+      expect { run_rake_task('gitlab:usage_data:dump_non_sql_in_json') }.to output(/.*"recorded_at":.*/).to_stdout
+    end
+  end
+
+  describe 'generate_sql_metrics_fixture' do
+    it 'generates fixture file correctly' do
+      run_rake_task('gitlab:usage_data:generate_sql_metrics_queries')
+
+      expect(Pathname.new(metrics_file)).to exist
+    end
+  end
+
+  describe 'generate_and_send' do
+    let(:service_ping_payload_url) do
+      File.join(ServicePing::SubmitService::STAGING_BASE_URL, ServicePing::SubmitService::USAGE_DATA_PATH)
+    end
+
+    let(:service_ping_metadata_url) do
+      File.join(ServicePing::SubmitService::STAGING_BASE_URL, ServicePing::SubmitService::METADATA_PATH)
+    end
+
+    let(:payload) { { recorded_at: Time.current } }
+
+    before do
+      allow_next_instance_of(ServicePing::BuildPayload) do |service|
+        allow(service).to receive(:execute).and_return(payload)
+      end
+      stub_response(body: payload.merge(conv_index: { usage_data_id: 123 }))
+      stub_response(body: nil, url: service_ping_metadata_url, status: 201)
+    end
+
+    it 'generates and sends Service Ping payload' do
+      expect { run_rake_task('gitlab:usage_data:generate_and_send') }.to output(/.*201.*/).to_stdout
+    end
+
+    private
+
+    def stub_response(url: service_ping_payload_url, body:, status: 201)
+      stub_full_request(url, method: :post)
+        .to_return(
+          headers: { 'Content-Type' => 'application/json' },
+          body: body.to_json,
+          status: status
+        )
+    end
+  end
 end