Implement scheduling and deleting a restart.

The scheduled restart API extends the already existing restart
endpoint by processing the parameters in the request body.

Only one scheduled restart at a time is support. DELETE method
on the /restart endpoint is used to remove an existing restart.

patroni/__init__.py

@@ -29,6 +29,7 @@ class Patroni(object):
         self.tags = self.get_tags()
         self.nap_time = self.config['loop_wait']
         self.next_run = time.time()
+        self.scheduled_restart = {}
 
         self._reload_config_scheduled = False
         self._received_sighup = False

patroni/api.py

@@ -7,6 +7,7 @@ import time
 import dateutil
 import datetime
 import pytz
+import re
 
 from patroni.exceptions import PostgresConnectionException
 from patroni.utils import deep_compare, patch_config, Retry, RetryFailedError
@@ -63,6 +64,8 @@ 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
         self._write_json_response(status_code, response)
 
     def do_GET(self, write_status_code_only=False):
@@ -86,7 +89,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
                 status_code = 503
         elif 'role' in response and response['role'] in path:
             status_code = 503 if response['role'] != 'master' and patroni.noloadbalance else 200
-        elif patroni.ha.restart_scheduled() and patroni.postgresql.role == 'master' and 'master' in path:
+        elif patroni.ha.immediate_restart_scheduled() and patroni.postgresql.role == 'master' and 'master' in path:
             # exceptional case for master node when the postgres is being restarted via API
             status_code = 200
         else:
@@ -112,9 +115,9 @@ 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 body_is_optional else None
         try:
             content_length = int(self.headers.get('content-length'))
             request = json.loads(self.rfile.read(content_length).decode('utf-8'))
@@ -162,17 +165,80 @@ 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()
+        if not request:  # unconditional restart
+            try:
+                status, data = self.server.patroni.ha.restart()
+                status_code = 200 if status else 503
+            except Exception:
+                logger.exception('Exception during restart')
+        else:
+            logger.info("received scheduled 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]?\.){2}[1-9][0-9]?$', request[k]):
+                        status_code = 400
+                        data = "PostgreSQL version should be in the first.major.minor format"
+                        break
+                else:
+                    status_code = 400
+                    data = "Unknown filter for the scheduled restart: {0}".format(k)
+                    break
+            else:
+                if 'schedule' not in request:
+                    data = "Schedule required for the scheduled restart"
+                    status_code = 400
+                else:
+                    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):
+        self.server.patroni.ha.delete_future_restart()
+        data = "scheduled restart deleted"
+        code = 200
+        self._write_response(code, data)
+
     @check_auth
     def do_POST_reinitialize(self):
         ha = self.server.patroni.ha
@@ -244,25 +310,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/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)'):
@@ -372,7 +377,19 @@ class Ha(object):
         with self._async_executor:
             return self._async_executor.schedule(action)
 
-    def restart_scheduled(self):
+    def schedule_future_restart(self, restart_data):
+        if isinstance(restart_data, dict):
+            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):
+        with self._async_executor:
+            self.patroni.scheduled_restart = {}
+
+    def immediate_restart_scheduled(self):
         return self._async_executor.scheduled_action == 'restart'
 
     def schedule_reinitialize(self):

tests/test_api.py

@@ -1,3 +1,4 @@
+
 import json
 import psycopg2
 import unittest
@@ -39,13 +40,21 @@ class MockHa(object):
         return (True, '')
 
     @staticmethod
-    def restart_scheduled():
+    def immediate_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 +66,7 @@ class MockPatroni(object):
     tags = {}
     version = '0.00'
     noloadbalance = Mock(return_value=False)
+    scheduled_restart = {'schedule': '2016-08-29 12:45TZ+1'}
 
     @staticmethod
     def sighup_handler():
@@ -104,7 +114,7 @@ class TestRestApiHandler(unittest.TestCase):
         MockPatroni.dcs.cluster = None
         with patch.object(RestApiHandler, 'get_postgresql_status', Mock(return_value={'role': 'master'})):
             MockRestApiServer(RestApiHandler, 'GET /master')
-        with patch.object(MockHa, 'restart_scheduled', Mock(return_value=True)):
+        with patch.object(MockHa, 'immediate_restart_scheduled', Mock(return_value=True)):
             MockRestApiServer(RestApiHandler, 'GET /master')
         self.assertIsNotNone(MockRestApiServer(RestApiHandler, 'GET /master'))
 
@@ -167,6 +177,33 @@ class TestRestApiHandler(unittest.TestCase):
         with patch.object(MockHa, 'restart', Mock(side_effect=Exception)):
             MockRestApiServer(RestApiHandler, request)
 
+        post = request + '\nContent-Length: '
+        # wrong version
+        request = post + '84\n\n{"schedule": "2016-08-20 12:45TZ+1", "role": "unknown", "postgres_version": "9.5.3"}'
+        MockRestApiServer(RestApiHandler, request)
+        # wrong role
+        request = post + '85\n\n{"schedule": "2016-08-20 12:45TZ+1", "role": "master", "postgres_version": "9.5.3.1"}'
+        MockRestApiServer(RestApiHandler, request)
+        # unknown filter
+        request = post + '55\n\n{"schedule": "2016-08-29 12:45TZ+1", "batman": "lives"}'
+        MockRestApiServer(RestApiHandler, request)
+        # incorrect schedule
+        request = post + '55\n\n{"schedule": "2016-08-42 12:45TZ+1", "role": "master"}'
+        MockRestApiServer(RestApiHandler, request)
+        # everything fine, but the schedule is missing
+        request = post + '47\n\n{"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 = post + '36\n\n{"schedule": "2016-08-29 12:45TZ+1"}'
+                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

@@ -1,4 +1,5 @@
 import datetime
+import dateutil
 import etcd
 import os
 import pytz
@@ -80,6 +81,7 @@ zookeeper:
         self.replicatefrom = None
         self.api.connection_string = 'http://127.0.0.1:8008'
         self.clonefrom = None
+        self.scheduled_restart = {'schedule': dateutil.parser.parse('2016-08-29 12:45TZ+1')}
 
 
 def run_async(func, args=()):
@@ -282,7 +284,7 @@ class TestHa(unittest.TestCase):
 
     def test_restart_in_progress(self):
         self.ha._async_executor.schedule('restart', True)
-        self.assertTrue(self.ha.restart_scheduled())
+        self.assertTrue(self.ha.immediate_restart_scheduled())
         self.assertEquals(self.ha.run_cycle(), 'not healthy enough for leader race')
 
         self.ha.cluster = get_cluster_initialized_with_leader()
@@ -395,3 +397,11 @@ 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 = {}
+        self.ha.schedule_future_restart({'schedule': '2016-08-30 12:45TZ+1"'})
+        self.ha.schedule_future_restart({'schedule': '2016-08-30 12:45TZ+1"'})
+
+    def test_delete_future_restarts(self):
+        self.ha.delete_future_restart()