diff --git a/helpers/postgresql.py b/helpers/postgresql.py index 7bb5b0ed..4616c325 100644 --- a/helpers/postgresql.py +++ b/helpers/postgresql.py @@ -61,10 +61,6 @@ class Postgresql: self.is_promoted = False self._pg_ctl = ['pg_ctl', '-w', '-D', self.data_dir] - self.wal_e = config.get('wal_e', None) - if self.wal_e: - self.wal_e_path = 'envdir {} wal-e --aws-instance-profile '.\ - format(self.wal_e.get('env_dir', '/home/postgres/etc/wal-e.d/env')) self.local_address = self.get_local_address() connect_address = config.get('connect_address', None) or self.local_address @@ -152,9 +148,13 @@ class Postgresql: def create_replica(self, master_connection, env): connstring = self.build_connstring(master_connection, master_connection) cmd = self.config['restore'] - ret = subprocess.call(shlex.split(os.path.abspath(cmd))+[self.scope, "replica", - self.data_dir, connstring], env=env) - self.delete_trigger_file() + try: + ret = subprocess.call(shlex.split(os.path.abspath(cmd))+[self.scope, "replica", + self.data_dir, connstring], env=env) + self.delete_trigger_file() + except Exception as e: + logger.error("Error when creating replica: {0}".format(e)) + return 1 return ret def is_leader(self, check_only=False): diff --git a/requirements-py2.txt b/requirements-py2.txt index 012d87a1..cb6203c5 100644 --- a/requirements-py2.txt +++ b/requirements-py2.txt @@ -1,6 +1,8 @@ boto dnspython +mock psycopg2 PyYAML requests +six >= 1.7 kazoo>=2.2.1 diff --git a/requirements-py3.txt b/requirements-py3.txt index bf481736..2c28b835 100644 --- a/requirements-py3.txt +++ b/requirements-py3.txt @@ -1,4 +1,5 @@ boto +mock dnspython3 psycopg2 PyYAML diff --git a/scripts/restore.py b/scripts/restore.py index 446ef134..2dca3775 100755 --- a/scripts/restore.py +++ b/scripts/restore.py @@ -17,17 +17,21 @@ 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): + 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() + self.env = os.environ.copy() if not env else env @staticmethod def parse_connstring(connstring): @@ -41,6 +45,9 @@ class Restore(object): result[key.strip()] = val.strip() return result + def setup(self): + pass + def replica_method(self): return self.create_replica_with_pg_basebackup @@ -50,26 +57,31 @@ class Restore(object): def run(self): """ creates a new replica using either pg_basebackup or WAL-E """ method_fn = self.replica_method() - ret = method_fn() + 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): - 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) + 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): - super(WALERestore, self).__init__(scope, role, datadir, connstring) + 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: @@ -105,9 +117,9 @@ class WALERestore(Restore): self.init_error = True def replica_method(self): - if self.should_use_s3_to_create_replica(self): + if self.should_use_s3_to_create_replica(): return self.create_replica_with_s3 - return 1 + return None def replica_fallback_method(self): return self.create_replica_with_pg_basebackup @@ -185,9 +197,11 @@ class WALERestore(Restore): def create_replica_with_s3(self): if self.init_error: return 1 - - ret = subprocess.call(self.wal_e.cmd + ' backup-fetch {} LATEST'.format(self.data_dir), env=self.env) - self.restore_configuration_files() + 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 @@ -195,5 +209,6 @@ 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/tests/test_restore.py b/tests/test_restore.py new file mode 100644 index 00000000..cb4f2a83 --- /dev/null +++ b/tests/test_restore.py @@ -0,0 +1,114 @@ +import unittest +from mock import MagicMock, patch +import os +from 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)