Merge pull request #227 from zalando/feature/scheduled_restarts

Implement scheduled restarts for Patroni via the API. 

Even normal (immediate) restarts will take advantage of additional modifiers you can supply to the restart endpoint:

- restart_pending: restart only if the pending restart flag is set (because of the configuration change)
- role: restart if the Postgres role is set to a certain value
- postgres_version (x.y.z) - restart if the current Postgres version is less than the one specified.

For the scheduled restart, the schedule parameter can be used the same way as it is currently used for the scheduled failovers. Particularly, we don't allow restarts in the past, and always require the timezone to be present in the request in order to avoid client/server TZ difference issues.

Unify the code that evaluates the schedule for the  scheduled restarts and scheduled failovers.
Use the RLock instead of Lock in the async_executor to avoid hanging if the thread takes the lock multiple times (mostly for the with blocks in the api).

features/patroni_api.feature

@@ -77,3 +77,16 @@ Scenario: check the scheduled failover
 	And postgres0 role is the primary after 5 seconds
 	And postgres1 role is the secondary after 10 seconds
 	And replication works from postgres0 to postgres1 after 25 seconds
+
+Scenario: check the scheduled restart
+	Given I issue a PATCH request to http://127.0.0.1:8008/config with {"postgresql": {"parameters": {"superuser_reserved_connections": "6"}}}
+	Then I receive a response code 200
+		And Response on GET http://127.0.0.1:8008/patroni contains pending_restart after 5 seconds
+	Given I issue a scheduled restart at http://127.0.0.1:8008 in 1 seconds with {"role": "replica"}
+	Then I receive a response code 202
+		And I sleep for 10 seconds
+		And Response on GET http://127.0.0.1:8008/patroni contains pending_restart after 10 seconds
+	Given I issue a scheduled restart at http://127.0.0.1:8008 in 1 seconds with {"restart_pending": "True"}
+	Then I receive a response code 202
+		And Response on GET http://127.0.0.1:8008/patroni does not contain pending_restart after 10 seconds
+

features/steps/patroni_api.py

