From 35efd36c5cae00cc5ff399b6c0779fac94953dd8 Mon Sep 17 00:00:00 2001 From: Oleksii Kliukin Date: Tue, 24 Nov 2015 15:20:29 +0100 Subject: [PATCH] Improve unittests and make minor bugfixes. In particular, remove restore.py in favor of wale_restore.py, fix minor bugs in the latter and add unit tests. --- patroni/postgresql.py | 8 +- patroni/scripts/restore.py | 216 -------------------------------- patroni/scripts/wale_restore.py | 20 +-- tests/test_postgresql.py | 10 ++ tests/test_restore.py | 111 ---------------- tests/test_wale_restore.py | 127 +++++++++++++++++++ 6 files changed, 147 insertions(+), 345 deletions(-) delete mode 100755 patroni/scripts/restore.py delete mode 100644 tests/test_restore.py create mode 100644 tests/test_wale_restore.py diff --git a/patroni/postgresql.py b/patroni/postgresql.py index 4ba9aab3..8a081b7d 100644 --- a/patroni/postgresql.py +++ b/patroni/postgresql.py @@ -229,8 +229,12 @@ class Postgresql: @staticmethod def build_connstring(conn): + """ + >>> Postgresql.build_connstring({'host': '127.0.0.1', 'port': '5432'}) == 'host=127.0.0.1 port=5432 ' + True + """ mconn = "" - for param, val in conn.items(): + for param, val in sorted(conn.items()): mconn = mconn + "{0}={1} ".format(param, val) return mconn @@ -282,7 +286,7 @@ class Postgresql: if ret == 0: break except Exception as e: - logger.exception('Error creating replica using method {0}: {1}'.format(replica_method, e.str)) + logger.exception('Error creating replica using method {0}: {1}'.format(replica_method, str(e))) ret = 1 return ret diff --git a/patroni/scripts/restore.py b/patroni/scripts/restore.py deleted file mode 100755 index 6b20e3e8..00000000 --- a/patroni/scripts/restore.py +++ /dev/null @@ -1,216 +0,0 @@ -#!/usr/bin/env python -# arguments are: -# - cluster scope -# - cluster role -# - master connection string - -# for the AWS, the folliowing environment variables should be defined: -# - WALE_ENV_DIR: directory where WAL-E environment is kept -# - WAL_S3_BUCKET: a name of the S3 bucket for WAL-E -# - WALE_BACKUP_THRESHOLD_MEGABYTES if WAL amount is above that - use pg_basebackup -# - WALE_BACKUP_THRESHOLD_PERCENTAGE if WAL size exceeds a certain percentage of the -# latest backup size -from collections import namedtuple -import logging -import os -import psycopg2 -import subprocess -import sys - - -if sys.hexversion >= 0x03000000: - long = int - -logger = logging.getLogger(__name__) - - -class Restore(object): - - def __init__(self, scope, role, datadir, connstring, env=None): - self.scope = scope - self.role = role - self.master_connection = Restore.parse_connstring(connstring) - self.data_dir = datadir - self.env = os.environ.copy() if not env else env - - @staticmethod - def parse_connstring(connstring): - # the connection string is in the form host= port= user= - # return the dictionary with all components as separare keys - result = {} - if connstring: - for x in connstring.split(): - if x and '=' in x: - key, val = x.split('=') - result[key.strip()] = val.strip() - return result - - def setup(self): - pass - - def replica_method(self): - return self.create_replica_with_pg_basebackup - - def replica_fallback_method(self): - return None - - def run(self): - """ creates a new replica using either pg_basebackup or WAL-E """ - method_fn = self.replica_method() - ret = method_fn() if method_fn else 1 - if ret != 0 and self.replica_fallback_method() is not None: - ret = (self.replica_fallback_method())() - return ret - - def create_replica_with_pg_basebackup(self): - try: - ret = subprocess.call(['pg_basebackup', '-R', '-D', - self.data_dir, '--host=' + self.master_connection['host'], - '--port=' + str(self.master_connection['port']), - '-U', self.master_connection['user']], - env=self.env) - except Exception as e: - logger.error('Error when fetching backup with pg_basebackup: {0}'.format(e)) - return 1 - return ret - - -class WALERestore(Restore): - - def __init__(self, scope, role, datadir, connstring, env=None): - super(WALERestore, self).__init__(scope, role, datadir, connstring, env) - # check the environment variables - self.init_error = False - - def setup(self): - if (self.env.get('WAL_S3_BUCKET') and - self.env.get('WALE_BACKUP_THRESHOLD_PERCENTAGE') and - self.env.get('WALE_BACKUP_THRESHOLD_MEGABYTES')) is None: - self.init_error = True - else: - self.wal_e = namedtuple('WALE', - 'threshold_megabytes threshold_backup_size_percentage s3_bucket cmd dir env_file') - - self.wal_e.dir = self.env.get('WALE_ENV_DIR', '/home/postgres/etc/wal-e.d/env') - self.wal_e.env_file = os.path.join(self.wal_e.dir, 'WALE_S3_PREFIX') - - self.wal_e.cmd = 'envdir {} wal-e --aws-instance-profile '.\ - format(self.wal_e.dir) - self.wal_e.s3_bucket = self.env['WAL_S3_BUCKET'] - self.wal_e.threshold_megabytes = self.env['WALE_BACKUP_THRESHOLD_MEGABYTES'] - self.wal_e.threshold_backup_size_percentage = self.env['WALE_BACKUP_THRESHOLD_PERCENTAGE'] - - # check that the env file exists, create it otherwise - try: - if not os.path.exists(self.wal_e.dir): - os.makedirs(self.wal_e.dir) - # if this is a directory - make sure we have full access there - elif not (os.path.isdir(self.wal_e.dir) and os.access(self.wal_e.dir, os.R_OK | os.W_OK | os.X_OK)): - logger.error("Unable to access {} or not a directory".format(self.wal_e.dir)) - self.init_error = True - # if WAL_S3_PREFIX is not there - create it and write the full path to bucket - if not self.init_error and not os.path.exists(self.wal_e.env_file): - with open(self.wal_e.env_file, 'w') as f: - f.write("s3://{0}/spilo/{1}/wal/\n".format(self.wal_e.s3_bucket, self.scope)) - - except (os.error, IOError) as e: - logger.error("{0}: WAL-e archiving is disabled".format(e)) - self.init_error = True - - def replica_method(self): - if self.should_use_s3_to_create_replica(): - return self.create_replica_with_s3 - return None - - def replica_fallback_method(self): - return self.create_replica_with_pg_basebackup - - def should_use_s3_to_create_replica(self): - """ determine whether it makes sense to use S3 and not pg_basebackup """ - if self.init_error: - return False - - threshold_megabytes = self.wal_e.threshold_megabytes - threshold_backup_size_percentage = self.wal_e.threshold_backup_size_percentage - - try: - latest_backup = subprocess.check_output(self.wal_e.cmd.split() + ['backup-list', '--detail', 'LATEST'], - env=self.env) - # 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.splitlines() if latest_backup else () - if len(backup_strings) != 2: - return False - - names = backup_strings[0].split() - vals = backup_strings[1].split() - if (len(names) != len(vals)) or (len(names) != 7): - return False - - backup_info = dict(zip(names, vals)) - except subprocess.CalledProcessError as e: - logger.error("could not query wal-e latest backup: {}".format(e)) - return False - - try: - 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 as e: - logger.error("unable to get some of S3 backup parameters: {}".format(e)) - return False - - # WAL filename is XXXXXXXXYYYYYYYY000000ZZ, where X - timeline, Y - LSN logical log file, - # ZZ - 2 high digits of LSN offset. The rest of the offset is the provided decimal offset, - # that we have to convert to hex and 'prepend' to the high offset digits. - - lsn_segment = backup_start_segment[8:16] - # first 2 characters of the result are 0x and the last one is L - lsn_offset = hex((long(backup_start_segment[16:32], 16) << 24) + long(backup_start_offset))[2:-1] - - # construct the LSN from the segment and offset - backup_start_lsn = '{}/{}'.format(lsn_segment, lsn_offset) - - conn = None - cursor = None - diff_in_bytes = long(backup_size) - try: - # get the difference in bytes between the current WAL location and the backup start offset - conn = psycopg2.connect(**(self.master_connection)) - conn.autocommit = True - cursor = conn.cursor() - cursor.execute("SELECT pg_xlog_location_diff(pg_current_xlog_location(), %s)", (backup_start_lsn,)) - diff_in_bytes = long(cursor.fetchone()[0]) - except psycopg2.Error as e: - logger.error('could not determine difference with the master location: {}'.format(e)) - return False - finally: - cursor and cursor.close() - conn and conn.close() - - # if the size of the accumulated WAL segments is more than a certan percentage of the backup size - # or exceeds the pre-determined size - pg_basebackup is chosen instead. - return (diff_in_bytes < long(threshold_megabytes) * 1048576) and\ - (diff_in_bytes < long(backup_size) * float(threshold_backup_size_percentage) / 100) - - def create_replica_with_s3(self): - if self.init_error: - return 1 - try: - ret = subprocess.call(self.wal_e.cmd + ' backup-fetch {} LATEST'.format(self.data_dir), env=self.env) - except Exception as e: - logger.error('Error when fetching backup with WAL-E: {0}'.format(e)) - return 1 - return ret - - -if __name__ == '__main__': - if len(sys.argv) == 5: - # scope, role, datadir, connstring - restore = WALERestore(*(sys.argv[1:])) - restore.setup() - sys.exit(restore.run()) - sys.exit("Usage: {0} scope role datadir connstring".format(sys.argv[0])) diff --git a/patroni/scripts/wale_restore.py b/patroni/scripts/wale_restore.py index 03ce000e..ec9a023f 100755 --- a/patroni/scripts/wale_restore.py +++ b/patroni/scripts/wale_restore.py @@ -50,27 +50,18 @@ class WALERestore(object): self.wal_e.dir = env_dir self.wal_e.threshold_mb = threshold_mb self.wal_e.threshold_pct = threshold_pct - if use_iam == 1: - self.wal_e.iam_string = ' --aws-instance-profile ' - else: - self.wal_e.iam_string = '' - if not os.path.exists(self.wal_e.dir): - self.init_error = True - else: - self.init_error = False - self.wal_e.cmd = 'envdir {0} wal-e {1} '.\ - format(self.wal_e.dir, self.wal_e.iam_string) + self.wal_e.iam_string = ' --aws-instance-profile ' if use_iam == 1 else '' + 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)) def run(self): """ creates a new replica using WAL-E """ - if self.should_use_s3_to_create_replica(): + if not self.init_error and self.should_use_s3_to_create_replica(): return self.create_replica_with_s3() return 2 def should_use_s3_to_create_replica(self): """ determine whether it makes sense to use S3 and not pg_basebackup """ - if self.init_error: - return False threshold_megabytes = self.wal_e.threshold_mb threshold_backup_size_percentage = self.wal_e.threshold_pct @@ -138,8 +129,6 @@ class WALERestore(object): (diff_in_bytes < long(backup_size) * float(threshold_backup_size_percentage) / 100) def create_replica_with_s3(self): - if self.init_error: - return 1 # 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']) @@ -151,7 +140,6 @@ class WALERestore(object): if __name__ == '__main__': - parser = argparse.ArgumentParser(description='Script to image replicas using WAL-E') parser.add_argument('--scope', required=True) parser.add_argument('--role', required=False) diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index ba108c28..e13e65fa 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -284,6 +284,16 @@ class TestPostgresql(unittest.TestCase): with patch('subprocess.call', Mock(side_effect=[Exception(), 0])): self.assertEquals(self.p.create_replica(self.leader, ''), 0) + self.p.config['create_replica_method'] = 'wale, basebackup' + self.p.config['wale'] = {'command': 'foo'} + with patch('subprocess.call', Mock(return_value=0)): + self.assertEquals(self.p.create_replica(self.leader, ''), 0) + del self.p.config['wale'] + self.assertEquals(self.p.create_replica(self.leader, ''), 0) + + with patch('subprocess.call', Mock(side_effect=Exception("foo"))): + self.assertEquals(self.p.create_replica(self.leader, ''), 1) + def test_create_connection_users(self): cfg = self.p.config cfg['superuser']['username'] = 'test' diff --git a/tests/test_restore.py b/tests/test_restore.py deleted file mode 100644 index 2ffd8a58..00000000 --- a/tests/test_restore.py +++ /dev/null @@ -1,111 +0,0 @@ -import unittest -from mock import MagicMock, patch -import os -from patroni.scripts.restore import Restore, WALERestore - - -def fake_cursor_fetchone(*args, **kwargs): - return ('16777216',) - - -def fake_call_fail_for_wal_e(*args, **kwargs): - if len(args) > 0 and 'backup-fetch' in args[0]: - return 1 - return 0 - - -def fake_call_fail_for_base_backup(*args, **kwargs): - if len(args) > 0 and 'backup-fetch' in args[0]: - return 0 - return 1 - - -def fake_backup_data(self, *args, **kwargs): - """ return the fake result of WAL-E backup-list""" - return """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 167772160 00000001000000000000007F 00000040 00000001000000000000007F 00000240 -""" - - -class TestRestore(unittest.TestCase): - - def setUp(self): - self.restore = Restore("batman", "master", "/data", "host=batman port=5432 user=batman") - pass - - def tearDown(self): - pass - - def test_parse_connstring(self): - self.assertDictEqual(self.restore.master_connection, {'host': 'batman', 'port': '5432', 'user': 'batman'}) - - @patch('subprocess.call', MagicMock(return_value=0)) - def test_run(self): - ret = self.restore.run() - self.assertEqual(ret, 0) - - @patch('subprocess.call', MagicMock(return_value=1)) - def test_run_fail(self): - ret = self.restore.run() - self.assertEqual(ret, 1) - - -@patch('os.access', MagicMock(return_value=True)) -@patch('os.makedirs', MagicMock(return_value=True)) -@patch('os.path.exists', MagicMock(return_value=True)) -@patch('os.path.isdir', MagicMock(return_value=True)) -@patch('psycopg2.extensions.cursor.fetchone', MagicMock(side_effect=fake_cursor_fetchone)) -@patch('psycopg2.extensions.cursor', MagicMock(autospec=True)) -@patch('psycopg2.extensions.connection', MagicMock(autospec=True)) -@patch('psycopg2.connect', MagicMock(autospec=True)) -@patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data)) -class TestWALERestore(unittest.TestCase): - - def setUp(self): - env = {} - env['WAL_S3_BUCKET'] = 'batman' - env['WALE_BACKUP_THRESHOLD_PERCENTAGE'] = 100 - env['WALE_BACKUP_THRESHOLD_MEGABYTES'] = 100 - self.wale_restore = WALERestore("batman", "master", "/data", "host=batman port=5432 user=batman", env=env) - - def tearDown(self): - pass - - def test_setup(self): - self.wale_restore.setup() - self.assertFalse(self.wale_restore.init_error) - - # have to redefine the class-level os.access mock inside the function - # since the class-level mock will be applied after the function level one. - @patch('os.access', return_value=False) - def test_setup_fail(self, mock_no_access): - os.access = mock_no_access - self.wale_restore.setup() - self.assertTrue(self.wale_restore.init_error) - - # The 3 tests above only differ with the mock function instead of a subprocess call - # in the first one, subprocess call should return success only for wal-e command, - # checking the primary use-case of restoring from WAL-E backup. - # In the second one, we test fallbacks by failing at WAL-E, but succeeding at - # pg_basebackup. - # Finally, the last use case is when all subprocess.call fails. resulting in a - # failure to restore from replica - @patch('subprocess.call', - MagicMock(side_effect=lambda *args, **kwargs: 0 if 'wal-e' in args[0] else 1)) - def test_run(self): - self.wale_restore.setup() - ret = self.wale_restore.run() - self.assertEqual(ret, 0) - - @patch('subprocess.call', - MagicMock(side_effect=lambda *args, **kwargs: 0 if 'pg_basebackup' in args[0] else 1)) - def test_run_fallback(self): - self.wale_restore.setup() - ret = self.wale_restore.run() - self.assertEqual(ret, 0) - - @patch('subprocess.call', MagicMock(return_value=1)) - def test_run_all_fail(self): - self.wale_restore.setup() - ret = self.wale_restore.run() - self.assertEqual(ret, 1) diff --git a/tests/test_wale_restore.py b/tests/test_wale_restore.py new file mode 100644 index 00000000..747fa5c2 --- /dev/null +++ b/tests/test_wale_restore.py @@ -0,0 +1,127 @@ +import unittest +from mock import MagicMock, patch, PropertyMock +import os +import psycopg2 +import subprocess +from patroni.scripts.wale_restore import WALERestore + + +def fake_cursor_fetchone(*args, **kwargs): + return ('16777216',) + + +def fake_call_fail_for_wal_e(*args, **kwargs): + if len(args) > 0 and 'backup-fetch' in args[0]: + return 1 + return 0 + + +def fake_call_fail_for_base_backup(*args, **kwargs): + if len(args) > 0 and 'backup-fetch' in args[0]: + return 0 + return 1 + + +def fake_backup_data(self, *args, **kwargs): + """ return the fake result of WAL-E backup-list""" + return """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 167772160 00000001000000000000007F 00000040 00000001000000000000007F 00000240 +""" + +def fake_backup_data_2(self, *args, **kwargs): + """ return the fake result of WAL-E backup-list""" + return """name last_modified expanded_size_bytes wal_segment_backup_start wal_segment_offset_backup_start wal_segment_backup_stop wal_segment_offset_backup_stop """ + +def fake_backup_data_3(self, *args, **kwargs): + """ return the fake result of WAL-E backup-list""" + return """name last_modified expanded_size_bytes wal_segment_backup_start wal_segment_offset_backup_start wal_segment_backup_stop +base_00000001000000000000007F_00000040 2015-05-18T10:13:25.000Z 167772160 00000001000000000000007F 00000040 00000001000000000000007F 00000240 +""" + +def fake_backup_data_4(self, *args, **kwargs): + """ return the fake result of WAL-E backup-list""" + return """name last_modified expanded_size_foo 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 167772160 00000001000000000000007F 00000040 00000001000000000000007F 00000240 +""" + + +@patch('os.access', MagicMock(return_value=True)) +@patch('os.makedirs', MagicMock(return_value=True)) +@patch('os.path.exists', MagicMock(return_value=True)) +@patch('os.path.isdir', MagicMock(return_value=True)) +@patch('psycopg2.extensions.cursor.fetchone', MagicMock(side_effect=fake_cursor_fetchone)) +@patch('psycopg2.extensions.cursor', MagicMock(autospec=True)) +@patch('psycopg2.extensions.connection', MagicMock(autospec=True)) +@patch('psycopg2.connect', MagicMock(autospec=True)) +@patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data)) +class TestWALERestore(unittest.TestCase): + + def setUp(self): + self.wale_restore = WALERestore("batman", "/data", + "host=batman port=5432 user=batman", "/etc", 100, 100, 1) + + def tearDown(self): + pass + + def test_should_use_s3_to_create_replica(self): + with patch('psycopg2.connect', MagicMock(side_effect=psycopg2.Error("foo"))): + self.assertFalse(self.wale_restore.should_use_s3_to_create_replica()) + with patch('subprocess.check_output', MagicMock(side_effect=subprocess.CalledProcessError(1, "cmd", "foo"))): + self.assertFalse(self.wale_restore.should_use_s3_to_create_replica()) + with patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data_2)): + self.assertFalse(self.wale_restore.should_use_s3_to_create_replica()) + with patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data_3)): + self.assertFalse(self.wale_restore.should_use_s3_to_create_replica()) + with patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data_4)): + self.assertFalse(self.wale_restore.should_use_s3_to_create_replica()) + + self.wale_restore.should_use_s3_to_create_replica() + + def test_create_replica_with_s3(self): + with patch('subprocess.call', MagicMock(return_value=0)): + self.assertEqual(self.wale_restore.create_replica_with_s3(), 0) + with patch('subprocess.call', MagicMock(side_effect=Exception("foo"))): + self.assertEqual(self.wale_restore.create_replica_with_s3(), 1) + + def test_run(self): + with patch.object(self.wale_restore, 'init_error', PropertyMock(return_value=True)): + self.assertEqual(self.wale_restore.run(), 2) + with patch.object(self.wale_restore, 'should_use_s3_to_create_replica', MagicMock(return_value=True)): + with patch.object(self.wale_restore, 'create_replica_with_s3', MagicMock(return_value=0)): + self.assertEqual(self.wale_restore.run(), 0) + + # with patch.object(self.wale_restore, 'init_error', PropertyMock(return_value=True)): + # self.assertFalse(self.wale_restore.create_replica_with_s3()) + + # @patch('subprocess.call', MagicMock(return_value=0)) + # def test_run(self): + # ret = self.wale_restore.run() + # self.assertEqual(ret, 0) + + + # The 3 tests above only differ with the mock function instead of a subprocess call + # in the first one, subprocess call should return success only for wal-e command, + # checking the primary use-case of restoring from WAL-E backup. + # In the second one, we test fallbacks by failing at WAL-E, but succeeding at + # pg_basebackup. + # Finally, the last use case is when all subprocess.call fails. resulting in a + # failure to restore from replica + # @patch('subprocess.call', + # MagicMock(side_effect=lambda *args, **kwargs: 0 if 'wal-e' in args[0] else 1)) + # def test_run(self): + # self.wale_restore.setup() + # ret = self.wale_restore.run() + # self.assertEqual(ret, 0) + + # @patch('subprocess.call', + # MagicMock(side_effect=lambda *args, **kwargs: 0 if 'pg_basebackup' in args[0] else 1)) + # def test_run_fallback(self): + # self.wale_restore.setup() + # ret = self.wale_restore.run() + # self.assertEqual(ret, 0) + + # @patch('subprocess.call', MagicMock(return_value=1)) + # def test_run_all_fail(self): + # self.wale_restore.setup() + # ret = self.wale_restore.run() + # self.assertEqual(ret, 1)