From 779fa4a7b7f2eef7910509df1787412844a3b3b8 Mon Sep 17 00:00:00 2001
From: Jacob Vosmaer <contact@jacobvosmaer.nl>
Date: Thu, 13 Mar 2014 18:32:30 +0100
Subject: [PATCH] Use batch_add_keys in `rake gitlab:shell:setup`

---
 CHANGELOG                   |  2 +-
 lib/gitlab/backend/shell.rb | 16 ++++++++++++++++
 lib/tasks/gitlab/shell.rake | 12 ++++++++----
 3 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG b/CHANGELOG
index 50cfc95ef33..8fd88bf1542 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -21,7 +21,7 @@ v 6.7.0
   - Reuse the GitLab LDAP connection within each request
   - Changed markdown new line behaviour to conform to markdown standards
   - Fix global search
-
+  - Faster authorized_keys rebuilding in `rake gitlab:shell:setup` (requires gitlab-shell 1.8.4)
 
 v 6.6.2
   - Fix 500 error on branch/tag create or remove via UI
diff --git a/lib/gitlab/backend/shell.rb b/lib/gitlab/backend/shell.rb
index 7121c8e40d2..b93800e235f 100644
--- a/lib/gitlab/backend/shell.rb
+++ b/lib/gitlab/backend/shell.rb
@@ -2,6 +2,12 @@ module Gitlab
   class Shell
     class AccessDenied < StandardError; end
 
+    class KeyAdder < Struct.new(:io)
+      def add_key(id, key)
+        io.puts("#{id}\t#{key.strip}")
+      end
+    end
+
     # Init new repository
     #
     # name - project path with namespace
@@ -130,6 +136,16 @@ module Gitlab
       system "#{gitlab_shell_path}/bin/gitlab-keys", "add-key", key_id, key_content
     end
 
+    # Batch-add keys to authorized_keys
+    #
+    # Ex.
+    #   batch_add_keys { |adder| adder.add_key("key-42", "sha-rsa ...") }
+    def batch_add_keys(&block)
+      IO.popen(%W(#{gitlab_shell_path}/bin/gitlab-keys batch-add-keys), 'w') do |io|
+        block.call(KeyAdder.new(io))
+      end
+    end
+
     # Remove ssh key from gitlab shell
     #
     # Ex.
diff --git a/lib/tasks/gitlab/shell.rake b/lib/tasks/gitlab/shell.rake
index 0d7a390bc92..08de0f2dd5d 100644
--- a/lib/tasks/gitlab/shell.rake
+++ b/lib/tasks/gitlab/shell.rake
@@ -34,14 +34,18 @@ namespace :gitlab do
 
     Gitlab::Shell.new.remove_all_keys
 
-    Key.find_each(batch_size: 1000) do |key|
-      if Gitlab::Shell.new.add_key(key.shell_id, key.key)
+    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 '.'
-      else
-        print 'F'
       end
     end
 
+    unless $?.success?
+      puts "Failed to add keys...".red
+      exit 1
+    end
+
   rescue Gitlab::TaskAbortedByUserError
     puts "Quitting...".red
     exit 1