Drop temporary tracking table when finished

6 files changed, 131 insertions, 92 deletions

diff --git a/db/post_migrate/20171103140253_track_untracked_uploads.rb b/db/post_migrate/20171103140253_track_untracked_uploads.rb
index 09ff21b103f..7a34abc85ee 100644
--- a/db/post_migrate/20171103140253_track_untracked_uploads.rb
+++ b/db/post_migrate/20171103140253_track_untracked_uploads.rb
@@ -10,6 +10,18 @@ class TrackUntrackedUploads < ActiveRecord::Migration
   MIGRATION = 'PrepareUntrackedUploads'
 
   def up
+    ensure_temporary_tracking_table_exists
+
+    BackgroundMigrationWorker.perform_async(MIGRATION)
+  end
+
+  def down
+    if table_exists?(:untracked_files_for_uploads)
+      drop_table :untracked_files_for_uploads
+    end
+  end
+
+  def ensure_temporary_tracking_table_exists
     unless table_exists?(:untracked_files_for_uploads)
       create_table :untracked_files_for_uploads do |t|
         t.string :path, limit: 600, null: false
@@ -25,13 +37,5 @@ class TrackUntrackedUploads < ActiveRecord::Migration
     unless index_exists?(:untracked_files_for_uploads, :tracked)
       add_index :untracked_files_for_uploads, :tracked
     end
-
-    BackgroundMigrationWorker.perform_async(MIGRATION)
-  end
-
-  def down
-    if table_exists?(:untracked_files_for_uploads)
-      drop_table :untracked_files_for_uploads
-    end
   end
 end
diff --git a/lib/gitlab/background_migration/populate_untracked_uploads.rb b/lib/gitlab/background_migration/populate_untracked_uploads.rb
index 8529e8d1d0b..f28892174bb 100644
--- a/lib/gitlab/background_migration/populate_untracked_uploads.rb
+++ b/lib/gitlab/background_migration/populate_untracked_uploads.rb
@@ -208,6 +208,8 @@ module Gitlab
             # to do.
           end
         end
+
+        drop_temp_table_if_finished
       end
 
       private
@@ -215,6 +217,10 @@ module Gitlab
       def migrate?
         UntrackedFile.table_exists? && Upload.table_exists?
       end
+
+      def drop_temp_table_if_finished
+        UntrackedFile.connection.drop_table(:untracked_files_for_uploads) if UntrackedFile.untracked.empty?
+      end
     end
   end
 end
diff --git a/spec/lib/gitlab/background_migration/populate_untracked_uploads_spec.rb b/spec/lib/gitlab/background_migration/populate_untracked_uploads_spec.rb
index c794a2f152b..52f57408bfa 100644
--- a/spec/lib/gitlab/background_migration/populate_untracked_uploads_spec.rb
+++ b/spec/lib/gitlab/background_migration/populate_untracked_uploads_spec.rb
@@ -1,11 +1,19 @@
 require 'spec_helper'
+require Rails.root.join('db', 'post_migrate', '20171103140253_track_untracked_uploads')
 
-describe Gitlab::BackgroundMigration::PopulateUntrackedUploads, :migration, :sidekiq, schema: 20171103140253 do
+describe Gitlab::BackgroundMigration::PopulateUntrackedUploads, :migration, :sidekiq, :temp_table_may_drop, schema: 20171103140253 do
   include TrackUntrackedUploadsHelpers
 
+  subject { described_class.new }
+
   let!(:untracked_files_for_uploads) { table(:untracked_files_for_uploads) }
   let!(:uploads) { table(:uploads) }
 
+  before do
+    # Prevent the TrackUntrackedUploads migration from running PrepareUntrackedUploads job
+    allow(BackgroundMigrationWorker).to receive(:perform_async).and_return(true)
+  end
+
   context 'with untracked files and tracked files in untracked_files_for_uploads' do
     let!(:appearance) { create(:appearance, logo: uploaded_file, header_logo: uploaded_file) }
     let!(:user1) { create(:user, :with_avatar) }
@@ -35,7 +43,7 @@ describe Gitlab::BackgroundMigration::PopulateUntrackedUploads, :migration, :sid
 
     it 'adds untracked files to the uploads table' do
       expect do
-        described_class.new.perform(1, 1000)
+        subject.perform(1, 1000)
       end.to change { uploads.count }.from(4).to(8)
 
       expect(user2.uploads.count).to eq(1)
@@ -44,13 +52,15 @@ describe Gitlab::BackgroundMigration::PopulateUntrackedUploads, :migration, :sid
     end
 
     it 'sets all added or confirmed tracked files to tracked' do
