From 3241ec2504224ef432893a7e4805935497edab8a Mon Sep 17 00:00:00 2001 From: Joar Wandborg Date: Thu, 27 Apr 2017 14:33:22 +0200 Subject: [PATCH] Use csv.DictReader when parsing wal-e backup-list (#436) wal-e outputs in CSV format using the 'excel-tab' dialect: https://github.com/wal-e/wal-e/blob/3164de68527e6ace269a2112291344b18b9ca6c5/wal_e/operator/backup.py#L63 The ISO date may be written with a space instead of'T' as delimiter between date and time, this causes the old parsing to fail. --- patroni/scripts/wale_restore.py | 153 +++++++++++++++++++++++--------- requirements.txt | 1 + tests/test_wale_restore.py | 151 +++++++++++++++++++++++++++---- 3 files changed, 248 insertions(+), 57 deletions(-) diff --git a/patroni/scripts/wale_restore.py b/patroni/scripts/wale_restore.py index e2b9099b..aa36dca3 100755 --- a/patroni/scripts/wale_restore.py +++ b/patroni/scripts/wale_restore.py @@ -23,10 +23,12 @@ # currently also requires that you configure the restore_command to use wal_e, example: # recovery_conf: # restore_command: envdir /etc/wal-e.d/env wal-e wal-fetch "%f" "%p" -p 1 - +import csv from collections import namedtuple import logging import os +from enum import IntEnum + import psycopg2 import subprocess import sys @@ -38,6 +40,19 @@ logger = logging.getLogger(__name__) RETRY_SLEEP_INTERVAL = 1 +class ExitCode(IntEnum): + """ + Gives meaningful names to the exit codes used by WALERestore + """ + + #: Succeeded + SUCCESS = 0 + #: External issue, retry later + RETRY_LATER = 1 + #: Don't try again unless configuration changes + FAIL = 2 + + # We need to know the current PG version in order to figure out the correct WAL directory name def get_major_version(data_dir): version_file = os.path.join(data_dir, 'PG_VERSION') @@ -50,33 +65,75 @@ def get_major_version(data_dir): return 0.0 +WALEConfig = namedtuple( + 'WALEConfig', + [ + 'env_dir', + 'threshold_mb', + 'threshold_pct', + 'cmd', + ] +) + + class WALERestore(object): - def __init__(self, scope, datadir, connstring, env_dir, threshold_mb, threshold_pct, use_iam, no_master, retries): + def __init__(self, scope, datadir, connstring, env_dir, threshold_mb, + threshold_pct, use_iam, no_master, retries): self.scope = scope self.master_connection = connstring self.data_dir = datadir - self.wal_e = namedtuple('wale', 'dir,threshold_mb,threshold_pct,iam_string,cmd') - self.wal_e.dir = env_dir - self.wal_e.threshold_mb = threshold_mb - self.wal_e.threshold_pct = threshold_pct - self.wal_e.iam_string = ' --aws-instance-profile ' if use_iam == 1 else '' self.no_master = no_master - self.wal_e.cmd = 'envdir {0} wal-e {1} '.format(self.wal_e.dir, self.wal_e.iam_string) - self.init_error = (not os.path.exists(self.wal_e.dir)) + + wale_cmd = [ + 'envdir', + env_dir, + 'wal-e', + ] + + if use_iam == 1: + wale_cmd += ['--aws-instance-profile'] + + self.wal_e = WALEConfig( + env_dir=env_dir, + threshold_mb=threshold_mb, + threshold_pct=threshold_pct, + cmd=wale_cmd, + ) + + self.init_error = (not os.path.exists(self.wal_e.env_dir)) self.retries = retries def run(self): - """ creates a new replica using WAL-E """ - if not self.init_error: - try: - ret = self.should_use_s3_to_create_replica() - if ret: - return self.create_replica_with_s3() - elif ret is None: # caught an exception, need to retry - return 1 - except Exception: - logger.exception("Exception when running WAL-E restore") - return 2 + """ + Creates a new replica using WAL-E + + Returns + ------- + ExitCode + 0 = Success + 1 = Error, try again + 2 = Error, don't try again + + """ + if self.init_error: + logger.error('init error: %r did not exist at initialization time', + self.wal_e.env_dir) + return ExitCode.FAIL + + try: + should_use_s3 = self.should_use_s3_to_create_replica() + if should_use_s3 is None: # Need to retry + return ExitCode.RETRY_LATER + elif should_use_s3: + return self.create_replica_with_s3() + elif not should_use_s3: + return ExitCode.FAIL + except Exception: + logger.exception("Unhandled exception when running WAL-E restore") + return ExitCode.FAIL + + logger.warning('Missing exit code', stack_info=True) + return ExitCode.FAIL def should_use_s3_to_create_replica(self): """ determine whether it makes sense to use S3 and not pg_basebackup """ @@ -85,22 +142,28 @@ class WALERestore(object): threshold_backup_size_percentage = self.wal_e.threshold_pct try: - latest_backup = subprocess.check_output(self.wal_e.cmd.split() + ['backup-list', '--detail', 'LATEST']) - # name last_modified expanded_size_bytes wal_segment_backup_start wal_segment_offset_backup_start - # wal_segment_backup_stop wal_segment_offset_backup_stop - # base_00000001000000000000007F_00000040 2015-05-18T10:13:25.000Z - # 20310671 00000001000000000000007F 00000040 - # 00000001000000000000007F 00000240 - backup_strings = latest_backup.decode('utf-8').splitlines() if latest_backup else () - if len(backup_strings) != 2: + cmd = self.wal_e.cmd + ['backup-list', '--detail', 'LATEST'] + + logger.debug('calling %r', cmd) + wale_output = subprocess.check_output(cmd) + + reader = csv.DictReader(wale_output.decode('utf-8').splitlines(), + dialect='excel-tab') + rows = list(reader) + if not len(rows): + logger.warning('wal-e did not find any backups') return False - names = backup_strings[0].split() - vals = backup_strings[1].split() - if (len(names) != len(vals)) or (len(names) != 7): + # This check might not add much, it was performed in the previous + # version of this code. since the old version rolled CSV parsing the + # check may have been part of the CSV parsing. + if len(rows) > 1: + logger.warning( + 'wal-e returned more than one row of backups: %r', + rows) return False - backup_info = dict(zip(names, vals)) + backup_info = rows[0] except subprocess.CalledProcessError: logger.exception("could not query wal-e latest backup") return None @@ -109,7 +172,7 @@ class WALERestore(object): backup_size = backup_info['expanded_size_bytes'] backup_start_segment = backup_info['wal_segment_backup_start'] backup_start_offset = backup_info['wal_segment_offset_backup_start'] - except Exception: + except KeyError: logger.exception("unable to get some of WALE backup parameters") return None @@ -187,15 +250,19 @@ class WALERestore(object): def create_replica_with_s3(self): # if we're set up, restore the replica using fetch latest try: - ret = subprocess.call(self.wal_e.cmd.split() + ['backup-fetch', '{}'.format(self.data_dir), 'LATEST']) + cmd = self.wal_e.cmd + ['backup-fetch', + '{}'.format(self.data_dir), + 'LATEST'] + logger.debug('calling: %r', cmd) + exit_code = subprocess.call(cmd) except Exception as e: logger.error('Error when fetching backup with WAL-E: {0}'.format(e)) - return 1 + return ExitCode.RETRY_LATER - if (ret == 0 and not + if (exit_code == 0 and not self.fix_subdirectory_path_if_broken('pg_xlog' if get_major_version(self.data_dir) < 10.0 else 'pg_wal')): - return 2 - return ret + return ExitCode.FAIL + return exit_code def main(): @@ -213,6 +280,9 @@ def main(): parser.add_argument('--no_master', type=int, default=0) args = parser.parse_args() + exit_code = None + assert args.retries >= 0 + # Retry cloning in a loop. We do separate retries for the master # connection attempt inside should_use_s3_to_create_replica, # because we need to differentiate between the last attempt and @@ -223,12 +293,13 @@ def main(): env_dir=args.envdir, threshold_mb=args.threshold_megabytes, threshold_pct=args.threshold_backup_size_percentage, use_iam=args.use_iam, no_master=args.no_master, retries=args.retries) - ret = restore.run() - if ret != 1: # only WAL-E failures lead to the retry + exit_code = restore.run() + if not exit_code == ExitCode.RETRY_LATER: # only WAL-E failures lead to the retry + logger.debug('exit_code is %r, not retrying', exit_code) break time.sleep(RETRY_SLEEP_INTERVAL) - return ret + return exit_code if __name__ == '__main__': diff --git a/requirements.txt b/requirements.txt index 7063351b..ba3fe65d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,3 +11,4 @@ click>=4.1 prettytable>=0.7 tzlocal python-dateutil +enum34 diff --git a/tests/test_wale_restore.py b/tests/test_wale_restore.py index 235ab822..39e9ad38 100644 --- a/tests/test_wale_restore.py +++ b/tests/test_wale_restore.py @@ -1,16 +1,132 @@ import psycopg2 import subprocess import unittest +import pytest from mock import Mock, MagicMock, patch, mock_open -from patroni.scripts.wale_restore import WALERestore, main as _main, get_major_version +from patroni.scripts import wale_restore +from patroni.scripts.wale_restore import WALERestore, main as _main, \ + get_major_version, ExitCode from six.moves import builtins -wale_output = b'name last_modified expanded_size_bytes wal_segment_backup_start ' +\ - b'wal_segment_offset_backup_start wal_segment_backup_stop wal_segment_offset_backup_stop\n' +\ - b'base_00000001000000000000007F_00000040 2015-05-18T10:13:25.000Z 167772160 ' +\ - b'00000001000000000000007F 00000040 00000001000000000000007F 00000240\n' +wale_output_header = ( + b'name\tlast_modified\t' + b'expanded_size_bytes\t' + b'wal_segment_backup_start\twal_segment_offset_backup_start\t' + b'wal_segment_backup_stop\twal_segment_offset_backup_stop\n' +) + +wale_output_values = ( + b'base_00000001000000000000007F_00000040\t2015-05-18T10:13:25.000Z\t' + b'167772160\t' + b'00000001000000000000007F\t00000040\t' + b'00000001000000000000007F\t00000240\n' +) + +wale_output = wale_output_header + wale_output_values + +wale_restore.RETRY_SLEEP_INTERVAL = 0.1 # Speed up retries +WALE_TEST_RETRIES = 2 + + +def make_wale_restore(): + return WALERestore( + scope="batman", + datadir="/data", + connstring="host=batman port=5432 user=batman", + env_dir="/etc", + threshold_mb=100, + threshold_pct=100, + use_iam=1, + no_master=0, + retries=WALE_TEST_RETRIES, + ) + + +@pytest.fixture(params=[ + # Nn space + wale_output, + # Space + wale_output.replace( + b'\t2015-05-18T10:13:25.000Z', + b'\t2015-05-18 10:13:25.000Z'), + ]) +def fx_wale_spaces(request): + return request.param + + +@pytest.fixture() +def fx_wale_restore(request): + patches = [ + patch('psycopg2.extensions.cursor', Mock(autospec=True)), + patch('psycopg2.extensions.connection', Mock(autospec=True)), + patch('psycopg2.connect', MagicMock(autospec=True)), + ] + for patch_ in patches: + patch_.start() + + def _finalize(): + for patch_ in patches: + patch_.stop() + + request.addfinalizer(_finalize) + + return make_wale_restore() + + +@pytest.mark.parametrize('exit_code_int,exit_code', [ + (0, ExitCode.SUCCESS), + (1, ExitCode.RETRY_LATER), + (2, ExitCode.FAIL), +]) +def test_exit_code_enum_members_are_int_compatible(exit_code_int, exit_code): + assert exit_code_int == exit_code + + +@pytest.mark.parametrize('mock,exit_code', [ + (Mock(return_value=True), ExitCode.SUCCESS), + (Mock(return_value=False), ExitCode.FAIL), + (Mock(return_value=None), ExitCode.RETRY_LATER), # Handled exception + (Mock(side_effect=Exception('Unhandled exception')), ExitCode.FAIL) +]) +def test_run_exit_codes_by_should_use_s3(mock, exit_code, fx_wale_restore): + """ + Verify that WALERestore.run() returns the correct values based on the + results of WALERestore.should_use_s3t_to_create_replica(). + """ + with patch.object(fx_wale_restore, 'should_use_s3_to_create_replica', + mock),\ + patch.object(fx_wale_restore, 'create_replica_with_s3', + Mock(return_value=ExitCode.SUCCESS)): + assert fx_wale_restore.run() == exit_code + + +def test_should_use_s3_too_many_rows(fx_wale_restore): + with patch('subprocess.check_output', + Mock(return_value=wale_output_header + + wale_output_values + + wale_output_values)): + assert not fx_wale_restore.should_use_s3_to_create_replica() + + +def test_should_use_s3_handles_space_in_date(fx_wale_restore, fx_wale_spaces): + with patch('subprocess.check_output', + Mock(return_value=fx_wale_spaces)): + + assert fx_wale_restore.should_use_s3_to_create_replica() + + +def test_should_use_s3_missing_unused_field(fx_wale_restore): + with patch('subprocess.check_output', + Mock(return_value=wale_output.replace(b'\twal_segment_offset_backup_stop', b''))): + assert fx_wale_restore.should_use_s3_to_create_replica() + + +def test_should_use_s3_missing_used_field(fx_wale_restore): + with patch('subprocess.check_output', + Mock(return_value=wale_output.replace(b'expanded_size_bytes', b'expanded_size_foo'))): + assert fx_wale_restore.should_use_s3_to_create_replica() is None @patch('os.access', Mock(return_value=True)) @@ -24,8 +140,7 @@ wale_output = b'name last_modified expanded_size_bytes wal_segment_backup_start class TestWALERestore(unittest.TestCase): def setUp(self): - self.wale_restore = WALERestore("batman", "/data", "host=batman port=5432 user=batman", - "/etc", 100, 100, 1, 0, 1) + self.wale_restore = make_wale_restore() def test_should_use_s3_to_create_replica(self): self.assertTrue(self.wale_restore.should_use_s3_to_create_replica()) @@ -35,8 +150,15 @@ class TestWALERestore(unittest.TestCase): save_master_connection = self.wale_restore.master_connection self.assertFalse(self.wale_restore.should_use_s3_to_create_replica()) - self.wale_restore.no_master = 1 - self.assertTrue(self.wale_restore.should_use_s3_to_create_replica()) # this would do 2 retries 1 sec each + + with patch('time.sleep', Mock(return_value=None)) as mock_sleep: + self.wale_restore.no_master = 1 + assert self.wale_restore.should_use_s3_to_create_replica() + # verify retries + mock_sleep.assert_has_calls( + [((wale_restore.RETRY_SLEEP_INTERVAL,),)] * WALE_TEST_RETRIES + ) + self.wale_restore.master_connection = '' self.assertTrue(self.wale_restore.should_use_s3_to_create_replica()) @@ -47,12 +169,6 @@ class TestWALERestore(unittest.TestCase): self.assertFalse(self.wale_restore.should_use_s3_to_create_replica()) with patch('subprocess.check_output', Mock(return_value=wale_output.split(b'\n')[0])): self.assertFalse(self.wale_restore.should_use_s3_to_create_replica()) - with patch('subprocess.check_output', - Mock(return_value=wale_output.replace(b' wal_segment_offset_backup_stop', b''))): - self.assertFalse(self.wale_restore.should_use_s3_to_create_replica()) - with patch('subprocess.check_output', - Mock(return_value=wale_output.replace(b'expanded_size_bytes', b'expanded_size_foo'))): - self.assertFalse(self.wale_restore.should_use_s3_to_create_replica()) def test_create_replica_with_s3(self): with patch('subprocess.call', Mock(return_value=0)): @@ -79,8 +195,11 @@ class TestWALERestore(unittest.TestCase): def test_main(self): with patch.object(WALERestore, 'run', Mock(return_value=0)): self.assertEqual(_main(), 0) - with patch.object(WALERestore, 'run', Mock(return_value=1)): + + with patch.object(WALERestore, 'run', Mock(return_value=1)), \ + patch('time.sleep', Mock(return_value=None)) as mock_sleep: self.assertEqual(_main(), 1) + assert mock_sleep.call_count == WALE_TEST_RETRIES @patch('os.path.isfile', Mock(return_value=True)) def test_get_major_version(self):