@@ -121,18 +121,30 @@ def scheduled_failover(context, from_host, to_host, in_seconds):
     """.format(from_host, to_host, datetime.now(pytz.utc) + timedelta(seconds=int(in_seconds))))
 
 
+@step('I issue a scheduled restart at {url:url} in {in_seconds:d} seconds with {data}')
+def scheduled_restart(context, url, in_seconds, data):
+    data = data and json.loads(data) or {}
+    data.update(schedule='{0}'.format((datetime.now(pytz.utc) + timedelta(seconds=int(in_seconds))).isoformat()))
+    context.execute_steps(u"""Given I issue a POST request to {0}/restart with {1}""".format(url, json.dumps(data)))
+
+
 @step('I add tag {tag:w} {value:w} to {pg_name:w} config')
 def add_tag_to_config(context, tag, value, pg_name):
     context.pctl.add_tag_to_config(pg_name, tag, value)
 
 
 @then('Response on GET {url} contains {value} after {timeout:d} seconds')
-def check_http_response(context, url, value, timeout):
+def check_http_response(context, url, value, timeout, negate=False):
     for _ in range(int(timeout)):
         r = requests.get(url)
-        if value in r.content.decode('utf-8'):
+        if (value in r.content.decode('utf-8')) != negate:
             break
         time.sleep(1)
     else:
         assert False,\
-            "Value {0} is not present in response after {1} seconds".format(value, timeout)
+            "Value {0} is {1} present in response after {2} seconds".format(value, "not" if not negate else "", timeout)
+
+
+@then('Response on GET {url} does not contain {value} after {timeout:d} seconds')
+def check_not_in_http_response(context, url, value, timeout):
+    check_http_response(context, url, value, timeout, negate=True)

patroni/__init__.py

@@ -32,6 +32,7 @@ class Patroni(object):
         self.tags = self.get_tags()
         self.nap_time = self.config['loop_wait']
         self.next_run = time.time()
+        self.scheduled_restart = {}
 
     def load_dynamic_configuration(self):
         while True:

patroni/api.py

@@ -7,6 +7,7 @@ import time
 import dateutil.parser
 import datetime
 import pytz
+import re
 
 from patroni.exceptions import PostgresConnectionException
 from patroni.utils import deep_compare, patch_config, Retry, RetryFailedError
@@ -63,6 +64,10 @@ class RestApiHandler(BaseHTTPRequestHandler):
         if patroni.postgresql.pending_restart:
             response['pending_restart'] = True
         response['patroni'] = {'version': patroni.version, 'scope': patroni.postgresql.scope}
+        if patroni.scheduled_restart and isinstance(patroni.scheduled_restart, dict):
+            response['scheduled_restart'] = patroni.scheduled_restart.copy()
+            del response['scheduled_restart']['postmaster_start_time']
+            response['scheduled_restart']['schedule'] = (response['scheduled_restart']['schedule']).isoformat()
         self._write_json_response(status_code, response)
 
     def do_GET(self, write_status_code_only=False):
@@ -112,13 +117,15 @@ class RestApiHandler(BaseHTTPRequestHandler):
         else:
             self.send_error(502)
 
-    def _read_json_content(self):
+    def _read_json_content(self, body_is_optional=False):
         if 'content-length' not in self.headers:
-            return self.send_error(411)
+            return self.send_error(411) if not body_is_optional else {}
         try:
             content_length = int(self.headers.get('content-length'))
+            if content_length == 0 and body_is_optional:
+                return {}
             request = json.loads(self.rfile.read(content_length).decode('utf-8'))
-            if isinstance(request, dict) and request:
+            if isinstance(request, dict) and (request or body_is_optional):
                 return request
         except Exception:
             logger.exception('Bad request')
@@ -162,17 +169,85 @@ class RestApiHandler(BaseHTTPRequestHandler):
             response = str(e)
         self._write_response(status_code, response)
 
+    @staticmethod
+    def parse_schedule(schedule, action):
+        """ parses the given schedule and validates at """
+        error = None
+        scheduled_at = None
+        try:
+            scheduled_at = dateutil.parser.parse(schedule)
+            if scheduled_at.tzinfo is None:
+                error = 'Timezone information is mandatory for the scheduled {0}'.format(action)
+                status_code = 400
+            elif scheduled_at < datetime.datetime.now(pytz.utc):
+                error = 'Cannot schedule {0} in the past'.format(action)
+                status_code = 422
+            else:
+                status_code = None
+        except (ValueError, TypeError):
+            logger.exception('Invalid scheduled %s time: %s', action, schedule)
+            error = 'Unable to parse scheduled timestamp. It should be in an unambiguous format, e.g. ISO 8601'
+            status_code = 422
+        return (status_code, error, scheduled_at)
+
     @check_auth
     def do_POST_restart(self):
         status_code = 500
         data = 'restart failed'
-        try:
-            status, data = self.server.patroni.ha.restart()
-            status_code = 200 if status else 503
-        except Exception:
-            logger.exception('Exception during restart')
+        request = self._read_json_content(body_is_optional=True)
+        if request is None:
+            # failed to parse the json
+            return
+        if request:
+            logger.debug("received restart request: {0}".format(request))
+        for k in request:
+            if k == 'schedule':
+                (_, data, request[k]) = self.parse_schedule(request[k], "restart")
+                if _:
+                    status_code = _
+                    break
+            elif k == 'role':
+                if request[k] not in ('master', 'replica'):
+                    status_code = 400
+                    data = "PostgreSQL role should be either master or replica"
+                    break
+            elif k == 'postgres_version':
+                if not re.match(r'[1-9][0-9]?(\.(0|([1-9][0-9]?))){2}$', request[k]):
+                    status_code = 400
+                    data = "PostgreSQL version should be in the first.major.minor format"
+                    break
+            elif k != 'restart_pending':
+                status_code = 400
+                data = "Unknown filter for the scheduled restart: {0}".format(k)
+                break
+        else:
+            if 'schedule' not in request:
+                try:
+                    status, data = self.server.patroni.ha.restart(request)
+                    status_code = 200 if status else 503
+                except Exception:
+                    logger.exception('Exception during restart')
+                    status_code = 400
+            else:
+                request['postmaster_start_time'] = self.server.patroni.ha.state_handler.postmaster_start_time()
+                if self.server.patroni.ha.schedule_future_restart(request):
+                    data = "Restart scheduled"
+                    status_code = 202
+                else:
+                    data = "Another restart is already scheduled"
+                    status_code = 409
         self._write_response(status_code, data)
 
+    @check_auth
+    def do_DELETE_restart(self):
+        if self.server.patroni.ha.delete_future_restart():
+            data = "scheduled restart deleted"
+            code = 200
+        else:
+            data = "no restarts are scheduled"
+            code = 404
+        self._write_response(code, data)
+
     @check_auth
     def do_POST_reinitialize(self):
         ha = self.server.patroni.ha
@@ -244,25 +319,16 @@ class RestApiHandler(BaseHTTPRequestHandler):
         data = ''
         if leader or candidate:
             if scheduled_at:
-                try:
-                    scheduled_at = dateutil.parser.parse(scheduled_at)
-                    if scheduled_at.tzinfo is None:
-                        data = 'Timezone information is mandatory for scheduled_at'
-                        status_code = 400
-                    elif scheduled_at < datetime.datetime.now(pytz.utc):
-                        data = 'Cannot schedule failover in the past'
-                        status_code = 422
-                    elif self.server.patroni.dcs.manual_failover(leader, candidate, scheduled_at=scheduled_at):
-                        self.server.patroni.dcs.event.set()
-                        data = 'Failover scheduled'
-                        status_code = 202
-                    else:
-                        data = 'failed to write failover key into DCS'
-                        status_code = 503
-                except (ValueError, TypeError):
-                    logger.exception('Invalid scheduled failover time: %s', request['scheduled_at'])
-                    data = 'Unable to parse scheduled timestamp. It should be in an unambiguous format, e.g. ISO 8601'
-                    status_code = 422
+                (_, data, scheduled_at) = self.parse_schedule(scheduled_at, "failover")
+                if _:
+                    status_code = _
+                elif self.server.patroni.dcs.manual_failover(leader, candidate, scheduled_at=scheduled_at):
+                    self.server.patroni.dcs.event.set()
+                    data = 'Failover scheduled'
+                    status_code = 202
+                else:
+                    data = 'failed to write failover key into DCS'
+                    status_code = 503
             else:
                 data = self.is_failover_possible(cluster, leader, candidate)
                 if not data:

patroni/async_executor.py

@@ -1,5 +1,5 @@
 import logging
-from threading import Lock, Thread
+from threading import RLock, Thread
 
 logger = logging.getLogger(__name__)
 
@@ -8,9 +8,9 @@ class AsyncExecutor(object):
 
     def __init__(self):
         self._busy = False
-        self._thread_lock = Lock()
+        self._thread_lock = RLock()
         self._scheduled_action = None
-        self._scheduled_action_lock = Lock()
+        self._scheduled_action_lock = RLock()
 
     @property
     def busy(self):
@@ -32,6 +32,7 @@ class AsyncExecutor(object):
     def reset_scheduled_action(self):
         with self._scheduled_action_lock:
             self._scheduled_action = None
+            self._busy = False
 
     def run(self, func, args=()):
         try:

patroni/ctl.py

@@ -127,7 +127,7 @@ def post_patroni(member, endpoint, content, headers=None):
         headers['Content-Type'] = 'application/json'
     return requests.post('{0}://{1}/{2}'.format(url.scheme, url.netloc, endpoint),
                          headers=headers,
-                         data=json.dumps(content), timeout=60)
+                         data=json.dumps(content) if content else None, timeout=60)
 
 
 def print_output(columns, rows=None, alignment=None, fmt='pretty', header=True, delimiter='\t'):

patroni/ha.py

@@ -66,6 +66,11 @@ class Ha(object):
                 data['xlog_location'] = self.state_handler.xlog_position()
             except:
                 pass
+        if self.patroni.scheduled_restart:
+            scheduled_restart_data = self.patroni.scheduled_restart.copy()
+            scheduled_restart_data['schedule'] = scheduled_restart_data['schedule'].isoformat()
+            data['scheduled_restart'] = scheduled_restart_data
+
         self.dcs.touch_member(json.dumps(data, separators=(',', ':')))
 
     def clone(self, clone_member=None, msg='(without leader)'):
@@ -168,7 +173,7 @@ class Ha(object):
             xlog_location = None if is_master else json['xlog']['replayed_location']
             return (member, True, not is_master, xlog_location, json.get('tags', {}))
         except:
-            logging.exception('request failed: GET %s', member.api_url)
+            logger.exception('request failed: GET %s', member.api_url)
         return (member, False, None, 0, {})
 
     def fetch_nodes_statuses(self, members):
@@ -279,31 +284,45 @@ class Ha(object):
         else:
             self.state_handler.follow(None, None)
 
-    def process_manual_failover_from_leader(self):
-        failover = self.cluster.failover
-
-        if failover.scheduled_at:
-            # If the failover is in the far future, we shouldn't do anything and just return.
-            # If the failover is in the past, we consider the value to be stale and we remove
+    def should_run_scheduled_action(self, action_name, scheduled_at, cleanup_fn):
+        if scheduled_at:
+            # If the scheduled action is in the far future, we shouldn't do anything and just return.
+            # If the scheduled action is in the past, we consider the value to be stale and we remove
             # the value.
-            # If the value is close to now, we initiate the failover
+            # If the value is close to now, we initiate the scheduled action
+            # Additionally, if the scheduled action cannot be executed altogether, i.e. there is an error
+            # or the action is in the past - we take care of cleaning it up.
             now = datetime.datetime.now(pytz.utc)
             try:
-                delta = (failover.scheduled_at - now).total_seconds()
+                delta = (scheduled_at - now).total_seconds()
 
                 if delta > self.patroni.nap_time:
-                    logging.info('Awaiting failover at %s (in %.0f seconds)', failover.scheduled_at.isoformat(), delta)
-                    return
+                    logger.info('Awaiting %s at %s (in %.0f seconds)',
+                                action_name, scheduled_at.isoformat(), delta)
+                    return False
                 elif delta < - int(self.patroni.nap_time * 1.5):
-                    logger.warning('Found a stale failover value, cleaning up: %s', failover.scheduled_at)
+                    logger.warning('Found a stale %s value, cleaning up: %s',
+                                   action_name, scheduled_at.isoformat())
+                    cleanup_fn()
                     self.dcs.manual_failover('', '', index=self.cluster.failover.index)
-                    return
+                    return False
 
                 # The value is very close to now
                 sleep(max(delta, 0))
-                logger.info('Manual scheduled failover at {}'.format(failover.scheduled_at.isoformat()))
+                logger.info('Manual scheduled {0} at %s'.format(action_name), scheduled_at.isoformat())
+                return True
             except TypeError:
-                logger.warning('Incorrect value in of scheduled_at: %s', failover.scheduled_at)
+                logger.warning('Incorrect value of scheduled_at: %s', scheduled_at)
+                cleanup_fn()
+        return False
+
+    def process_manual_failover_from_leader(self):
+        failover = self.cluster.failover
+
+        if (failover.scheduled_at and not
+            self.should_run_scheduled_action("failover", failover.scheduled_at, lambda:
+                                             self.dcs.manual_failover('', '', index=self.cluster.failover.index))):
+            return
 
         if not failover.leader or failover.leader == self.state_handler.name:
             if not failover.candidate or failover.candidate != self.state_handler.name:
@@ -361,12 +380,71 @@ class Ha(object):
         return self.follow('demoting self because i do not have the lock and i was a leader',
                            'no action.  i am a secondary and i am following a leader', False)
 
-    def schedule(self, action):
-        with self._async_executor:
-            return self._async_executor.schedule(action)
+    def evaluate_scheduled_restart(self):
+        # restart if we need to
+        restart_data = self.future_restart_scheduled()
+        if restart_data:
+            recent_time = self.state_handler.postmaster_start_time()
+            request_time = restart_data['postmaster_start_time']
+            # check if postmaster start time has changed since the last restart
+            if recent_time and request_time and recent_time != request_time:
+                logger.info("Cancelling scheduled restart: postgres restart has already happened at %s", recent_time)
+                self.delete_future_restart()
+                return None
 
-    def restart_scheduled(self):
-        return self._async_executor.scheduled_action == 'restart'
+        if (restart_data and
+           self.should_run_scheduled_action('restart', restart_data['schedule'], self.delete_future_restart)):
+            try:
+                ret, message = self.restart(restart_data, run_async=True)
+                if not ret:
+                    logger.warning("Scheduled restart: %s", message)
+                    return None
+                return message
+            finally:
+                self.delete_future_restart()
+
+    def restart_matches(self, role, postgres_version, pending_restart):
+        reason_to_cancel = ""
+        # checking the restart filters here seem to be less ugly than moving them into the
+        # run_scheduled_action.
+        if role and role != self.state_handler.role:
+            reason_to_cancel = "host role mismatch"
+
+        if (postgres_version and
+           self.state_handler.postgres_version_to_int(postgres_version) <= int(self.state_handler.server_version)):
+            reason_to_cancel = "postgres version mismatch"
+
+        if pending_restart and not self.state_handler.pending_restart:
+            reason_to_cancel = "pending restart flag is not set"
+
+        if not reason_to_cancel:
+            return True
+        else:
+            logger.info("not proceeding with the restart: %s", reason_to_cancel)
+        return False
+
+    def schedule(self, action, immediate=False):
+        with self._async_executor:
+            return self._async_executor.schedule(action, immediate)
+
+    def schedule_future_restart(self, restart_data):
+        with self._async_executor:
+            if not self.patroni.scheduled_restart:
+                self.patroni.scheduled_restart = restart_data
+                return True
+        return False
+
+    def delete_future_restart(self):
+        ret = False
+        with self._async_executor:
+            if self.patroni.scheduled_restart:
+                self.patroni.scheduled_restart = {}
+                ret = True
+        return ret
+
+    def future_restart_scheduled(self):
+        return self.patroni.scheduled_restart.copy() if (self.patroni.scheduled_restart and
+                                                         isinstance(self.patroni.scheduled_restart, dict)) else None
 
     def schedule_reinitialize(self):
         return self.schedule('reinitialize')
@@ -374,15 +452,32 @@ class Ha(object):
     def reinitialize_scheduled(self):
         return self._async_executor.scheduled_action == 'reinitialize'
 
-    def restart(self):
+    def schedule_restart(self, immediate=False):
+        return self.schedule('restart', immediate)
+
+    def restart_scheduled(self):
+        return self._async_executor.scheduled_action == 'restart'
+
+    def restart(self, restart_data=None, run_async=False):
+        """ conditional and unconditional restart """
+        if (restart_data and isinstance(restart_data, dict) and
+            not self.restart_matches(restart_data.get('role'),
+                                     restart_data.get('postgres_version'),
+                                     ('restart_pending' in restart_data))):
+            return (False, "restart conditions are not satisfied")
+
         with self._async_executor:
-            prev = self._async_executor.schedule('restart', True)
+            prev = self.schedule_restart(immediate=(not run_async))
             if prev is not None:
                 return (False, prev + ' already in progress')
-        if self._async_executor.run(self.state_handler.restart):
-            return (True, 'restarted successfully')
-        else:
-            return (False, 'restart failed')
+            if not run_async:
+                if self._async_executor.run(self.state_handler.restart):
+                    return (True, 'restarted successfully')
+                else:
+                    return (False, 'restart failed')
+            else:
+                self._async_executor.run_async(self.state_handler.restart)
+                return (True, "restart initiated")
 
     def reinitialize(self, cluster):
         self.state_handler.stop('immediate')
@@ -482,6 +577,9 @@ class Ha(object):
                 if self.cluster.is_unlocked():
                     return self.process_unhealthy_cluster()
                 else:
+                    msg = self.evaluate_scheduled_restart()
+                    if msg is not None:
+                        return msg
                     return self.process_healthy_cluster()
             finally:
                 # we might not have a valid PostgreSQL connection here if another thread

patroni/postgresql.py

@@ -858,6 +858,13 @@ $$""".format(name, ' '.join(options)), name, password, password)
             self._replication_slots = [r[0] for r in cursor]
             self._schedule_load_slots = False
 