+      expect(subject).to receive(:drop_temp_table_if_finished) # Don't drop the table so we can look at it
+
       expect do
-        described_class.new.perform(1, 1000)
+        subject.perform(1, 1000)
       end.to change { untracked_files_for_uploads.where(tracked: true).count }.from(0).to(8)
     end
 
     it 'does not create duplicate uploads of already tracked files' do
-      described_class.new.perform(1, 1000)
+      subject.perform(1, 1000)
 
       expect(user1.uploads.count).to eq(1)
       expect(project1.uploads.count).to eq(2)
@@ -62,7 +72,7 @@ describe Gitlab::BackgroundMigration::PopulateUntrackedUploads, :migration, :sid
       end_id = untracked_files_for_uploads.all.to_a[3].id
 
       expect do
-        described_class.new.perform(start_id, end_id)
+        subject.perform(start_id, end_id)
       end.to change { uploads.count }.from(4).to(6)
 
       expect(user1.uploads.count).to eq(1)
@@ -80,7 +90,7 @@ describe Gitlab::BackgroundMigration::PopulateUntrackedUploads, :migration, :sid
       end_id = untracked_files_for_uploads.all.to_a[7].id
 
       expect do
-        described_class.new.perform(start_id, end_id)
+        subject.perform(start_id, end_id)
       end.to change { uploads.count }.from(4).to(6)
 
       expect(user1.uploads.count).to eq(1)
@@ -92,12 +102,24 @@ describe Gitlab::BackgroundMigration::PopulateUntrackedUploads, :migration, :sid
       # Only 4 have been either confirmed or added to uploads
       expect(untracked_files_for_uploads.where(tracked: true).count).to eq(4)
     end
+
+    it 'does not drop the temporary tracking table after processing the batch, if there are still untracked rows' do
+      subject.perform(1, untracked_files_for_uploads.last.id - 1)
+
+      expect(table_exists?(:untracked_files_for_uploads)).to be_truthy
+    end
+
+    it 'drops the temporary tracking table after processing the batch, if there are no untracked rows left' do
+      subject.perform(1, untracked_files_for_uploads.last.id)
+
+      expect(table_exists?(:untracked_files_for_uploads)).to be_falsey
+    end
   end
 
   context 'with no untracked files' do
     it 'does not add to the uploads table (and does not raise error)' do
       expect do
-        described_class.new.perform(1, 1000)
+        subject.perform(1, 1000)
       end.not_to change { uploads.count }.from(0)
     end
   end
diff --git a/spec/lib/gitlab/background_migration/prepare_untracked_uploads_spec.rb b/spec/lib/gitlab/background_migration/prepare_untracked_uploads_spec.rb
index d61135912dd..8fd20fd0bb3 100644
--- a/spec/lib/gitlab/background_migration/prepare_untracked_uploads_spec.rb
+++ b/spec/lib/gitlab/background_migration/prepare_untracked_uploads_spec.rb
@@ -17,6 +17,13 @@ describe Gitlab::BackgroundMigration::PrepareUntrackedUploads, :migration, :side
     end
   end
 
+  around do |example|
+    # Especially important so the follow-up migration does not get run
+    Sidekiq::Testing.fake! do
+      example.run
+    end
+  end
+
   context 'when files were uploaded before and after hashed storage was enabled' do
     let!(:appearance) { create(:appearance, logo: uploaded_file, header_logo: uploaded_file) }
     let!(:user) { create(:user, :with_avatar) }
@@ -34,38 +41,30 @@ describe Gitlab::BackgroundMigration::PrepareUntrackedUploads, :migration, :side
     end
 
     it 'adds unhashed files to the untracked_files_for_uploads table' do
-      Sidekiq::Testing.fake! do
-        expect do
-          described_class.new.perform
-        end.to change { untracked_files_for_uploads.count }.from(0).to(5)
-      end
+      expect do
+        described_class.new.perform
+      end.to change { untracked_files_for_uploads.count }.from(0).to(5)
     end
 
     it 'adds files with paths relative to CarrierWave.root' do
-      Sidekiq::Testing.fake! do
-        described_class.new.perform
-        untracked_files_for_uploads.all.each do |file|
-          expect(file.path.start_with?('uploads/')).to be_truthy
-        end
+      described_class.new.perform
+      untracked_files_for_uploads.all.each do |file|
+        expect(file.path.start_with?('uploads/')).to be_truthy
       end
     end
 
     it 'does not add hashed files to the untracked_files_for_uploads table' do
