Various bugfixes and improvements: (#346)

* Replace pytz.UTC with dateutil.tz.tzutc, it helps to reduce memory by more than 4Mb...

* fix check of python version: 0x0300000 => 0x3000000

* Update leader key before restart and demote

features/steps/patroni_api.py

@@ -1,7 +1,6 @@
 import json
 import os
 import parse
-import pytz
 import requests
 import shlex
 import subprocess
@@ -9,8 +8,11 @@ import time
 import yaml
 
 from behave import register_type, step, then
+from dateutil import tz
 from datetime import datetime, timedelta
 
+tzutc = tz.tzutc()
+
 
 @parse.with_pattern(r'https?://(?:\w|\.|:|/)+')
 def parse_url(text):
@@ -123,13 +125,13 @@ def check_response(context, component, data):
 def scheduled_failover(context, from_host, to_host, in_seconds):
     context.execute_steps(u"""
         Given I run patronictl.py failover batman --master {0} --candidate {1} --scheduled "{2}" --force
-    """.format(from_host, to_host, datetime.now(pytz.utc) + timedelta(seconds=int(in_seconds))))
+    """.format(from_host, to_host, datetime.now(tzutc) + 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()))
+    data.update(schedule='{0}'.format((datetime.now(tzutc) + 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)))
 
 

patroni/api.py

@@ -6,10 +6,9 @@ import psycopg2
 import time
 import dateutil.parser
 import datetime
-import pytz
 
 from patroni.exceptions import PostgresConnectionException
-from patroni.utils import deep_compare, patch_config, Retry, RetryFailedError, is_valid_pg_version
+from patroni.utils import deep_compare, patch_config, Retry, RetryFailedError, is_valid_pg_version, tzutc
 from six.moves.BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
 from six.moves.socketserver import ThreadingMixIn
 from threading import Thread
@@ -180,7 +179,7 @@ class RestApiHandler(BaseHTTPRequestHandler):
             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):
+            elif scheduled_at < datetime.datetime.now(tzutc):
                 error = 'Cannot schedule {0} in the past'.format(action)
                 status_code = 422
             else:

patroni/ha.py

@@ -5,14 +5,13 @@ import psycopg2
 import requests
 import sys
 import datetime
-import pytz
 from threading import RLock
 
 from multiprocessing.pool import ThreadPool
 from patroni.async_executor import AsyncExecutor
 from patroni.exceptions import DCSError, PostgresConnectionException
 from patroni.postgresql import ACTION_ON_START
-from patroni.utils import polling_loop, sleep
+from patroni.utils import polling_loop, sleep, tzutc
 
 logger = logging.getLogger(__name__)
 
@@ -445,7 +444,7 @@ class Ha(object):
             # 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)
+            now = datetime.datetime.now(tzutc)
             try:
                 delta = (scheduled_at - now).total_seconds()
 
@@ -469,7 +468,14 @@ class Ha(object):
         return False
 
     def process_manual_failover_from_leader(self):
+        """Checks if manual failover is requested and takes action if appropriate.
+
+        Cleans up failover key if failover conditions are not matched.
+
+        :returns: action message if demote was initiated, None if no action was taken"""
         failover = self.cluster.failover
+        if not failover or (self.is_paused() and not self.state_handler.is_leader()):
+            return
 
         if (failover.scheduled_at and not
             self.should_run_scheduled_action("failover", failover.scheduled_at, lambda:
@@ -534,11 +540,6 @@ class Ha(object):
 
     def process_healthy_cluster(self):
         if self.has_lock():
-            if self.cluster.failover and (not self.is_paused() or self.state_handler.is_leader()):
-                msg = self.process_manual_failover_from_leader()
-                if msg is not None:
-                    return msg
-
             if self.is_paused() and not self.state_handler.is_leader():
                 if self.cluster.failover and self.cluster.failover.candidate == self.state_handler.name:
                     return 'waiting to become master after promote...'
@@ -548,6 +549,10 @@ class Ha(object):
                 return 'removed leader lock because postgres is not running as master'
 
             if self.update_lock(True):
+                msg = self.process_manual_failover_from_leader()
+                if msg is not None:
+                    return msg
+
                 return self.enforce_master_role('no action.  i am the leader with the lock',
                                                 'promoted self to leader because i had the session lock')
             else:
@@ -561,6 +566,9 @@ class Ha(object):
                            'no action.  i am a secondary and i am following a leader', False)
 
     def evaluate_scheduled_restart(self):
+        if self._async_executor.busy:  # Restart already in progress
+            return None
+
         # restart if we need to
         restart_data = self.future_restart_scheduled()
         if restart_data:
@@ -758,10 +766,8 @@ 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()
+                    msg = self.process_healthy_cluster()
+                    return self.evaluate_scheduled_restart() or msg
             finally:
                 # we might not have a valid PostgreSQL connection here if another thread
                 # stops PostgreSQL, therefore, we only reload replication slots if no

patroni/scripts/wale_restore.py

@@ -32,8 +32,7 @@ import subprocess
 import sys
 import argparse
 
-
-if sys.hexversion >= 0x0300000:
+if sys.hexversion >= 0x3000000:
     long = int
 
 logger = logging.getLogger(__name__)

patroni/utils.py

@@ -4,11 +4,13 @@ import sys
 import time
 import re
 
+from dateutil import tz
 from patroni.exceptions import PatroniException
 
-if sys.hexversion >= 0x0300000:
+if sys.hexversion >= 0x3000000:
     long = int
 
+tzutc = tz.tzutc()
 __interrupted_sleep = False
 __reap_children = False
 

tests/test_api.py

@@ -1,19 +1,19 @@
 import datetime
 import json
 import psycopg2
-import pytz
 import unittest
 
 from mock import Mock, patch
 from patroni.api import RestApiHandler, RestApiServer
 from patroni.dcs import ClusterConfig, Member
+from patroni.utils import tzutc
 from six import BytesIO as IO
 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)
+future_restart_time = datetime.datetime.now(tzutc) + datetime.timedelta(days=5)
+postmaster_start_time = datetime.datetime.now(tzutc)
 
 
 class MockPostgresql(object):

tests/test_consul.py

@@ -157,6 +157,8 @@ class TestConsul(unittest.TestCase):
     def test_set_retry_timeout(self):
         self.c.set_retry_timeout(10)
 
+    @patch.object(consul.Consul.KV, 'delete', Mock(return_value=True))
+    @patch.object(consul.Consul.KV, 'put', Mock(return_value=True))
     def test_sync_state(self):
-        self.assertFalse(self.c.set_sync_state_value('{}'))
-        self.assertFalse(self.c.delete_sync_state())
+        self.assertTrue(self.c.set_sync_state_value('{}'))
+        self.assertTrue(self.c.delete_sync_state())

tests/test_ctl.py

@@ -73,6 +73,7 @@ class TestCtl(unittest.TestCase):
     def test_failover(self, mock_get_dcs):
         mock_get_dcs.return_value = self.e
         mock_get_dcs.return_value.get_cluster = get_cluster_initialized_with_leader
+        mock_get_dcs.return_value.set_failover_value = Mock()
         result = self.runner.invoke(ctl, ['failover', 'dummy'], input='leader\nother\n\ny')
         assert 'leader' in result.output
 
@@ -360,6 +361,7 @@ class TestCtl(unittest.TestCase):
         mock_get_dcs.return_value.initialize = Mock(return_value=True)
         mock_get_dcs.return_value.touch_member = Mock(return_value=True)
         mock_get_dcs.return_value.attempt_to_acquire_leader = Mock(return_value=True)
+        mock_get_dcs.return_value.delete_cluster = Mock()
 
         with patch.object(self.e, 'initialize', return_value=False):
             result = self.runner.invoke(ctl, ['scaffold', 'alpha'])

tests/test_ha.py

@@ -1,7 +1,6 @@
 import datetime
 import etcd
 import os
-import pytz
 import unittest
 
 from mock import Mock, MagicMock, PropertyMock, patch
@@ -11,6 +10,7 @@ from patroni.dcs.etcd import Client
 from patroni.exceptions import DCSError, PostgresException
 from patroni.ha import Ha
 from patroni.postgresql import Postgresql
+from patroni.utils import tzutc
 from test_etcd import socket_getaddrinfo, etcd_read, etcd_write, requests_get
 from test_postgresql import psycopg2_connect
 
@@ -53,8 +53,8 @@ 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, None)
 
-future_restart_time = datetime.datetime.now(pytz.utc) + datetime.timedelta(days=5)
-postmaster_start_time = datetime.datetime.now(pytz.utc)
+future_restart_time = datetime.datetime.now(tzutc) + datetime.timedelta(days=5)
+postmaster_start_time = datetime.datetime.now(tzutc)
 
 
 class MockPatroni(object):
@@ -343,7 +343,7 @@ class TestHa(unittest.TestCase):
         self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, 'blabla', self.p.name, scheduled))
         self.ha.run_cycle()
 
-        scheduled = datetime.datetime.utcnow().replace(tzinfo=pytz.UTC)
+        scheduled = datetime.datetime.utcnow().replace(tzinfo=tzutc)
         self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, 'blabla', self.p.name, scheduled))
         self.assertEquals('no action.  i am the leader with the lock', self.ha.run_cycle())