+    def postmaster_start_time(self):
+        try:
+            cursor = self.query("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ')""")
+            return cursor.fetchone()[0]
+        except psycopg2.Error:
+            return None
+
     def sync_replication_slots(self, cluster):
         if self.use_slots:
             try:
@@ -964,3 +971,37 @@ $$""".format(name, ' '.join(options)), name, password, password)
                 time.sleep(5)
 
         return ret
+
+    @staticmethod
+    def postgres_version_to_int(pg_version):
+        """ Convert the server_version to integer
+
+        >>> Postgresql.postgres_version_to_int('9.5.3')
+        90503
+        >>> Postgresql.postgres_version_to_int('9.3.13')
+        90313
+        >>> Postgresql.postgres_version_to_int('10.1')
+        100001
+        >>> Postgresql.postgres_version_to_int('10')
+        Traceback (most recent call last):
+            ...
+        Exception: Invalid PostgreSQL format: X.Y or X.Y.Z is accepted: 10
+        >>> Postgresql.postgres_version_to_int('a.b.c')
+        Traceback (most recent call last):
+            ...
+        Exception: Invalid PostgreSQL version: a.b.c
+        """
+        components = pg_version.split('.')
+
+        result = []
+        if len(components) < 2 or len(components) > 3:
+            raise Exception("Invalid PostgreSQL format: X.Y or X.Y.Z is accepted: {0}".format(pg_version))
+        if len(components) == 2:
+            # new style verion numbers, i.e. 10.1 becomes 100001
+            components.insert(1, '0')
+        try:
+            result = [c if int(c) > 10 else '0{0}'.format(c) for c in components]
+            result = int(''.join(result))
+        except ValueError:
+            raise Exception("Invalid PostgreSQL version: {0}".format(pg_version))
+        return result

tests/test_api.py

@@ -1,5 +1,7 @@
+import datetime
 import json
 import psycopg2
+import pytz
 import unittest
 
 from mock import Mock, patch
@@ -10,6 +12,10 @@ from six.moves import BaseHTTPServer
 from test_postgresql import psycopg2_connect, MockCursor
 
 
+future_restart_time = datetime.datetime.now(pytz.utc) + datetime.timedelta(days=5)
+postmaster_start_time = datetime.datetime.now(pytz.utc)
+
+
 class MockPostgresql(object):
 
     name = 'test'
@@ -24,6 +30,10 @@ class MockPostgresql(object):
     def connection():
         return psycopg2_connect()
 
+    @staticmethod
+    def postmaster_start_time():
+        return str(postmaster_start_time)
+
 
 class MockHa(object):
 
@@ -42,10 +52,18 @@ class MockHa(object):
     def restart_scheduled():
         return False
 
+    @staticmethod
+    def delete_future_restart():
+        return True
+
     @staticmethod
     def fetch_nodes_statuses(members):
         return [[None, True, None, None, {}]]
 
+    @staticmethod
+    def schedule_future_restart(data):
+        return True
+
 
 class MockPatroni(object):
 
@@ -57,6 +75,8 @@ class MockPatroni(object):
     tags = {}
     version = '0.00'
     noloadbalance = Mock(return_value=False)
+    scheduled_restart = {'schedule': future_restart_time,
+                         'postmaster_start_time': postgresql.postmaster_start_time()}
 
     @staticmethod
     def sighup_handler():
@@ -167,6 +187,46 @@ class TestRestApiHandler(unittest.TestCase):
         with patch.object(MockHa, 'restart', Mock(side_effect=Exception)):
             MockRestApiServer(RestApiHandler, request)
 
+        post = request + '\nContent-Length: '
+
+        def make_request(request):
+            return '{0}{1}\n\n{2}'.format(post, len(request), request)
+
+        # empty request
+        request = make_request('')
+        MockRestApiServer(RestApiHandler, request)
+        # invalid request
+        request = make_request('foobar=baz')
+        MockRestApiServer(RestApiHandler, request)
+        # wrong role
+        request = make_request('{"schedule": "2016-08-20 12:45TZ+1", "role": "unknown", "postgres_version": "9.5.3"}')
+        MockRestApiServer(RestApiHandler, request)
+        # wrong version
+        request = make_request('{"schedule": "2016-08-20 12:45TZ+1", "role": "master", "postgres_version": "9.5.3.1"}')
+        MockRestApiServer(RestApiHandler, request)
+        # unknown filter
+        request = make_request('{"schedule": "2016-08-29 12:45TZ+1", "batman": "lives"}')
+        MockRestApiServer(RestApiHandler, request)
+        # incorrect schedule
+        request = make_request('{"schedule": "2016-08-42 12:45TZ+1", "role": "master"}')
+        MockRestApiServer(RestApiHandler, request)
+        # everything fine, but the schedule is missing
+        request = make_request('{"role": "master", "postgres_version": "9.5.2"}')
+        MockRestApiServer(RestApiHandler, request)
+        for retval in (True, False):
+            with patch.object(MockHa, 'schedule_future_restart', Mock(return_value=retval)):
+                request = make_request('{"schedule": "2016-08-29 12:45TZ+1"}')
+                MockRestApiServer(RestApiHandler, request)
+            with patch.object(MockHa, 'restart', Mock(return_value=(retval, "foo"))):
+                request = make_request('{"role": "master", "postgres_version": "9.5.2"}')
+                MockRestApiServer(RestApiHandler, request)
+
+    def test_do_DELETE_restart(self):
+        for retval in (True, False):
+            with patch.object(MockHa, 'delete_future_restart', Mock(return_value=retval)):
+                request = 'DELETE /restart HTTP/1.0' + self._authorization
+                self.assertIsNotNone(MockRestApiServer(RestApiHandler, request))
+
     @patch.object(MockHa, 'dcs')
     def test_do_POST_reinitialize(self, dcs):
         cluster = dcs.get_cluster.return_value

tests/test_ha.py

@@ -46,6 +46,9 @@ def get_cluster_initialized_with_only_leader(failover=None):
     l = get_cluster_initialized_without_leader(leader=True, failover=failover).leader
     return get_cluster(True, l, [l], failover)
 
+future_restart_time = datetime.datetime.now(pytz.utc) + datetime.timedelta(days=5)
+postmaster_start_time = datetime.datetime.now(pytz.utc)
+
 
 class MockPatroni(object):
 
@@ -80,6 +83,8 @@ zookeeper:
         self.replicatefrom = None
         self.api.connection_string = 'http://127.0.0.1:8008'
         self.clonefrom = None
+        self.scheduled_restart = {'schedule': future_restart_time,
+                                  'postmaster_start_time': str(postmaster_start_time)}
 
 
 def run_async(func, args=()):
@@ -117,6 +122,7 @@ class TestHa(unittest.TestCase):
                                                 'hot_standby': 'on', 'max_wal_senders': 5, 'wal_keep_segments': 8}})
             self.p.set_state('running')
             self.p.set_role('replica')
+            self.p.postmaster_start_time = MagicMock(return_value=str(postmaster_start_time))
             self.p.check_replication_lag = true
             self.p.can_create_replica_without_replication_connection = MagicMock(return_value=False)
             self.e = get_dcs({'etcd': {'ttl': 30, 'host': 'ok:2379', 'scope': 'test',
@@ -279,6 +285,8 @@ class TestHa(unittest.TestCase):
         self.assertEquals(self.ha.restart(), (False, 'restart failed'))
         self.ha.schedule_reinitialize()
         self.assertEquals(self.ha.restart(), (False, 'reinitialize already in progress'))
+        with patch.object(self.ha, "restart_matches", return_value=False):
+            self.assertEquals(self.ha.restart({'foo': 'bar'}), (False, "restart conditions are not satisfied"))
 
     def test_restart_in_progress(self):
         self.ha._async_executor.schedule('restart', True)
@@ -395,3 +403,48 @@ class TestHa(unittest.TestCase):
         self.assertEqual(self.ha.post_recover(), 'failed to start postgres')
         self.p.is_running = true
         self.assertIsNone(self.ha.post_recover())
+
+    def test_schedule_future_restart(self):
+        self.ha.patroni.scheduled_restart = {}
+        # do the restart 2 times. The first one should succeed, the second one should fail
+        self.assertTrue(self.ha.schedule_future_restart({'schedule': str(future_restart_time)}))
+        self.assertFalse(self.ha.schedule_future_restart({'schedule': str(future_restart_time)}))
+
+    def test_delete_future_restarts(self):
+        self.ha.delete_future_restart()
+
+    def test_evaluate_scheduled_restart(self):
+        self.p.postmaster_start_time = Mock(return_value=str(postmaster_start_time))
+        # restart while the postmaster has been already restarted, fails
+        with patch.object(self.ha,
+                          'future_restart_scheduled',
+                          Mock(return_value={'postmaster_start_time':
+                                             str(postmaster_start_time - datetime.timedelta(days=1)),
+                                             'schedule': str(future_restart_time)})):
+            self.assertIsNone(self.ha.evaluate_scheduled_restart())
+        with patch.object(self.ha,
+                          'future_restart_scheduled',
+                          Mock(return_value={'postmaster_start_time': str(postmaster_start_time),
+                                             'schedule': str(future_restart_time)})):
+            with patch.object(self.ha,
+                              'should_run_scheduled_action', Mock(return_value=True)):
+                # restart in the future, ok
+                self.assertIsNotNone(self.ha.evaluate_scheduled_restart())
+                with patch.object(self.ha, 'restart', Mock(return_value=(False, "Test"))):
+                    # restart in the future, bit the actual restart failed
+                    self.assertIsNone(self.ha.evaluate_scheduled_restart())
+
+    def test_scheduled_restart(self):
+        self.ha.cluster = get_cluster_initialized_with_leader()
+        with patch.object(self.ha, "evaluate_scheduled_restart", Mock(return_value="restart scheduled")):
+            self.assertEquals(self.ha.run_cycle(), "restart scheduled")
+
+    def test_restart_matches(self):
+        self.p._role = 'replica'
+        self.p.server_version = 90500
+        self.p._pending_restart = True
+        self.assertFalse(self.ha.restart_matches("master", "9.5.0", True))
+        self.assertFalse(self.ha.restart_matches("replica", "9.4.3", True))
+        self.p._pending_restart = False
+        self.assertFalse(self.ha.restart_matches("replica", "9.5.2", True))
+        self.assertTrue(self.ha.restart_matches("replica", "9.5.2", False))

tests/test_postgresql.py

@@ -537,3 +537,9 @@ class TestPostgresql(unittest.TestCase):
             self.assertEquals(self.p.get_major_version(), 9.4)
         with patch.object(builtins, 'open', Mock(side_effect=Exception)):
             self.assertEquals(self.p.get_major_version(), 0.0)
+
+    def test_postmaster_start_time(self):
+        with patch.object(MockCursor, "fetchone", Mock(return_value=('foo', True, '', '', '', '', False))):
+            self.assertEqual(self.p.postmaster_start_time(), 'foo')
+        with patch.object(MockCursor, "execute", side_effect=psycopg2.Error):
+            self.assertIsNone(self.p.postmaster_start_time())