-      Sidekiq::Testing.fake! do
-        described_class.new.perform
+      described_class.new.perform
 
-        hashed_file_path = project2.uploads.where(uploader: 'FileUploader').first.path
-        expect(untracked_files_for_uploads.where("path like '%#{hashed_file_path}%'").exists?).to be_falsey
-      end
+      hashed_file_path = project2.uploads.where(uploader: 'FileUploader').first.path
+      expect(untracked_files_for_uploads.where("path like '%#{hashed_file_path}%'").exists?).to be_falsey
     end
 
     it 'correctly schedules the follow-up background migration jobs' do
-      Sidekiq::Testing.fake! do
-        described_class.new.perform
+      described_class.new.perform
 
-        expect(described_class::FOLLOW_UP_MIGRATION).to be_scheduled_migration(1, 5)
-        expect(BackgroundMigrationWorker.jobs.size).to eq(1)
-      end
+      expect(described_class::FOLLOW_UP_MIGRATION).to be_scheduled_migration(1, 5)
+      expect(BackgroundMigrationWorker.jobs.size).to eq(1)
     end
 
     # E.g. from a previous failed run of this background migration
@@ -75,11 +74,9 @@ describe Gitlab::BackgroundMigration::PrepareUntrackedUploads, :migration, :side
       end
 
       it 'does not error or produce duplicates of existing data' do
-        Sidekiq::Testing.fake! do
-          expect do
-            described_class.new.perform
-          end.not_to change { untracked_files_for_uploads.count }.from(5)
-        end
+        expect do
+          described_class.new.perform
+        end.not_to change { untracked_files_for_uploads.count }.from(5)
       end
     end
 
@@ -97,11 +94,9 @@ describe Gitlab::BackgroundMigration::PrepareUntrackedUploads, :migration, :side
       end
 
       it 'does not add files from /uploads/tmp' do
-        Sidekiq::Testing.fake! do
-          expect do
-            described_class.new.perform
-          end.to change { untracked_files_for_uploads.count }.from(0).to(5)
-        end
+        expect do
+          described_class.new.perform
+        end.to change { untracked_files_for_uploads.count }.from(0).to(5)
       end
     end
   end
@@ -110,11 +105,9 @@ describe Gitlab::BackgroundMigration::PrepareUntrackedUploads, :migration, :side
   # may not have an upload directory because they have no uploads.
   context 'when no files were ever uploaded' do
     it 'does not add to the untracked_files_for_uploads table (and does not raise error)' do
-      Sidekiq::Testing.fake! do
-        expect do
-          described_class.new.perform
-        end.not_to change { untracked_files_for_uploads.count }.from(0)
-      end
+      expect do
+        described_class.new.perform
+      end.not_to change { untracked_files_for_uploads.count }.from(0)
     end
   end
 end
diff --git a/spec/migrations/track_untracked_uploads_spec.rb b/spec/migrations/track_untracked_uploads_spec.rb
index a6e880279b6..a17251ba052 100644
--- a/spec/migrations/track_untracked_uploads_spec.rb
+++ b/spec/migrations/track_untracked_uploads_spec.rb
@@ -1,7 +1,7 @@
 require 'spec_helper'
 require Rails.root.join('db', 'post_migrate', '20171103140253_track_untracked_uploads')
 
-describe TrackUntrackedUploads, :migration, :sidekiq do
+describe TrackUntrackedUploads, :migration, :sidekiq, :temp_table_may_drop do
   include TrackUntrackedUploadsHelpers
 
   let(:untracked_files_for_uploads) { table(:untracked_files_for_uploads) }
@@ -18,34 +18,41 @@ describe TrackUntrackedUploads, :migration, :sidekiq do
     end
   end
 
-  it 'correctly schedules the follow-up background migration' do
-    Sidekiq::Testing.fake! do
+  context 'before running the background migration' do
+    around do |example|
+      # Especially important so the follow-up migration does not get run
+      Sidekiq::Testing.fake! do
+        example.run
+      end
+    end
+
+    it 'correctly schedules the follow-up background migration' do
       migrate!
 
       expect(described_class::MIGRATION).to be_scheduled_migration
       expect(BackgroundMigrationWorker.jobs.size).to eq(1)
     end
-  end
 
-  it 'ensures the untracked_files_for_uploads table exists' do
-    expect do
-      migrate!
-    end.to change { table_exists?(:untracked_files_for_uploads) }.from(false).to(true)
-  end
+    it 'ensures the untracked_files_for_uploads table exists' do
+      expect do
+        migrate!
+      end.to change { table_exists?(:untracked_files_for_uploads) }.from(false).to(true)
+    end
 
