Commit 8d6e72f3 authored by Michael Kozono's avatar Michael Kozono

Move batch-add-logic to background migration

Do the work synchronously to avoid multiple workers attempting to add batches of keys at the same time.

Also, make the delete portion wait until after adding is done.
parent 12d0a64d
......@@ -4,19 +4,6 @@ class GitlabShellWorker
include DedicatedSidekiqQueue
def perform(action, *arg)
if action.to_s == 'batch_add_keys_in_db_starting_from'
batch_add_keys_in_db_starting_from(arg.first)
else
gitlab_shell.send(action, *arg)
end
end
# Not added to Gitlab::Shell because I don't expect this to be used again
def batch_add_keys_in_db_starting_from(start_id)
gitlab_shell.batch_add_keys do |adder|
Key.find_each(start: start_id, batch_size: 1000) do |key|
adder.add_key(key.shell_id, key.key)
end
end
gitlab_shell.send(action, *arg)
end
end
......@@ -3,6 +3,10 @@ module Gitlab
class UpdateAuthorizedKeysFileSince
class Key < ActiveRecord::Base
self.table_name = 'keys'
def shell_id
"key-#{id}"
end
end
def perform(cutoff_datetime)
......@@ -14,13 +18,26 @@ module Gitlab
def add_keys_since(cutoff_datetime)
start_key = Key.select(:id).where("created_at >= ?", cutoff_datetime).take
if start_key
GitlabShellWorker.perform_async(:batch_add_keys_in_db_starting_from, start_key.id)
batch_add_keys_in_db_starting_from(start_key.id)
end
end
def remove_keys_not_found_in_db
GitlabShellWorker.perform_async(:remove_keys_not_found_in_db)
end
# Not added to Gitlab::Shell because I don't expect this to be used again
def batch_add_keys_in_db_starting_from(start_id)
gitlab_shell.batch_add_keys do |adder|
Key.find_each(start: start_id, batch_size: 1000) do |key|
adder.add_key(key.shell_id, key.key)
end
end
end
def gitlab_shell
@gitlab_shell ||= Gitlab::Shell.new
end
end
end
end
require 'spec_helper'
describe Gitlab::BackgroundMigration::UpdateAuthorizedKeysFileSince do
let(:background_migration) { described_class.new }
describe '#perform' do
let!(:cutoff_datetime) { DateTime.now }
subject { described_class.new.perform(cutoff_datetime) }
subject { background_migration.perform(cutoff_datetime) }
context 'when an SSH key was created after the cutoff datetime' do
before do
......@@ -19,9 +21,8 @@ describe Gitlab::BackgroundMigration::UpdateAuthorizedKeysFileSince do
@key = create(:key)
end
it 'queues a batch_add_keys_from call to GitlabShellWorker, including the start key ID' do
expect(GitlabShellWorker).to receive(:perform_async).with(:batch_add_keys_in_db_starting_from, @key.id)
allow(GitlabShellWorker).to receive(:perform_async).with(:remove_keys_not_found_in_db)
it 'calls batch_add_keys_in_db_starting_from with the start key ID' do
expect(background_migration).to receive(:batch_add_keys_in_db_starting_from).with(@key.id)
subject
end
end
......@@ -35,7 +36,7 @@ describe Gitlab::BackgroundMigration::UpdateAuthorizedKeysFileSince do
describe '#add_keys_since' do
let!(:cutoff_datetime) { DateTime.now }
subject { described_class.new.add_keys_since(cutoff_datetime) }
subject { background_migration.add_keys_since(cutoff_datetime) }
before do
Timecop.freeze
......@@ -50,15 +51,15 @@ describe Gitlab::BackgroundMigration::UpdateAuthorizedKeysFileSince do
@key = create(:key)
end
it 'queues a batch_add_keys_from call to GitlabShellWorker, including the start key ID' do
expect(GitlabShellWorker).to receive(:perform_async).with(:batch_add_keys_in_db_starting_from, @key.id)
it 'calls batch_add_keys_in_db_starting_from with the start key ID' do
expect(background_migration).to receive(:batch_add_keys_in_db_starting_from).with(@key.id)
subject
end
end
context 'when an SSH key was not created after the cutoff datetime' do
it 'does not use GitlabShellWorker' do
expect(GitlabShellWorker).not_to receive(:perform_async)
it 'does not call batch_add_keys_in_db_starting_from' do
expect(background_migration).not_to receive(:batch_add_keys_in_db_starting_from)
subject
end
end
......@@ -67,7 +68,32 @@ describe Gitlab::BackgroundMigration::UpdateAuthorizedKeysFileSince do
describe '#remove_keys_not_found_in_db' do
it 'queues a rm_keys_not_in_db call to GitlabShellWorker' do
expect(GitlabShellWorker).to receive(:perform_async).with(:remove_keys_not_found_in_db)
described_class.new.remove_keys_not_found_in_db
background_migration.remove_keys_not_found_in_db
end
end
describe '#batch_add_keys_in_db_starting_from' do
context 'when there are many keys in the DB' do
before do
@keys = []
10.times do
@keys << create(:key)
end
end
it 'adds all the keys in the DB, starting from the given ID, to the authorized_keys file' do
Gitlab::Shell.new.remove_all_keys
background_migration.batch_add_keys_in_db_starting_from(@keys[3].id)
file = File.read(Rails.root.join('tmp/tests/.ssh/authorized_keys'))
expect(file.scan(/ssh-rsa/).count).to eq(7)
expect(file).to_not include(Gitlab::Shell.strip_key(@keys[0].key))
expect(file).to_not include(Gitlab::Shell.strip_key(@keys[2].key))
expect(file).to include(Gitlab::Shell.strip_key(@keys[3].key))
expect(file).to include(Gitlab::Shell.strip_key(@keys[9].key))
end
end
end
end
......@@ -9,29 +9,4 @@ describe GitlabShellWorker do
worker.perform(:add_key, 'foo', 'bar')
end
end
describe '#perform with batch_add_keys_in_db_starting_from' do
context 'when there are many keys in the DB' do
before do
@keys = []
10.times do
@keys << create(:key)
end
end
it 'adds all the keys in the DB, starting from the given ID, to the authorized_keys file' do
Gitlab::Shell.new.remove_all_keys
worker.perform(:batch_add_keys_in_db_starting_from, @keys[3].id)
file = File.read(Rails.root.join('tmp/tests/.ssh/authorized_keys'))
expect(file.scan(/ssh-rsa/).count).to eq(7)
expect(file).to_not include(Gitlab::Shell.strip_key(@keys[0].key))
expect(file).to_not include(Gitlab::Shell.strip_key(@keys[2].key))
expect(file).to include(Gitlab::Shell.strip_key(@keys[3].key))
expect(file).to include(Gitlab::Shell.strip_key(@keys[9].key))
end
end
end
end
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment