Add pause/resume replication commands to gitlab-ctl

- Update POC code to also call rake task that calls primary api to pause
- Add changelog

Relates to https://gitlab.com/gitlab-org/omnibus-gitlab/issues/35913

changelogs/unreleased/geo-pause-resume-commands.yml

+---
+title: Add Commands to Pause/Resume Replication
+merge_request: 4331
+author:
+type: changed

files/gitlab-ctl-commands-ee/geo_replication.rb

+require "#{base_path}/embedded/service/omnibus-ctl-ee/lib/geo/replication_toggle_command"
+
+add_command_under_category('geo-replication-pause', 'gitlab-geo', 'Replication Process', 2) do |_cmd_name, *args|
+  Geo::ReplicationToggleCommand.new(self, 'pause', ARGV).execute!
+end
+
+add_command_under_category('geo-replication-resume', 'gitlab-geo', 'Replication Process', 2) do |_cmd_name, *args|
+  Geo::ReplicationToggleCommand.new(self, 'resume', ARGV).execute!
+end

files/gitlab-ctl-commands-ee/lib/geo/replication_process.rb

-require 'io/console'
 require 'rainbow/ext/string'
-# For testing purposes, if the first path cannot be found load the second
-begin
-  require_relative '../../../omnibus-ctl/lib/postgresql'
-rescue LoadError
-  require_relative '../../../gitlab-ctl-commands/lib/postgresql'
-end
-
 module Geo
   class ReplicationProcess
     attr_accessor :base_path, :ctl
@@ -21,25 +13,49 @@ module Geo
     def pause
       puts '* Pausing replication'.color(:green)
-      run_query('SELECT pg_wal_replay_pause();')
+
+      task = run_task('geo:replication:pause')
+      # This isn't an error because theoretically if the primary is down you
+      # may want to use this to pause the WAL anyway.
+      puts task.stdout.strip.color(:red) if task.error?
+
+      query = run_query('SELECT pg_wal_replay_pause();')
+      raise PsqlError, "Unable to pause postgres replication #{query.stdout.strip}" if query.error?
+
+      puts '* Replication paused'.color(:green)
     end
     def resume
       puts '* Resume replication'.color(:green)
-      run_query('pg_wal_replay_resume();')
+
+      query = run_query('pg_wal_replay_resume();')
+      raise PsqlError, "Unable to resume postgres replication #{query.stdout.strip}" if query.error?
+
+      task = run_task('geo:replication:resume')
+      raise RakeError, "Unable to resume replication from primary #{task.stdout.strip}" if task.error?
+
+      puts '* Replication resumed'.color(:green)
     end
     private
     def run_query(query)
