1 files changed, 4 insertions, 11 deletions
diff --git a/lib/tasks/gitlab/shell.rake b/lib/tasks/gitlab/shell.rake
index 0ebc6f00793..ee3ef9dad6e 100644
--- a/lib/tasks/gitlab/shell.rake
+++ b/lib/tasks/gitlab/shell.rake
@@ -103,19 +103,12 @@ namespace :gitlab do
 
     Gitlab::Shell.new.remove_all_keys
 
-    Gitlab::Shell.new.batch_add_keys do |adder|
-      Key.find_each(batch_size: 1000) do |key|
-        adder.add_key(key.shell_id, key.key)
-        print '.'
+    Key.find_in_batches(batch_size: 1000) do |keys|
+      unless Gitlab::Shell.new.batch_add_keys(keys)
+        puts "Failed to add keys...".color(:red)
+        exit 1
       end
     end
-    puts ""
-
-    unless $?.success?
-      puts "Failed to add keys...".color(:red)
-      exit 1
-    end
-
   rescue Gitlab::TaskAbortedByUserError
     puts "Quitting...".color(:red)
     exit 1