From 0e01bb33bb51d97063e89f6cca818efe4c54eba3 Mon Sep 17 00:00:00 2001 From: Alexander Kukushkin Date: Thu, 4 Jan 2018 10:31:44 +0100 Subject: [PATCH] Improve patronictl reinit (#576) Make it possible to cancel a running task if you want to reinitialize replica. There are two possible ways to trigger it: 1. patronictl will ask whether you want to cancel already running task if an attempt to trigger reinitialize has failed 2. if you are using `--force` argument with `patronictl reinit` --- patroni/api.py | 11 ++- patroni/async_executor.py | 32 ++++++++- patroni/ctl.py | 18 +++-- patroni/ha.py | 10 ++- patroni/postgresql.py | 101 ++++++++++++++++++++------- tests/test_api.py | 4 +- tests/test_async_executor.py | 8 ++- tests/test_ctl.py | 2 +- tests/test_etcd.py | 11 ++- tests/test_ha.py | 3 +- tests/test_postgresql.py | 129 ++++++++++++++++++++++------------- 11 files changed, 239 insertions(+), 90 deletions(-) diff --git a/patroni/api.py b/patroni/api.py index 807de770..2bb43ad4 100644 --- a/patroni/api.py +++ b/patroni/api.py @@ -8,7 +8,7 @@ import dateutil.parser import datetime from patroni.postgresql import PostgresConnectionException, PostgresException, Postgresql -from patroni.utils import deep_compare, patch_config, Retry, RetryFailedError, parse_int, tzutc +from patroni.utils import deep_compare, parse_bool, patch_config, Retry, RetryFailedError, parse_int, tzutc from six.moves.BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer from six.moves.socketserver import ThreadingMixIn from threading import Thread @@ -266,7 +266,14 @@ class RestApiHandler(BaseHTTPRequestHandler): @check_auth def do_POST_reinitialize(self): - data = self.server.patroni.ha.reinitialize() + request = self._read_json_content(body_is_optional=True) + + if request: + logger.debug('received reinitialize request: %s', request) + + force = isinstance(request, dict) and parse_bool(request.get('force')) or False + + data = self.server.patroni.ha.reinitialize(force) if data is None: status_code = 200 data = 'reinitialize started' diff --git a/patroni/async_executor.py b/patroni/async_executor.py index f1de6490..730ebdeb 100644 --- a/patroni/async_executor.py +++ b/patroni/async_executor.py @@ -1,5 +1,5 @@ import logging -from threading import Lock, RLock, Thread +from threading import Event, Lock, RLock, Thread logger = logging.getLogger(__name__) @@ -52,22 +52,27 @@ class CriticalTask(object): class AsyncExecutor(object): - def __init__(self, ha_wakeup): + def __init__(self, state_handler, ha_wakeup): + self.state_handler = state_handler self._ha_wakeup = ha_wakeup self._thread_lock = RLock() self._scheduled_action = None self._scheduled_action_lock = RLock() + self._is_cancelled = False + self._finish_event = Event() self.critical_task = CriticalTask() @property def busy(self): return self.scheduled_action is not None - def schedule(self, action, immediately=False): + def schedule(self, action): with self._scheduled_action_lock: if self._scheduled_action is not None: return self._scheduled_action self._scheduled_action = action + self._is_cancelled = False + self._finish_event.set() return None @property @@ -82,6 +87,12 @@ class AsyncExecutor(object): def run(self, func, args=()): wakeup = False try: + with self: + if self._is_cancelled: + return + self._finish_event.clear() + + self.state_handler.reset_is_cancelled() # if the func returned something (not None) - wake up main HA loop wakeup = func(*args) if args else func() return wakeup @@ -90,6 +101,7 @@ class AsyncExecutor(object): finally: with self: self.reset_scheduled_action() + self._finish_event.set() with self.critical_task: self.critical_task.reset() if wakeup is not None: @@ -98,6 +110,20 @@ class AsyncExecutor(object): def run_async(self, func, args=()): Thread(target=self.run, args=(func, args)).start() + def cancel(self): + with self: + with self._scheduled_action_lock: + if self._scheduled_action is None: + return + logger.warning('Cancelling long running task %s', self._scheduled_action) + self._is_cancelled = True + + self.state_handler.cancel() + self._finish_event.wait() + + with self: + self.reset_scheduled_action() + def __enter__(self): self._thread_lock.acquire() diff --git a/patroni/ctl.py b/patroni/ctl.py index b5643bdb..80ac4d2a 100644 --- a/patroni/ctl.py +++ b/patroni/ctl.py @@ -257,9 +257,9 @@ def get_members(cluster, cluster_name, member_names, role, force, action): member_names = [click.prompt('Which member do you want to {0} [{1}]?'.format(action, ', '.join(candidates.keys())), type=str, default='')] - for mn in member_names: - if mn not in candidates: - raise PatroniCtlException('{0} is not a member of cluster'.format(mn)) + for member_name in member_names: + if member_name not in candidates: + raise PatroniCtlException('{0} is not a member of cluster'.format(member_name)) if not force: confirm = click.confirm('Are you sure you want to {0} members {1}?'.format(action, ', '.join(member_names))) @@ -419,8 +419,10 @@ def check_response(response, member_name, action_name, silent_success=False): click.echo('Failed: {0} for member {1}, status code={2}, ({3})'.format( action_name, member_name, response.status_code, response.text )) + return False elif not silent_success: click.echo('Success: {0} for member {1}'.format(action_name, member_name)) + return True def parse_scheduled(scheduled): @@ -517,8 +519,14 @@ def reinit(obj, cluster_name, member_names, force): members = get_members(cluster, cluster_name, member_names, None, force, 'reinitialize') for member in members: - r = request_patroni(member, 'post', 'reinitialize', headers=auth_header(obj)) - check_response(r, member.name, 'reinitialize') + body = {'force': force} + while True: + r = request_patroni(member, 'post', 'reinitialize', body, auth_header(obj)) + if not check_response(r, member.name, 'reinitialize') and r.text.endswith(' already in progress') \ + and not force and click.confirm('Do you want to cancel it and reinitialize anyway?'): + body['force'] = True + continue + break @ctl.command('failover', help='Failover to a replica') diff --git a/patroni/ha.py b/patroni/ha.py index 9971c170..73428ab0 100644 --- a/patroni/ha.py +++ b/patroni/ha.py @@ -61,7 +61,7 @@ class Ha(object): self._post_bootstrap_task = None self._crash_recovery_executed = False self._start_timeout = None - self._async_executor = AsyncExecutor(self.wakeup) + self._async_executor = AsyncExecutor(self.state_handler, self.wakeup) self.watchdog = patroni.watchdog # Each member publishes various pieces of information to the DCS using touch_member. This lock protects @@ -853,7 +853,7 @@ class Ha(object): member_role = 'leader' if clone_member == self.cluster.leader else 'replica' return self.clone(clone_member, "from {0} '{1}'".format(member_role, clone_member.name)) - def reinitialize(self): + def reinitialize(self, force=False): with self._async_executor: self.load_cluster_from_dcs() @@ -863,7 +863,11 @@ class Ha(object): if self.cluster.leader.name == self.state_handler.name: return 'I am the leader, can not reinitialize' - action = self._async_executor.schedule('reinitialize', immediately=True) + if force: + self._async_executor.cancel() + + with self._async_executor: + action = self._async_executor.schedule('reinitialize') if action is not None: return '{0} already in progress'.format(action) diff --git a/patroni/postgresql.py b/patroni/postgresql.py index 581d4bff..b1c65c56 100644 --- a/patroni/postgresql.py +++ b/patroni/postgresql.py @@ -137,6 +137,10 @@ class Postgresql(object): self._trigger_file = config.get('recovery_conf', {}).get('trigger_file') or 'promote' self._trigger_file = os.path.abspath(os.path.join(self._data_dir, self._trigger_file)) + self._is_cancelled = False + self._cancellable = None + self._cancellable_lock = Lock() + self._connection_lock = Lock() self._connection = None self._cursor_holder = None @@ -536,7 +540,7 @@ class Postgresql(object): params = ['--scope=' + self.scope, '--datadir=' + self._data_dir] try: logger.info('Running custom bootstrap script: %s', config['command']) - if subprocess.call(shlex.split(config['command']) + params) != 0: + if self.cancellable_subprocess_call(shlex.split(config['command']) + params) != 0: self.set_state('custom bootstrap failed') return False except Exception: @@ -582,7 +586,7 @@ class Postgresql(object): env = self.write_pgpass(r) if 'password' in r else None try: - ret = subprocess.call(shlex.split(cmd) + [connstring], env=env) + ret = self.cancellable_subprocess_call(shlex.split(cmd) + [connstring], env=env) except OSError: logger.error('post_init script %s failed', cmd) return False @@ -646,6 +650,9 @@ class Postgresql(object): # go through them in priority order ret = 1 for replica_method in replica_methods: + with self._cancellable_lock: + if self._is_cancelled: + break # if the method is basebackup, then use the built-in if replica_method == "basebackup": ret = self.basebackup(connstring, env) @@ -675,7 +682,7 @@ class Postgresql(object): params = ["--{0}={1}".format(arg, val) for arg, val in method_config.items()] try: # call script with the full set of parameters - ret = subprocess.call(shlex.split(cmd) + params, env=env) + ret = self.cancellable_subprocess_call(shlex.split(cmd) + params, env=env) # if we succeeded, stop if ret == 0: logger.info('replica has been created using %s', replica_method) @@ -768,6 +775,10 @@ class Postgresql(object): def wait_for_port_open(self, postmaster, timeout): """Waits until PostgreSQL opens ports.""" for _ in polling_loop(timeout): + with self._cancellable_lock: + if self._is_cancelled: + return False + if not postmaster.is_running(): logger.error('postmaster is not running') self.set_state('start failed') @@ -814,6 +825,10 @@ class Postgresql(object): options = ['--{0}={1}'.format(p, self._server_parameters[p]) for p in self.CMDLINE_OPTIONS if p in self._server_parameters and p != 'wal_keep_segments'] + with self._cancellable_lock: + if self._is_cancelled: + return False + with task or null_context(): if task and task.is_cancelled: logger.info("PostgreSQL start cancelled.") @@ -823,6 +838,7 @@ class Postgresql(object): self._data_dir, self._postgresql_conf, options) + if task: task.complete(self._postmaster_proc) @@ -988,6 +1004,9 @@ class Postgresql(object): logger.warning("wait_for_startup() called when not in starting state") while not self.check_startup_state_changed(): + with self._cancellable_lock: + if self._is_cancelled: + return None if timeout and self.time_in_state() > timeout: return None time.sleep(1) @@ -1108,10 +1127,9 @@ class Postgresql(object): dsn = " ".join("{0}={1}".format(k, v) for k, v in dsn_attrs if v is not None) logger.info('running pg_rewind from %s', dsn) try: - return subprocess.call([self._pgcommand('pg_rewind'), - '-D', self._data_dir, - '--source-server', dsn, - ], env=env) == 0 + return self.cancellable_subprocess_call([self._pgcommand('pg_rewind'), + '-D', self._data_dir, + '--source-server', dsn], env=env) == 0 except OSError: return False @@ -1531,11 +1549,6 @@ $$""".format(name, ' '.join(options)), name, password, password) self.move_data_directory() def basebackup(self, conn_url, env): - # save environ to restore it later - old_env = os.environ.copy() - os.environ.clear() - os.environ.update(env) - # creates a replica data dir using pg_basebackup. # this is the default, built-in create_replica_method # tries twice, then returns failure (as 1) @@ -1543,12 +1556,15 @@ $$""".format(name, ' '.join(options)), name, password, password) maxfailures = 2 ret = 1 for bbfailures in range(0, maxfailures): + with self._cancellable_lock: + if self._is_cancelled: + break if not self.data_directory_empty(): self.remove_data_directory() try: - ret = subprocess.call([self._pgcommand('pg_basebackup'), '--pgdata=' + self._data_dir, - '-X', 'stream', '--dbname=' + conn_url]) + ret = self.cancellable_subprocess_call([self._pgcommand('pg_basebackup'), '--pgdata=' + self._data_dir, + '-X', 'stream', '--dbname=' + conn_url], env=env) if ret == 0: break else: @@ -1561,10 +1577,6 @@ $$""".format(name, ' '.join(options)), name, password, password) logger.warning('Trying again in 5 seconds') time.sleep(5) - # restore environ - os.environ.clear() - os.environ.update(old_env) - return ret def pick_synchronous_standby(self, cluster): @@ -1684,13 +1696,7 @@ $$""".format(name, ' '.join(options)), name, password, password) cmd.extend(['-c', '{0}={1}'.format(opt, val)]) # need a database name to connect cmd.append(self._database) - p = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=open(os.devnull, 'w'), stderr=subprocess.STDOUT) - if p: - if command: - p.communicate('{0}\n'.format(command)) - p.stdin.close() - return p.wait() - return 1 + return self.cancellable_subprocess_call(cmd, communicate_input=command) def cleanup_archive_status(self): status_dir = os.path.join(self._data_dir, 'pg_' + self.wal_name, 'archive_status') @@ -1716,3 +1722,48 @@ $$""".format(name, ' '.join(options)), name, password, password) if os.path.isfile(self._recovery_conf) or os.path.islink(self._recovery_conf): os.unlink(self._recovery_conf) return self.single_user_mode(options=opts) == 0 or None + + def cancellable_subprocess_call(self, *args, **kwargs): + communicate_input = kwargs.pop('communicate_input', None) + for s in ('stdin', 'stdout', 'stderr'): + kwargs.pop(s, None) + + try: + with self._cancellable_lock: + if self._is_cancelled: + raise PostgresException('cancelled') + + self._is_cancelled = False + self._cancellable = subprocess.Popen(*args, **kwargs) + + if communicate_input: + kwargs['stdin'] = subprocess.PIPE + if communicate_input[-1] != '\n': + communicate_input += '\n' + self._cancellable.communicate(communicate_input + '\n') + self._cancellable.stdin.close() + + return self._cancellable.wait() + finally: + with self._cancellable_lock: + self._cancellable = None + + def reset_is_cancelled(self): + with self._cancellable_lock: + self._is_cancelled = False + + def cancel(self): + with self._cancellable_lock: + self._is_cancelled = True + if self._cancellable is None or self._cancellable.returncode is not None: + return + self._cancellable.terminate() + + for _ in polling_loop(10): + with self._cancellable_lock: + if self._cancellable is None or self._cancellable.returncode is not None: + return + + with self._cancellable_lock: + if self._cancellable is not None and self._cancellable.returncode is None: + self._cancellable.kill() diff --git a/tests/test_api.py b/tests/test_api.py index 653958d8..90d84f9d 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -48,7 +48,7 @@ class MockHa(object): watchdog = MockWatchdog() @staticmethod - def reinitialize(): + def reinitialize(_): return 'reinitialize' @staticmethod @@ -264,7 +264,7 @@ class TestRestApiHandler(unittest.TestCase): def test_do_POST_reinitialize(self, mock_dcs): cluster = mock_dcs.get_cluster.return_value cluster.is_paused.return_value = False - request = 'POST /reinitialize HTTP/1.0' + self._authorization + request = 'POST /reinitialize HTTP/1.0' + self._authorization + '\nContent-Length: 15\n\n{"force": true}' MockRestApiServer(RestApiHandler, request) with patch.object(MockHa, 'reinitialize', Mock(return_value=None)): MockRestApiServer(RestApiHandler, request) diff --git a/tests/test_async_executor.py b/tests/test_async_executor.py index 2c726c0e..e5dea105 100644 --- a/tests/test_async_executor.py +++ b/tests/test_async_executor.py @@ -8,7 +8,7 @@ from threading import Thread class TestAsyncExecutor(unittest.TestCase): def setUp(self): - self.a = AsyncExecutor(Mock()) + self.a = AsyncExecutor(Mock(), Mock()) @patch.object(Thread, 'start', Mock()) def test_run_async(self): @@ -17,6 +17,12 @@ class TestAsyncExecutor(unittest.TestCase): def test_run(self): self.a.run(Mock(side_effect=Exception())) + def test_cancel(self): + self.a.cancel() + self.a.schedule('foo') + self.a.cancel() + self.a.run(Mock()) + class TestCriticalTask(unittest.TestCase): diff --git a/tests/test_ctl.py b/tests/test_ctl.py index ca030129..2189279e 100644 --- a/tests/test_ctl.py +++ b/tests/test_ctl.py @@ -222,7 +222,7 @@ class TestCtl(unittest.TestCase): assert result.exit_code == 1 # successful reinit - result = self.runner.invoke(ctl, ['reinit', 'alpha', 'other'], input='y') + result = self.runner.invoke(ctl, ['reinit', 'alpha', 'other'], input='y\ny') assert result.exit_code == 0 # Aborted restart diff --git a/tests/test_etcd.py b/tests/test_etcd.py index 03c4dd5e..b3ed44c2 100644 --- a/tests/test_etcd.py +++ b/tests/test_etcd.py @@ -18,7 +18,6 @@ class MockResponse(object): self.status_code = status_code self.content = '{}' self.ok = True - self.text = '' def json(self): return json.loads(self.content) @@ -27,6 +26,10 @@ class MockResponse(object): def data(self): return self.content.encode('utf-8') + @property + def text(self): + return self.content + @property def status(self): return self.status_code @@ -48,6 +51,12 @@ def requests_get(url, **kwargs): response.content = '[{}]' if url.startswith('http://error') else members elif url.startswith('http://exhibitor'): response.content = '{"servers":["127.0.0.1","127.0.0.2","127.0.0.3"],"port":2181}' + elif url.endswith(':8011/reinitialize'): + data = kwargs.get('data', '') + if ' false}' in data: + response.status_code = 503 + response.ok = False + response.content = 'restarting after failure already in progress' else: response.status_code = 404 response.ok = False diff --git a/tests/test_ha.py b/tests/test_ha.py index 68014fa3..175c2672 100644 --- a/tests/test_ha.py +++ b/tests/test_ha.py @@ -125,6 +125,7 @@ def run_async(self, func, args=()): @patch.object(Postgresql, 'query', Mock()) @patch.object(Postgresql, 'checkpoint', Mock()) @patch.object(Postgresql, 'call_nowait', Mock()) +@patch.object(Postgresql, 'cancellable_subprocess_call', Mock(return_value=0)) @patch.object(etcd.Client, 'write', etcd_write) @patch.object(etcd.Client, 'read', etcd_read) @patch.object(etcd.Client, 'delete', Mock(side_effect=etcd.EtcdException)) @@ -356,7 +357,7 @@ class TestHa(unittest.TestCase): self.assertIsNotNone(self.ha.reinitialize()) self.ha.cluster = get_cluster_initialized_with_leader() - self.assertIsNone(self.ha.reinitialize()) + self.assertIsNone(self.ha.reinitialize(True)) self.assertIsNotNone(self.ha.reinitialize()) diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index a3928434..dbfac543 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -8,7 +8,7 @@ import unittest from mock import Mock, MagicMock, PropertyMock, patch, mock_open from patroni.async_executor import CriticalTask from patroni.dcs import Cluster, Leader, Member, SyncState -from patroni.exceptions import PostgresConnectionException +from patroni.exceptions import PostgresConnectionException, PostgresException from patroni.postgresql import Postgresql, STATE_REJECT, STATE_NO_RESPONSE from patroni.postmaster import PostmasterProcess from patroni.utils import RetryFailedError @@ -248,6 +248,9 @@ class TestPostgresql(unittest.TestCase): task.cancel() self.assertFalse(self.p.start(task=task)) + self.p.cancel() + self.assertFalse(self.p.start()) + @patch.object(Postgresql, 'pg_isready') @patch('patroni.postgresql.polling_loop', Mock(return_value=range(1))) def test_wait_for_port_open(self, mock_pg_isready): @@ -266,6 +269,10 @@ class TestPostgresql(unittest.TestCase): mock_pg_isready.return_value = 'garbage' self.assertTrue(self.p.wait_for_port_open(mock_postmaster, 1)) + # cancelled + self.p.cancel() + self.assertFalse(self.p.wait_for_port_open(mock_postmaster, 1)) + @patch('time.sleep', Mock()) @patch.object(Postgresql, 'is_running') @patch.object(Postgresql, '_wait_for_connection_close', Mock()) @@ -310,12 +317,13 @@ class TestPostgresql(unittest.TestCase): self.assertIsNone(self.p.checkpoint()) self.assertEquals(self.p.checkpoint(), 'not accessible or not healty') - @patch('subprocess.call', side_effect=OSError) + @patch.object(Postgresql, 'cancellable_subprocess_call') @patch('patroni.postgresql.Postgresql.write_pgpass', MagicMock(return_value=dict())) - def test_pg_rewind(self, mock_call): + def test_pg_rewind(self, mock_cancellable_subprocess_call): r = {'user': '', 'host': '', 'port': '', 'database': '', 'password': ''} + mock_cancellable_subprocess_call.return_value = 0 self.assertTrue(self.p.pg_rewind(r)) - subprocess.call = mock_call + mock_cancellable_subprocess_call.side_effect = OSError self.assertFalse(self.p.pg_rewind(r)) def test_check_recovery_conf(self): @@ -369,6 +377,7 @@ class TestPostgresql(unittest.TestCase): self.p.check_leader_is_not_in_recovery() self.p.check_leader_is_not_in_recovery() + @patch.object(Postgresql, 'cancellable_subprocess_call', Mock(return_value=0)) @patch.object(Postgresql, 'checkpoint', side_effect=['', '1']) @patch.object(Postgresql, 'stop', Mock(return_value=False)) @patch.object(Postgresql, 'start', Mock()) @@ -405,26 +414,36 @@ class TestPostgresql(unittest.TestCase): self.assertFalse(self.p.can_rewind) @patch('time.sleep', Mock()) + @patch.object(Postgresql, 'cancellable_subprocess_call') @patch.object(Postgresql, 'remove_data_directory', Mock(return_value=True)) - def test_create_replica(self): + def test_create_replica(self, mock_cancellable_subprocess_call): self.p.delete_trigger_file = Mock(side_effect=OSError) - with patch('subprocess.call', Mock(side_effect=[1, 0])): - self.assertEquals(self.p.create_replica(self.leader), 0) - 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) + mock_cancellable_subprocess_call.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) + mock_cancellable_subprocess_call.return_value = 1 + self.assertEquals(self.p.create_replica(self.leader), 1) - with patch('subprocess.call', Mock(return_value=1)): - self.assertEquals(self.p.create_replica(self.leader), 1) + mock_cancellable_subprocess_call.side_effect = Exception('foo') + self.assertEquals(self.p.create_replica(self.leader), 1) + + mock_cancellable_subprocess_call.side_effect = [1, 0] + self.assertEquals(self.p.create_replica(self.leader), 0) + + mock_cancellable_subprocess_call.side_effect = [Exception(), 0] + self.assertEquals(self.p.create_replica(self.leader), 0) + + self.p.cancel() + self.assertEquals(self.p.create_replica(self.leader), 1) + + def test_basebackup(self): + self.p.cancel() + self.p.basebackup(None, None) @patch.object(Postgresql, 'is_running', Mock(return_value=True)) def test_sync_replication_slots(self): @@ -543,14 +562,15 @@ class TestPostgresql(unittest.TestCase): lines = f.readlines() self.assertTrue('host replication replicator 127.0.0.1/32 md5\n' in lines) - def test_custom_bootstrap(self): + @patch.object(Postgresql, 'cancellable_subprocess_call') + def test_custom_bootstrap(self, mock_cancellable_subprocess_call): config = {'method': 'foo', 'foo': {'command': 'bar'}} - with patch('subprocess.call', Mock(return_value=1)): - self.assertFalse(self.p.bootstrap(config)) - with patch('subprocess.call', Mock(side_effect=Exception)): - self.assertFalse(self.p.bootstrap(config)) - with patch('subprocess.call', Mock(return_value=0)),\ - patch('subprocess.Popen', Mock(side_effect=Exception("42"))),\ + + mock_cancellable_subprocess_call.return_value = 1 + self.assertFalse(self.p.bootstrap(config)) + + mock_cancellable_subprocess_call.return_value = 0 + with patch('subprocess.Popen', Mock(side_effect=Exception("42"))),\ patch('os.path.isfile', Mock(return_value=True)),\ patch('os.unlink', Mock()),\ patch.object(Postgresql, 'save_configuration_files', Mock()),\ @@ -566,6 +586,9 @@ class TestPostgresql(unittest.TestCase): self.p.bootstrap(config) self.assertEqual(str(e.exception), '42') + mock_cancellable_subprocess_call.side_effect = Exception + self.assertFalse(self.p.bootstrap(config)) + @patch('time.sleep', Mock()) @patch('os.unlink', Mock()) @patch.object(Postgresql, 'run_bootstrap_post_init', Mock(return_value=True)) @@ -593,26 +616,27 @@ class TestPostgresql(unittest.TestCase): self.p.post_bootstrap({}, task) mock_restart.assert_called_once() - def test_run_bootstrap_post_init(self): - with patch('subprocess.call', Mock(return_value=1)): - self.assertFalse(self.p.run_bootstrap_post_init({'post_init': '/bin/false'})) + @patch.object(Postgresql, 'cancellable_subprocess_call') + def test_run_bootstrap_post_init(self, mock_cancellable_subprocess_call): + mock_cancellable_subprocess_call.return_value = 1 + self.assertFalse(self.p.run_bootstrap_post_init({'post_init': '/bin/false'})) - with patch('subprocess.call', Mock(side_effect=OSError)): - self.assertFalse(self.p.run_bootstrap_post_init({'post_init': '/bin/false'})) + mock_cancellable_subprocess_call.return_value = 0 + self.p._superuser.pop('username') + self.assertTrue(self.p.run_bootstrap_post_init({'post_init': '/bin/false'})) + mock_cancellable_subprocess_call.assert_called() + args, kwargs = mock_cancellable_subprocess_call.call_args + self.assertTrue('PGPASSFILE' in kwargs['env']) + self.assertEquals(args[0], ['/bin/false', 'postgres://127.0.0.2:5432/postgres']) - with patch('subprocess.call', Mock(return_value=0)) as mock_method: - self.p._superuser.pop('username') - self.assertTrue(self.p.run_bootstrap_post_init({'post_init': '/bin/false'})) - mock_method.assert_called() - args, kwargs = mock_method.call_args - self.assertTrue('PGPASSFILE' in kwargs['env']) - self.assertEquals(args[0], ['/bin/false', 'postgres://127.0.0.2:5432/postgres']) + mock_cancellable_subprocess_call.reset_mock() + self.p._local_address.pop('host') + self.assertTrue(self.p.run_bootstrap_post_init({'post_init': '/bin/false'})) + mock_cancellable_subprocess_call.assert_called() + self.assertEquals(mock_cancellable_subprocess_call.call_args[0][0], ['/bin/false', 'postgres://:5432/postgres']) - mock_method.reset_mock() - self.p._local_address.pop('host') - self.assertTrue(self.p.run_bootstrap_post_init({'post_init': '/bin/false'})) - mock_method.assert_called() - self.assertEquals(mock_method.call_args[0][0], ['/bin/false', 'postgres://:5432/postgres']) + mock_cancellable_subprocess_call.side_effect = OSError + self.assertFalse(self.p.run_bootstrap_post_init({'post_init': '/bin/false'})) @patch('patroni.postgresql.Postgresql.create_replica', Mock(return_value=0)) def test_clone(self): @@ -775,6 +799,11 @@ class TestPostgresql(unittest.TestCase): self.assertFalse(self.p.wait_for_startup(timeout=2)) self.assertEquals(state['sleeps'], 3) + with patch.object(Postgresql, 'check_startup_state_changed', Mock(return_value=False)): + self.p.cancel() + self.p._state = 'starting' + self.assertIsNone(self.p.wait_for_startup()) + def test_read_pid_file(self): pidfile = os.path.join(self.data_dir, 'postmaster.pid') if os.path.exists(pidfile): @@ -886,13 +915,9 @@ class TestPostgresql(unittest.TestCase): self.assertEqual(data, dict()) @patch('subprocess.Popen') - @patch.object(builtins, 'open', Mock(return_value=42)) def test_single_user_mode(self, subprocess_popen_mock): subprocess_popen_mock.return_value.wait.return_value = 0 - self.assertEquals(self.p.single_user_mode(command="CHECKPOINT"), 0) - subprocess_popen_mock.return_value = None - self.assertEquals(self.p.single_user_mode(), 1) - self.assertEquals(self.p.single_user_mode(options={'archive_mode': 'on'}), 1) + self.assertEquals(self.p.single_user_mode('CHECKPOINT', {'archive_mode': 'on'}), 0) @patch('os.listdir', Mock(side_effect=[OSError, ['a', 'b']])) @patch('os.unlink', Mock(side_effect=OSError)) @@ -908,3 +933,15 @@ class TestPostgresql(unittest.TestCase): @patch.object(Postgresql, 'single_user_mode', Mock(return_value=0)) def test_fix_cluster_state(self): self.assertTrue(self.p.fix_cluster_state()) + + def test_cancellable_subprocess_call(self): + self.p.cancel() + self.assertRaises(PostgresException, self.p.cancellable_subprocess_call) + + @patch('patroni.postgresql.polling_loop', Mock(return_value=[0, 0])) + def test_cancel(self): + self.p._cancellable = Mock() + self.p._cancellable.returncode = None + self.p.cancel() + type(self.p._cancellable).returncode = PropertyMock(side_effect=[None, -15]) + self.p.cancel()