diff --git a/patroni/scripts/wale_restore.py b/patroni/scripts/wale_restore.py index d4a45bf0..156e38ba 100755 --- a/patroni/scripts/wale_restore.py +++ b/patroni/scripts/wale_restore.py @@ -30,6 +30,7 @@ import os import psycopg2 import subprocess import sys +import time import argparse if sys.hexversion >= 0x3000000: @@ -37,10 +38,23 @@ if sys.hexversion >= 0x3000000: logger = logging.getLogger(__name__) +RETRY_SLEEP_INTERVAL = 1 + + +# 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') + if os.path.isfile(version_file): # version file exists + try: + with open(version_file) as f: + return float(f.read()) + except Exception: + logger.exception('Failed to read PG_VERSION from %s', data_dir) + return 0.0 + class WALERestore(object): - - def __init__(self, scope, datadir, connstring, env_dir, threshold_mb, threshold_pct, use_iam, no_master): + 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 @@ -52,11 +66,19 @@ class WALERestore(object): 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)) + self.retries = retries def run(self): """ creates a new replica using WAL-E """ - if not self.init_error and self.should_use_s3_to_create_replica(): - return self.create_replica_with_s3() + 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 def should_use_s3_to_create_replica(self): @@ -82,17 +104,17 @@ class WALERestore(object): 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 + except subprocess.CalledProcessError: + logger.exception("could not query wal-e latest backup") + return None 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 WALE backup parameters: {}".format(e)) - return False + except Exception: + logger.exception("unable to get some of WALE backup parameters") + return None # 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, @@ -106,26 +128,57 @@ class WALERestore(object): backup_start_lsn = '{0}/{1}'.format(lsn_segment, lsn_offset) diff_in_bytes = long(backup_size) - if not self.no_master: - try: - # get the difference in bytes between the current WAL location and the backup start offset - with psycopg2.connect(self.master_connection) as con: - con.autocommit = True - with con.cursor() as cur: - cur.execute("SELECT pg_xlog_location_diff(pg_current_xlog_location(), %s)", (backup_start_lsn,)) - diff_in_bytes = long(cur.fetchone()[0]) - except psycopg2.Error as e: - logger.error('could not determine difference with the master location: %s', e) - return False - else: - # always try to use WAL-E if base backup is available - diff_in_bytes = 0 + attempts_no = 0 + while True: + if self.master_connection: + try: + # get the difference in bytes between the current WAL location and the backup start offset + with psycopg2.connect(self.master_connection) as con: + con.autocommit = True + with con.cursor() as cur: + cur.execute("SELECT pg_xlog_location_diff(pg_current_xlog_location(), %s)", + (backup_start_lsn,)) + diff_in_bytes = long(cur.fetchone()[0]) + except psycopg2.Error: + logger.exception('could not determine difference with the master location') + if attempts_no < self.retries: # retry in case of a temporarily connection issue + attempts_no = attempts_no + 1 + time.sleep(RETRY_SLEEP_INTERVAL) + continue + else: + if not self.no_master: + return False # do no more retries on the outer level + logger.info("continue with base backup from S3 since master is not available") + diff_in_bytes = 0 + break + else: + # always try to use WAL-E if master connection string is not available + diff_in_bytes = 0 + break # 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 fix_subdirectory_path_if_broken(self, dirname): + # in case it is a symlink pointing to a non-existing location, remove it and create the actual directory + path = os.path.join(self.data_dir, dirname) + if not os.path.exists(path): + if os.path.islink(path): # broken xlog symlink, to remove + try: + os.remove(path) + except OSError: + logger.exception("could not remove broken %s symlink pointing to %s", + dirname, os.readlink(path)) + return False + try: + os.mkdir(path) + except OSError: + logger.exception("coud not create missing %s directory path", dirname) + return False + return True + def create_replica_with_s3(self): # if we're set up, restore the replica using fetch latest try: @@ -134,6 +187,9 @@ class WALERestore(object): logger.error('Error when fetching backup with WAL-E: {0}'.format(e)) return 1 + if (ret == 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 @@ -152,17 +208,23 @@ def main(): parser.add_argument('--no_master', type=int, default=0) args = parser.parse_args() - # retry cloning in a loop + # 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 + # the rest and make a decision when the last attempt fails on + # whether to use WAL-E or not depending on the no_master flag. for _ in range(0, args.retries + 1): restore = WALERestore(scope=args.scope, datadir=args.datadir, connstring=args.connstring, 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) + no_master=args.no_master, retries=args.retries) ret = restore.run() - if ret == 0: + if ret != 1: # only WAL-E failures lead to the retry break + time.sleep(RETRY_SLEEP_INTERVAL) + + return ret - sys.exit(ret) if __name__ == '__main__': - main() + sys.exit(main()) diff --git a/tests/test_wale_restore.py b/tests/test_wale_restore.py index 4e7080e0..7e06487b 100644 --- a/tests/test_wale_restore.py +++ b/tests/test_wale_restore.py @@ -2,8 +2,9 @@ import psycopg2 import subprocess import unittest -from mock import MagicMock, patch, PropertyMock -from patroni.scripts.wale_restore import WALERestore, main as _main +from mock import Mock, MagicMock, patch, mock_open +from patroni.scripts.wale_restore import WALERestore, main as _main, get_major_version +from six.moves import builtins wale_output = b'name last_modified expanded_size_bytes wal_segment_backup_start ' +\ @@ -12,51 +13,88 @@ wale_output = b'name last_modified expanded_size_bytes wal_segment_backup_start b'00000001000000000000007F 00000040 00000001000000000000007F 00000240\n' -@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', MagicMock(autospec=True)) -@patch('psycopg2.extensions.connection', MagicMock(autospec=True)) +@patch('os.access', Mock(return_value=True)) +@patch('os.makedirs', Mock(return_value=True)) +@patch('os.path.exists', Mock(return_value=True)) +@patch('os.path.isdir', Mock(return_value=True)) +@patch('psycopg2.extensions.cursor', Mock(autospec=True)) +@patch('psycopg2.extensions.connection', Mock(autospec=True)) @patch('psycopg2.connect', MagicMock(autospec=True)) -@patch('subprocess.check_output', MagicMock(return_value=wale_output)) +@patch('subprocess.check_output', Mock(return_value=wale_output)) class TestWALERestore(unittest.TestCase): def setUp(self): - self.wale_restore = WALERestore("batman", "/data", "host=batman port=5432 user=batman", "/etc", 100, 100, 1, 0) + self.wale_restore = WALERestore("batman", "/data", "host=batman port=5432 user=batman", "/etc", 100, 100, 1, 0, 1) 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(return_value=wale_output.split(b'\n')[0])): - self.assertFalse(self.wale_restore.should_use_s3_to_create_replica()) - with patch('subprocess.check_output', - MagicMock(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', - MagicMock(return_value=wale_output.replace(b'expanded_size_bytes', b'expanded_size_foo'))): - self.assertFalse(self.wale_restore.should_use_s3_to_create_replica()) - - 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()) + with patch('psycopg2.connect', Mock(side_effect=psycopg2.Error("foo"))): + save_no_master = self.wale_restore.no_master + 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 + self.wale_restore.master_connection = '' + self.assertTrue(self.wale_restore.should_use_s3_to_create_replica()) + + self.wale_restore.no_master = save_no_master + self.wale_restore.master_connection = save_master_connection + + with patch('subprocess.check_output', Mock(side_effect=subprocess.CalledProcessError(1, "cmd", "foo"))): + 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', MagicMock(return_value=0)): + with patch('subprocess.call', Mock(return_value=0)): self.assertEqual(self.wale_restore.create_replica_with_s3(), 0) - with patch('subprocess.call', MagicMock(side_effect=Exception("foo"))): + with patch.object(self.wale_restore, 'fix_subdirectory_path_if_broken', Mock(return_value=False)): + self.assertEqual(self.wale_restore.create_replica_with_s3(), 2) + + with patch('subprocess.call', Mock(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.wale_restore.init_error = True + self.assertEqual(self.wale_restore.run(), 2) # this would do 2 retries 1 sec each + self.wale_restore.init_error = False + with patch.object(self.wale_restore, 'should_use_s3_to_create_replica', Mock(return_value=True)): + with patch.object(self.wale_restore, 'create_replica_with_s3', Mock(return_value=0)): self.assertEqual(self.wale_restore.run(), 0) + with patch.object(self.wale_restore, 'should_use_s3_to_create_replica', Mock(return_value=None)): + self.assertEqual(self.wale_restore.run(), 1) + with patch.object(self.wale_restore, 'should_use_s3_to_create_replica', Mock(side_effect=Exception)): + self.assertEqual(self.wale_restore.run(), 2) - @patch('sys.exit', MagicMock()) - @patch.object(WALERestore, 'run', MagicMock(return_value=0)) + @patch('sys.exit', Mock()) def test_main(self): - self.assertEqual(_main(), None) + with patch.object(WALERestore, 'run', Mock(return_value=0)): + self.assertEqual(_main(), 0) + with patch.object(WALERestore, 'run', Mock(return_value=1)): + self.assertEqual(_main(), 1) + + @patch('os.path.isfile', Mock(return_value=True)) + def test_get_major_version(self): + with patch.object(builtins, 'open', mock_open(read_data='9.4')): + self.assertEqual(get_major_version("data"), 9.4) + with patch.object(builtins, 'open', side_effect=OSError): + self.assertEqual(get_major_version("data"), 0.0) + + @patch('os.path.islink', Mock(return_value=True)) + @patch('os.readlink', Mock(return_value="foo")) + @patch('os.remove', Mock()) + @patch('os.mkdir', Mock()) + def test_fix_subdirectory_path_if_broken(self): + with patch('os.path.exists', Mock(return_value=False)): # overriding the class-wide mock + self.assertTrue(self.wale_restore.fix_subdirectory_path_if_broken("data1")) + for fn in ('os.remove', 'os.mkdir'): + with patch(fn, side_effect=OSError): + self.assertFalse(self.wale_restore.fix_subdirectory_path_if_broken("data3"))