-  it 'has a path field long enough for really long paths' do
-    migrate!
+    it 'has a path field long enough for really long paths' do
+      migrate!
 
-    component = 'a'*255
+      component = 'a'*255
 
-    long_path = [
-      'uploads',
-      component, # project.full_path
-      component  # filename
-    ].flatten.join('/')
+      long_path = [
+        'uploads',
+        component, # project.full_path
+        component  # filename
+      ].flatten.join('/')
 
-    record = untracked_files_for_uploads.create!(path: long_path)
-    expect(record.reload.path.size).to eq(519)
+      record = untracked_files_for_uploads.create!(path: long_path)
+      expect(record.reload.path.size).to eq(519)
+    end
   end
 
   context 'with tracked and untracked uploads' do
@@ -77,36 +84,29 @@ describe TrackUntrackedUploads, :migration, :sidekiq do
     end
 
     it 'tracks untracked uploads' do
-      Sidekiq::Testing.inline! do
-        expect do
-          migrate!
-        end.to change { uploads.count }.from(4).to(8)
-
-        expect(appearance.reload.uploads.where("path like '%/header_logo/%'").first.attributes).to include(@appearance_header_logo_attributes)
-        expect(user2.reload.uploads.first.attributes).to include(@user2_avatar_attributes)
-        expect(project2.reload.uploads.first.attributes).to include(@project2_avatar_attributes)
-        expect(project2.uploads.last.attributes).to include(@project2_markdown_attributes)
-      end
+      expect do
+        migrate!
+      end.to change { uploads.count }.from(4).to(8)
+
+      expect(appearance.reload.uploads.where("path like '%/header_logo/%'").first.attributes).to include(@appearance_header_logo_attributes)
+      expect(user2.reload.uploads.first.attributes).to include(@user2_avatar_attributes)
+      expect(project2.reload.uploads.first.attributes).to include(@project2_avatar_attributes)
+      expect(project2.uploads.last.attributes).to include(@project2_markdown_attributes)
     end
 
     it 'ignores already-tracked uploads' do
-      Sidekiq::Testing.inline! do
-        migrate!
+      migrate!
 
-        expect(appearance.reload.uploads.where("path like '%/logo/%'").first.attributes).to include(@appearance_logo_attributes)
-        expect(user1.reload.uploads.first.attributes).to include(@user1_avatar_attributes)
-        expect(project1.reload.uploads.first.attributes).to include(@project1_avatar_attributes)
-        expect(project1.uploads.last.attributes).to include(@project1_markdown_attributes)
-      end
+      expect(appearance.reload.uploads.where("path like '%/logo/%'").first.attributes).to include(@appearance_logo_attributes)
+      expect(user1.reload.uploads.first.attributes).to include(@user1_avatar_attributes)
+      expect(project1.reload.uploads.first.attributes).to include(@project1_avatar_attributes)
+      expect(project1.uploads.last.attributes).to include(@project1_markdown_attributes)
     end
 
-    it 'all untracked_files_for_uploads records are marked as tracked' do
-      Sidekiq::Testing.inline! do
-        migrate!
+    it 'the temporary table untracked_files_for_uploads no longer exists' do
+      migrate!
 
-        expect(untracked_files_for_uploads.count).to eq(8)
-        expect(untracked_files_for_uploads.count).to eq(untracked_files_for_uploads.where(tracked: true).count)
-      end
+      expect(table_exists?(:untracked_files_for_uploads)).to be_falsey
     end
   end
 end
diff --git a/spec/support/track_untracked_uploads_helpers.rb b/spec/support/track_untracked_uploads_helpers.rb
index 5b832929602..bb700bc53f1 100644
--- a/spec/support/track_untracked_uploads_helpers.rb
+++ b/spec/support/track_untracked_uploads_helpers.rb
@@ -3,4 +3,18 @@ module TrackUntrackedUploadsHelpers
     fixture_path = Rails.root.join('spec', 'fixtures', 'rails_sample.jpg')
     fixture_file_upload(fixture_path)
   end
+
+  def recreate_temp_table_if_dropped
+    TrackUntrackedUploads.new.ensure_temporary_tracking_table_exists
+  end
+
+  RSpec.configure do |config|
+    config.after(:each, :temp_table_may_drop) do
+      recreate_temp_table_if_dropped
+    end
+
+    config.after(:context, :temp_table_may_drop) do
+      recreate_temp_table_if_dropped
+    end
+  end
 end