From 2ab3031bd35213802e508fef6eebceaaf40cee9b Mon Sep 17 00:00:00 2001
From: Michael Kozono <mkozono@gmail.com>
Date: Wed, 8 Nov 2017 15:05:08 -0800
Subject: Refactor, no change in behavior

---
 .../populate_untracked_uploads.rb                  | 10 +++----
 .../prepare_unhashed_uploads.rb                    | 35 +++++++++++++---------
 2 files changed, 25 insertions(+), 20 deletions(-)

(limited to 'lib')

diff --git a/lib/gitlab/background_migration/populate_untracked_uploads.rb b/lib/gitlab/background_migration/populate_untracked_uploads.rb
index 42ad28400f4..e63220c8001 100644
--- a/lib/gitlab/background_migration/populate_untracked_uploads.rb
+++ b/lib/gitlab/background_migration/populate_untracked_uploads.rb
@@ -55,9 +55,7 @@ module Gitlab
         def ensure_tracked!
           return if persisted? && tracked?
 
-          unless in_uploads?
-            add_to_uploads
-          end
+          add_to_uploads unless in_uploads?
 
           mark_as_tracked
         end
@@ -82,8 +80,7 @@ module Gitlab
         end
 
         def mark_as_tracked
-          self.tracked = true
-          self.save!
+          update!(tracked: true)
         end
 
         def upload_path
@@ -121,7 +118,8 @@ module Gitlab
 
         # Not including a leading slash
         def path_relative_to_upload_dir
-          @path_relative_to_upload_dir ||= path.sub(%r{\A#{Gitlab::BackgroundMigration::PrepareUnhashedUploads::UPLOAD_DIR}/}, '')
+          base = %r{\A#{Regexp.escape(Gitlab::BackgroundMigration::PrepareUnhashedUploads::UPLOAD_DIR)}/}
+          @path_relative_to_upload_dir ||= path.sub(base, '')
         end
 
         # Not including a leading slash
diff --git a/lib/gitlab/background_migration/prepare_unhashed_uploads.rb b/lib/gitlab/background_migration/prepare_unhashed_uploads.rb
index ce488542df9..8033f994959 100644
--- a/lib/gitlab/background_migration/prepare_unhashed_uploads.rb
+++ b/lib/gitlab/background_migration/prepare_unhashed_uploads.rb
@@ -35,29 +35,36 @@ module Gitlab
       def store_unhashed_upload_file_paths
         return unless Dir.exist?(UPLOAD_DIR)
 
-        file_paths = []
-        each_file_path(UPLOAD_DIR) do |file_path|
-          file_paths << file_path
-
-          if file_paths.size >= FILE_PATH_BATCH_SIZE
-            insert_file_paths(file_paths)
-            file_paths = []
-          end
+        each_file_batch(UPLOAD_DIR, FILE_PATH_BATCH_SIZE) do |file_paths|
+          insert_file_paths(file_paths)
         end
-
-        insert_file_paths(file_paths) if file_paths.any?
       end
 
-      def each_file_path(search_dir, &block)
+      def each_file_batch(search_dir, batch_size, &block)
         cmd = build_find_command(search_dir)
+
         Open3.popen2(*cmd) do |stdin, stdout, status_thread|
-          stdout.each_line("\0") do |line|
-            yield(line.chomp("\0"))
-          end
+          yield_paths_in_batches(stdout, batch_size, &block)
+
           raise "Find command failed" unless status_thread.value.success?
         end
       end
 
+      def yield_paths_in_batches(stdout, batch_size, &block)
+        paths = []
+
+        stdout.each_line("\0") do |line|
+          paths << line.chomp("\0")
+
+          if paths.size >= batch_size
+            yield(paths)
+            paths = []
+          end
+        end
+
+        yield(paths)
+      end
+
       def build_find_command(search_dir)
         hashed_path = "#{UPLOAD_DIR}/@hashed/*"
         tmp_path = "#{UPLOAD_DIR}/tmp/*"
-- 
cgit v1.2.1