Fix the WAL-E restore (#359)

* Fix broken WAL directory symlinks after WAL-E restore.

* Add unit-tests for wale_restore.

* Reduce the amount of MagicMock to the one (for psycopg2.connect)

* Make WAL-E restore process more robuts.

Allow retries only on WAL-E failures.
Sleep after each attempt

* Update the tests.

* Change WAL-E behavior when master is absent, tests.

- Challenge the use of WAL-E even when 'no_master' flag is set. This flag in
  fact does not indicate that the master is absent. In order to check the master
  absense the script looks whether the connection string is not empty.

- Retry on a failure to fetch current xlog position from the master. The reason
  it has to be separate from retries in the main loop is that we don't just
  retry the connection attempt, but also make a decision when either it was
  successfull or all attempts are exhausted.

- Remove wrong usages of ProperyMocks from the tests.

* Avoid redundant output of the exception message in logger.exception

* Address issues uncovered by flake8
This commit is contained in:
Oleksii Kliukin
2016-12-06 17:01:40 +01:00
committed by GitHub
parent b299b12f58
commit b38d98a6a3
2 changed files with 163 additions and 63 deletions
+91 -29
View File
@@ -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())
+72 -34
View File
@@ -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"))