-      status = GitlabCtl::Util.run_command(
+      GitlabCtl::Util.run_command(
         "#{base_path}/bin/gitlab-psql -d #{db_name} -c '#{query}' -q -t"
       )
-      status.error? ? false : status.stdout.strip
+    end
+
+    def run_task(task)
+      GitlabCtl::Util.run_command(
+        "#{base_path}/bin/gitlab-rake #{task}"
+      )
     end
     def db_name
       @options[:db_name]
     end
   end
+
+  PsqlError = Class.new(StandardError)
+  RakeError = Class.new(StandardError)
 end

files/gitlab-ctl-commands-ee/lib/geo/replication_toggle_command.rb

+require_relative "./replication_process"
+require 'optparse'
+require 'English'
+
+module Geo
+  class ReplicationToggleCommand
+    def initialize(ctl, action, args)
+      @ctl = ctl
+      @args = args
+      @action = action
+
+      @options = {
+        db_name: 'gitlabhq_production'
+      }
+      parse_options!
+
+      @replication_process = Geo::ReplicationProcess.new(@ctl, @options)
+    end
+
+    def execute!
+      @replication_process.send(@action.to_sym)
+    rescue Geo::PsqlError => e
+      puts "Postgres encountered an error: #{e.message}"
+      exit 1
+    rescue Geo::RakeError => e
+      puts "Rake encountered an error: #{e.message}"
+      exit 1
+    end
+
+    def arguments
+      @args.dup
+    end
+
+    private
+
+    def parse_options!
+      opts_parser = OptionParser.new do |opts|
+        opts.banner = "Usage: gitlab-ctl replication-process-#{@action} [options]"
+
+        opts.separator ''
+        opts.separator 'Specific options:'
+
+        opts.on('--db_name=gitlabhq_production', 'Specify the database name') do |db_name|
+          @options[:db_name] = db_name
+        end
+
+        opts.on_tail('-h', '--help', 'Show this message') do
+          puts opts
+          exit
+        end
+      end
+
+      opts_parser.parse!(arguments)
+    end
+  end
+end

files/gitlab-ctl-commands-ee/replication_process.rb

-require "#{base_path}/embedded/service/omnibus-ctl-ee/lib/geo/replication_process"
-require 'optparse'
-require 'English'
-
-add_command_under_category('replication-process-pause', 'gitlab-geo', 'Replication Process', 2) do |_cmd_name, *args|
-  ReplicationProcessCommand.new(self, ARGV).execute!
-end
-
-class ReplicationProcessCommand
-  def initialize(ctl, args)
-    @ctl = ctl
-    @args = args
-
-    @options = {
-      db_name: 'gitlabhq_production'
-    }
-
-    parse_options!
-  end
-
-  def execute!
-    Geo::ReplicationProcess.new(@ctl, @options).pause
-  end
-
-  def arguments
-    @args.dup
-  end
-
-  private
-
-  def parse_options!
-    opts_parser = OptionParser.new do |opts|
-      opts.banner = 'Usage: gitlab-ctl replication-process-pause [options]'
-
-      opts.separator ''
-      opts.separator 'Specific @options:'
-
-      opts.on('--db_name=gitlabhq_production', 'Specify the database name') do |db_name|
-        @options[:db_name] = db_name
-      end
-
-      opts.on_tail('-h', '--help', 'Show this message') do
-        puts opts
-        exit
-      end
-    end
-
-    opts_parser.parse!(arguments)
-  end
-end

spec/gitlab-ctl-commands-ee/geo_replication_spec.rb

+require 'spec_helper'
+require 'omnibus-ctl'
+
+describe 'gitlab-ctl geo-replication' do
+  let(:toggle_command) { double(Geo::ReplicationToggleCommand) }
+
+  subject { Omnibus::Ctl.new('testing-ctl') }
+
+  before do
+    allow_any_instance_of(Omnibus::Ctl).to receive(:require).and_call_original
+    allow_any_instance_of(Omnibus::Ctl).to receive(:require).with(
+      '/opt/testing-ctl/embedded/service/omnibus-ctl-ee/lib/geo/replication_process'
+    ) do
+      require_relative('../../files/gitlab-ctl-commands-ee/lib/geo/replication_process')
+    end
+    allow_any_instance_of(Omnibus::Ctl).to receive(:require).with(
+      '/opt/testing-ctl/embedded/service/omnibus-ctl-ee/lib/geo/replication_toggle_command'
+    ) do
+      require_relative('../../files/gitlab-ctl-commands-ee/lib/geo/replication_toggle_command')
+    end
+
+    subject.load_file('files/gitlab-ctl-commands-ee/geo_replication.rb')
+  end
+
+  it 'appends the geo replication pause and resume commands' do
+    expect(subject.get_all_commands_hash).to include('geo-replication-pause')
+    expect(subject.get_all_commands_hash).to include('geo-replication-resume')
+  end
+
+  describe 'pause' do
+    it 'calls pause' do
+      expect(Geo::ReplicationToggleCommand).to receive(:new).with(anything, 'pause', anything).and_return(toggle_command)
+      expect(toggle_command).to receive(:execute!)
+
+      subject.geo_replication_pause
+    end
+  end
+
+  describe 'resume' do
+    it 'calls resume' do
+      expect(Geo::ReplicationToggleCommand).to receive(:new).with(anything, 'resume', anything).and_return(toggle_command)
+      expect(toggle_command).to receive(:execute!)
+
+      subject.geo_replication_resume
+    end
+  end
+end

spec/gitlab-ctl-commands-ee/lib/geo/replication_process_spec.rb

+require 'spec_helper'
+
+$LOAD_PATH << './files/gitlab-ctl-commands-ee/lib'
+$LOAD_PATH << './files/gitlab-ctl-commands/lib'
+
+require 'geo/replication_process'
+require 'gitlab_ctl/util'
+
+describe Geo::ReplicationProcess do
+  let(:error_text) { 'AN ERROR' }
+  let(:good_status) { double('Command status', error?: false) }
+  let(:bad_status) { double('Command status', error?: true, stdout: error_text) }
+  let(:instance) { double(base_path: '/opt/gitlab/embedded', data_path: '/var/opt/gitlab/postgresql/data') }
+  let(:db_name) { 'gitlab_db_name' }
+  let(:options) do
+    {
+      db_name: db_name
+    }
+  end
+
+  subject { described_class.new(instance, options) }
+
+  describe '#pause' do
+    it 'logs a message if the rake task throws an error' do
+      expect(GitlabCtl::Util).to receive(:run_command).with(/gitlab-rake geo:replication:pause/).and_return(bad_status)
+      expect(GitlabCtl::Util).to receive(:run_command).with(/gitlab-psql/).and_return(good_status)
+
+      expect do
+        subject.pause
+      end.to output(/#{error_text}/).to_stdout
+    end
+
+    it 'raises an exception if unable to pause replication' do
+      expect(GitlabCtl::Util).to receive(:run_command).with(/gitlab-rake geo:replication:pause/).and_return(good_status)
+      expect(GitlabCtl::Util).to receive(:run_command).with(/gitlab-psql/).and_return(bad_status)
+
+      expect do
+        subject.pause
+      end.to raise_error(/Unable to pause postgres replication/)
+    end
+
+    it 'provides the requested database from the options' do
+      expect(GitlabCtl::Util).to receive(:run_command).with(/gitlab-rake geo:replication:pause/).and_return(good_status)
+      expect(GitlabCtl::Util).to receive(:run_command).with(/gitlab-psql -d #{db_name}/).and_return(good_status)
+
+      subject.pause
+    end
+  end
+
+  describe '#resume' do
+    it 'raises an exception if unable to pause pg replication' do
+      expect(GitlabCtl::Util).to receive(:run_command).with(/gitlab-psql/).and_return(bad_status)
+
+      expect do
+        subject.resume
+      end.to raise_error(/Unable to resume postgres replication/)
+    end
+
+    it 'raises an error if rake task to resume fails' do
+      expect(GitlabCtl::Util).to receive(:run_command).with(/gitlab-rake geo:replication:resume/).and_return(bad_status)
+      expect(GitlabCtl::Util).to receive(:run_command).with(/gitlab-psql/).and_return(good_status)
+
+      expect do
+        subject.resume
+      end.to raise_error(/Unable to resume replication from primary/)
+    end
+
+    it 'provides the requested database from the options' do
+      expect(GitlabCtl::Util).to receive(:run_command).with(/gitlab-rake geo:replication:resume/).and_return(good_status)
+      expect(GitlabCtl::Util).to receive(:run_command).with(/gitlab-psql -d #{db_name}/).and_return(good_status)
+
+      subject.resume
+    end
+  end
+end

spec/gitlab-ctl-commands-ee/lib/geo/replication_toggle_command_spec.rb

+require 'spec_helper'
+
+$LOAD_PATH << './files/gitlab-ctl-commands-ee/lib'
+$LOAD_PATH << './files/gitlab-ctl-commands/lib'
+
+require 'geo/replication_toggle_command'
+
+describe Geo::ReplicationToggleCommand do
+  let(:status) { double('Command status', error?: false) }
+  let(:arguments) { [] }
+  let(:ctl_instance) { double('gitlab-ctl instance', base_path: '') }
+
+  describe 'pause' do
+    subject { described_class.new(ctl_instance, 'pause', arguments) }
+
+    it 'calls pause' do
+      expect_any_instance_of(Geo::ReplicationProcess).to receive(:pause)
+
+      subject.execute!
+    end
+
+    it 'rescues and exits if postgres has an error' do
+      expect_any_instance_of(Geo::ReplicationProcess).to receive(:pause).and_raise(Geo::PsqlError, "Oh nose!")
+
+      expect do
+        expect { subject.execute! }.to raise_error(SystemExit)
+      end.to output(/Postgres encountered an error: Oh nose!/).to_stdout
+    end
+
+    context 'database specified' do
+      let(:arguments) { %w(--db_name=database_i_want) }
+
+      it 'uses the specified database' do
+        expect(Geo::ReplicationProcess).to receive(:new).with(any_args, { db_name: 'database_i_want' }).and_call_original
+        expect_any_instance_of(Geo::ReplicationProcess).to receive(:pause)
+
+        subject.execute!
+      end
+    end
+  end
+
+  describe 'resume' do
+    subject { described_class.new(ctl_instance, 'resume', arguments) }
+
+    it 'calls resume' do
+      expect_any_instance_of(Geo::ReplicationProcess).to receive(:resume)
+
+      subject.execute!
+    end
+
+    it 'rescues and exits if postgres has an error' do
+      expect_any_instance_of(Geo::ReplicationProcess).to receive(:resume).and_raise(Geo::PsqlError, "Oh nose!")
+
+      expect do
+        expect { subject.execute! }.to raise_error(SystemExit)
+      end.to output(/Postgres encountered an error: Oh nose!/).to_stdout
+    end
+
+    it 'rescues and exits if rake has an error' do
+      expect_any_instance_of(Geo::ReplicationProcess).to receive(:resume).and_raise(Geo::RakeError, "Oh nose!")
+
+      expect do
+        expect { subject.execute! }.to raise_error(SystemExit)
+      end.to output(/Rake encountered an error: Oh nose!/).to_stdout
+    end
+
+    context 'database specified' do
+      let(:arguments) { %w(--db_name=database_i_want) }
+
+      it 'uses the specified database' do
+        expect(Geo::ReplicationProcess).to receive(:new).with(any_args, { db_name: 'database_i_want' }).and_call_original
+        expect_any_instance_of(Geo::ReplicationProcess).to receive(:resume)
+
+        subject.execute!
+      end
+    end
+  end
+end