diff --git a/patroni/postgresql/__init__.py b/patroni/postgresql/__init__.py index ce69c7fa..ad3541a4 100644 --- a/patroni/postgresql/__init__.py +++ b/patroni/postgresql/__init__.py @@ -137,6 +137,10 @@ class Postgresql(object): def callback(self): return self.config.get('callbacks') or {} + @property + def wal_dir(self): + return os.path.join(self._data_dir, 'pg_' + self.wal_name) + @property def wal_name(self): return 'wal' if self._major_version >= 100000 else 'xlog' @@ -743,7 +747,7 @@ class Postgresql(object): return self._cluster_info_state_get('timeline') def get_history(self, timeline): - history_path = os.path.join(self._data_dir, 'pg_' + self.wal_name, '{0:08X}.history'.format(timeline)) + history_path = os.path.join(self.wal_dir, '{0:08X}.history'.format(timeline)) history_mtime = mtime(history_path) if history_mtime: try: diff --git a/patroni/postgresql/cancellable.py b/patroni/postgresql/cancellable.py index ffb7040e..fc734d03 100644 --- a/patroni/postgresql/cancellable.py +++ b/patroni/postgresql/cancellable.py @@ -1,11 +1,9 @@ import logging -import os import psutil import subprocess from patroni.exceptions import PostgresException from patroni.utils import polling_loop -from six import string_types from threading import Lock logger = logging.getLogger(__name__) @@ -75,16 +73,16 @@ class CancellableSubprocess(CancellableExecutor): for s in ('stdin', 'stdout', 'stderr'): kwargs.pop(s, None) - communicate_input = 'communicate_input' in kwargs - if communicate_input: - input_data = kwargs.pop('communicate_input', None) - if not isinstance(input_data, string_types): - input_data = '' - if input_data and input_data[-1] != '\n': - input_data += '\n' + communicate = kwargs.pop('communicate', None) + if isinstance(communicate, dict): + input_data = communicate.get('input') + if input_data: + if input_data[-1] != '\n': + input_data += '\n' + input_data = input_data.encode('utf-8') kwargs['stdin'] = subprocess.PIPE - kwargs['stdout'] = open(os.devnull, 'w') - kwargs['stderr'] = subprocess.STDOUT + kwargs['stdout'] = subprocess.PIPE + kwargs['stderr'] = subprocess.PIPE try: with self._lock: @@ -95,10 +93,8 @@ class CancellableSubprocess(CancellableExecutor): started = self._start_process(*args, **kwargs) if started: - if communicate_input: - if input_data: - self._process.communicate(input_data) - self._process.stdin.close() + if isinstance(communicate, dict): + communicate['stdout'], communicate['stderr'] = self._process.communicate(input_data) return self._process.wait() finally: with self._lock: diff --git a/patroni/postgresql/rewind.py b/patroni/postgresql/rewind.py index b90cf038..937cdae2 100644 --- a/patroni/postgresql/rewind.py +++ b/patroni/postgresql/rewind.py @@ -250,21 +250,82 @@ class Rewind(object): def checkpoint_after_promote(self): return self._state == REWIND_STATUS.CHECKPOINT + def _fetch_missing_wal(self, restore_command, wal_filename): + cmd = '' + length = len(restore_command) + i = 0 + while i < length: + if restore_command[i] == '%' and i + 1 < length: + i += 1 + if restore_command[i] == 'p': + cmd += os.path.join(self._postgresql.wal_dir, wal_filename) + elif restore_command[i] == 'f': + cmd += wal_filename + elif restore_command[i] == 'r': + cmd += '000000010000000000000001' + elif restore_command[i] == '%': + cmd += '%' + else: + cmd += '%' + i -= 1 + else: + cmd += restore_command[i] + i += 1 + + logger.info('Trying to fetch the missing wal: %s', cmd) + return self._postgresql.cancellable.call(cmd, shell=True) == 0 + + def _find_missing_wal(self, data): + # could not open file "$PGDATA/pg_wal/0000000A00006AA100000068": No such file or directory + pattern = 'could not open file "' + for line in data.decode('utf-8').split('\n'): + b = line.find(pattern) + if b > -1: + b += len(pattern) + e = line.find('": ', b) + if e > -1: + waldir, wal_filename = os.path.split(line[b:e]) + if waldir.endswith(os.path.sep + 'pg_' + self._postgresql.wal_name) and len(wal_filename) == 24: + return wal_filename + def pg_rewind(self, r): # prepare pg_rewind connection env = self._postgresql.config.write_pgpass(r) - env['PGOPTIONS'] = '-c statement_timeout=0' + env.update(LANG='C', LC_ALL='C', PGOPTIONS='-c statement_timeout=0') dsn = self._postgresql.config.format_dsn(r, True) logger.info('running pg_rewind from %s', dsn) + restore_command = self._postgresql.config.get('recovery_conf', {}).get('restore_command') \ + if self._postgresql.major_version < 120000 else self._postgresql.get_guc_value('restore_command') + cmd = [self._postgresql.pgcommand('pg_rewind')] - if self._postgresql.major_version >= 130000 and self._postgresql.get_guc_value('restore_command'): + if self._postgresql.major_version >= 130000 and restore_command: cmd.append('--restore-target-wal') cmd.extend(['-D', self._postgresql.data_dir, '--source-server', dsn]) - try: - return self._postgresql.cancellable.call(cmd, env=env) == 0 - except OSError: - return False + + while True: + results = {} + ret = self._postgresql.cancellable.call(cmd, env=env, communicate=results) + + logger.info('pg_rewind exit code=%s', ret) + if ret is None: + return False + + logger.info(' stdout=%s', results['stdout'].decode('utf-8')) + logger.info(' stderr=%s', results['stderr'].decode('utf-8')) + if ret == 0: + return True + + if not restore_command or self._postgresql.major_version >= 130000: + return False + + missing_wal = self._find_missing_wal(results['stderr']) or self._find_missing_wal(results['stdout']) + if not missing_wal: + return False + + if not self._fetch_missing_wal(restore_command, missing_wal): + logger.info('Failed to fetch WAL segment %s required for pg_rewind', missing_wal) + return False def execute(self, leader): if self._postgresql.is_running() and not self._postgresql.stop(checkpoint=False): @@ -335,17 +396,17 @@ class Rewind(object): logger.exception('Error when reading postmaster.opts') return result - def single_user_mode(self, command=None, options=None): + def single_user_mode(self, communicate=None, options=None): """run a given command in a single-user mode. If the command is empty - then just start and stop""" cmd = [self._postgresql.pgcommand('postgres'), '--single', '-D', self._postgresql.data_dir] for opt, val in sorted((options or {}).items()): cmd.extend(['-c', '{0}={1}'.format(opt, val)]) # need a database name to connect cmd.append('template1') - return self._postgresql.cancellable.call(cmd, communicate_input=command) + return self._postgresql.cancellable.call(cmd, communicate=communicate) def cleanup_archive_status(self): - status_dir = os.path.join(self._postgresql.data_dir, 'pg_' + self._postgresql.wal_name, 'archive_status') + status_dir = os.path.join(self._postgresql.wal_dir, 'archive_status') try: for f in os.listdir(status_dir): path = os.path.join(status_dir, f) diff --git a/tests/test_cancellable.py b/tests/test_cancellable.py index 7a0e24b7..b555472b 100644 --- a/tests/test_cancellable.py +++ b/tests/test_cancellable.py @@ -13,7 +13,7 @@ class TestCancellableSubprocess(unittest.TestCase): def test_call(self): self.c.cancel() - self.assertRaises(PostgresException, self.c.call, communicate_input=None) + self.assertRaises(PostgresException, self.c.call) def test__kill_children(self): self.c._process_children = [Mock()] diff --git a/tests/test_rewind.py b/tests/test_rewind.py index 0f5c2805..db4bce39 100644 --- a/tests/test_rewind.py +++ b/tests/test_rewind.py @@ -18,6 +18,28 @@ class MockThread(object): self._target(*self._args) +def mock_cancellable_call(*args, **kwargs): + communicate = kwargs.pop('communicate', None) + if isinstance(communicate, dict): + communicate.update(stdout=b'', stderr=b'pg_rewind: error: could not open file ' + + b'"data/postgresql0/pg_xlog/000000010000000000000003": No such file') + return 1 + + +def mock_cancellable_call0(*args, **kwargs): + communicate = kwargs.pop('communicate', None) + if isinstance(communicate, dict): + communicate.update(stdout=b'', stderr=b'') + return 0 + + +def mock_cancellable_call1(*args, **kwargs): + communicate = kwargs.pop('communicate', None) + if isinstance(communicate, dict): + communicate.update(stdout=b'', stderr=b'') + return 1 + + @patch('subprocess.call', Mock(return_value=0)) @patch('psycopg2.connect', psycopg2_connect) class TestRewind(BaseTestPostgresql): @@ -36,16 +58,23 @@ class TestRewind(BaseTestPostgresql): self.p.config._config['use_pg_rewind'] = False self.assertFalse(self.r.can_rewind) - @patch.object(Postgresql, 'major_version', PropertyMock(return_value=130000)) - @patch.object(CancellableSubprocess, 'call') - def test_pg_rewind(self, mock_cancellable_subprocess_call): + def test_pg_rewind(self): r = {'user': '', 'host': '', 'port': '', 'database': '', 'password': ''} - mock_cancellable_subprocess_call.return_value = 0 - with patch('subprocess.check_output', Mock(return_value=b'foo')): - self.assertTrue(self.r.pg_rewind(r)) - mock_cancellable_subprocess_call.side_effect = OSError - with patch('subprocess.check_output', Mock(side_effect=Exception)): - self.assertFalse(self.r.pg_rewind(r)) + with patch.object(Postgresql, 'major_version', PropertyMock(return_value=130000)),\ + patch.object(CancellableSubprocess, 'call', Mock(return_value=None)): + with patch('subprocess.check_output', Mock(return_value=b'boo')): + self.assertFalse(self.r.pg_rewind(r)) + with patch('subprocess.check_output', Mock(side_effect=Exception)): + self.assertFalse(self.r.pg_rewind(r)) + + with patch.object(Postgresql, 'major_version', PropertyMock(return_value=120000)),\ + patch('subprocess.check_output', Mock(return_value=b'foo %f %p %r %% % %')): + with patch.object(CancellableSubprocess, 'call', mock_cancellable_call): + self.assertFalse(self.r.pg_rewind(r)) + with patch.object(CancellableSubprocess, 'call', mock_cancellable_call0): + self.assertTrue(self.r.pg_rewind(r)) + with patch.object(CancellableSubprocess, 'call', mock_cancellable_call1): + self.assertFalse(self.r.pg_rewind(r)) @patch.object(Rewind, 'can_rewind', PropertyMock(return_value=True)) def test__get_local_timeline_lsn(self): @@ -61,7 +90,7 @@ class TestRewind(BaseTestPostgresql): with patch.object(MockCursor, 'fetchone', Mock(side_effect=[(0, 0, 1, 1,), Exception])): self.r.rewind_or_reinitialize_needed_and_possible(self.leader) - @patch.object(CancellableSubprocess, 'call', Mock(return_value=0)) + @patch.object(CancellableSubprocess, 'call', mock_cancellable_call) @patch.object(Postgresql, 'checkpoint', side_effect=['', '1'],) @patch.object(Postgresql, 'stop', Mock(return_value=False)) @patch.object(Postgresql, 'start', Mock()) @@ -161,7 +190,8 @@ class TestRewind(BaseTestPostgresql): @patch('psutil.Popen') def test_single_user_mode(self, subprocess_popen_mock): subprocess_popen_mock.return_value.wait.return_value = 0 - self.assertEqual(self.r.single_user_mode('CHECKPOINT', {'archive_mode': 'on'}), 0) + subprocess_popen_mock.return_value.communicate.return_value = ('', '') + self.assertEqual(self.r.single_user_mode({'input': 'CHECKPOINT'}, {'archive_mode': 'on'}), 0) @patch('os.listdir', Mock(side_effect=[OSError, ['a', 'b']])) @patch('os.unlink', Mock(side_effect=OSError))