Merge pull request #64 from zalando/bugfix/pg_controldata_changes

Work around the differences in pg_controldata names.

Dockerfile

@@ -8,7 +8,7 @@ RUN apt-get update -y && apt-get install curl -y
 
 # Add PGDG repositories
 RUN echo "deb http://apt.postgresql.org/pub/repos/apt/ $(lsb_release -cs)-pgdg main" > /etc/apt/sources.list.d/pgdg.list
-RUN curl https://www.postgresql.org/media/keys/ACCC4CF8.asc | sudo apt-key add -
+RUN curl https://www.postgresql.org/media/keys/ACCC4CF8.asc | apt-key add -
 RUN apt-get update -y
 RUN apt-get upgrade -y
 
@@ -18,16 +18,13 @@ RUN pip install python-etcd
 
 ENV PATH /usr/lib/postgresql/${PGVERSION}/bin:$PATH
 
-RUN mkdir -p /patroni/helpers
-RUN mkdir -p /patroni/scripts
-ADD patroni.py /patroni/patroni.py
-ADD helpers /patroni/helpers
-ADD scripts /patroni/scripts
+ADD patroni.py /patroni.py
+ADD patroni/ /patroni
 
 ENV ETCDVERSION 2.0.13
 RUN curl -L https://github.com/coreos/etcd/releases/download/v${ETCDVERSION}/etcd-v${ETCDVERSION}-linux-amd64.tar.gz | tar xz -C /bin --strip=1 --wildcards --no-anchored etcd etcdctl
 
-## Setting up a simple script that will serve as an entrypoint
+### Setting up a simple script that will serve as an entrypoint
 RUN mkdir /data/ && touch /var/log/etcd.log /var/log/etcd.err /pgpass /patroni/postgres.yml
 RUN chown postgres:postgres -R /patroni/ /data/ /pgpass /var/log/etcd.* /patroni/postgres.yml
 ADD docker/entrypoint.sh /entrypoint.sh

MANIFEST.in

@@ -0,0 +1,3 @@
+include requirements*
+include *.rst
+recursive-include patroni *.py

README.md

@@ -1,129 +0,0 @@
-[![Build Status](https://travis-ci.org/zalando/patroni.svg?branch=master)](https://travis-ci.org/zalando/patroni)
-[![Coverage Status](https://coveralls.io/repos/zalando/patroni/badge.svg?branch=master)](https://coveralls.io/r/zalando/patroni?branch=master)
-# Patroni: A Template for PostgreSQL HA with ZooKeeper or etcd
-
-Patroni was previously known as Governor.
-
-*There are many ways to run high availability with PostgreSQL; here we present a template for you to create your own custom fit high availability solution using python and distributed configuration store (like ZooKeeper or etcd) for maximum accessibility.*
-
-## Getting Started
-To get started, do the following from different terminals:
-
-```
-> etcd --data-dir=data/etcd
-> ./patroni.py postgres0.yml
-> ./patroni.py postgres1.yml
-```
-
-From there, you will see a high-availability cluster start up. Test
-different settings in the YAML files to see how behavior changes.  Kill
-some of the different components to see how the system behaves.
-
-Add more `postgres*.yml` files to create an even larger cluster.
-
-We provide a haproxy configuration, which will give your application a single endpoint for connecting to the cluster's leader.  To configure, run:
-
-```
-> haproxy -f haproxy.cfg
-```
-
-```
-> psql --host 127.0.0.1 --port 5000 postgres
-```
-
-## How Patroni works
-
-For a diagram of the high availability decision loop, see the included a PDF: [postgres-ha.pdf](https://github.com/zalando/patroni/blob/master/postgres-ha.pdf)
-
-## YAML Configuration
-
-For an example file, see `postgres0.yml`.  Below is an explanation of settings:
-
-* *ttl*: the TTL to acquire the leader lock.  Think of it as the length of time before automatic failover process is initiated.
-* *loop_wait*: the number of seconds the loop will sleep
-
-* *restapi*
-  * *listen*: ip address + port that Patroni will listen to provide health-check information for haproxy.
-  * *connect_address*: ip address + port through which restapi is accessible.
-
-* *etcd*
-  * *scope*: the relative path used on etcd's http api for this deployment, thus you can run multiple HA deployments from a single etcd
-  * *ttl*: the TTL to acquire the leader lock.  Think of it as the length of time before automatic failover process is initiated.
-  * *host*: the host:port for the etcd endpoint
-
-* *zookeeper*
-  * *scope*: the relative path used on etcd's http api for this deployment, thus you can run multiple HA deployments from a single etcd
-  * *session_timeout*: the TTL to acquire the leader lock. Think of it as the length of time before automatic failover process is initiated.
-  * *reconnects_timeout*: how long we should try to reconnect to ZooKeeper after connection loss. After this timeout we assume that we don't have lock anymore and will restart in read-only mode.
-  * *hosts*: list of ZooKeeper cluster members in format: [ 'host1:port1', 'host2:port2', 'etc...']
-  * *exhibitor*: if you are running ZooKeeper cluster under Exhibitor supervisory the following section could be interesting for you
-    * *poll_interval*: how often list of ZooKeeper and Exhibitor nodes should be updated from Exhibitor
-    * *port*: Exhibitor port
-    * *hosts*: initial list of Exhibitor (ZooKeeper) nodes in format: [ 'host1', 'host2', 'etc...' ]. This list would be updated automatically when Exhibitor (ZooKeeper) cluster topology changes.
-
-* *postgresql*
-  * *name*: the name of the Postgres host, must be unique for the cluster
-  * *listen*: ip address + port that Postgres listening. Must be accessible from other nodes in the cluster if using streaming replication.
-  * *connect_address*: ip address + port through which Postgres is accessible from other nodes and applications.
-  * *data_dir*: file path to initialize and store Postgres data files
-  * *maximum_lag_on_failover*: the maximum bytes a follower may lag before it is not eligible become leader
-  * *pg_hba*: list of lines which should be added to pg_hba.conf
-    * *- host all all 0.0.0.0/0 md5*
-  * *replication*
-    * *username*: replication username, user will be created during initialization
-    * *password*: replication password, user will be created during initialization
-    * *network*: network setting for replication in pg_hba.conf
-  * *callbacks* callback scripts to run on certain actions. Patroni will pass current action, role and cluster name. See scripts/aws.py as an example on how to write them.
-    * *on_start*: a script to run when the cluster starts
-    * *on_stop*: a script to run when the cluster stops
-    * *on_restart*: a script to run when the cluster restarts
-    * *on_reload*: a script to run when configuration reload is triggered
-    * *on_role_change*: a script to run when the cluster is being promoted or demoted
-  * *superuser*
-    * *password*: password for postgres user. It would be set during initialization
-  * *admin*:
-    * *username*: admin username, user will be created during initialization. It would have CREATEDB and CREATEROLE privileges
-    * *password*: admin password, user will be created during initialization.
-  * *recovery_conf*: configuration settings written to recovery.conf when configuring follower
-  * *parameters*: list of configuration settings for Postgres
-
-## Replication choices
-
-Patroni uses Postgres' streaming replication.  By default, this replication is asynchronous.  For more information, see the [Postgres documentation on streaming replication](http://www.postgresql.org/docs/current/static/warm-standby.html#STREAMING-REPLICATION). 
-
-Patroni's asynchronous replication configuration allows for `maximum_lag_on_failover` settings. This setting ensures failover will not occur if a follower is more than a certain number of bytes behind the follower.  This setting should be increased or decreased based on business requirements.
-
-When asynchronous replication is not best for your use-case, investigate how Postgres's [synchronous replication](http://www.postgresql.org/docs/current/static/warm-standby.html#SYNCHRONOUS-REPLICATION) works.  Synchronous replication ensures consistency across a cluster by confirming that writes are written to a secondary before returning to the connecting client with a success.  The cost of synchronous replication will be reduced throughput on writes.  This throughput will be entirely based on network performance.  In hosted datacenter environments (like AWS, Rackspace, or any network you do not control), synchrous replication increases the variability of write performance significantly.  If followers become inaccessible from the leader, the leader will becomes effectively readonly.
-
-To enable a simple synchronous replication test, add the follow lines to the `parameters` section of your YAML configuration files.
-
-```YAML
-    synchronous_commit: "on"
-    synchronous_standby_names: "*"
-```
-
-When using synchronous replication, use at least a 3-Postgres data nodes to ensure write availability if one host fails.
-
-Choosing your replication schema is dependent on the many business decisions.  Investigate both async and sync replication, as well as other HA solutions, to determine which solution is best for you.
-
-## Applications should not use superusers
-
-When connecting from an application, always use a non-superuser. Patroni requires access to the database to function properly.  By using a superuser from application, you can potentially use the entire connection pool, including the connections reserved for superusers with the `superuser_reserved_connections` setting. If Patroni cannot access the Primary, because the connection pool is full, behavior will be undesireable.
-
-## Requirements on a Mac
-
-Run the following on a Mac to install requirements:
-
-```
-brew install postgresql etcd haproxy libyaml python
-pip install psycopg2 pyyaml
-```
-
-## Notice
-
-There are many different ways to do HA with PostgreSQL, see [the
-PostgreSQL documentation](https://wiki.postgresql.org/wiki/Replication,_Clustering,_and_Connection_Pooling) for a complete list.
-
-We call this project a "template" because it is far from a one-size fits
-all, or a plug-and-play replication system.  It will have it's own
-caveats.  Use wisely.

README.rst

@@ -0,0 +1,190 @@
+|Build Status| |Coverage Status|
+
+Patroni: A Template for PostgreSQL HA with ZooKeeper or etcd
+------------------------------------------------------------
+
+Patroni was previously known as Governor.
+
+*There are many ways to run high availability with PostgreSQL; here we
+present a template for you to create your own custom fit high
+availability solution using python and distributed configuration store
+(like ZooKeeper or etcd) for maximum accessibility.*
+
+Getting Started
+---------------
+
+To get started, do the following from different terminals:
+
+::
+
+    > etcd --data-dir=data/etcd
+    > ./patroni.py postgres0.yml
+    > ./patroni.py postgres1.yml
+
+From there, you will see a high-availability cluster start up. Test
+different settings in the YAML files to see how behavior changes. Kill
+some of the different components to see how the system behaves.
+
+Add more ``postgres*.yml`` files to create an even larger cluster.
+
+We provide a haproxy configuration, which will give your application a
+single endpoint for connecting to the cluster's leader. To configure,
+run:
+
+::
+
+    > haproxy -f haproxy.cfg
+
+::
+
+    > psql --host 127.0.0.1 --port 5000 postgres
+
+How Patroni works
+-----------------
+
+For a diagram of the high availability decision loop, see the included a
+PDF:
+`postgres-ha.pdf <https://github.com/zalando/patroni/blob/master/postgres-ha.pdf>`__
+
+YAML Configuration
+------------------
+
+For an example file, see ``postgres0.yml``. Below is an explanation of
+settings:
+
+-  *ttl*: the TTL to acquire the leader lock. Think of it as the length of time before automatic failover process is initiated.
+-  *loop\_wait*: the number of seconds the loop will sleep
+
+-  *restapi*:
+    -  *listen*: ip address + port that Patroni will listen to provide health-check information for haproxy.
+    -  *connect\_address*: ip address + port through which restapi is accessible.
+    -  *auth*: (optional) 'username:password' to protect some dangerous REST API endpoints.
+    -  *certfile*: (optional) Specifies a file with the certificate in the PEM format. If certfile is not specified or empty API server will work without SSL.
+    -  *keyfile*: (optional) Specifies a file with the secret key in the PEM format.
+
+-  *etcd*:
+    -  *scope*: the relative path used on etcd's http api for this deployment, thus you can run multiple HA deployments from a single etcd
+    -  *ttl*: the TTL to acquire the leader lock. Think of it as the length of time before automatic failover process is initiated.
+    -  *host*: the host:port for the etcd endpoint
+
+-  *zookeeper*:
+    -  *scope*: the relative path used on etcd's http api for this deployment, thus you can run multiple HA deployments from a single etcd
+    -  *session\_timeout*: the TTL to acquire the leader lock. Think of it as the length of time before automatic failover process is initiated.
+    -  *reconnect\_timeout*: how long we should try to reconnect to ZooKeeper after connection loss. After this timeout we assume that we don't have lock anymore and will restart in read-only mode.
+    -  *hosts*: list of ZooKeeper cluster members in format: ['host1:port1', 'host2:port2', 'etc...']
+    -  *exhibitor*: if you are running ZooKeeper cluster under Exhibitor supervisory the following section could be interesting for you
+        -  *poll\_interval*: how often list of ZooKeeper and Exhibitor nodes should be updated from Exhibitor
+        -  *port*: Exhibitor port
+        -  *hosts*: initial list of Exhibitor (ZooKeeper) nodes in format: ['host1', 'host2', 'etc...' ]. This list would be updated automatically when Exhibitor (ZooKeeper) cluster topology changes.
+
+-  *postgresql*:
+    -  *name*: the name of the Postgres host, must be unique for the cluster
+    -  *listen*: ip address + port that Postgres listening. Must be accessible from other nodes in the cluster if using streaming replication.
+    -  *connect\_address*: ip address + port through which Postgres is accessible from other nodes and applications.
+    -  *data\_dir*: file path to initialize and store Postgres data files
+    -  *maximum\_lag\_on\_failover*: the maximum bytes a follower may lag
+    -  *use\_slots*: whether or not to use replication_slots.  Must be False for PostgreSQL 9.3, and you should comment out max_replication_slots. before it is not eligible become leader
+    -  *pg\_hba*: list of lines which should be added to pg\_hba.conf
+        -  *- host all all 0.0.0.0/0 md5*
+
+    -  *replication*:
+        -  *username*: replication username, user will be created during initialization
+        -  *password*: replication password, user will be created during initialization
+        -  *network*: network setting for replication in pg\_hba.conf
+
+    -  *callbacks* callback scripts to run on certain actions. Patroni will pass current action, role and cluster name. See scripts/aws.py as an example on how to write them.
+        -  *on\_start*: a script to run when the cluster starts
+        -  *on\_stop*: a script to run when the cluster stops
+        -  *on\_restart*: a script to run when the cluster restarts
+        -  *on\_reload*: a script to run when configuration reload is triggered
+        -  *on\_role\_change*: a script to run when the cluster is being promoted or demoted
+
+    -  *superuser*:
+        -  *password*: password for postgres user. It would be set during initialization
+
+    -  *admin*:
+        -  *username*: admin username, user will be created during initialization. It would have CREATEDB and CREATEROLE privileges
+        -  *password*: admin password, user will be created during initialization.
+
+    -  *recovery\_conf*: additional configuration settings written to recovery.conf when configuring follower
+        -  *parameters*: list of configuration settings for Postgres.  Many of these are required for replication to work.
+
+Replication choices
+-------------------
+
+Patroni uses Postgres' streaming replication. By default, this
+replication is asynchronous. For more information, see the `Postgres
+documentation on streaming
+replication <http://www.postgresql.org/docs/current/static/warm-standby.html#STREAMING-REPLICATION>`__.
+
+Patroni's asynchronous replication configuration allows for
+``maximum_lag_on_failover`` settings. This setting ensures failover will
+not occur if a follower is more than a certain number of bytes behind
+the follower. This setting should be increased or decreased based on
+business requirements.
+
+When asynchronous replication is not best for your use-case, investigate
+how Postgres's `synchronous
+replication <http://www.postgresql.org/docs/current/static/warm-standby.html#SYNCHRONOUS-REPLICATION>`__
+works. Synchronous replication ensures consistency across a cluster by
+confirming that writes are written to a secondary before returning to
+the connecting client with a success. The cost of synchronous
+replication will be reduced throughput on writes. This throughput will
+be entirely based on network performance. In hosted datacenter
+environments (like AWS, Rackspace, or any network you do not control),
+synchrous replication increases the variability of write performance
+significantly. If followers become inaccessible from the leader, the
+leader will becomes effectively readonly.
+
+To enable a simple synchronous replication test, add the follow lines to
+the ``parameters`` section of your YAML configuration files.
+
+.. code:: YAML
+
+        synchronous_commit: "on"
+        synchronous_standby_names: "*"
+
+When using synchronous replication, use at least a 3-Postgres data nodes
+to ensure write availability if one host fails.
+
+Choosing your replication schema is dependent on the many business
+decisions. Investigate both async and sync replication, as well as other
+HA solutions, to determine which solution is best for you.
+
+Applications should not use superusers
+--------------------------------------
+
+When connecting from an application, always use a non-superuser. Patroni
+requires access to the database to function properly. By using a
+superuser from application, you can potentially use the entire
+connection pool, including the connections reserved for superusers with
+the ``superuser_reserved_connections`` setting. If Patroni cannot access
+the Primary, because the connection pool is full, behavior will be
+undesireable.
+
+Requirements on a Mac
+---------------------
+
+Run the following on a Mac to install requirements:
+
+::
+
+    brew install postgresql etcd haproxy libyaml python
+    pip install psycopg2 pyyaml
+
+Notice
+------
+
+There are many different ways to do HA with PostgreSQL, see `the
+PostgreSQL
+documentation <https://wiki.postgresql.org/wiki/Replication,_Clustering,_and_Connection_Pooling>`__
+for a complete list.
+
+We call this project a "template" because it is far from a one-size fits
+all, or a plug-and-play replication system. It will have it's own
+caveats. Use wisely.
+
+.. |Build Status| image:: https://travis-ci.org/zalando/patroni.svg?branch=master
+   :target: https://travis-ci.org/zalando/patroni
+.. |Coverage Status| image:: https://coveralls.io/repos/zalando/patroni/badge.svg?branch=master
+   :target: https://coveralls.io/r/zalando/patroni?branch=master

docker/README.md

@@ -1,6 +1,7 @@
 # Patroni Dockerfile
-You can run Patroni in a docker container using this Dockerfile, or by using the Docker image at
-    https://os-registry.stups.zalan.do/acid/patroni-1.0-SNAPSHOT
+You can run Patroni in a docker container using this Dockerfile, or by using one of the Docker image at
+
+    https://registry.opensource.zalan.do/v1/repositories/acid/patroni/tags
 
 This Dockerfile is meant in aiding development of Patroni and quick testing of features. It is not a production-worthy
 Dockerfile
@@ -9,7 +10,7 @@ Dockerfile
 
 ## Standalone Patroni
 
-    docker run -d os-registry.stups.zalan.do/acid/patroni:1.0-SNAPSHOT
+    docker run -d registry.opensource.zalan.do/acid/patroni:1.0-SNAPSHOT
 
 ## Multiple Patroni's communicating with a standalone etcd inside Docker
 
@@ -35,12 +36,12 @@ To automate this you can run the following script:
 
 Example session:
 
-    $ ./dev_patroni_cluster.sh --image os-registry.stups.zalan.do/acid/patroni:1.0-SNAPSHOT --members=2 --name=bravo
+    $ ./dev_patroni_cluster.sh --image registry.opensource.zalan.do/acid/patroni:1.0-SNAPSHOT --members=2 --name=bravo
     The etcd container is 6be871a11cb373406ca5ea1c6b39e1.0-SNAPSHOTfdde9fb1d6177212d6ad0c0d1bd9b563, ip=172.17.1.24
     Started Patroni container 67e611f2eca7c40f9e6e0e24a4a8f2cba7e3e56d22a420e15ab9240a37a9d7a4, ip=172.17.1.25
     Started Patroni container 47dd12ae635ab83b039f5889e250048b606ed5e48e3650b69e365e7e1d4acbcf, ip=172.17.1.26
     $ docker ps
     CONTAINER ID        IMAGE                                         COMMAND                CREATED             STATUS              PORTS                          NAMES
-    47dd12ae635a        os-registry.stups.zalan.do/acid/patroni:1.0-SNAPSHOT   "/bin/bash /entrypoi   10 seconds ago      Up 8 seconds        4001/tcp, 5432/tcp, 2380/tcp   bravo_OR64g8bx
-    67e611f2eca7        os-registry.stups.zalan.do/acid/patroni:1.0-SNAPSHOT   "/bin/bash /entrypoi   11 seconds ago      Up 10 seconds       2380/tcp, 4001/tcp, 5432/tcp   bravo_si9no8iz
-    6be871a11cb3        os-registry.stups.zalan.do/acid/patroni:1.0-SNAPSHOT   "/bin/bash /entrypoi   12 seconds ago      Up 10 seconds       4001/tcp, 5432/tcp, 2380/tcp   bravo_etcd
+    47dd12ae635a        registry.opensource.zalan.do/acid/patroni:1.0-SNAPSHOT   "/bin/bash /entrypoi   10 seconds ago      Up 8 seconds        4001/tcp, 5432/tcp, 2380/tcp   bravo_OR64g8bx
+    67e611f2eca7        registry.opensource.zalan.do/acid/patroni:1.0-SNAPSHOT   "/bin/bash /entrypoi   11 seconds ago      Up 10 seconds       2380/tcp, 4001/tcp, 5432/tcp   bravo_si9no8iz
+    6be871a11cb3        registry.opensource.zalan.do/acid/patroni:1.0-SNAPSHOT   "/bin/bash /entrypoi   12 seconds ago      Up 10 seconds       4001/tcp, 5432/tcp, 2380/tcp   bravo_etcd

docker/entrypoint.sh

@@ -133,5 +133,5 @@ then
         sleep 60
     done
 else
-    exec /patroni/patroni.py /patroni/postgres.yml
+    exec python /patroni.py /patroni/postgres.yml
 fi

extras/startup-scripts/README.md

@@ -0,0 +1,10 @@
+# startup scripts for Patroni
+
+This directory contains sample startup scripts for various OSes 
+and management tools for Patroni.
+
+Scripts supplied:
+
+### patroni.upstart.conf
+
+Upstart job for Ubuntu 12.04 or 14.04.  Requires Upstart > 1.4. Intended for systems where Patroni has been installed on a base system, rather than in Docker.

extras/startup-scripts/patroni.upstart.conf

@@ -0,0 +1,36 @@
+# patroni - patroni daemon
+#
+# controls startup/shutdown of postgres
+# you should disable any postgres start jobs
+#
+# assumes that patroni has been installed into the
+# pythonpath by using setup.py install
+
+description    "patroni start daemon"
+
+start on net-device-up
+stop on runlevel [06]
+
+respawn
+respawn limit 5 10
+
+# set location of patroni
+env PATRONI=/usr/local/bin/patroni
+# virtualenv example
+# env PATRONI=/var/lib/postgresql/patronienv/bin/patroni
+
+# set location of config file
+env PATRONICONF=/etc/patroni/patroni.yml
+
+# set log dir for patroni logs
+# postgres user must have write permission
+env POSTGRESLOGDIR=/var/log/postgresql
+
+setuid postgres
+setgid postgres
+
+script
+    exec start-stop-daemon --start \
+        --exec $PATRONI -- $PATRONICONF \
+        >> $POSTGRESLOGDIR/patroni.log 2>&1
+end script

helpers/api.py

@@ -1,73 +0,0 @@
-import fcntl
-import json
-import logging
-import psycopg2
-
-from six.moves.BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
-from six.moves.socketserver import ThreadingMixIn
-from threading import Thread
-
-logger = logging.getLogger(__name__)
-
-
-class RestApiHandler(BaseHTTPRequestHandler):
-
-    def do_GET(self):
-        response = self.get_postgresql_status()
-
-        path = '/master' if self.path == '/' else self.path
-        status_code = 200 if response['running'] and 'role' in response and response['role'] in path else 503
-
-        self.send_response(status_code)
-        self.send_header('Content-Type', 'application/json')
-        self.end_headers()
-        self.wfile.write(json.dumps(response).encode('utf-8'))
-
-    def get_postgresql_status(self):
-        try:
-            row = self.server.query("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'),
-                                              pg_is_in_recovery(),
-                                              CASE WHEN pg_is_in_recovery()
-                                                   THEN null
-                                                   ELSE pg_current_xlog_location() END,
-                                              pg_last_xlog_receive_location(),
-                                              pg_last_xlog_replay_location(),
-                                              pg_is_in_recovery() AND pg_is_xlog_replay_paused()""")[0]
-            return {
-                'running': True,
-                'postmaster_start_time': row[0],
-                'role': 'slave' if row[1] else 'master',
-                'xlog': ({
-                    'received_location': row[3],
-                    'replayed_location': row[4],
-                    'paused': row[5]} if row[1] else {
-                    'location': row[2]
-                })
-            }
-        except (psycopg2.OperationalError, psycopg2.InterfaceError):
-            logger.exception('get_postgresql_status')
-            return {'running': self.server.patroni.postgresql.is_running()}
-
-
-class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
-
-    def __init__(self, patroni, config):
-        self.connection_string = 'http://{}/patroni'.format(config.get('connect_address', None) or config['listen'])
-        host, port = config['listen'].split(':')
-        HTTPServer.__init__(self, (host, int(port)), RestApiHandler)
-        Thread.__init__(self, target=self.serve_forever)
-        self._set_fd_cloexec(self.socket)
-        self.patroni = patroni
-        self.daemon = True
-
-    def query(self, sql, *params):
-        cursor = self.patroni.postgresql.connection().cursor()
-        cursor.execute(sql, params)
-        ret = [r for r in cursor]
-        cursor.close()
-        return ret
-
-    @staticmethod
-    def _set_fd_cloexec(fd):
-        flags = fcntl.fcntl(fd, fcntl.F_GETFD)
-        fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)

helpers/dcs.py

@@ -1,148 +0,0 @@
-import abc
-
-from collections import namedtuple
-from helpers.utils import calculate_ttl, sleep
-from six.moves.urllib_parse import urlparse, urlunparse, parse_qsl
-
-
-def parse_connection_string(value):
-    """Original Governor stores connection strings for each cluster members if a following format:
-        postgres://{username}:{password}@{connect_address}/postgres
-    Since each of our patroni instances provides own REST API endpoint it's good to store this information
-    in DCS among with postgresql connection string. In order to not introduce new keys and be compatible with
-    original Governor we decided to extend original connection string in a following way:
-        postgres://{username}:{password}@{connect_address}/postgres?application_name={api_url}
-    This way original Governor could use such connection string as it is, because of feature of `libpq` library.
-
-    This method is able to split connection string stored in DCS into two parts, `conn_url` and `api_url`"""
-
-    scheme, netloc, path, params, query, fragment = urlparse(value)
-    conn_url = urlunparse((scheme, netloc, path, params, '', fragment))
-    api_url = ([v for n, v in parse_qsl(query) if n == 'application_name'] or [None])[0]
-    return conn_url, api_url
-
-
-class DCSError(Exception):
-    """Parent class for all kind of exceptions related to selected distributed configuration store"""
-
-    def __init__(self, value):
-        self.value = value
-
-    def __str__(self):
-        """
-        >>> str(DCSError('foo'))
-        "'foo'"
-        """
-        return repr(self.value)
-
-
-class Member(namedtuple('Member', 'index,name,conn_url,api_url,expiration,ttl')):
-    """Immutable object (namedtuple) which represents single member of PostgreSQL cluster.
-    Consists of the following fields:
-    :param index: modification index of a given member key in DCS
-    :param name: name of PostgreSQL cluster member
-    :param conn_url: connection string containing host, user and password which could be used to access this member.
-    :param api_url: REST API url of patroni instance
-    :param expiration: expiration time of given member key
-    :param ttl: ttl of given member key in seconds"""
-
-    def real_ttl(self):
-        return calculate_ttl(self.expiration) or -1
-
-
-class Cluster(namedtuple('Cluster', 'initialize,leader,last_leader_operation,members')):
-    """Immutable object (namedtuple) which represents PostgreSQL cluster.
-    Consists of the following fields:
-    :param initialize: boolean, shows whether this cluster has initialization key stored in DC or not.
-    :param leader: `Member` object which represents current leader of the cluster
-    :param last_leader_operation: int or long object containing position of last known leader operation.
-        This value is stored in `/optime/leader` key
-    :param members: list of Member object, all PostgreSQL cluster members including leader"""
-
-    def is_unlocked(self):
-        return not (self.leader and self.leader.name)
-
-
-class AbstractDCS:
-
-    __metaclass__ = abc.ABCMeta
-
-    def __init__(self, name, config):
-        """
-        :param name: name of current instance (the same value as `~Postgresql.name`)
-        :param config: dict, reference to config section of selected DCS.
-            i.e.: `zookeeper` for zookeeper, `etcd` for etcd, etc...
-        """
-        self._name = name
-        self._base_path = '/service/' + config['scope']
-
-    def client_path(self, path):
-        return self._base_path + path
-
-    @abc.abstractmethod
-    def get_cluster(self):
-        """:returns: `Cluster` object which represent current state and topology of the cluster
-        raise `~DCSError` in case of communication or other problems with DCS. If current instance was
-            running as a master and exception raised instance would be demoted."""
-
-    @abc.abstractmethod
-    def update_leader(self, state_handler):
-        """Update leader key (or session) ttl and `/optime/leader` key in DCS.
-
-        :param state_handler: reference to `Postgresql` object
-        :returns: `!True` if leader key (or session) has been updated successfully.
-            If not, `!False` must be returned and current instance would be demoted.
-
-        If you failed to update `/optime/leader` this error is not critical and you can return `!True`
-        You have to use CAS (Compare And Swap) operation in order to update leader key,
-        for example for etcd `prevValue` parameter must be used."""
-
-    @abc.abstractmethod
-    def attempt_to_acquire_leader(self):
-        """Attempt to acquire leader lock
-        This method should create `/leader` key with value=`~self._name`
-        :returns: `!True` if key has been created successfully.
-
-        Key must be created atomically. In case if key already exists it should not be
-        overwritten and `!False` must be returned"""
-
-    def current_leader(self):
-        try:
-            cluster = self.get_cluster()
-            return None if cluster.is_unlocked() else cluster.leader
-        except DCSError:
-            return None
-
-    @abc.abstractmethod
-    def touch_member(self, connection_string, ttl=None):
-        """Update member key in DCS.
-        This method should create or update key with the name = '/members/' + `~self._name`
-        and value = connection_string in a given DCS.
-
-        :param connection_string: how this instance can be accessed by other instances
-        :param ttl: ttl for member key, optional parameter. If it is None `~self.member_ttl will be used`
-        :returns: `!True` on success otherwise `!False`
-        """
-
-    @abc.abstractmethod
-    def take_leader(self):
-        """This method should create leader key with value = `~self._name` and ttl=`~self.ttl`
-        Since it could be called only on initial cluster bootstrap it could create this key regardless,
-        overwriting the key if necessary."""
-
-    @abc.abstractmethod
-    def race(self, path):
-        """Race for cluster initialization.
-        :param path: usually this is just '/initialize'
-        :returns: `!True` if key has been created successfully.
-
-        this method should create atomically `path` key and return `!True`
-        otherwise it should return `!False`"""
-
-    @abc.abstractmethod
-    def delete_leader(self):
-        """Voluntarily remove leader key from DCS
-        This method should remove leader key if current instance is the leader"""
-
-    def sleep(self, timeout):
-        sleep(timeout)

helpers/ha.py

@@ -1,101 +0,0 @@
-import logging
-
-from helpers.dcs import DCSError
-from psycopg2 import InterfaceError, OperationalError
-
-logger = logging.getLogger(__name__)
-
-
-class Ha:
-
-    def __init__(self, state_handler, etcd):
-        self.state_handler = state_handler
-        self.dcs = etcd
-        self.cluster = None
-        self.old_cluster = None
-
-    def load_cluster_from_dcs(self):
-        cluster = self.dcs.get_cluster()
-
-        # We want to keep the state of cluster when it was healhy
-        if cluster.is_unlocked() and self.cluster and not self.cluster.is_unlocked():
-            self.old_cluster = self.cluster
-        if not self.old_cluster:
-            self.old_cluster = cluster
-        self.cluster = cluster
-
-    def acquire_lock(self):
-        return self.dcs.attempt_to_acquire_leader()
-
-    def update_lock(self):
-        return self.dcs.update_leader(self.state_handler)
-
-    def has_lock(self):
-        lock_owner = self.cluster.leader and self.cluster.leader.name
-        logger.info('Lock owner: %s; I am %s', lock_owner, self.state_handler.name)
-        return lock_owner == self.state_handler.name
-
-    def demote(self):
-        return self.state_handler.demote(self.cluster.leader)
-
-    def follow_the_leader(self):
-        return self.state_handler.follow_the_leader(self.cluster.leader)
-
-    def run_cycle(self):
-        try:
-            self.load_cluster_from_dcs()
-            if not self.state_handler.is_healthy():
-                has_lock = self.has_lock()
-                self.state_handler.write_recovery_conf(None if has_lock else self.cluster.leader)
-                self.state_handler.start()
-                if not has_lock:
-                    return 'started as a secondary'
-                logger.info('started as readonly because i had the session lock')
-                self.load_cluster_from_dcs()
-
-            if self.cluster.is_unlocked():
-                if self.state_handler.is_healthiest_node(self.old_cluster):
-                    if self.acquire_lock():
-                        if self.state_handler.is_leader() or self.state_handler.is_promoted:
-                            return 'acquired session lock as a leader'
-                        else:
-                            self.state_handler.promote()
-                            return 'promoted self to leader by acquiring session lock'
-                    else:
-                        self.load_cluster_from_dcs()
-                        if self.state_handler.is_leader():
-                            self.demote()
-                            return 'demoted self due after trying and failing to obtain lock'
-                        else:
-                            self.follow_the_leader()
-                            return 'following new leader after trying and failing to obtain lock'
-                else:
-                    self.load_cluster_from_dcs()
-                    if self.state_handler.is_leader():
-                        self.demote()
-                        return 'demoting self because i am not the healthiest node'
-                    else:
-                        self.follow_the_leader()
-                        return 'following a different leader because i am not the healthiest node'
-            else:
-                if self.has_lock() and self.update_lock():
-                    if self.state_handler.is_leader() or self.state_handler.is_promoted:
-                        return 'no action.  i am the leader with the lock'
-                    else:
-                        self.state_handler.promote()
-                        return 'promoted self to leader because i had the session lock'
-                else:
-                    logger.info('does not have lock')
-                    if self.state_handler.is_leader():
-                        self.demote()
-                        return 'demoting self because i do not have the lock and i was a leader'
-                    else:
-                        self.follow_the_leader()
-                        return 'no action.  i am a secondary and i am following a leader'
-        except DCSError:
-            logger.error('Error communicating with DCS')
-            if self.state_handler.is_leader():
-                self.state_handler.demote(None)
-                return 'demoted self because DCS is not accessible and i was a leader'
-        except (InterfaceError, OperationalError):
-            logger.error('Error communicating with Postgresql.  Will try again')

helpers/postgresql.py

@@ -1,393 +0,0 @@
-import logging
-import os
-import psycopg2
-import shlex
-import shutil
-import subprocess
-import six
-
-from helpers.utils import sleep
-from six.moves.urllib_parse import urlparse
-
-if six.PY3:
-    long = int
-
-logger = logging.getLogger(__name__)
-
-ACTION_ON_START = "on_start"
-ACTION_ON_STOP = "on_stop"
-ACTION_ON_RESTART = "on_restart"
-ACTION_ON_RELOAD = "on_reload"
-ACTION_ON_ROLE_CHANGE = "on_role_change"
-
-
-def parseurl(url):
-    r = urlparse(url)
-    ret = {
-        'host': r.hostname,
-        'port': r.port or 5432,
-        'database': r.path[1:],
-        'fallback_application_name': 'Patroni',
-        'connect_timeout': 3,
-        'options': '-c statement_timeout=2000',
-    }
-    if r.username:
-        ret['user'] = r.username
-    if r.password:
-        ret['password'] = r.password
-    return ret
-
-
-class Postgresql:
-
-    def __init__(self, config):
-        self.config = config
-        self.name = config['name']
-        self.scope = config['scope']
-        self.listen_addresses, self.port = config['listen'].split(':')
-        self.data_dir = config['data_dir']
-        self.replication = config['replication']
-        self.superuser = config['superuser']
-        self.admin = config['admin']
-        self.callback = config.get('callbacks', {})
-        self.recovery_conf = os.path.join(self.data_dir, 'recovery.conf')
-        self.configuration_to_save = (os.path.join(self.data_dir, 'pg_hba.conf'),
-                                      os.path.join(self.data_dir, 'postgresql.conf'))
-        self.postmaster_pid = os.path.join(self.data_dir, 'postmaster.pid')
-        self.trigger_file = config.get('recovery_conf', {}).get('trigger_file', None) or 'promote'
-        self.trigger_file = os.path.abspath(os.path.join(self.data_dir, self.trigger_file))
-        self.is_promoted = False
-
-        self._pg_ctl = ['pg_ctl', '-w', '-D', self.data_dir]
-
-        self.local_address = self.get_local_address()
-        connect_address = config.get('connect_address', None) or self.local_address
-        self.connection_string = 'postgres://{username}:{password}@{connect_address}/postgres'.format(
-            connect_address=connect_address, **self.replication)
-
-        self._connection = None
-        self._cursor_holder = None
-        self.members = []  # list of already existing replication slots
-
-    def get_local_address(self):
-        listen_addresses = self.listen_addresses.split(',')
-        local_address = listen_addresses[0].strip()  # take first address from listen_addresses
-
-        for la in listen_addresses:
-            if la.strip() in ['*', '0.0.0.0']:  # we are listening on *
-                local_address = 'localhost'  # connection via localhost is preferred
-                break
-        return local_address + ':' + self.port
-
-    def connection(self):
-        if not self._connection or self._connection.closed != 0:
-            r = parseurl('postgres://{}/postgres'.format(self.local_address))
-            self._connection = psycopg2.connect(**r)
-            self._connection.autocommit = True
-        return self._connection
-
-    def _cursor(self):
-        if not self._cursor_holder or self._cursor_holder.closed:
-            self._cursor_holder = self.connection().cursor()
-        return self._cursor_holder
-
-    def disconnect(self):
-        self._connection and self._connection.close()
-        self._connection = self._cursor_holder = None
-
-    def query(self, sql, *params):
-        max_attempts = 0
-        while True:
-            ex = None
-            try:
-                cursor = self._cursor()
-                cursor.execute(sql, params)
-                return cursor
-            except psycopg2.InterfaceError as e:
-                ex = e
-            except psycopg2.OperationalError as e:
-                if self._connection and self._connection.closed == 0:
-                    raise e
-                ex = e
-            if ex:
-                self.disconnect()
-                max_attempts += 1
-                if max_attempts >= 3:
-                    raise ex
-                sleep(5)
-
-    def data_directory_empty(self):
-        return not os.path.exists(self.data_dir) or os.listdir(self.data_dir) == []
-
-    def initialize(self):
-        ret = subprocess.call(self._pg_ctl + ['initdb', '-o', '--encoding=UTF8']) == 0
-        ret and self.write_pg_hba()
-        return ret
-
-    def delete_trigger_file(self):
-        os.path.exists(self.trigger_file) and os.unlink(self.trigger_file)
-
-    def sync_from_leader(self, leader):
-        r = parseurl(leader.conn_url)
-
-        pgpass = 'pgpass'
-        with open(pgpass, 'w') as f:
-            os.fchmod(f.fileno(), 0o600)
-            f.write('{host}:{port}:*:{user}:{password}\n'.format(**r))
-
-        env = os.environ.copy()
-        env['PGPASSFILE'] = pgpass
-        return self.create_replica(r, env) == 0
-
-    @staticmethod
-    def build_connstring(conn):
-        return "host={host} port={port} user={user}".format(**conn)
-
-    def create_replica(self, master_connection, env):
-        connstring = self.build_connstring(master_connection)
-        cmd = self.config['restore']
-        try:
-            ret = subprocess.call(shlex.split(cmd) + [self.scope, "replica", self.data_dir, connstring], env=env)
-            self.delete_trigger_file()
-        except:
-            logger.exception('Error when creating replica')
-            return 1
-        return ret
-
-    def is_leader(self, check_only=False):
-        ret = not self.query('SELECT pg_is_in_recovery()').fetchone()[0]
-        if ret and self.is_promoted and not check_only:
-            self.delete_trigger_file()
-            self.is_promoted = False
-        return ret
-
-    def is_running(self):
-        return subprocess.call(' '.join(self._pg_ctl) + ' status > /dev/null', shell=True) == 0
-
-    def call_nowait(self, cb_name, is_leader=None):
-        """ pick a callback command and call it without waiting for it to finish """
-        if not self.callback or cb_name not in self.callback:
-            return False
-        cmd = self.callback[cb_name]
-        if is_leader is None:
-            try:
-                is_leader = self.is_leader(check_only=True)
-            except psycopg2.OperationalError as e:
-                logger.warning("unable to perform {0} action, cannot obtain the cluster role: {1}".format(cb_name, e))
-                return False
-        try:
-            role = "master" if is_leader else "replica"
-            subprocess.Popen(shlex.split(cmd) + [cb_name, role, self.scope])
-        except:
-            logger.exception('callback %s %s %s %s failed', cmd, cb_name, role, self.scope)
-            return False
-        return True
-
-    def start(self):
-        if self.is_running():
-            self.load_replication_slots()
-            logger.error('Cannot start PostgreSQL because one is already running.')
-            return False
-
-        if os.path.exists(self.postmaster_pid):
-            os.remove(self.postmaster_pid)
-            logger.info('Removed %s', self.postmaster_pid)
-
-        ret = subprocess.call(self._pg_ctl + ['start', '-o', self.server_options()]) == 0
-        ret and self.load_replication_slots()
-        self.save_configuration_files()
-        if ret and ACTION_ON_START in self.callback:
-            self.call_nowait(ACTION_ON_START)
-        return ret
-
-    def stop(self):
-        try:
-            is_leader = self.is_leader(check_only=True)
-        except:
-            is_leader = None
-            pass
-        ret = subprocess.call(self._pg_ctl + ['stop', '-m', 'fast'])
-        if ret == 0 and ACTION_ON_STOP in self.callback:
-            self.call_nowait(ACTION_ON_STOP, is_leader=is_leader)
-        return ret == 0
-
-    def reload(self):
-        ret = subprocess.call(self._pg_ctl + ['reload'])
-        if ret == 0 and ACTION_ON_RELOAD in self.callback:
-            self.call_nowait(ACTION_ON_RELOAD)
-        return ret == 0
-
-    def restart(self):
-        try:
-            is_leader = self.is_leader(check_only=True)
-        except:
-            is_leader = None
-            pass
-        ret = subprocess.call(self._pg_ctl + ['restart', '-m', 'fast'])
-        if ret == 0 and ACTION_ON_RESTART in self.callback:
-            self.call_nowait(ACTION_ON_RESTART, is_leader=is_leader)
-        return ret == 0
-
-    def server_options(self):
-        options = "--listen_addresses='{}' --port={}".format(self.listen_addresses, self.port)
-        for setting, value in self.config['parameters'].items():
-            options += " --{}='{}'".format(setting, value)
-        return options
-
-    def is_healthy(self):
-        if not self.is_running():
-            logger.warning('Postgresql is not running.')
-            return False
-        return True
-
-    def is_healthiest_node(self, cluster):
-        if self.is_leader():
-            return True
-
-        if cluster.last_leader_operation - self.xlog_position() > self.config.get('maximum_lag_on_failover', 0):
-            return False
-
-        for member in cluster.members:
-            if member.name == self.name:
-                continue
-            try:
-                r = parseurl(member.conn_url)
-                member_conn = psycopg2.connect(**r)
-                member_conn.autocommit = True
-                member_cursor = member_conn.cursor()
-                member_cursor.execute(
-                    "SELECT pg_is_in_recovery(), %s - (pg_last_xlog_replay_location() - '0/0000000'::pg_lsn)",
-                    (self.xlog_position(), ))
-                row = member_cursor.fetchone()
-                member_cursor.close()
-                member_conn.close()
-                logger.error([self.name, member.name, row])
-                if not row[0]:
-                    logger.warning('Master (%s) is still alive', member.name)
-                    return False
-                if row[1] < 0:
-                    return False
-            except psycopg2.Error:
-                continue
-        return True
-
-    def write_pg_hba(self):
-        with open(os.path.join(self.data_dir, 'pg_hba.conf'), 'a') as f:
-            f.write('\nhost replication {username} {network} md5\n'.format(**self.replication))
-            for line in self.config.get('pg_hba', []):
-                if line.split()[0].strip() == 'hostssl' and self.config['parameters'].get('ssl', 'off').lower() != 'on':
-                    continue
-                f.write(line + '\n')
-
-    @staticmethod
-    def primary_conninfo(leader_url):
-        r = parseurl(leader_url)
-        return 'user={user} password={password} host={host} port={port} sslmode=prefer sslcompression=1'.format(**r)
-
-    def check_recovery_conf(self, leader):
-        if not os.path.isfile(self.recovery_conf):
-            return False
-
-        pattern = leader and leader.conn_url and self.primary_conninfo(leader.conn_url)
-
-        with open(self.recovery_conf, 'r') as f:
-            for line in f:
-                if line.startswith('primary_conninfo'):
-                    if not pattern:
-                        return False
-                    return pattern in line
-
-        return not pattern
-
-    def write_recovery_conf(self, leader):
-        with open(self.recovery_conf, 'w') as f:
-            f.write("""standby_mode = 'on'
-recovery_target_timeline = 'latest'
-""")
-            if leader and leader.conn_url:
-                f.write("""
-primary_slot_name = '{}'
-primary_conninfo = '{}'
-""".format(self.name, self.primary_conninfo(leader.conn_url)))
-                for name, value in self.config.get('recovery_conf', {}).items():
-                    f.write("{} = '{}'\n".format(name, value))
-
-    def follow_the_leader(self, leader):
-        if not self.check_recovery_conf(leader):
-            self.write_recovery_conf(leader)
-            self.restart()
-            if ACTION_ON_ROLE_CHANGE in self.callback:
-                self.call_nowait(ACTION_ON_ROLE_CHANGE)
-
-    def save_configuration_files(self):
-        """
-            copy postgresql.conf to postgresql.conf.backup to preserve it in the WAL-e backup.
-            see http://comments.gmane.org/gmane.comp.db.postgresql.wal-e/239
-        """
-        for f in self.configuration_to_save:
-            shutil.copy(f, f + '.backup')
-
-    def restore_configuration_files(self):
-        """ restore a previously saved postgresql.conf """
-        try:
-            for f in self.configuration_to_save:
-                shutil.copy(f + '.backup', f)
-        except:
-            logger.exception('unable to restore configuration from WAL-E backup')
-
-    def promote(self):
-        self.is_promoted = subprocess.call(self._pg_ctl + ['promote']) == 0
-        if self.is_promoted and ACTION_ON_ROLE_CHANGE in self.callback:
-            self.call_nowait(ACTION_ON_ROLE_CHANGE)
-        return self.is_promoted
-
-    def demote(self, leader):
-        self.follow_the_leader(leader)
-
-    def create_replication_user(self):
-        self.query('CREATE USER "{}" WITH REPLICATION ENCRYPTED PASSWORD %s'.format(
-            self.replication['username']), self.replication['password'])
-
-    def create_connection_users(self):
-        if self.superuser:
-            if 'username' in self.superuser:
-                self.query('CREATE ROLE "{0}" WITH LOGIN SUPERUSER PASSWORD %s'.format(
-                    self.superuser['username']), self.superuser['password'])
-            else:
-                rolsuper = self.query("""SELECT rolname FROM pg_authid WHERE rolsuper = 't'""").fetchone()[0]
-                self.query('ALTER ROLE "{0}" WITH PASSWORD %s'.format(rolsuper), self.superuser['password'])
-        if self.admin:
-            self.query('CREATE ROLE "{0}" WITH LOGIN CREATEDB CREATEROLE PASSWORD %s'.format(
-                self.admin['username']), self.admin['password'])
-
-    def xlog_position(self):
-        return self.query("""SELECT CASE WHEN pg_is_in_recovery()
-                                         THEN pg_last_xlog_replay_location() - '0/0000000'::pg_lsn
-                                         ELSE pg_current_xlog_location() - '0/00000'::pg_lsn END""").fetchone()[0]
-
-    def load_replication_slots(self):
-        cursor = self.query("SELECT slot_name FROM pg_replication_slots WHERE slot_type='physical'")
-        self.members = [r[0] for r in cursor]
-
-    def sync_replication_slots(self, members):
-        # drop unused slots
-        for slot in set(self.members) - set(members):
-            self.query("""SELECT pg_drop_replication_slot(%s)
-                           WHERE EXISTS(SELECT 1 FROM pg_replication_slots
-                           WHERE slot_name = %s)""", slot, slot)
-
-        # create new slots
-        for slot in set(members) - set(self.members):
-            self.query("""SELECT pg_create_physical_replication_slot(%s)
-                           WHERE NOT EXISTS (SELECT 1 FROM pg_replication_slots
-                           WHERE slot_name = %s)""", slot, slot)
-        self.members = members
-
-    def create_replication_slots(self, cluster):
-        self.sync_replication_slots([m.name for m in cluster.members if m.name != self.name])
-
-    def drop_replication_slots(self):
-        self.sync_replication_slots([])
-
-    def last_operation(self):
-        return str(self.xlog_position())

helpers/utils.py

@@ -1,109 +0,0 @@
-import datetime
-import os
-import re
-import signal
-import sys
-import time
-
-interrupted_sleep = False
-reap_children = False
-
-_DATE_TIME_RE = re.compile(r'''^
-(?P<year>\d{4})\-(?P<month>\d{2})\-(?P<day>\d{2})  # date
-T
-(?P<hour>\d{2}):(?P<minute>\d{2}):(?P<second>\d{2})\.(?P<microsecond>\d{6})  # time
-\d*Z$''', re.X)
-
-
-def parse_datetime(time_str):
-    """
-    >>> parse_datetime('2015-06-10T12:56:30.552539016Z')
-    datetime.datetime(2015, 6, 10, 12, 56, 30, 552539)
-    >>> parse_datetime('2015-06-10 12:56:30.552539016Z')
-    """
-    m = _DATE_TIME_RE.match(time_str)
-    if not m:
-        return None
-    p = dict((n, int(m.group(n))) for n in 'year month day hour minute second microsecond'.split(' '))
-    return datetime.datetime(**p)
-
-
-def calculate_ttl(expiration):
-    """
-    >>> calculate_ttl(None)
-    >>> calculate_ttl('2015-06-10 12:56:30.552539016Z')
-    """
-    if not expiration:
-        return None
-    expiration = parse_datetime(expiration)
-    if not expiration:
-        return None
-    now = datetime.datetime.utcnow()
-    return int((expiration - now).total_seconds())
-
-
-def lsn_to_bytes(value):
-    """
-    >>> lsn_to_bytes('1/66000060')
-    6006243424
-    >>> lsn_to_bytes('j/66000060')
-    0
-    """
-    try:
-        e = value.split('/')
-        if len(e) == 2 and len(e[0]) > 0 and len(e[1]) > 0:
-            return (int(e[0], 16) << 32) | int(e[1], 16)
-    except ValueError:
-        pass
-    return 0
-
-
-def bytes_to_lsn(value):
-    """
-    >>> bytes_to_lsn(6006243424)
-    '1/66000060'
-    """
-    id = value >> 32
-    off = value & 0xffffffff
-    return '%x/%x' % (id, off)
-
-
-def sigterm_handler(signo, stack_frame):
-    sys.exit()
-
-
-def sigchld_handler(signo, stack_frame):
-    global interrupted_sleep, reap_children
-    reap_children = interrupted_sleep = True
-
-
-def sleep(interval):
-    global interrupted_sleep
-    current_time = time.time()
-    end_time = current_time + interval
-    while current_time < end_time:
-        interrupted_sleep = False
-        time.sleep(end_time - current_time)
-        if not interrupted_sleep:  # we will ignore only sigchld
-            break
-        current_time = time.time()
-    interrupted_sleep = False
-
-
-def setup_signal_handlers():
-    signal.signal(signal.SIGTERM, sigterm_handler)
-    signal.signal(signal.SIGCHLD, sigchld_handler)
-
-
-def reap_children():
-    global reap_children
-    if reap_children:
-        try:
-            while True:
-                ret = os.waitpid(-1, os.WNOHANG)
-                if ret == (0, 0):
-                    break
-        except OSError:
-            pass
-        finally:
-            reap_children = False

patroni.py

@@ -1,123 +1,5 @@
 #!/usr/bin/env python
-import logging
-import os
-import sys
-import time
-import yaml
-
-from helpers.api import RestApiServer
-from helpers.etcd import Etcd
-from helpers.ha import Ha
-from helpers.postgresql import Postgresql
-from helpers.utils import setup_signal_handlers, sleep, reap_children
-from helpers.zookeeper import ZooKeeper
-
-logger = logging.getLogger(__name__)
-
-
-class Patroni:
-
-    def __init__(self, config):
-        self.nap_time = config['loop_wait']
-        self.postgresql = Postgresql(config['postgresql'])
-        self.ha = Ha(self.postgresql, self.get_dcs(self.postgresql.name, config))
-        host, port = config['restapi']['listen'].split(':')
-        self.api = RestApiServer(self, config['restapi'])
-        self.next_run = time.time()
-        self.shutdown_member_ttl = 300
-
-    @staticmethod
-    def get_dcs(name, config):
-        if 'etcd' in config:
-            return Etcd(name, config['etcd'])
-        if 'zookeeper' in config:
-            return ZooKeeper(name, config['zookeeper'])
-        raise Exception('Can not find sutable configuration of distributed configuration store')
-
-    def touch_member(self, ttl=None):
-        connection_string = self.postgresql.connection_string + '?application_name=' + self.api.connection_string
-        if self.ha.cluster:
-            for m in self.ha.cluster.members:
-                # Do not update member TTL when it is far from being expired
-                if m.name == self.postgresql.name and m.real_ttl() > self.shutdown_member_ttl:
-                    return True
-        return self.ha.dcs.touch_member(connection_string, ttl)
-
-    def initialize(self):
-        # wait for etcd to be available
-        while not self.touch_member():
-            logger.info('waiting on DCS')
-            sleep(5)
-
-        # is data directory empty?
-        if self.postgresql.data_directory_empty():
-            # racing to initialize
-            if self.ha.dcs.race('/initialize'):
-                self.postgresql.initialize()
-                self.ha.dcs.take_leader()
-                self.postgresql.start()
-                self.postgresql.create_replication_user()
-                self.postgresql.create_connection_users()
-            else:
-                while True:
-                    leader = self.ha.dcs.current_leader()
-                    if leader and self.postgresql.sync_from_leader(leader):
-                        self.postgresql.write_recovery_conf(leader)
-                        self.postgresql.start()
-                        break
-                    sleep(5)
-        elif self.postgresql.is_running():
-            self.postgresql.load_replication_slots()
-
-    def schedule_next_run(self):
-        self.next_run += self.nap_time
-        current_time = time.time()
-        nap_time = self.next_run - current_time
-        if nap_time <= 0:
-            self.next_run = current_time
-        else:
-            self.ha.dcs.sleep(nap_time)
-
-    def run(self):
-        self.api.start()
-        self.next_run = time.time()
-
-        while True:
-            self.touch_member()
-            logger.info(self.ha.run_cycle())
-            try:
-                if self.ha.state_handler.is_leader():
-                    self.ha.cluster and self.ha.state_handler.create_replication_slots(self.ha.cluster)
-                else:
-                    self.ha.state_handler.drop_replication_slots()
-            except:
-                logger.exception('Exception when changing replication slots')
-            reap_children()
-            self.schedule_next_run()
-
-
-def main():
-    logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO)
-    logging.getLogger('requests').setLevel(logging.WARNING)
-    setup_signal_handlers()
-
-    if len(sys.argv) < 2 or not os.path.isfile(sys.argv[1]):
-        print('Usage: {} config.yml'.format(sys.argv[0]))
-        return
-
-    with open(sys.argv[1], 'r') as f:
-        config = yaml.load(f)
-
-    patroni = Patroni(config)
-    try:
-        patroni.initialize()
-        patroni.run()
-    except KeyboardInterrupt:
-        pass
-    finally:
-        patroni.touch_member(patroni.shutdown_member_ttl)  # schedule member removal
-        patroni.postgresql.stop()
-        patroni.ha.dcs.delete_leader()
+from patroni import main
 
 
 if __name__ == '__main__':

patroni/__init__.py

@@ -0,0 +1,75 @@
+import logging
+import os
+import sys
+import time
+import yaml
+
+from patroni.api import RestApiServer
+from patroni.etcd import Etcd
+from patroni.ha import Ha
+from patroni.postgresql import Postgresql
+from patroni.utils import setup_signal_handlers, reap_children
+from patroni.zookeeper import ZooKeeper
+
+logger = logging.getLogger(__name__)
+
+
+class Patroni:
+
+    def __init__(self, config):
+        self.nap_time = config['loop_wait']
+        self.postgresql = Postgresql(config['postgresql'])
+        self.dcs = self.get_dcs(self.postgresql.name, config)
+        host, port = config['restapi']['listen'].split(':')
+        self.api = RestApiServer(self, config['restapi'])
+        self.ha = Ha(self)
+        self.next_run = time.time()
+
+    @staticmethod
+    def get_dcs(name, config):
+        if 'etcd' in config:
+            return Etcd(name, config['etcd'])
+        if 'zookeeper' in config:
+            return ZooKeeper(name, config['zookeeper'])
+        raise Exception('Can not find sutable configuration of distributed configuration store')
+
+    def schedule_next_run(self):
+        self.next_run += self.nap_time
+        current_time = time.time()
+        nap_time = self.next_run - current_time
+        if nap_time <= 0:
+            self.next_run = current_time
+        elif self.dcs.watch(nap_time):
+            self.next_run = time.time()
+
+    def run(self):
+        self.api.start()
+        self.next_run = time.time()
+
+        while True:
+            logger.info(self.ha.run_cycle())
+            reap_children()
+            self.schedule_next_run()
+
+
+def main():
+    logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO)
+    logging.getLogger('requests').setLevel(logging.WARNING)
+    setup_signal_handlers()
+
+    if len(sys.argv) < 2 or not os.path.isfile(sys.argv[1]):
+        print('Usage: {} config.yml'.format(sys.argv[0]))
+        return
+
+    with open(sys.argv[1], 'r') as f:
+        config = yaml.load(f)
+
+    patroni = Patroni(config)
+    try:
+        patroni.run()
+    except KeyboardInterrupt:
+        pass
+    finally:
+        patroni.api.shutdown()
+        patroni.postgresql.stop()
+        patroni.dcs.delete_leader()

patroni/__main__.py

@@ -0,0 +1,5 @@
+from patroni import main
+
+
+if __name__ == '__main__':
+    main()

patroni/api.py

@@ -0,0 +1,226 @@
+import base64
+import fcntl
+import json
+import logging
+import psycopg2
+
+from patroni.exceptions import PostgresConnectionException
+from patroni.utils import Retry, RetryFailedError
+from six.moves.BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
+from six.moves.socketserver import ThreadingMixIn
+from threading import Thread
+
+logger = logging.getLogger(__name__)
+
+
+def check_auth(func):
+    """Decorator function to check authorization header.
+
+    Usage example:
+    @check_auth
+    def do_PUT_foo():
+        pass
+    """
+    def wrapper(handler):
+        if handler.check_auth_header():
+            return func(handler)
+    return wrapper
+
+
+class RestApiHandler(BaseHTTPRequestHandler):
+
+    def send_auth_request(self, body):
+        self.send_response(401)
+        self.send_header('WWW-Authenticate', 'Basic realm=\"Patroni\"')
+        self.send_header('Content-type', 'text/html')
+        self.end_headers()
+        self.wfile.write(body.encode('utf-8'))
+
+    def check_auth_header(self):
+        auth_header = self.headers.get('Authorization')
+        status = self.server.check_auth_header(auth_header)
+        return not status or self.send_auth_request(status)
+
+    def do_GET(self):
+        """Default method for processing all GET requests which can not be routed to other methods"""
+
+        path = '/master' if self.path == '/' else self.path
+        response = self.get_postgresql_status()
+
+        patroni = self.server.patroni
+        cluster = patroni.dcs.cluster
+        if cluster:  # dcs available
+            if cluster.leader and cluster.leader.name == patroni.postgresql.name:  # is_leader
+                status_code = 200 if 'master' in path else 503
+            elif 'role' not in response:
+                status_code = 503
+            elif response['role'] == 'master':  # running as master but without leader lock!!!!
+                status_code = 503
+            elif response['role'] in path:
+                status_code = 200
+            else:
+                status_code = 503
+        elif 'role' in response and response['role'] in path:
+            status_code = 200
+        elif patroni.ha.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:
+            status_code = 503
+
+        self.send_response(status_code)
+        self.send_header('Content-Type', 'application/json')
+        self.end_headers()
+        self.wfile.write(json.dumps(response).encode('utf-8'))
+
+    def do_GET_patroni(self):
+        response = self.get_postgresql_status(True)
+
+        self.send_response(200)
+        self.send_header('Content-Type', 'application/json')
+        self.end_headers()
+        self.wfile.write(json.dumps(response).encode('utf-8'))
+
+    @check_auth
+    def do_POST_restart(self):
+        status_code = 503
+        data = b'restart failed'
+        try:
+            status, msg = self.server.patroni.ha.restart()
+            status_code = 200 if status else 503
+            data = msg.encode('utf-8')
+        except:
+            logger.exception('Exception during restart')
+
+        self.send_response(status_code)
+        self.send_header('Content-Type', 'text/html')
+        self.end_headers()
+        self.wfile.write(data)
+
+    @check_auth
+    def do_POST_reinitialize(self):
+        ha = self.server.patroni.ha
+        cluster = ha.dcs.get_cluster()
+        if cluster.is_unlocked():
+            status_code = 503
+            data = b'Cluster has no leader, can not reinitialize'
+        elif cluster.leader.name == ha.state_handler.name:
+            status_code = 503
+            data = b'I am the leader, can not reinitialize'
+        else:
+            action = ha.schedule_reinitialize()
+            if action is not None:
+                status_code = 503
+                data = (action + ' already in progress').encode('utf-8')
+            else:
+                status_code = 200
+                data = b'reinitialize scheduled'
+
+        self.send_response(status_code)
+        self.send_header('Content-Type', 'text/html')
+        self.end_headers()
+        self.wfile.write(data)
+
+    def parse_request(self):
+        """Override parse_request method to enrich basic functionality of `BaseHTTPRequestHandler` class
+
+        Original class can only invoke do_GET, do_POST, do_PUT, etc method implementations if they are defined.
+        But we would like to have at least some simple routing mechanism, i.e.:
+        GET /uri1/part2 request should invoke `do_GET_uri1()`
+        POST /other should invoke `do_POST_other()`
+
+        If the `do_<REQUEST_METHOD>_<first_part_url>` method does not exists we'll fallback to original behavior."""
+
+        ret = BaseHTTPRequestHandler.parse_request(self)
+        if ret:
+            mname = self.path.lstrip('/').split('/')[0]
+            mname = self.command + ('_' + mname if mname else '')
+            if hasattr(self, 'do_' + mname):
+                self.command = mname
+        return ret
+
+    def query(self, sql, *params, **kwargs):
+        if not kwargs.get('retry', False):
+            return self.server.query(sql, *params)
+        retry = Retry(delay=1, retry_exceptions=PostgresConnectionException)
+        return retry(self.server.query, sql, *params)
+
+    def get_postgresql_status(self, retry=False):
+        try:
+            row = self.query("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'),
+                                       pg_is_in_recovery(),
+                                       CASE WHEN pg_is_in_recovery()
+                                            THEN 0
+                                            ELSE pg_xlog_location_diff(pg_current_xlog_location(), '0/0')::bigint
+                                       END,
+                                       pg_xlog_location_diff(pg_last_xlog_receive_location(), '0/0')::bigint,
+                                       pg_xlog_location_diff(pg_last_xlog_replay_location(), '0/0')::bigint,
+                                       pg_is_in_recovery() AND pg_is_xlog_replay_paused()""", retry=retry)[0]
+            return {
+                'state': self.server.patroni.postgresql.state,
+                'postmaster_start_time': row[0],
+                'role': 'replica' if row[1] else 'master',
+                'xlog': ({
+                    'received_location': row[3],
+                    'replayed_location': row[4],
+                    'paused': row[5]} if row[1] else {
+                    'location': row[2]
+                })
+            }
+        except (psycopg2.Error, RetryFailedError, PostgresConnectionException):
+            state = self.server.patroni.postgresql.state
+            if state in ['stopped', 'starting', 'stopping', 'restarting', 'running']:
+                logger.exception('get_postgresql_status')
+                state = 'unknown' if state == 'running' else state
+            return {'state': state}
+
+
+class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
+
+    def __init__(self, patroni, config):
+        self._auth_key = base64.b64encode(config['auth'].encode('utf-8')).decode('utf-8') if 'auth' in config else None
+        host, port = config['listen'].split(':')
+        HTTPServer.__init__(self, (host, int(port)), RestApiHandler)
+        Thread.__init__(self, target=self.serve_forever)
+        self._set_fd_cloexec(self.socket)
+
+        protocol = 'http'
+
+        # wrap socket with ssl if 'certfile' is defined in a config.yaml
+        # Sometime it's also needed to pass reference to a 'keyfile'.
+        options = {option: config[option] for option in ['certfile', 'keyfile'] if option in config}
+        if options.get('certfile', None):
+            import ssl
+            self.socket = ssl.wrap_socket(self.socket, server_side=True, **options)
+            protocol = 'https'
+
+        self.connection_string = '{}://{}/patroni'.format(protocol, config.get('connect_address', config['listen']))
+
+        self.patroni = patroni
+        self.daemon = True
+
+    def query(self, sql, *params):
+        cursor = None
+        try:
+            with self.patroni.postgresql.connection().cursor() as cursor:
+                cursor.execute(sql, params)
+                return [r for r in cursor]
+        except psycopg2.Error as e:
+            if cursor and cursor.connection.closed == 0:
+                raise e
+            raise PostgresConnectionException('connection problems')
+
+    @staticmethod
+    def _set_fd_cloexec(fd):
+        flags = fcntl.fcntl(fd, fcntl.F_GETFD)
+        fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)
+
+    def check_basic_auth_key(self, key):
+        return self._auth_key == key
+
+    def check_auth_header(self, auth_header):
+        if self._auth_key:
+            if auth_header is None:
+                return 'no auth header received'
+            if not auth_header.startswith('Basic ') or not self.check_basic_auth_key(auth_header[6:]):
+                return 'not authenticated'

patroni/async_executor.py

@@ -0,0 +1,55 @@
+import logging
+from threading import Lock, Thread
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncExecutor:
+
+    def __init__(self):
+        Lock.__init__(self)
+        self._busy = False
+        self._thread_lock = Lock()
+        self._scheduled_action = None
+        self._scheduled_action_lock = Lock()
+
+    @property
+    def busy(self):
+        return self._busy
+
+    def schedule(self, action, immediately=False):
+        with self._scheduled_action_lock:
+            if self._scheduled_action is not None:
+                return self._scheduled_action
+            self._scheduled_action = action
+            self._busy = immediately
+        return None
+
+    @property
+    def scheduled_action(self):
+        with self._scheduled_action_lock:
+            return self._scheduled_action
+
+    def reset_scheduled_action(self):
+        with self._scheduled_action_lock:
+            self._scheduled_action = None
+
+    def run(self, func, args=()):
+        try:
+            return func(*args) if args else func()
+        except:
+            logger.exception('Exception during execution of long running task %s', self.scheduled_action)
+        finally:
+            with self:
+                self._busy = False
+                self.reset_scheduled_action()
+
+    def run_async(self, func, args=()):
+        self._busy = True
+        Thread(target=self.run, args=(func, args)).start()
+
+    def __enter__(self):
+        self._thread_lock.acquire()
+
+    def __exit__(self, type, value, traceback):
+        self._thread_lock.release()

patroni/dcs.py

@@ -0,0 +1,267 @@
+import abc
+import json
+
+from collections import namedtuple
+from patroni.exceptions import DCSError
+from six.moves.urllib_parse import urlparse, urlunparse, parse_qsl
+from threading import Event, Lock
+
+
+def parse_connection_string(value):
+    """Original Governor stores connection strings for each cluster members if a following format:
+        postgres://{username}:{password}@{connect_address}/postgres
+    Since each of our patroni instances provides own REST API endpoint it's good to store this information
+    in DCS among with postgresql connection string. In order to not introduce new keys and be compatible with
+    original Governor we decided to extend original connection string in a following way:
+        postgres://{username}:{password}@{connect_address}/postgres?application_name={api_url}
+    This way original Governor could use such connection string as it is, because of feature of `libpq` library.
+
+    This method is able to split connection string stored in DCS into two parts, `conn_url` and `api_url`"""
+
+    scheme, netloc, path, params, query, fragment = urlparse(value)
+    conn_url = urlunparse((scheme, netloc, path, params, '', fragment))
+    api_url = ([v for n, v in parse_qsl(query) if n == 'application_name'] or [None])[0]
+    return conn_url, api_url
+
+
+class Member(namedtuple('Member', 'index,name,session,data')):
+
+    """Immutable object (namedtuple) which represents single member of PostgreSQL cluster.
+    Consists of the following fields:
+    :param index: modification index of a given member key in a Configuration Store
+    :param name: name of PostgreSQL cluster member
+    :param session: either session id or just ttl in seconds
+    :param data: arbitrary data i.e. conn_url, api_url, xlog location, state, role, tags, etc...
+
+    There are two mandatory keys in a data:
+    conn_url: connection string containing host, user and password which could be used to access this member.
+    api_url: REST API url of patroni instance"""
+
+    @staticmethod
+    def from_node(index, name, session, data):
+        """
+        >>> Member.from_node(-1, '', '', '{"conn_url": "postgres://foo@bar/postgres"}') is not None
+        True
+        >>> Member.from_node(-1, '', '', '{')
+        Member(index=-1, name='', session='', data={})
+        """
+        if data.startswith('postgres'):
+            conn_url, api_url = parse_connection_string(data)
+            data = {'conn_url': conn_url, 'api_url': api_url}
+        else:
+            try:
+                data = json.loads(data)
+            except:
+                data = {}
+        return Member(index, name, session, data)
+
+    @property
+    def conn_url(self):
+        return self.data.get('conn_url', None)
+
+    @property
+    def api_url(self):
+        return self.data.get('api_url', None)
+
+
+class Leader(namedtuple('Leader', 'index,session,member')):
+
+    """Immutable object (namedtuple) which represents leader key.
+    Consists of the following fields:
+    :param index: modification index of a leader key in a Configuration Store
+    :param session: either session id or just ttl in seconds
+    :param member: reference to a `Member` object which represents current leader (see `Cluster.members`)"""
+
+    @property
+    def name(self):
+        return self.member.name
+
+    @property
+    def conn_url(self):
+        return self.member.conn_url
+
+
+class Failover(namedtuple('Failover', 'index,leader,member')):
+
+    @staticmethod
+    def from_node(index, value):
+        t = [a.strip() for a in value.split(':')] + ['']
+        return Failover(index, t[0], t[1]) if t[0] or t[1] else None
+
+
+class Cluster(namedtuple('Cluster', 'initialize,leader,last_leader_operation,members,failover')):
+
+    """Immutable object (namedtuple) which represents PostgreSQL cluster.
+    Consists of the following fields:
+    :param initialize: boolean, shows whether this cluster has initialization key stored in DC or not.
+    :param leader: `Leader` object which represents current leader of the cluster
+    :param last_leader_operation: int or long object containing position of last known leader operation.
+        This value is stored in `/optime/leader` key
+    :param members: list of Member object, all PostgreSQL cluster members including leader
+    :param failover: reference to `Failover` object"""
+
+    def is_unlocked(self):
+        return not (self.leader and self.leader.name)
+
+
+class AbstractDCS:
+
+    __metaclass__ = abc.ABCMeta
+
+    _INITIALIZE = 'initialize'
+    _LEADER = 'leader'
+    _FAILOVER = 'failover'
+    _MEMBERS = 'members/'
+    _OPTIME = 'optime'
+    _LEADER_OPTIME = _OPTIME + '/' + _LEADER
+
+    def __init__(self, name, config):
+        """
+        :param name: name of current instance (the same value as `~Postgresql.name`)
+        :param config: dict, reference to config section of selected DCS.
+            i.e.: `zookeeper` for zookeeper, `etcd` for etcd, etc...
+        """
+        self._name = name
+        self._scope = config['scope']
+        self._base_path = '/service/' + self._scope
+
+        self._cluster = None
+        self._cluster_thread_lock = Lock()
+        self.event = Event()
+
+    def client_path(self, path):
+        return '/'.join([self._base_path, path.lstrip('/')])
+
+    @property
+    def initialize_path(self):
+        return self.client_path(self._INITIALIZE)
+
+    @property
+    def members_path(self):
+        return self.client_path(self._MEMBERS)
+
+    @property
+    def member_path(self):
+        return self.client_path(self._MEMBERS + self._name)
+
+    @property
+    def leader_path(self):
+        return self.client_path(self._LEADER)
+
+    @property
+    def failover_path(self):
+        return self.client_path(self._FAILOVER)
+
+    @property
+    def leader_optime_path(self):
+        return self.client_path(self._LEADER_OPTIME)
+
+    @abc.abstractmethod
+    def _load_cluster(self):
+        """Internally this method should build  `Cluster` object which
+           represents current state and topology of the cluster in DCS.
+           this method supposed to be called only by `get_cluster` method.
+
+           raise `~DCSError` in case of communication or other problems with DCS.
+           If the current node was running as a master and exception raised,
+           instance would be demoted."""
+
+    def get_cluster(self):
+        with self._cluster_thread_lock:
+            try:
+                self._load_cluster()
+            except:
+                self._cluster = None
+                raise
+            return self._cluster
+
+    @property
+    def cluster(self):
+        with self._cluster_thread_lock:
+            return self._cluster
+
+    def reset_cluster(self):
+        with self._cluster_thread_lock:
+            self._cluster = None
+
+    @abc.abstractmethod
+    def write_leader_optime(self, last_operation):
+        """write current xlog location into `/optime/leader` key in DCS
+        :param last_operation: absolute xlog location in bytes"""
+
+    @abc.abstractmethod
+    def update_leader(self):
+        """Update leader key (or session) ttl
+
+        :returns: `!True` if leader key (or session) has been updated successfully.
+            If not, `!False` must be returned and current instance would be demoted.
+
+        You have to use CAS (Compare And Swap) operation in order to update leader key,
+        for example for etcd `prevValue` parameter must be used."""
+
+    @abc.abstractmethod
+    def attempt_to_acquire_leader(self):
+        """Attempt to acquire leader lock
+        This method should create `/leader` key with value=`~self._name`
+        :returns: `!True` if key has been created successfully.
+
+        Key must be created atomically. In case if key already exists it should not be
+        overwritten and `!False` must be returned"""
+
+    @abc.abstractmethod
+    def set_failover_value(self, value, index=None):
+        """Create or update `/failover` key"""
+
+    def manual_failover(self, leader, member, index=None):
+        return self.set_failover_value(leader + (':' + member if member else ''), index)
+
+    def current_leader(self):
+        try:
+            cluster = self.get_cluster()
+            return None if cluster.is_unlocked() else cluster.leader
+        except DCSError:
+            return None
+
+    @abc.abstractmethod
+    def touch_member(self, connection_string, ttl=None):
+        """Update member key in DCS.
+        This method should create or update key with the name = '/members/' + `~self._name`
+        and value = connection_string in a given DCS.
+
+        :param connection_string: how this instance can be accessed by other instances
+        :param ttl: ttl for member key, optional parameter. If it is None `~self.member_ttl will be used`
+        :returns: `!True` on success otherwise `!False`
+        """
+
+    @abc.abstractmethod
+    def take_leader(self):
+        """This method should create leader key with value = `~self._name` and ttl=`~self.ttl`
+        Since it could be called only on initial cluster bootstrap it could create this key regardless,
+        overwriting the key if necessary."""
+
+    @abc.abstractmethod
+    def initialize(self):
+        """Race for cluster initialization.
+        :returns: `!True` if key has been created successfully.
+
+        this method should create atomically initialize key and return `!True`
+        otherwise it should return `!False`"""
+
+    @abc.abstractmethod
+    def delete_leader(self):
+        """Voluntarily remove leader key from DCS
+        This method should remove leader key if current instance is the leader"""
+
+    @abc.abstractmethod
+    def cancel_initialization(self):
+        """ Removes the initialize key for a cluster """
+
+    def watch(self, timeout):
+        """If the current node is a master it should just sleep.
+        Any other node should watch for changes of leader key with a given timeout
+
+        :param timeout: timeout in seconds
+        :returns: `!True` if you would like to reschedule the next run of ha cycle"""
+
+        self.event.wait(timeout)
+        return self.event.isSet()

patroni/etcd.py

@@ -5,11 +5,14 @@ import os
 import random
 import requests
 import socket
+import time
+import urllib3
 
 from dns.exception import DNSException
 from dns import resolver
-from helpers.dcs import AbstractDCS, Cluster, DCSError, Member, parse_connection_string
-from helpers.utils import sleep
+from patroni.dcs import AbstractDCS, Cluster, Failover, Leader, Member
+from patroni.exceptions import DCSError
+from patroni.utils import Retry, RetryFailedError, sleep
 from requests.exceptions import RequestException
 
 logger = logging.getLogger(__name__)
@@ -59,6 +62,16 @@ class Client(etcd.Client):
             logger.exception('Can not resolve SRV for %s', host)
         return []
 
+    # try to workarond bug in python-etcd: https://github.com/jplana/python-etcd/issues/81
+    def _result_from_response(self, response):
+        try:
+            response.data.decode('utf-8')
+        except urllib3.exceptions.TimeoutError:
+            raise
+        except Exception as e:
+            raise etcd.EtcdException('Unable to decode server response: %s' % e)
+        return super(Client, self)._result_from_response(response)
+
     def _get_machines_cache_from_srv(self, discovery_srv):
         """Fetch list of etcd-cluster member by resolving _etcd-server._tcp. SRV record.
         This record should contain list of host and peer ports which could be used to run
@@ -68,7 +81,7 @@ class Client(etcd.Client):
         for host, port in self.get_srv_record(discovery_srv):
             url = '{}://{}:{}/members'.format(self._protocol, host, port)
             try:
-                response = requests.get(url)
+                response = requests.get(url, timeout=5)
                 if response.ok:
                     for member in response.json():
                         ret.extend(member['clientURLs'])
@@ -124,7 +137,7 @@ def catch_etcd_errors(func):
     def wrapper(*args, **kwargs):
         try:
             return not func(*args, **kwargs) is None
-        except etcd.EtcdException:
+        except (RetryFailedError, etcd.EtcdException):
             return False
     return wrapper
 
@@ -134,9 +147,16 @@ class Etcd(AbstractDCS):
     def __init__(self, name, config):
         super(Etcd, self).__init__(name, config)
         self.ttl = config['ttl']
-        self.member_ttl = config.get('member_ttl', 3600)
+        self._retry = Retry(deadline=10, max_delay=1, max_tries=-1,
+                            retry_exceptions=(etcd.EtcdConnectionFailed,
+                                              etcd.EtcdLeaderElectionInProgress,
+                                              etcd.EtcdWatcherCleared,
+                                              etcd.EtcdEventIndexCleared))
         self.client = self.get_etcd_client(config)
 
+    def retry(self, *args, **kwargs):
+        return self._retry.copy()(*args, **kwargs)
+
     def get_etcd_client(self, config):
         client = None
         while not client:
@@ -149,66 +169,105 @@ class Etcd(AbstractDCS):
 
     @staticmethod
     def member(node):
-        conn_url, api_url = parse_connection_string(node.value)
-        return Member(node.modifiedIndex, os.path.basename(node.key), conn_url, api_url, node.expiration, node.ttl)
+        return Member.from_node(node.modifiedIndex, os.path.basename(node.key), node.ttl, node.value)
 
-    def get_cluster(self):
+    def _load_cluster(self):
         try:
-            result = self.client.read(self.client_path(''), recursive=True)
+            result = self.retry(self.client.read, self.client_path(''), recursive=True)
             nodes = {os.path.relpath(node.key, result.key): node for node in result.leaves}
 
             # get initialize flag
-            initialize = bool(nodes.get('initialize', False))
+            initialize = bool(nodes.get(self._INITIALIZE, False))
 
             # get last leader operation
-            last_leader_operation = nodes.get('optime/leader', None)
+            last_leader_operation = nodes.get(self._LEADER_OPTIME, None)
             last_leader_operation = 0 if last_leader_operation is None else int(last_leader_operation.value)
 
             # get list of members
-            members = [self.member(n) for k, n in nodes.items() if k.startswith('members/') and len(k.split('/')) == 2]
+            members = [self.member(n) for k, n in nodes.items() if k.startswith(self._MEMBERS) and k.count('/') == 1]
 
             # get leader
-            leader = nodes.get('leader', None)
+            leader = nodes.get(self._LEADER, None)
             if leader:
-                leader = Member(-1, leader.value, None, None, None, None)
-                leader = ([m for m in members if m.name == leader.name] or [leader])[0]
+                member = Member(-1, leader.value, None, {})
+                member = ([m for m in members if m.name == leader.value] or [member])[0]
+                leader = Leader(leader.modifiedIndex, leader.ttl, member)
 
-            return Cluster(initialize, leader, last_leader_operation, members)
+            # failover key
+            failover = nodes.get(self._FAILOVER, None)
+            if failover:
+                failover = Failover.from_node(failover.modifiedIndex, failover.value)
+
+            self._cluster = Cluster(initialize, leader, last_leader_operation, members, failover)
         except etcd.EtcdKeyNotFound:
-            return Cluster(False, None, None, [])
+            self._cluster = Cluster(False, None, None, [], None)
         except:
             logger.exception('get_cluster')
-
-        raise EtcdError('Etcd is not responding properly')
+            raise EtcdError('Etcd is not responding properly')
 
     @catch_etcd_errors
     def touch_member(self, connection_string, ttl=None):
-        return self.client.set(self.client_path('/members/' + self._name), connection_string, ttl or self.member_ttl)
+        return self.retry(self.client.set, self.member_path, connection_string, ttl or self.ttl)
 
     @catch_etcd_errors
     def take_leader(self):
-        return self.client.set(self.client_path('/leader'), self._name, self.ttl)
+        return self.retry(self.client.set, self.leader_path, self._name, self.ttl)
 
-    @catch_etcd_errors
     def attempt_to_acquire_leader(self):
-        ret = self.client.write(self.client_path('/leader'), self._name, ttl=self.ttl, prevExist=False)
-        ret or logger.info('Could not take out TTL lock')
-        return ret
+        try:
+            return bool(self.retry(self.client.write, self.leader_path, self._name, ttl=self.ttl, prevExist=False))
+        except etcd.EtcdAlreadyExist:
+            logger.info('Could not take out TTL lock')
+        except (RetryFailedError, etcd.EtcdException):
+            pass
+        return False
 
     @catch_etcd_errors
-    def write_leader_optime(self, state_handler):
-        return self.client.set(self.client_path('/optime/leader'), state_handler.last_operation())
+    def set_failover_value(self, value, index=None):
+        return self.client.write(self.failover_path, value, prevIndex=index or 0)
 
     @catch_etcd_errors
-    def update_leader(self, state_handler):
-        ret = self.client.test_and_set(self.client_path('/leader'), self._name, self._name, self.ttl)
-        ret and self.write_leader_optime(state_handler)
-        return ret
+    def write_leader_optime(self, last_operation):
+        return self.client.set(self.leader_optime_path, last_operation)
 
     @catch_etcd_errors
-    def race(self, path):
-        return self.client.write(self.client_path(path), self._name, prevExist=False)
+    def update_leader(self):
+        return self.retry(self.client.test_and_set, self.leader_path, self._name, self._name, self.ttl)
+
+    @catch_etcd_errors
+    def initialize(self):
+        return self.retry(self.client.write, self.initialize_path, self._name, prevExist=False)
 
     @catch_etcd_errors
     def delete_leader(self):
-        return self.client.delete(self.client_path('/leader'), prevValue=self._name)
+        return self.client.delete(self.leader_path, prevValue=self._name)
+
+    @catch_etcd_errors
+    def cancel_initialization(self):
+        return self.retry(self.client.delete, self.initialize_path, prevValue=self._name)
+
+    def watch(self, timeout):
+        cluster = self.cluster
+        # watch on leader key changes if it is defined and current node is not lock owner
+        if cluster and cluster.leader and cluster.leader.name != self._name:
+            end_time = time.time() + timeout
+            index = cluster.leader.index
+
+            while index and timeout >= 1:  # when timeout is too small urllib3 doesn't have enough time to connect
+                try:
+                    self.client.watch(self.leader_path, index=index + 1, timeout=timeout + 0.5)
+                    # Synchronous work of all cluster members with etcd is less expensive
+                    # than reestablishing http connection every time from every replica.
+                    return True
+                except urllib3.exceptions.TimeoutError:
+                    self.client.http.clear()
+                    return False
+                except etcd.EtcdException:
+                    logging.exception('watch')
+
+                timeout = end_time - time.time()
+
+        try:
+            return super(Etcd, self).watch(timeout)
+        finally:
+            self.event.clear()

patroni/exceptions.py

@@ -0,0 +1,25 @@
+class PatroniException(Exception):
+
+    """Parent class for all kind of exceptions related to selected distributed configuration store"""
+
+    def __init__(self, value):
+        self.value = value
+
+    def __str__(self):
+        """
+        >>> str(PatroniException('foo'))
+        "'foo'"
+        """
+        return repr(self.value)
+
+
+class PostgresException(PatroniException):
+    pass
+
+
+class DCSError(PatroniException):
+    pass
+
+
+class PostgresConnectionException(PostgresException):
+    pass

patroni/ha.py

@@ -0,0 +1,401 @@
+import json
+import logging
+import psycopg2
+import requests
+
+from patroni.async_executor import AsyncExecutor
+from patroni.exceptions import DCSError, PostgresConnectionException
+from multiprocessing.pool import ThreadPool
+
+logger = logging.getLogger(__name__)
+
+
+class Ha:
+
+    def __init__(self, patroni):
+        self.patroni = patroni
+        self.state_handler = patroni.postgresql
+        self.dcs = patroni.dcs
+        self.cluster = None
+        self.old_cluster = None
+        self._async_executor = AsyncExecutor()
+
+    def load_cluster_from_dcs(self):
+        cluster = self.dcs.get_cluster()
+
+        # We want to keep the state of cluster when it was healhy
+        if not cluster.is_unlocked() or not self.old_cluster:
+            self.old_cluster = cluster
+        self.cluster = cluster
+
+    def acquire_lock(self):
+        return self.dcs.attempt_to_acquire_leader()
+
+    def update_lock(self):
+        ret = self.dcs.update_leader()
+        if ret:
+            try:
+                self.dcs.write_leader_optime(self.state_handler.last_operation())
+            except:
+                pass
+        return ret
+
+    def has_lock(self):
+        lock_owner = self.cluster.leader and self.cluster.leader.name
+        logger.info('Lock owner: %s; I am %s', lock_owner, self.state_handler.name)
+        return lock_owner == self.state_handler.name
+
+    def touch_member(self):
+        data = {
+            'conn_url': self.state_handler.connection_string,
+            'api_url': self.patroni.api.connection_string,
+            'state': self.state_handler.state,
+            'role': self.state_handler.role
+        }
+        if data['state'] in ['running', 'restarting', 'starting']:
+            try:
+                data['xlog_location'] = self.state_handler.xlog_position()
+            except:
+                pass
+        self.dcs.touch_member(json.dumps(data, separators=(',', ':')))
+
+    def copy_backup_from_leader(self, leader):
+        if self.state_handler.bootstrap(leader):
+            logger.info('bootstrapped from leader')
+        else:
+            self.state_handler.stop('immediate')
+            self.state_handler.remove_data_directory()
+            logger.error('failed to bootstrap from leader')
+
+    def bootstrap(self):
+        if not self.cluster.is_unlocked():  # cluster already has leader
+            self._async_executor.schedule('bootstrap from leader')
+            self._async_executor.run_async(self.copy_backup_from_leader, args=(self.cluster.leader, ))
+            return 'trying to bootstrap from leader'
+        elif not self.cluster.initialize:  # no initialize key
+            if self.dcs.initialize():  # race for initialization
+                try:
+                    self.state_handler.bootstrap()
+                except:  # initdb or start failed
+                    # remove initialization key and give a chance to other members
+                    logger.info("removing initialize key after failed attempt to initialize the cluster")
+                    self.dcs.cancel_initialization()
+                    self.state_handler.stop('immediate')
+                    self.state_handler.move_data_directory()
+                    raise
+                self.dcs.take_leader()
+                return 'initialized a new cluster'
+            else:
+                return 'failed to acquire initialize lock'
+        else:
+            return 'waiting for leader to bootstrap'
+
+    def recover(self):
+        has_lock = self.has_lock()
+
+        # try to see if we are the former master that crashed. If so - we likely need to run pg_rewind
+        # in order to join the former standby being promoted.
+        pg_controldata = self.state_handler.controldata()
+        if not has_lock and pg_controldata and\
+                pg_controldata.get('Database cluster state', '') == 'in production':  # crashed master
+            self.state_handler.require_rewind()
+
+        # XXX: follow the leader calls stop, which might take quite some time.
+        # perhaps we should run sync asynchronously
+        # (we still need the exit code from follow_the_leader)
+        ret = self.state_handler.follow_the_leader(None if has_lock else self.cluster.leader, recovery=True)
+        if not ret:
+            if not has_lock:
+                return 'failed to start postgres'
+            self.dcs.delete_leader()
+            self.dcs.reset_cluster()
+            return 'removed leader key after trying and failing to start postgres'
+        if not has_lock:
+            return 'started as a secondary'
+        logger.info('started as readonly because i had the session lock')
+        self.load_cluster_from_dcs()
+
+    def follow_the_leader(self, demote_reason, follow_reason, refresh=True):
+        refresh and self.load_cluster_from_dcs()
+        ret = demote_reason if self.state_handler.is_leader() else follow_reason
+        leader = self.cluster.leader
+        leader = None if (leader and leader.name) == self.state_handler.name else leader
+        if not self.state_handler.check_recovery_conf(leader):
+            self._async_executor.schedule('changing primary_conninfo and restarting')
+            self._async_executor.run_async(self.state_handler.follow_the_leader, (leader, ))
+        return ret
+
+    def enforce_master_role(self, message, promote_message):
+        if self.state_handler.is_leader() or self.state_handler.role == 'master':
+            return message
+        else:
+            self.state_handler.promote()
+            return promote_message
+
+    @staticmethod
+    def fetch_node_status(member):
+        """This function perform http get request on member.api_url and fetches its status
+        :returns: tuple(`member`, reachable, in_recovery, xlog_location)
+
+        reachable - `!False` if the node is not reachable or is not responding with correct JSON
+        in_recovery - `!True` if pg_is_in_recovery() == true
+        xlog_location - value of `replayed_location` or `location` from JSON, dependin on its role."""
+
+        try:
+            response = requests.get(member.api_url, timeout=2, verify=False)
+            logger.info('Got response from %s %s: %s', member.name, member.api_url, response.content)
+            json = response.json()
+            is_master = json['role'] == 'master'
+            xlog_location = json['xlog']['location' if is_master else 'replayed_location']
+            return (member, True, not is_master, xlog_location)
+        except:
+            logging.exception('request failed: GET %s', member.api_url)
+        return (member, False, None, 0)
+
+    def fetch_nodes_statuses(self, members):
+        pool = ThreadPool(len(members))
+        results = pool.map(self.fetch_node_status, members)  # Run API calls on members in parallel
+        pool.close()
+        pool.join()
+        return results
+
+    def _is_healthiest_node(self, members, check_replication_lag=True):
+        """This method tries to determine whether I am healthy enough to became a new leader candidate or not."""
+
+        if self.state_handler.is_leader():
+            return True
+
+        if check_replication_lag and not self.state_handler.check_replication_lag(self.cluster.last_leader_operation):
+            return False  # Too far behind last reported xlog location on master
+
+        # Prepare list of nodes to run check against
+        members = [m for m in members if m.name != self.state_handler.name and m.api_url]
+
+        if members:
+            my_xlog_location = self.state_handler.xlog_position()
+            for member, reachable, in_recovery, xlog_location in self.fetch_nodes_statuses(members):
+                if reachable:  # If the node is unreachable it's not healhy
+                    if not in_recovery:
+                        logger.warning('Master (%s) is still alive', member.name)
+                        return False
+                    if my_xlog_location < xlog_location:
+                        return False
+        return True
+
+    def is_failover_possible(self, members):
+        ret = False
+        members = [m for m in members if m.name != self.state_handler.name and m.api_url]
+        if members:
+            for member, reachable, in_recovery, xlog_location in self.fetch_nodes_statuses(members):
+                if reachable:
+                    ret = True  # TODO: check xlog_location
+                else:
+                    logger.info('Member %s is not reachable', member.name)
+        else:
+            logger.warning('manual failover: members list is empty')
+        return ret
+
+    def manual_failover_process_no_leader(self):
+        failover = self.cluster.failover
+        if failover.member:  # manual failover to specific member
+            if failover.member == self.state_handler.name:  # manual failover to me
+                return True
+
+            # find specific node and check that it is healthy
+            members = [m for m in self.cluster.members if m.name == failover.member]
+            if members:
+                member, reachable, in_recovery, xlog_location = self.fetch_node_status(members[0])
+                if reachable:  # node is healthy
+                    logger.info('manual failover: to %s, i am %s', member.name, self.state_handler.name)
+                    return False
+                # we wanted to failover to specific member but it is not healthy
+                logger.warning('manual failover: member %s is unhealthy', member.name)
+
+            # at this point we should consider all members as a candidates for failover
+            # i.e. we assume that failover.member is None
+
+        # try to pick some other members to failover and check that they are healthy
+        if failover.leader:
+            if self.state_handler.name == failover.leader:  # I was the leader
+                # exclude me and desired member which is unhealthy (failover.member can be None)
+                members = [m for m in self.cluster.members if m.name != failover.member]
+                if self.is_failover_possible(members):  # check that there are healthy members
+                    return False
+                else:  # I was the leader and it looks like currently I am the only healthy member
+                    return True
+
+            # at this point we assume that our node is a candidate for a failover among all nodes except former leader
+
+        # exclude former leader from the list (failover.leader can be None)
+        members = [m for m in self.cluster.members if m.name != failover.leader]
+        return self._is_healthiest_node(members, check_replication_lag=False)
+
+    def is_healthiest_node(self):
+        if self.cluster.failover:
+            return self.manual_failover_process_no_leader()
+
+        # run usual health check
+        members = {m.name: m for m in self.cluster.members + self.old_cluster.members}
+        return self._is_healthiest_node(members.values())
+
+    def demote(self, delete_leader=True):
+        if delete_leader:
+            self.state_handler.stop()
+            self.dcs.delete_leader()
+            self.dcs.reset_cluster()
+        self.state_handler.follow_the_leader(None)
+
+    def process_manual_failover_from_leader(self):
+        failover = self.cluster.failover
+        if not failover.leader or failover.leader == self.state_handler.name:
+            if not failover.member or failover.member != self.state_handler.name:
+                members = [m for m in self.cluster.members if not failover.member or m.name == failover.member]
+                if self.is_failover_possible(members):  # check that there are healthy members
+                    self._async_executor.schedule('manual failover: demote')
+                    self._async_executor.run_async(self.demote)
+                    return 'manual failover: demoting myself'
+                else:
+                    logger.warning('manual failover: no healthy members found, failover is not possible')
+            else:
+                logger.warning('manual failover: I am already the leader, no need to failover')
+        else:
+            logger.warning('manual failover: leader name does not match: %s != %s',
+                           self.cluster.failover.leader, self.state_handler.name)
+
+        logger.info('Trying to clean up failover key')
+        self.dcs.manual_failover('', '', self.cluster.failover.index)
+
+    def process_unhealthy_cluster(self):
+        if self.is_healthiest_node():
+            if self.acquire_lock():
+                if self.cluster.failover:
+                    logger.info('Cleanning up failover key after acquiring leader lock...')
+                    self.dcs.manual_failover('', '')
+                return self.enforce_master_role('acquired session lock as a leader',
+                                                'promoted self to leader by acquiring session lock')
+            else:
+                return self.follow_the_leader('demoted self due after trying and failing to obtain lock',
+                                              'following new leader after trying and failing to obtain lock')
+        else:
+            return self.follow_the_leader('demoting self because i am not the healthiest node',
+                                          'following a different leader because i am not the healthiest node')
+
+    def process_healthy_cluster(self):
+        if self.has_lock():
+            if self.cluster.failover:
+                msg = self.process_manual_failover_from_leader()
+                if msg is not None:
+                    return msg
+
+            if self.update_lock():
+                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:
+                # Either there is no connection to DCS or someone else acquired the lock
+                logger.error('failed to update leader lock')
+                self.load_cluster_from_dcs()
+        else:
+            logger.info('does not have lock')
+        return self.follow_the_leader('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 restart_scheduled(self):
+        return self._async_executor.scheduled_action == 'restart'
+
+    def schedule_reinitialize(self):
+        return self.schedule('reinitialize')
+
+    def reinitialize_scheduled(self):
+        return self._async_executor.scheduled_action == 'reinitialize'
+
+    def restart(self):
+        with self._async_executor:
+            prev = self._async_executor.schedule('restart', True)
+            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')
+
+    def reinitialize(self, cluster):
+        self.state_handler.stop('immediate')
+        self.state_handler.remove_data_directory()
+        self.copy_backup_from_leader(cluster.leader)
+
+    def process_scheduled_action(self):
+        if self.reinitialize_scheduled():
+            if self.cluster.is_unlocked():
+                logger.error('Cluster has no leader, can not reinitialize')
+                self._async_executor.reset_scheduled_action()
+            elif self.has_lock():
+                logger.error('I am the leader, can not reinitialize')
+                self._async_executor.reset_scheduled_action()
+            else:
+                self._async_executor.run_async(self.reinitialize, args=(self.cluster, ))
+                return 'reinitialize started'
+
+    def handle_long_action_in_progress(self):
+        if self.has_lock():
+            if self.update_lock():
+                return 'updated leader lock during ' + self._async_executor.scheduled_action
+            else:
+                return 'failed to update leader lock during ' + self._async_executor.scheduled_action
+        elif self.cluster.is_unlocked():
+            return 'not healthy enough for leader race'
+        else:
+            return self._async_executor.scheduled_action + ' in progress'
+
+    def _run_cycle(self):
+        try:
+            self.load_cluster_from_dcs()
+
+            self.touch_member()
+
+            # cluster has leader key but not initialize key
+            if not self.cluster.is_unlocked() and not self.cluster.initialize:
+                self.dcs.initialize()  # fix it
+
+            if self._async_executor.busy:
+                return self.handle_long_action_in_progress()
+
+            # currently it can trigger only reinitialize
+            msg = self.process_scheduled_action()
+            if msg is not None:
+                return msg
+
+            # is data directory empty?
+            if self.state_handler.data_directory_empty():
+                return self.bootstrap()  # new node
+            # "bootstrap", but data directory is not empty
+            elif not self.cluster.initialize and self.cluster.is_unlocked():
+                self.dcs.initialize()
+
+            # try to start dead postgres
+            if not self.state_handler.is_healthy():
+                msg = self.recover()
+                if msg is not None:
+                    return msg
+
+            try:
+                if self.cluster.is_unlocked():
+                    return self.process_unhealthy_cluster()
+                else:
+                    return self.process_healthy_cluster()
+            finally:
+                self.state_handler.sync_replication_slots(self.cluster)
+        except DCSError:
+            logger.error('Error communicating with DCS')
+            if self.state_handler.is_running() and self.state_handler.is_leader():
+                self.demote(delete_leader=False)
+                return 'demoted self because DCS is not accessible and i was a leader'
+        except (psycopg2.Error, PostgresConnectionException):
+            logger.exception('Error communicating with Postgresql. Will try again later')
+
+    def run_cycle(self):
+        with self._async_executor:
+            return self._run_cycle()

patroni/postgresql.py

@@ -0,0 +1,601 @@
+import logging
+import os
+import psycopg2
+import shlex
+import shutil
+import subprocess
+import time
+
+from patroni.exceptions import PostgresConnectionException, PostgresException
+from patroni.utils import Retry, RetryFailedError
+from six.moves.urllib_parse import urlparse
+from threading import Lock
+
+logger = logging.getLogger(__name__)
+
+ACTION_ON_START = "on_start"
+ACTION_ON_STOP = "on_stop"
+ACTION_ON_RESTART = "on_restart"
+ACTION_ON_RELOAD = "on_reload"
+ACTION_ON_ROLE_CHANGE = "on_role_change"
+
+
+def parseurl(url):
+    r = urlparse(url)
+    ret = {
+        'host': r.hostname,
+        'port': r.port or 5432,
+        'database': r.path[1:],
+        'fallback_application_name': 'Patroni',
+        'connect_timeout': 3,
+        'options': '-c statement_timeout=2000',
+    }
+    if r.username:
+        ret['user'] = r.username
+    if r.password:
+        ret['password'] = r.password
+    return ret
+
+
+class Postgresql:
+
+    def __init__(self, config):
+        self.config = config
+        self.name = config['name']
+        self.scope = config['scope']
+        self.listen_addresses, self.port = config['listen'].split(':')
+        self.data_dir = config['data_dir']
+        self.replication = config['replication']
+        self.superuser = config['superuser']
+        self.admin = config['admin']
+        self.pg_rewind = config.get('pg_rewind', {})
+        self.callback = config.get('callbacks', {})
+        self.use_slots = config.get('use_slots', True)
+        self.schedule_load_slots = self.use_slots
+        self.recovery_conf = os.path.join(self.data_dir, 'recovery.conf')
+        self.configuration_to_save = (os.path.join(self.data_dir, 'pg_hba.conf'),
+                                      os.path.join(self.data_dir, 'postgresql.conf'))
+        self.postmaster_pid = os.path.join(self.data_dir, 'postmaster.pid')
+        self.trigger_file = config.get('recovery_conf', {}).get('trigger_file', None) or 'promote'
+        self.trigger_file = os.path.abspath(os.path.join(self.data_dir, self.trigger_file))
+
+        self._pg_ctl = ['pg_ctl', '-w', '-D', self.data_dir]
+
+        self.local_address = self.get_local_address()
+        connect_address = config.get('connect_address', None) or self.local_address
+        self.connection_string = 'postgres://{username}:{password}@{connect_address}/postgres'.format(
+            connect_address=connect_address, **self.replication)
+
+        self._connection = None
+        self._cursor_holder = None
+        self._need_rewind = False
+        self.replication_slots = []  # list of already existing replication slots
+        self.retry = Retry(max_tries=-1, deadline=5, max_delay=1, retry_exceptions=PostgresConnectionException)
+
+        self._state = 'stopped'
+        self._state_lock = Lock()
+        self._role = 'replica'
+        self._role_lock = Lock()
+
+        if self.is_running():
+            self._state = 'running'
+            self._role = 'master' if self.is_leader() else 'replica'
+
+    @property
+    def can_rewind(self):
+        """ check if pg_rewind executable is there and that pg_controldata indicates
+            we have either wal_log_hints or checksums turned on
+        """
+        # low-hanging fruit: check if pg_rewind configuration is there
+        if not self.pg_rewind or\
+           not (self.pg_rewind.get('username', '') and self.pg_rewind.get('password', '')):
+            return False
+
+        cmd = ['pg_rewind', '--help']
+        try:
+            ret = subprocess.call(cmd, stdout=open(os.devnull, 'w'), stderr=subprocess.STDOUT)
+            if ret != 0:  # pg_rewind is not there, close up the shop and go home
+                return False
+        except OSError:
+            return False
+        # check if the cluster's configuration permits pg_rewind
+        data = self.controldata()
+        if data:
+            return data.get('wal_log_hints setting', 'off') == 'on' or\
+                data.get('Data page checksum version', '0') != '0'
+        return False
+
+    def require_rewind(self):
+        self._need_rewind = True
+
+    def get_local_address(self):
+        listen_addresses = self.listen_addresses.split(',')
+        local_address = listen_addresses[0].strip()  # take first address from listen_addresses
+
+        for la in listen_addresses:
+            if la.strip() in ['*', '0.0.0.0']:  # we are listening on *
+                local_address = 'localhost'  # connection via localhost is preferred
+                break
+        return local_address + ':' + self.port
+
+    def connection(self):
+        if not self._connection or self._connection.closed != 0:
+            r = parseurl('postgres://{}/postgres'.format(self.local_address))
+            self._connection = psycopg2.connect(**r)
+            self._connection.autocommit = True
+        return self._connection
+
+    def _cursor(self):
+        if not self._cursor_holder or self._cursor_holder.closed or self._cursor_holder.connection.closed != 0:
+            self._cursor_holder = self.connection().cursor()
+        return self._cursor_holder
+
+    def _query(self, sql, *params):
+        cursor = None
+        try:
+            cursor = self._cursor()
+            cursor.execute(sql, params)
+            return cursor
+        except psycopg2.Error as e:
+            if cursor and cursor.connection.closed == 0:
+                raise e
+            if self.state == 'restarting':
+                raise RetryFailedError('cluster is being restarted')
+            raise PostgresConnectionException('connection problems')
+
+    def query(self, sql, *params):
+        try:
+            return self.retry(self._query, sql, *params)
+        except RetryFailedError as e:
+            raise PostgresConnectionException(str(e))
+
+    def data_directory_empty(self):
+        return not os.path.exists(self.data_dir) or os.listdir(self.data_dir) == []
+
+    def initialize(self):
+        self.set_state('initalizing new cluster')
+        ret = subprocess.call(self._pg_ctl + ['initdb', '-o', '--encoding=UTF8']) == 0
+        if ret:
+            self.write_pg_hba()
+        else:
+            self.set_state('initdb failed')
+        return ret
+
+    def delete_trigger_file(self):
+        os.path.exists(self.trigger_file) and os.unlink(self.trigger_file)
+
+    def write_pgpass(self, record):
+        pgpass = 'pgpass'
+        with open(pgpass, 'w') as f:
+            os.fchmod(f.fileno(), 0o600)
+            f.write('{host}:{port}:*:{user}:{password}\n'.format(**record))
+        env = os.environ.copy()
+        env['PGPASSFILE'] = pgpass
+        return env
+
+    def sync_from_leader(self, leader):
+        r = parseurl(leader.conn_url)
+
+        env = self.write_pgpass(r)
+        return self.create_replica(r, env) == 0
+
+    @staticmethod
+    def build_connstring(conn):
+        return "host={host} port={port} user={user}".format(**conn)
+
+    def create_replica(self, master_connection, env):
+        self.set_state('building replica from {host}:{port}'.format(**master_connection))
+        connstring = self.build_connstring(master_connection)
+        cmd = self.config['restore']
+        try:
+            ret = subprocess.call(shlex.split(cmd) + [self.scope, "replica", self.data_dir, connstring], env=env)
+            self.delete_trigger_file()
+        except:
+            logger.exception('Error when creating replica')
+            ret = 1
+        if ret != 0:
+            self.set_state('failed to build replica from {host}:{port}'.format(**master_connection))
+        return ret
+
+    def is_leader(self):
+        return not self.query('SELECT pg_is_in_recovery()').fetchone()[0]
+
+    def is_running(self):
+        return subprocess.call(' '.join(self._pg_ctl) + ' status > /dev/null 2>&1', shell=True) == 0
+
+    def call_nowait(self, cb_name):
+        """ pick a callback command and call it without waiting for it to finish """
+        if not self.callback or cb_name not in self.callback:
+            return False
+        cmd = self.callback[cb_name]
+        try:
+            subprocess.Popen(shlex.split(cmd) + [cb_name, self.role, self.scope])
+        except:
+            logger.exception('callback %s %s %s %s failed', cmd, cb_name, self.role, self.scope)
+            return False
+        return True
+
+    @property
+    def role(self):
+        with self._role_lock:
+            return self._role
+
+    def set_role(self, value):
+        with self._role_lock:
+            self._role = value
+
+    @property
+    def state(self):
+        with self._state_lock:
+            return self._state
+
+    def set_state(self, value):
+        with self._state_lock:
+            self._state = value
+
+    def start(self, block_callbacks=False):
+        if self.is_running():
+            logger.error('Cannot start PostgreSQL because one is already running.')
+            return True
+
+        self.set_role('replica' if os.path.exists(self.recovery_conf) else 'master')
+        if os.path.exists(self.postmaster_pid):
+            os.remove(self.postmaster_pid)
+            logger.info('Removed %s', self.postmaster_pid)
+
+        if not block_callbacks:
+            self.set_state('starting')
+
+        ret = subprocess.call(self._pg_ctl + ['start', '-o', self.server_options()]) == 0
+
+        self.set_state('running' if ret else 'start failed')
+
+        self.schedule_load_slots = ret and self.use_slots
+        self.save_configuration_files()
+        # block_callbacks is used during restart to avoid
+        # running start/stop callbacks in addition to restart ones
+        ret and not block_callbacks and self.call_nowait(ACTION_ON_START)
+        return ret
+
+    def checkpoint(self):
+        try:
+            r = parseurl('postgres://{}/postgres'.format(self.local_address))
+            r['options'] = '-c statement_timeout=0'
+            with psycopg2.connect(**r) as conn:
+                conn.autocommit = True
+                with conn.cursor() as cur:
+                    cur.execute('CHECKPOINT')
+        except:
+            logging.exception('Exception during CHECKPOINT')
+
+    def stop(self, mode='fast', block_callbacks=False):
+        if not self.is_running():
+            if not block_callbacks:
+                self.set_state('stopped')
+            return True
+
+        if block_callbacks:
+            self.checkpoint()
+        else:
+            self.set_state('stopping')
+
+        ret = subprocess.call(self._pg_ctl + ['stop', '-m', mode]) == 0
+        # block_callbacks is used during restart to avoid
+        # running start/stop callbacks in addition to restart ones
+        if not ret:
+            self.set_state('stop failed')
+        elif not block_callbacks:
+            self.set_state('stopped')
+            self.call_nowait(ACTION_ON_STOP)
+        return ret
+
+    def reload(self):
+        ret = subprocess.call(self._pg_ctl + ['reload']) == 0
+        ret and self.call_nowait(ACTION_ON_RELOAD)
+        return ret
+
+    def restart(self):
+        self.set_state('restarting')
+        ret = self.stop(block_callbacks=True) and self.start(block_callbacks=True)
+        if ret:
+            self.call_nowait(ACTION_ON_RESTART)
+        else:
+            self.set_state('restart failed ({})'.format(self.state))
+        return ret
+
+    def server_options(self):
+        options = "--listen_addresses='{}' --port={}".format(self.listen_addresses, self.port)
+        for setting, value in self.config['parameters'].items():
+            options += " --{}='{}'".format(setting, value)
+        return options
+
+    def is_healthy(self):
+        if not self.is_running():
+            logger.warning('Postgresql is not running.')
+            return False
+        return True
+
+    def check_replication_lag(self, last_leader_operation):
+        return last_leader_operation - self.xlog_position() <= self.config.get('maximum_lag_on_failover', 0)
+
+    def write_pg_hba(self):
+        with open(os.path.join(self.data_dir, 'pg_hba.conf'), 'a') as f:
+            f.write('\nhost replication {username} {network} md5\n'.format(**self.replication))
+            for line in self.config.get('pg_hba', []):
+                if line.split()[0].strip() == 'hostssl' and self.config['parameters'].get('ssl', 'off').lower() != 'on':
+                    continue
+                f.write(line + '\n')
+
+    @staticmethod
+    def primary_conninfo(leader_url):
+        r = parseurl(leader_url)
+        return 'user={user} password={password} host={host} port={port} sslmode=prefer sslcompression=1'.format(**r)
+
+    def check_recovery_conf(self, leader):
+        if not os.path.isfile(self.recovery_conf):
+            return False
+
+        pattern = leader and leader.conn_url and self.primary_conninfo(leader.conn_url)
+
+        with open(self.recovery_conf, 'r') as f:
+            for line in f:
+                if line.startswith('primary_conninfo'):
+                    return pattern and (pattern in line)
+        return not pattern
+
+    def write_recovery_conf(self, leader):
+        with open(self.recovery_conf, 'w') as f:
+            f.write("""standby_mode = 'on'
+recovery_target_timeline = 'latest'
+""")
+            if leader and leader.conn_url:
+                f.write("""primary_conninfo = '{}'\n""".format(self.primary_conninfo(leader.conn_url)))
+                if self.use_slots:
+                    f.write("""primary_slot_name = '{}'\n""".format(self.name))
+                for name, value in self.config.get('recovery_conf', {}).items():
+                    f.write("{} = '{}'\n".format(name, value))
+
+    def rewind(self, leader):
+        # prepare pg_rewind connection
+        r = parseurl(leader.conn_url)
+        r.update(self.pg_rewind)
+        r['user'] = r['username']
+        env = self.write_pgpass(r)
+        pc = "user={user} host={host} port={port} dbname=postgres sslmode=prefer sslcompression=1".format(**r)
+        logger.info("running pg_rewind from {}".format(pc))
+        pg_rewind = ['pg_rewind', '-D', self.data_dir, '--source-server', pc]
+        try:
+            ret = (subprocess.call(pg_rewind, env=env) == 0)
+        except:
+            ret = False
+        if ret:
+            self.write_recovery_conf(leader)
+        return ret
+
+    def controldata(self):
+        """ return the contents of pg_controldata, or non-True value if pg_controldata call failed """
+        result = {}
+        try:
+            data = subprocess.check_output(['pg_controldata', self.data_dir])
+            if data:
+                data = data.splitlines()
+                result = {l.split(':')[0].replace('Current ', '', 1): l.split(':')[1].strip() for l in data if l}
+        except subprocess.CalledProcessError:
+            logger.exception("Error when calling pg_controldata")
+        finally:
+            return result
+
+    def read_postmaster_opts(self):
+        """ returns the list of option names/values from postgres.opts, Empty dict if read failed or no file """
+        result = {}
+        try:
+            with open(os.path.join(self.data_dir, "postmaster.opts")) as f:
+                data = f.read()
+                opts = [opt.strip('"\n') for opt in data.split(' "')]
+                for opt in opts:
+                    if '=' in opt and opt.startswith('--'):
+                        name, val = opt.split('=', 1)
+                        name = name.strip('-')
+                        result[name] = val
+        except IOError:
+            logger.exception('Error when reading postmaster.opts')
+        finally:
+            return result
+
+    def single_user_mode(self, command=None, options={}):
+        """ run a given command in a single-user mode. If the command is empty - then just start and stop """
+        cmd = ['postgres', '--single', '-D', self.data_dir]
+        for opt in sorted(options):
+            cmd.extend(['-c', '{0}={1}'.format(opt, options[opt])])
+        # need a database name to connect
+        cmd.append('postgres')
+        p = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=open(os.devnull, 'w'), stderr=subprocess.STDOUT)
+        if p:
+            command and p.communicate('{}\n'.format(command))
+            p.stdin.close()
+            return p.wait()
+        return 1
+
+    def cleanup_archive_status(self):
+        status_dir = os.path.join(self.data_dir, 'pg_xlog', 'archive_status')
+        if os.path.isdir(status_dir):
+            for f in os.listdir(status_dir):
+                path = os.path.join(status_dir, f)
+                try:
+                    if os.path.islink(path):
+                        os.unlink(path)
+                    elif os.path.isfile(path):
+                        os.remove(path)
+                except:
+                    logger.exception("Unable to remove {}".format(path))
+
+    def follow_the_leader(self, leader, recovery=False):
+        if not self.check_recovery_conf(leader) or recovery:
+            change_role = (self.role == 'master')
+
+            self._need_rewind = (self._need_rewind or change_role) and self.can_rewind
+            if self._need_rewind:
+                logger.info("set the rewind flag after demote")
+            self.write_recovery_conf(leader)
+            if not leader or not self._need_rewind:  # do not rewind until the leader becomes available
+                ret = self.restart()
+            else:  # we have a leader and need to rewind
+                if self.is_running():
+                    self.stop()
+                # at present, pg_rewind only runs when the cluster is shut down cleanly
+                # and not shutdown in recovery. We have to remove the recovery.conf if present
+                # and start/shutdown in a single user mode to emulate this.
+                # XXX: if recovery.conf is linked, it will be written anew as a normal file.
+                if os.path.islink(self.recovery_conf):
+                    os.unlink(self.recovery_conf)
+                else:
+                    os.remove(self.recovery_conf)
+                # Archived segments might be useful to pg_rewind,
+                # clean the flags that tell we should remove them.
+                self.cleanup_archive_status()
+                # Start in a single user mode and stop to produce a clean shutdown
+                opts = self.read_postmaster_opts()
+                opts['archive_mode'] = 'on'
+                opts['archive_command'] = 'false'
+                self.single_user_mode(options=opts)
+                if self.rewind(leader):
+                    ret = self.start()
+                else:
+                    logger.error("unable to rewind the former master")
+                    self.remove_data_directory()
+                    ret = True
+                self._need_rewind = False
+            change_role and ret and self.call_nowait(ACTION_ON_ROLE_CHANGE)
+            return ret
+        else:
+            return True
+
+    def save_configuration_files(self):
+        """
+            copy postgresql.conf to postgresql.conf.backup to preserve it in the WAL-e backup.
+            see http://comments.gmane.org/gmane.comp.db.postgresql.wal-e/239
+        """
+        for f in self.configuration_to_save:
+            shutil.copy(f, f + '.backup')
+
+    def restore_configuration_files(self):
+        """ restore a previously saved postgresql.conf """
+        try:
+            for f in self.configuration_to_save:
+                shutil.copy(f + '.backup', f)
+        except:
+            logger.exception('unable to restore configuration from WAL-E backup')
+
+    def promote(self):
+        if self.role == 'master':
+            return True
+        ret = subprocess.call(self._pg_ctl + ['promote']) == 0
+        if ret:
+            self.set_role('master')
+            logger.info("cleared rewind flag after becoming the leader")
+            self._need_rewind = False
+            self.call_nowait(ACTION_ON_ROLE_CHANGE)
+        return ret
+
+    def demote(self):
+        self.follow_the_leader(None)
+
+    def create_replication_user(self):
+        self.query('CREATE USER "{}" WITH REPLICATION ENCRYPTED PASSWORD %s'.format(
+            self.replication['username']), self.replication['password'])
+
+    def create_connection_users(self):
+        if self.superuser:
+            if 'username' in self.superuser:
+                self.query('CREATE ROLE "{0}" WITH LOGIN SUPERUSER PASSWORD %s'.format(
+                    self.superuser['username']), self.superuser['password'])
+            else:
+                rolsuper = self.query("""SELECT rolname FROM pg_authid WHERE rolsuper = 't'""").fetchone()[0]
+                self.query('ALTER ROLE "{0}" WITH PASSWORD %s'.format(rolsuper), self.superuser['password'])
+        if self.admin:
+            self.query('CREATE ROLE "{0}" WITH LOGIN CREATEDB CREATEROLE PASSWORD %s'.format(
+                self.admin['username']), self.admin['password'])
+
+    def xlog_position(self):
+        return self.query("""SELECT pg_xlog_location_diff(CASE WHEN pg_is_in_recovery()
+                                                               THEN pg_last_xlog_replay_location()
+                                                               ELSE pg_current_xlog_location()
+                                                          END, '0/0')::bigint""").fetchone()[0]
+
+    def load_replication_slots(self):
+        if self.use_slots and self.schedule_load_slots:
+            cursor = self.query("SELECT slot_name FROM pg_replication_slots WHERE slot_type='physical'")
+            self.replication_slots = [r[0] for r in cursor]
+            self.schedule_load_slots = False
+
+    def sync_replication_slots(self, cluster):
+        if self.use_slots:
+            try:
+                self.load_replication_slots()
+                slots = [m.name for m in cluster.members if m.name != self.name] if self.role == 'master' else []
+                # drop unused slots
+                for slot in set(self.replication_slots) - set(slots):
+                    self.query("""SELECT pg_drop_replication_slot(%s)
+                                   WHERE EXISTS(SELECT 1 FROM pg_replication_slots
+                                   WHERE slot_name = %s)""", slot, slot)
+
+                # create new slots
+                for slot in set(slots) - set(self.replication_slots):
+                    self.query("""SELECT pg_create_physical_replication_slot(%s)
+                                   WHERE NOT EXISTS (SELECT 1 FROM pg_replication_slots
+                                   WHERE slot_name = %s)""", slot, slot)
+
+                self.replication_slots = slots
+            except:
+                logger.exception('Exception when changing replication slots')
+
+    def last_operation(self):
+        return str(self.xlog_position())
+
+    def bootstrap(self, current_leader=None):
+        """
+            Initially bootstrap PostgreSQL, either by creating a data
+            directory with initdb, or by initalizing a replica from an
+            exiting leader. Failure in the first case always leads to
+            exception, since there is no point in continuing if initdb failed.
+            In the second case, however, a False is returned on failure, since
+            it is normal for the replica to retry a failed attempt to initialize
+            from the master.
+        """
+        ret = False
+        if not current_leader:
+            ret = self.initialize() and self.start()
+            if ret:
+                self.create_replication_user()
+                self.create_connection_users()
+            else:
+                raise PostgresException("Could not bootstrap master PostgreSQL")
+        else:
+            if self.sync_from_leader(current_leader):
+                self.write_recovery_conf(current_leader)
+                ret = self.start()
+        return ret
+
+    def move_data_directory(self):
+        if os.path.isdir(self.data_dir) and not self.is_running():
+            try:
+                new_name = '{0}_{1}'.format(self.data_dir, time.strftime('%Y-%m-%d-%H-%M-%S'))
+                logger.info('renaming data directory to %s', new_name)
+                os.rename(self.data_dir, new_name)
+            except:
+                logger.exception("Could not rename data directory %s", self.data_dir)
+
+    def remove_data_directory(self):
+        logger.info('Removing data directory: %s', self.data_dir)
+        try:
+            if os.path.islink(self.data_dir):
+                os.unlink(self.data_dir)
+            elif not os.path.exists(self.data_dir):
+                return
+            elif os.path.isfile(self.data_dir):
+                os.remove(self.data_dir)
+            elif os.path.isdir(self.data_dir):
+                shutil.rmtree(self.data_dir)
+        except:
+            logger.exception('Could not remove data directory %s', self.data_dir)
+            self.move_data_directory()

patroni/scripts/aws.py

@@ -1,4 +1,4 @@
-#!/usr/bin/python
+#!/usr/bin/env python
 
 import logging
 import requests

patroni/scripts/restore.py

@@ -1,4 +1,4 @@
-#!/usr/bin/python
+#!/usr/bin/env python
 # arguments are:
 #   - cluster scope
 #   - cluster role

patroni/utils.py

@@ -0,0 +1,165 @@
+import datetime
+import os
+import random
+import re
+import signal
+import sys
+import time
+
+from patroni.exceptions import PatroniException
+
+ignore_sigterm = False
+interrupted_sleep = False
+reap_children = False
+
+_DATE_TIME_RE = re.compile(r'''^
+(?P<year>\d{4})\-(?P<month>\d{2})\-(?P<day>\d{2})  # date
+T
+(?P<hour>\d{2}):(?P<minute>\d{2}):(?P<second>\d{2})\.(?P<microsecond>\d{6})  # time
+\d*Z$''', re.X)
+
+
+def parse_datetime(time_str):
+    """
+    >>> parse_datetime('2015-06-10T12:56:30.552539016Z')
+    datetime.datetime(2015, 6, 10, 12, 56, 30, 552539)
+    >>> parse_datetime('2015-06-10 12:56:30.552539016Z')
+    """
+    m = _DATE_TIME_RE.match(time_str)
+    if not m:
+        return None
+    p = dict((n, int(m.group(n))) for n in 'year month day hour minute second microsecond'.split(' '))
+    return datetime.datetime(**p)
+
+
+def calculate_ttl(expiration):
+    """
+    >>> calculate_ttl(None)
+    >>> calculate_ttl('2015-06-10 12:56:30.552539016Z')
+    >>> calculate_ttl('2015-06-10T12:56:30.552539016Z') < 0
+    True
+    """
+    if not expiration:
+        return None
+    expiration = parse_datetime(expiration)
+    if not expiration:
+        return None
+    now = datetime.datetime.utcnow()
+    return int((expiration - now).total_seconds())
+
+
+def sigterm_handler(signo, stack_frame):
+    global ignore_sigterm
+    if not ignore_sigterm:
+        ignore_sigterm = True
+        sys.exit()
+
+
+def sigchld_handler(signo, stack_frame):
+    global interrupted_sleep, reap_children
+    reap_children = interrupted_sleep = True
+
+
+def sleep(interval):
+    global interrupted_sleep
+    current_time = time.time()
+    end_time = current_time + interval
+    while current_time < end_time:
+        interrupted_sleep = False
+        time.sleep(end_time - current_time)
+        if not interrupted_sleep:  # we will ignore only sigchld
+            break
+        current_time = time.time()
+    interrupted_sleep = False
+
+
+def setup_signal_handlers():
+    signal.signal(signal.SIGTERM, sigterm_handler)
+    signal.signal(signal.SIGCHLD, sigchld_handler)
+
+
+def reap_children():
+    global reap_children
+    if reap_children:
+        try:
+            while True:
+                ret = os.waitpid(-1, os.WNOHANG)
+                if ret == (0, 0):
+                    break
+        except OSError:
+            pass
+        finally:
+            reap_children = False
+
+
+class RetryFailedError(PatroniException):
+
+    """Raised when retrying an operation ultimately failed, after retrying the maximum number of attempts."""
+
+
+class Retry:
+
+    """Helper for retrying a method in the face of retry-able exceptions"""
+
+    def __init__(self, max_tries=1, delay=0.1, backoff=2, max_jitter=0.8, max_delay=3600,
+                 sleep_func=sleep, deadline=None, retry_exceptions=PatroniException):
+        """Create a :class:`Retry` instance for retrying function calls
+
+        :param max_tries: How many times to retry the command. -1 means infinite tries.
+        :param delay: Initial delay between retry attempts.
+        :param backoff: Backoff multiplier between retry attempts. Defaults to 2 for exponential backoff.
+        :param max_jitter: Additional max jitter period to wait between retry attempts to avoid slamming the server.
+        :param max_delay: Maximum delay in seconds, regardless of other backoff settings. Defaults to one hour.
+        :param retry_exceptions: single exception or tuple"""
+
+        self.max_tries = max_tries
+        self.delay = delay
+        self.backoff = backoff
+        self.max_jitter = int(max_jitter * 100)
+        self.max_delay = float(max_delay)
+        self._attempts = 0
+        self._cur_delay = delay
+        self.deadline = deadline
+        self._cur_stoptime = None
+        self.sleep_func = sleep_func
+        self.retry_exceptions = retry_exceptions
+
+    def reset(self):
+        """Reset the attempt counter"""
+        self._attempts = 0
+        self._cur_delay = self.delay
+        self._cur_stoptime = None
+
+    def copy(self):
+        """Return a clone of this retry manager"""
+        return Retry(max_tries=self.max_tries, delay=self.delay, backoff=self.backoff,
+                     max_jitter=self.max_jitter / 100.0, max_delay=self.max_delay, sleep_func=self.sleep_func,
+                     deadline=self.deadline, retry_exceptions=self.retry_exceptions)
+
+    def __call__(self, func, *args, **kwargs):
+        """Call a function with arguments until it completes without throwing a `retry_exceptions`
+
+        :param func: Function to call
+        :param args: Positional arguments to call the function with
+        :params kwargs: Keyword arguments to call the function with
+
+        The function will be called until it doesn't throw one of the retryable exceptions"""
+        self.reset()
+
+        while True:
+            try:
+                if self.deadline is not None and self._cur_stoptime is None:
+                    self._cur_stoptime = time.time() + self.deadline
+                return func(*args, **kwargs)
+            except self.retry_exceptions:
+                # Note: max_tries == -1 means infinite tries.
+                if self._attempts == self.max_tries:
+                    raise RetryFailedError("Too many retry attempts")
+                self._attempts += 1
+                sleeptime = self._cur_delay + (random.randint(0, self.max_jitter) / 100.0)
+
+                if self._cur_stoptime is not None and time.time() + sleeptime >= self._cur_stoptime:
+                    raise RetryFailedError("Exceeded retry deadline")
+                else:
+                    self.sleep_func(sleeptime)
+                self._cur_delay = min(self._cur_delay * self.backoff, self.max_delay)

patroni/version.py

@@ -0,0 +1 @@
+__version__ = '0.2'

patroni/zookeeper.py

@@ -3,10 +3,11 @@ import random
 import requests
 import time
 
-from helpers.dcs import AbstractDCS, Cluster, DCSError, Member, parse_connection_string
-from helpers.utils import sleep
 from kazoo.client import KazooClient, KazooState
 from kazoo.exceptions import NoNodeError, NodeExistsError
+from patroni.dcs import AbstractDCS, Cluster, Failover, Leader, Member
+from patroni.exceptions import DCSError
+from patroni.utils import sleep
 from requests.exceptions import RequestException
 
 logger = logging.getLogger(__name__)
@@ -90,11 +91,9 @@ class ZooKeeper(AbstractDCS):
                                       'max_tries': -1},
                                   connection_retry={'max_delay': 1, 'max_tries': -1})
         self.client.add_listener(self.session_listener)
-        self.cluster_event = self.client.handler.event_object()
 
+        self._my_member_data = None
         self.fetch_cluster = True
-        self.members = []
-        self.leader = None
         self.last_leader_operation = 0
 
         self.client.start(None)
@@ -105,57 +104,72 @@ class ZooKeeper(AbstractDCS):
 
     def cluster_watcher(self, event):
         self.fetch_cluster = True
-        self.cluster_event.set()
+        self.event.set()
 
-    def get_node(self, name, watch=None):
+    def get_node(self, key, watch=None):
         try:
-            return self.client.get(self.client_path(name), watch)
+            ret = self.client.get(key, watch)
+            return (ret[0].decode('utf-8'), ret[1])
         except NoNodeError:
-            pass
-        except:
-            logger.exception('get_node')
-        return None
+            return None
 
     @staticmethod
     def member(name, value, znode):
-        conn_url, api_url = parse_connection_string(value)
-        return Member(znode.mzxid, name, conn_url, api_url, None, None)
+        return Member.from_node(znode.version, name, znode.ephemeralOwner, value)
+
+    def get_children(self, key, watch=None):
+        try:
+            return self.client.get_children(key, watch)
+        except NoNodeError:
+            return []
 
     def load_members(self):
         members = []
-        for member in self.client.get_children(self.client_path('/members'), self.cluster_watcher):
-            data = self.get_node('/members/' + member)
+        for member in self.get_children(self.members_path, self.cluster_watcher):
+            data = self.get_node(self.members_path + member)
             if data is not None:
                 members.append(self.member(member, *data))
         return members
 
     def _inner_load_cluster(self):
-        self.cluster_event.clear()
-        leader = self.get_node('/leader', self.cluster_watcher)
-        self.members = self.load_members()
+        self.fetch_cluster = False
+        self.event.clear()
+        nodes = set(self.get_children(self.client_path(''), self.cluster_watcher))
+        if not nodes:
+            self.fetch_cluster = True
+
+        # get initialize flag
+        initialize = self._INITIALIZE in nodes
+
+        # get list of members
+        members = self.load_members() if self._MEMBERS[:-1] in nodes else []
+
+        # get leader
+        leader = self.get_node(self.leader_path) if self._LEADER in nodes else None
         if leader:
-            if leader[0] == self._name:
-                client_id = self.client.client_id
-                if client_id is not None and client_id[0] != leader[1].ephemeralOwner:
-                    logger.info('I am leader but not owner of the session. Removing leader node')
-                    self.client.delete(self.client_path('/leader'))
-                    leader = None
+            client_id = self.client.client_id
+            if leader[0] == self._name and client_id is not None and client_id[0] != leader[1].ephemeralOwner:
+                logger.info('I am leader but not owner of the session. Removing leader node')
+                self.client.delete(self.leader_path)
+                leader = None
 
             if leader:
-                for member in self.members:
-                    if member.name == leader[0]:
-                        leader = member
-                        self.fetch_cluster = False
-                        break
-            if not isinstance(leader, Member):
-                leader = Member(-1, leader, None, None, None, None)
-        self.leader = leader
-        if self.fetch_cluster:
-            last_leader_operation = self.get_node('/optime/leader')
-            if last_leader_operation:
-                self.last_leader_operation = int(last_leader_operation[0])
+                member = Member(-1, leader[0], None, {})
+                member = ([m for m in members if m.name == leader[0]] or [member])[0]
+                leader = Leader(leader[1].version, leader[1].ephemeralOwner, member)
+                self.fetch_cluster = member.index == -1
 
-    def get_cluster(self):
+        # failover key
+        failover = self.get_node(self.failover_path, watch=self.cluster_watcher) if self._FAILOVER in nodes else None
+        if failover:
+            failover = Failover.from_node(failover[1].version, failover[0])
+
+        # get last leader operation
+        optime = self.get_node(self.leader_optime_path) if self._OPTIME in nodes and self.fetch_cluster else None
+        self.last_leader_operation = 0 if optime is None else int(optime[0])
+        self._cluster = Cluster(initialize, leader, self.last_leader_operation, members, failover)
+
+    def _load_cluster(self):
         if self.exhibitor and self.exhibitor.poll():
             self.client.set_hosts(self.exhibitor.zookeeper_hosts)
 
@@ -166,48 +180,76 @@ class ZooKeeper(AbstractDCS):
                 logger.exception('get_cluster')
                 self.session_listener(KazooState.LOST)
                 raise ZooKeeperError('ZooKeeper in not responding properly')
-        return Cluster(True, self.leader, self.last_leader_operation, self.members)
 
     def _create(self, path, value, **kwargs):
         try:
-            self.client.retry(self.client.create, self.client_path(path), value, **kwargs)
+            self.client.retry(self.client.create, path, value.encode('utf-8'), **kwargs)
             return True
         except:
             return False
 
     def attempt_to_acquire_leader(self):
-        ret = self._create('/leader', self._name, makepath=True, ephemeral=True)
+        ret = self._create(self.leader_path, self._name, makepath=True, ephemeral=True)
         ret or logger.info('Could not take out TTL lock')
         return ret
 
-    def race(self, path):
-        return self._create(path, self._name, makepath=True)
-
-    def touch_member(self, connection_string, ttl=None):
-        for m in self.members:
-            if m.name == self._name:
-                return True
-        path = self.client_path('/members/' + self._name)
+    def set_failover_value(self, value, index=None):
         try:
-            self.client.retry(self.client.create, path, connection_string, makepath=True, ephemeral=True)
+            self.client.retry(self.client.set, self.failover_path, value.encode('utf-8'), version=index or -1)
+            return True
+        except NoNodeError:
+            return value == '' or (not index and self._create(self.failover_path, value.encode('utf-8')))
+        except:
+            logging.exception('set_failover_value')
+            return False
+
+    def initialize(self):
+        return self._create(self.initialize_path, self._name, makepath=True)
+
+    def touch_member(self, data, ttl=None):
+        cluster = self.cluster
+        me = cluster and ([m for m in cluster.members if m.name == self._name] or [None])[0]
+        path = self.member_path
+        data = data.encode('utf-8')
+        create = not me
+        if me and self.client.client_id is not None and me.session != self.client.client_id[0]:
+            try:
+                self.client.retry(self.client.delete, path)
+            except NoNodeError:
+                pass
+            except:
+                return False
+            create = True
+
+        if not create and data == self._my_member_data:
+            return True
+
+        try:
+            if create:
+                self.client.retry(self.client.create, path, data, makepath=True, ephemeral=True)
+            else:
+                self.client.retry(self.client.set, path, data)
+            self._my_member_data = data
             return True
         except NodeExistsError:
             try:
-                self.client.retry(self.client.delete, path)
-                self.client.retry(self.client.create, path, connection_string, makepath=True, ephemeral=True)
+                self.client.retry(self.client.set, path, data)
+                self._my_member_data = data
                 return True
             except:
                 logger.exception('touch_member')
+        except:
+            logger.exception('touch_member')
         return False
 
     def take_leader(self):
         return self.attempt_to_acquire_leader()
 
-    def update_leader(self, state_handler):
-        last_operation = state_handler.last_operation()
+    def write_leader_optime(self, last_operation):
+        last_operation = last_operation.encode('utf-8')
         if last_operation != self.last_leader_operation:
             self.last_leader_operation = last_operation
-            path = self.client_path('/optime/leader')
+            path = self.leader_optime_path
             try:
                 self.client.retry(self.client.set, path, last_operation)
             except NoNodeError:
@@ -217,13 +259,27 @@ class ZooKeeper(AbstractDCS):
                     logger.exception('Failed to create %s', path)
             except:
                 logger.exception('Failed to update %s', path)
+
+    def update_leader(self):
         return True
 
     def delete_leader(self):
-        if isinstance(self.leader, Member) and self.leader.name == self._name:
-            self.client.delete(self.client_path('/leader'))
+        self.client.restart()
+        self._my_member_data = None
+        return True
 
-    def sleep(self, timeout):
-        self.cluster_event.wait(timeout)
-        if self.cluster_event.isSet():
+    def _cancel_initialization(self):
+        node = self.get_node(self.initialize_path)
+        if node and node[0] == self._name:
+            self.client.delete(self.initialize_path, version=node[1].version)
+
+    def cancel_initialization(self):
+        try:
+            self.client.retry(self._cancel_initialization)
+        except:
+            logger.exception("Unable to delete initialize key")
+
+    def watch(self, timeout):
+        if super(ZooKeeper, self).watch(timeout):
             self.fetch_cluster = True
+        return self.fetch_cluster

postgres0.yml

@@ -4,6 +4,9 @@ scope: &scope batman
 restapi:
   listen: 127.0.0.1:8008
   connect_address: 127.0.0.1:8008
+  auth: 'username:password'
+#  certfile: /etc/ssl/certs/ssl-cert-snakeoil.pem
+#  keyfile: /etc/ssl/private/ssl-cert-snakeoil.key
 etcd:
   scope: *scope
   ttl: *ttl
@@ -30,6 +33,10 @@ postgresql:
   connect_address: 127.0.0.1:5432
   data_dir: data/postgresql0
   maximum_lag_on_failover: 1048576 # 1 megabyte in bytes
+  use_slots: True
+  pg_rewind:
+    username: postgres
+    password: zalando
   pg_hba:
   - host all all 0.0.0.0/0 md5
   - hostssl all all 0.0.0.0/0 md5
@@ -38,6 +45,7 @@ postgresql:
     password: rep-pass
     network:  127.0.0.1/32
   superuser:
+    username: postgres
     password: zalando
   admin:
     username: admin
@@ -46,7 +54,7 @@ postgresql:
     env_dir: /home/postgres/etc/wal-e.d/env
     threshold_megabytes: 10240
     threshold_backup_size_percentage: 30
-  restore: scripts/restore.py
+  restore: patroni/scripts/restore.py
   #recovery_conf:
     #restore_command: cp ../wal_archive/%f %p
   parameters:
@@ -58,3 +66,4 @@ postgresql:
     archive_timeout: 1800s
     max_replication_slots: 5
     hot_standby: "on"
+    wal_log_hints: "on"

postgres1.yml

@@ -4,6 +4,9 @@ scope: &scope batman
 restapi:
   listen: 127.0.0.1:8009
   connect_address: 127.0.0.1:8009
+  auth: 'username:password'
+#  certfile: /etc/ssl/certs/ssl-cert-snakeoil.pem
+#  keyfile: /etc/ssl/private/ssl-cert-snakeoil.key
 etcd:
   scope: *scope
   ttl: *ttl
@@ -30,6 +33,10 @@ postgresql:
   connect_address: 127.0.0.1:5433
   data_dir: data/postgresql1
   maximum_lag_on_failover: 1048576 # 1 megabyte in bytes
+  use_slots: True
+  pg_rewind:
+    username: postgres
+    password: zalando
   pg_hba:
   - host all all 0.0.0.0/0 md5
   - hostssl all all 0.0.0.0/0 md5
@@ -38,6 +45,7 @@ postgresql:
     password: rep-pass
     network: 127.0.0.1/32
   superuser:
+    user: postgres
     password: zalando
   admin:
     username: admin
@@ -48,7 +56,7 @@ postgresql:
     env_dir: /home/postgres/etc/wal-e.d/env
     threshold_megabytes: 10240
     threshold_backup_size_percentage: 30
-  restore: scripts/restore.py
+  restore: patroni/scripts/restore.py
   parameters:
     archive_mode: "on"
     wal_level: hot_standby
@@ -58,3 +66,4 @@ postgresql:
     archive_timeout: 1800s
     max_replication_slots: 5
     hot_standby: "on"
+    wal_log_hints: "on"

release.sh

@@ -0,0 +1,31 @@
+#!/bin/sh
+
+if [ $# -ne 1 ]; then
+    >&2 echo "usage: $0 <version>"
+    exit 1
+fi
+
+readonly VERSIONFILE="patroni/version.py"
+
+## Bail out on any non-zero exitcode from the called processes
+set -xe
+
+python3 --version
+git --version
+
+version=$1
+
+sed -i "s/__version__ = .*/__version__ = '${version}'/"  "${VERSIONFILE}"
+python3 setup.py clean
+python3 setup.py test
+python3 setup.py flake8
+
+git add "${VERSIONFILE}"
+
+git commit -m "Bumped version to $version"
+git push
+
+python3 setup.py sdist bdist_wheel upload
+
+git tag v${version}
+git push --tags

setup.py

@@ -19,13 +19,22 @@ if sys.version_info < (2, 7, 0):
 __location__ = os.path.join(os.getcwd(), os.path.dirname(inspect.getfile(inspect.currentframe())))
 
 
+def read_version(package):
+    data = {}
+    with open(os.path.join(package, 'version.py'), 'r') as fd:
+        exec(fd.read(), data)
+    return data['__version__']
+
 NAME = 'patroni'
-MAIN_PACKAGE = 'patroni.py'
-HELPERS = 'helpers'
+MAIN_PACKAGE = NAME
 SCRIPTS = 'scripts'
-VERSION = '0.1'
-DESCRIPTION = 'A Template for PostgreSQL HA with etcd'
+VERSION = read_version(MAIN_PACKAGE)
+DESCRIPTION = 'PostgreSQL High-Available orchestrator and CLI'
 LICENSE = 'The MIT License'
+URL = 'https://github.com/zalando/patroni'
+AUTHOR = 'Alexander Kukushkin, Alexey Klyukin, Feike Steenbergen'
+AUTHOR_EMAIL = '[email protected], [email protected], [email protected]'
+KEYWORDS = 'etcd governor patroni postgresql postgres ha zookeeper streaming replication'
 
 COVERAGE_XML = True
 COVERAGE_HTML = False
@@ -38,7 +47,7 @@ CLASSIFIERS = [
     'Environment :: Console',
     'Intended Audience :: Developers',
     'Intended Audience :: System Administrators',
-    'License :: OSI Approved :: The MIT License',
+    'License :: OSI Approved :: MIT License',
     'Operating System :: POSIX :: Linux',
     'Programming Language :: Python',
     'Programming Language :: Python :: 2.7',
@@ -47,6 +56,8 @@ CLASSIFIERS = [
     'Programming Language :: Python :: Implementation :: CPython',
 ]
 
+CONSOLE_SCRIPTS = ['patroni = patroni:main']
+
 
 class PyTest(TestCommand):
 
@@ -62,8 +73,7 @@ class PyTest(TestCommand):
     def finalize_options(self):
         TestCommand.finalize_options(self)
         if self.cov_xml or self.cov_html:
-            self.cov = ['--cov', MAIN_PACKAGE, '--cov', HELPERS, '--cov', SCRIPTS, '--cov-report',
-                        'term-missing']
+            self.cov = ['--cov', MAIN_PACKAGE, '--cov', MAIN_PACKAGE, '--cov-report', 'term-missing']
             if self.cov_xml:
                 self.cov.extend(['--cov-report', 'xml'])
             if self.cov_html:
@@ -79,10 +89,9 @@ class PyTest(TestCommand):
         params = {'args': self.test_args}
         if self.cov:
             params['args'] += self.cov
-            params['plugins'] = ['cov']
         if self.junitxml:
             params['args'] += self.junitxml
-        params['args'] += ['--doctest-modules', HELPERS, '--doctest-modules', SCRIPTS, '-s']
+        params['args'] += ['--doctest-modules', MAIN_PACKAGE, '-s', '-vv']
         errno = pytest.main(**params)
         sys.exit(errno)
 
@@ -118,10 +127,13 @@ def setup_package():
     setup(
         name=NAME,
         version=version,
+        url=URL,
+        author=AUTHOR,
+        author_email=AUTHOR_EMAIL,
         description=DESCRIPTION,
         license=LICENSE,
-        keywords='etcd governor patroni postgresql postgres ha zookeeper',
-        long_description=read('README.md'),
+        keywords=KEYWORDS,
+        long_description=read('README.rst'),
         classifiers=CLASSIFIERS,
         test_suite='tests',
         packages=setuptools.find_packages(exclude=['tests', 'tests.*']),
@@ -131,6 +143,7 @@ def setup_package():
         cmdclass=cmdclass,
         tests_require=['pytest-cov', 'pytest'],
         command_options=command_options,
+        entry_points={'console_scripts': CONSOLE_SCRIPTS},
     )
 
 

tests/test_api.py

@@ -1,16 +1,18 @@
 import psycopg2
 import unittest
 
-from helpers.api import RestApiHandler, RestApiServer
+from mock import Mock, patch
+from patroni.api import RestApiHandler, RestApiServer
 from six import BytesIO as IO
-from test_postgresql import psycopg2_connect
+from six.moves import BaseHTTPServer
+from test_postgresql import psycopg2_connect, MockCursor
 
 
-def throws(*args, **kwargs):
-    raise psycopg2.OperationalError()
+class MockPostgresql(Mock):
 
-
-class MockPostgresql:
+    name = 'test'
+    state = 'running'
+    role = 'master'
 
     def connection(self):
         return psycopg2_connect()
@@ -19,10 +21,29 @@ class MockPostgresql:
         return True
 
 
+class MockHa(Mock):
+
+    dcs = Mock()
+    state_handler = MockPostgresql()
+
+    def schedule_restart(self):
+        return 'restart'
+
+    def schedule_reinitialize(self):
+        return 'reinitialize'
+
+    def restart(self):
+        return (True, '')
+
+    def restart_scheduled(self):
+        return False
+
+
 class MockPatroni:
 
-    def __init__(self):
-        self.postgresql = MockPostgresql()
+    postgresql = MockPostgresql()
+    ha = MockHa()
+    dcs = Mock()
 
 
 class MockRequest:
@@ -36,18 +57,63 @@ class MockRequest:
 
 class MockRestApiServer(RestApiServer):
 
-    def __init__(self, Handler, path, *args):
-        self.patroni = MockPatroni()
-        if len(args) > 0:
-            self.query = args[0]
+    def __init__(self, Handler, path):
+        self.socket = 0
+        BaseHTTPServer.HTTPServer.__init__ = Mock()
+        MockRestApiServer._BaseServer__is_shut_down = Mock()
+        MockRestApiServer._BaseServer__shutdown_request = True
+        config = {'listen': '127.0.0.1:8008', 'auth': 'test:test', 'certfile': 'dumb'}
+        super(MockRestApiServer, self).__init__(MockPatroni(), config)
         Handler(MockRequest(path), ('0.0.0.0', 8080), self)
 
 
+@patch('ssl.wrap_socket', Mock(return_value=0))
 class TestRestApiHandler(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        super(TestRestApiHandler, self).__init__(method_name)
-
     def test_do_GET(self):
-        MockRestApiServer(RestApiHandler, b'GET /')
-        MockRestApiServer(RestApiHandler, b'GET /', throws)
+        MockRestApiServer(RestApiHandler, b'GET /replica')
+        with patch.object(RestApiHandler, 'get_postgresql_status', Mock(return_value={})):
+            MockRestApiServer(RestApiHandler, b'GET /replica')
+        with patch.object(RestApiHandler, 'get_postgresql_status', Mock(return_value={'role': 'master'})):
+            MockRestApiServer(RestApiHandler, b'GET /replica')
+        MockRestApiServer(RestApiHandler, b'GET /master')
+        MockPatroni.dcs.cluster.leader.name = MockPostgresql.name
+        MockRestApiServer(RestApiHandler, b'GET /replica')
+        MockPatroni.dcs.cluster = None
+        with patch.object(RestApiHandler, 'get_postgresql_status', Mock(return_value={'role': 'master'})):
+            MockRestApiServer(RestApiHandler, b'GET /master')
+        with patch.object(MockHa, 'restart_scheduled', Mock(return_value=True)):
+            MockRestApiServer(RestApiHandler, b'GET /master')
+        MockRestApiServer(RestApiHandler, b'GET /master')
+
+    def test_do_GET_patroni(self):
+        MockRestApiServer(RestApiHandler, b'GET /patroni')
+
+    def test_basicauth(self):
+        MockRestApiServer(RestApiHandler, b'POST /restart HTTP/1.0')
+        MockRestApiServer(RestApiHandler, b'POST /restart HTTP/1.0\nAuthorization:')
+
+    def test_do_POST_restart(self):
+        request = b'POST /restart HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0'
+        MockRestApiServer(RestApiHandler, request)
+        with patch.object(MockHa, 'restart', Mock(side_effect=Exception)):
+            MockRestApiServer(RestApiHandler, request)
+
+    @patch.object(MockHa, 'dcs')
+    def test_do_POST_reinitialize(self, dcs):
+        cluster = dcs.get_cluster.return_value
+        request = b'POST /reinitialize HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0'
+        MockRestApiServer(RestApiHandler, request)
+        cluster.is_unlocked.return_value = False
+        MockRestApiServer(RestApiHandler, request)
+        with patch.object(MockHa, 'schedule_reinitialize', Mock(return_value=None)):
+            MockRestApiServer(RestApiHandler, request)
+        cluster.leader.name = 'test'
+        MockRestApiServer(RestApiHandler, request)
+
+    @patch('time.sleep', Mock())
+    def test_RestApiServer_query(self):
+        with patch.object(MockCursor, 'execute', Mock(side_effect=psycopg2.OperationalError)):
+            MockRestApiServer(RestApiHandler, b'GET /patroni')
+        with patch.object(MockPostgresql, 'connection', Mock(side_effect=psycopg2.OperationalError)):
+            MockRestApiServer(RestApiHandler, b'GET /patroni')

tests/test_async_executor.py

@@ -0,0 +1,18 @@
+import unittest
+
+from mock import Mock, patch
+from patroni.async_executor import AsyncExecutor
+from threading import Thread
+
+
+class TestAsyncExecutor(unittest.TestCase):
+
+    def setUp(self):
+        self.a = AsyncExecutor()
+
+    @patch.object(Thread, 'start', Mock())
+    def test_run_async(self):
+        self.a.run_async(Mock(return_value=True))
+
+    def test_run(self):
+        self.a.run(Mock(side_effect=Exception()))

tests/test_aws.py

@@ -2,7 +2,7 @@ import unittest
 import requests
 import boto.ec2
 from collections import namedtuple
-from scripts.aws import AWSConnection
+from patroni.scripts.aws import AWSConnection
 from requests.exceptions import RequestException
 
 

tests/test_etcd.py

@@ -1,16 +1,14 @@
-import datetime
-import dns.resolver
 import etcd
 import json
 import requests
+import urllib3
 import socket
-import time
 import unittest
 
 from dns.exception import DNSException
-from helpers.dcs import Cluster, DCSError, Member
-from helpers.etcd import Client, Etcd
 from mock import Mock, patch
+from patroni.dcs import Cluster, DCSError, Leader
+from patroni.etcd import Client, Etcd
 
 
 class MockResponse:
@@ -25,6 +23,10 @@ class MockResponse:
 
     @property
     def data(self):
+        if self.content == 'TimeoutError':
+            raise urllib3.exceptions.TimeoutError
+        if self.content == 'Exception':
+            raise Exception
         return self.content
 
     @property
@@ -35,11 +37,10 @@ class MockResponse:
         return ''
 
 
-class MockPostgresql:
-    name = ''
+class MockPostgresql(Mock):
 
     def last_operation(self):
-        return 0
+        return '0'
 
 
 def requests_get(url, **kwargs):
@@ -48,6 +49,8 @@ def requests_get(url, **kwargs):
     response = MockResponse()
     if url.startswith('http://local'):
         raise requests.exceptions.RequestException()
+    elif ':8011/patroni' in url:
+        response.content = '{"role": "replica", "xlog": {"replayed_location": 0}}'
     elif url.endswith('/members'):
         if url.startswith('http://error'):
             response.content = '[{}]'
@@ -61,24 +64,37 @@ def requests_get(url, **kwargs):
     return response
 
 
+def etcd_watch(key, index=None, timeout=None, recursive=None):
+    if timeout == 2.0:
+        raise urllib3.exceptions.TimeoutError
+    elif timeout == 5.0:
+        return etcd.EtcdResult('delete', {})
+    elif timeout == 10.0:
+        raise etcd.EtcdException
+    elif index == 20729:
+        return etcd.EtcdResult('set', {'value': 'postgresql1', 'modifiedIndex': index + 1})
+    elif index == 20731:
+        return etcd.EtcdResult('set', {'value': 'postgresql2', 'modifiedIndex': index + 1})
+
+
 def etcd_write(key, value, **kwargs):
+    if key == '/service/exists/leader':
+        raise etcd.EtcdAlreadyExist
     if key == '/service/test/leader':
         if kwargs.get('prevValue', None) == 'foo' or not kwargs.get('prevExist', True):
             return True
     raise etcd.EtcdException
 
 
-def etcd_delete(key, **kwargs):
-    raise etcd.EtcdException
-
-
 def etcd_read(key, **kwargs):
-    if key == '/service/noleader':
+    if key == '/service/noleader/':
         raise DCSError('noleader')
-    elif key == '/service/nocluster':
+    elif key == '/service/nocluster/':
         raise etcd.EtcdKeyNotFound
 
     response = {"action": "get", "node": {"key": "/service/batman5", "dir": True, "nodes": [
+                {"key": "/service/batman5/failover", "value": "",
+                 "modifiedIndex": 1582, "createdIndex": 1582},
                 {"key": "/service/batman5/initialize", "value": "postgresql0",
                  "modifiedIndex": 1582, "createdIndex": 1582},
                 {"key": "/service/batman5/leader", "value": "postgresql1",
@@ -103,14 +119,10 @@ def etcd_read(key, **kwargs):
     return etcd.EtcdResult(**response)
 
 
-def time_sleep(_):
+class SleepException(Exception):
     pass
 
 
-def time_sleep_exception(_):
-    raise Exception()
-
-
 class MockSRV:
     port = 2380
     target = '127.0.0.1'
@@ -127,7 +139,7 @@ def dns_query(name, type):
 def socket_getaddrinfo(*args):
     if args[0] == 'ok':
         return [(2, 1, 6, '', ('127.0.0.1', 2379)), (2, 1, 6, '', ('127.0.0.1', 2379))]
-    raise socket.error()
+    raise socket.error
 
 
 def http_request(method, url, **kwargs):
@@ -136,28 +148,14 @@ def http_request(method, url, **kwargs):
     raise socket.error
 
 
-class TestMember(unittest.TestCase):
-
-    def __init__(self, method_name='runTest'):
-        super(TestMember, self).__init__(method_name)
-
-    def test_real_ttl(self):
-        now = datetime.datetime.utcnow()
-        member = Member(0, 'a', 'b', 'c', (now + datetime.timedelta(seconds=2)).strftime('%Y-%m-%dT%H:%M:%S.%fZ'), None)
-        self.assertLess(member.real_ttl(), 2)
-        self.assertEquals(Member(0, 'a', 'b', 'c', '', None).real_ttl(), -1)
-
-
+@patch('dns.resolver.query', dns_query)
+@patch('socket.getaddrinfo', socket_getaddrinfo)
+@patch('requests.get', requests_get)
 class TestClient(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestClient, self).__init__(method_name)
-
-    def set_up(self):
-        socket.getaddrinfo = socket_getaddrinfo
-        requests.get = requests_get
-        dns.resolver.query = dns_query
+    @patch('dns.resolver.query', dns_query)
+    @patch('requests.get', requests_get)
+    def setUp(self):
         with patch.object(etcd.Client, 'machines') as mock_machines:
             mock_machines.__get__ = Mock(return_value=['http://localhost:2379', 'http://localhost:4001'])
             self.client = Client({'discovery_srv': 'test'})
@@ -172,12 +170,21 @@ class TestClient(unittest.TestCase):
         self.assertEquals(self.client.get_srv_record('blabla'), [])
         self.assertEquals(self.client.get_srv_record('exception'), [])
 
+    def test__result_from_response(self):
+        response = MockResponse()
+        response.content = 'TimeoutError'
+        self.assertRaises(urllib3.exceptions.TimeoutError, self.client._result_from_response, response)
+        response.content = 'Exception'
+        self.assertRaises(etcd.EtcdException, self.client._result_from_response, response)
+        response.content = b'{}'
+        self.assertRaises(etcd.EtcdException, self.client._result_from_response, response)
+
     def test__get_machines_cache_from_srv(self):
-        self.client.get_srv_record = lambda e: [('localhost', 2380)]
+        self.client.get_srv_record = Mock(return_value=[('localhost', 2380)])
         self.client._get_machines_cache_from_srv('blabla')
 
     def test__get_machines_cache_from_dns(self):
-        self.client._get_machines_cache_from_dns('ok:2379')
+        self.client._get_machines_cache_from_dns('error:2379')
 
     def test__load_machines_cache(self):
         self.client._config = {}
@@ -186,25 +193,23 @@ class TestClient(unittest.TestCase):
         self.assertRaises(etcd.EtcdException, self.client._load_machines_cache)
 
 
+@patch('requests.get', requests_get)
 class TestEtcd(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestEtcd, self).__init__(method_name)
-
-    def set_up(self):
-        time.sleep = time_sleep
+    def setUp(self):
         with patch.object(Client, 'machines') as mock_machines:
             mock_machines.__get__ = Mock(return_value=['http://localhost:2379', 'http://localhost:4001'])
             self.etcd = Etcd('foo', {'ttl': 30, 'host': 'localhost:2379', 'scope': 'test'})
             self.etcd.client.write = etcd_write
             self.etcd.client.read = etcd_read
+            self.etcd.client.delete = Mock(side_effect=etcd.EtcdException())
 
+    @patch('dns.resolver.query', dns_query)
     def test_get_etcd_client(self):
-        time.sleep = time_sleep_exception
         with patch.object(etcd.Client, 'machines') as mock_machines:
             mock_machines.__get__ = Mock(side_effect=etcd.EtcdException)
-            self.assertRaises(Exception, self.etcd.get_etcd_client, {'discovery_srv': 'test'})
+            with patch('time.sleep', Mock(side_effect=SleepException())):
+                self.assertRaises(SleepException, self.etcd.get_etcd_client, {'discovery_srv': 'test'})
 
     def test_get_cluster(self):
         self.assertIsInstance(self.etcd.get_cluster(), Cluster)
@@ -214,7 +219,7 @@ class TestEtcd(unittest.TestCase):
         self.assertIsNone(cluster.leader)
 
     def test_current_leader(self):
-        self.assertIsInstance(self.etcd.current_leader(), Member)
+        self.assertIsInstance(self.etcd.current_leader(), Leader)
         self.etcd._base_path = '/service/noleader'
         self.assertIsNone(self.etcd.current_leader())
 
@@ -224,12 +229,32 @@ class TestEtcd(unittest.TestCase):
     def test_take_leader(self):
         self.assertFalse(self.etcd.take_leader())
 
-    def test_update_leader(self):
-        self.assertTrue(self.etcd.update_leader(MockPostgresql()))
+    def test_attempt_to_acquire_leader(self):
+        self.etcd._base_path = '/service/exists'
+        self.assertFalse(self.etcd.attempt_to_acquire_leader())
+        self.etcd._base_path = '/service/failed'
+        self.assertFalse(self.etcd.attempt_to_acquire_leader())
 
-    def test_race(self):
-        self.assertFalse(self.etcd.race(''))
+    def test_write_leader_optime(self):
+        self.etcd.write_leader_optime('0')
+
+    def test_update_leader(self):
+        self.assertTrue(self.etcd.update_leader())
+
+    def test_initialize(self):
+        self.assertFalse(self.etcd.initialize())
+
+    def test_cancel_initializion(self):
+        self.assertFalse(self.etcd.cancel_initialization())
 
     def test_delete_leader(self):
-        self.etcd.client.delete = etcd_delete
         self.assertFalse(self.etcd.delete_leader())
+
+    def test_watch(self):
+        self.etcd.client.watch = etcd_watch
+        self.etcd.watch(0)
+        self.etcd.get_cluster()
+        self.etcd.watch(1.5)
+        self.etcd.watch(4.5)
+        self.etcd.watch(9.5)
+        self.etcd.watch(100)

tests/test_ha.py

@@ -1,10 +1,11 @@
 import unittest
 
-from helpers.dcs import Cluster, DCSError
-from helpers.etcd import Client, Etcd
-from helpers.ha import Ha
 from mock import Mock, patch
-from test_etcd import etcd_read, etcd_write
+from patroni.dcs import Cluster, Failover, Leader, Member
+from patroni.etcd import Client, Etcd
+from patroni.exceptions import DCSError, PostgresException
+from patroni.ha import Ha
+from test_etcd import socket_getaddrinfo, etcd_read, etcd_write, requests_get
 
 
 def true(*args, **kwargs):
@@ -15,18 +16,37 @@ def false(*args, **kwargs):
     return False
 
 
-class MockPostgresql:
+def get_cluster(initialize, leader, members, failover):
+    return Cluster(initialize, leader, None, members, failover)
 
-    def __init__(self):
-        self.name = 'postgresql0'
-        self.is_promoted = False
+
+def get_cluster_not_initialized_without_leader():
+    return get_cluster(None, None, [], None)
+
+
+def get_cluster_initialized_without_leader(leader=False, failover=None):
+    m = Member(0, 'leader', 28, {'conn_url': 'postgres://replicator:[email protected]:5435/postgres',
+                                 'api_url': 'http://127.0.0.1:8008/patroni'})
+    l = Leader(0, 0, m) if leader else None
+    o = Member(0, 'other', 28, {'conn_url': 'postgres://replicator:[email protected]:5436/postgres',
+                                'api_url': 'http://127.0.0.1:8011/patroni'})
+    return get_cluster(True, l, [m, o], failover)
+
+
+def get_cluster_initialized_with_leader(failover=None):
+    return get_cluster_initialized_without_leader(leader=True, failover=failover)
+
+
+class MockPostgresql(Mock):
+
+    name = 'postgresql0'
+    role = 'replica'
+    state = 'running'
+    connection_string = 'postgres://foo@bar/postgres'
 
     def is_healthy(self):
         return True
 
-    def write_recovery_conf(self, _):
-        return True
-
     def start(self):
         return True
 
@@ -36,64 +56,80 @@ class MockPostgresql:
     def is_leader(self):
         return True
 
-    def promote(self):
-        return True
-
-    def demote(self, _):
-        return True
-
-    def follow_the_leader(self, _):
-        return True
-
-    def create_replication_slots(self, _):
-        return True
+    def xlog_position(self):
+        return 0
 
     def last_operation(self):
         return 0
 
+    def data_directory_empty(self):
+        return False
 
-def nop(*args, **kwargs):
-    pass
+    def bootstrap(self, *args, **kwargs):
+        return True
+
+    def check_replication_lag(self, last_leader_operation):
+        return True
+
+    def check_recovery_conf(self, leader):
+        return False
 
 
-def dead_etcd():
-    raise DCSError('Etcd is not responding properly')
+class MockPatroni:
+
+    def __init__(self, p, d):
+        self.postgresql = p
+        self.dcs = d
+        self.api = Mock()
+        self.api.connection_string = 'http://127.0.0.1:8008'
 
 
-def get_unlocked_cluster():
-    return Cluster(False, None, None, [])
+def run_async(func, args=()):
+    func(*args) if args else func()
 
 
 class TestHa(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestHa, self).__init__(method_name)
-
-    def set_up(self):
+    @patch('socket.getaddrinfo', socket_getaddrinfo)
+    @patch.object(Client, 'machines')
+    def setUp(self, mock_machines):
+        mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
         self.p = MockPostgresql()
-        with patch.object(Client, 'machines') as mock_machines:
-            mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
-            self.e = Etcd('foo', {'ttl': 30, 'host': 'remotehost:2379', 'scope': 'test'})
-            self.e.client.read = etcd_read
-            self.e.client.write = etcd_write
-            self.ha = Ha(self.p, self.e)
-            self.ha.load_cluster_from_dcs()
-            self.ha.cluster = get_unlocked_cluster()
-            self.ha.load_cluster_from_dcs = nop
+        self.e = Etcd('foo', {'ttl': 30, 'host': 'ok:2379', 'scope': 'test'})
+        self.e.client.read = etcd_read
+        self.e.client.write = etcd_write
+        self.ha = Ha(MockPatroni(self.p, self.e))
+        self.ha._async_executor.run_async = run_async
+        self.ha.old_cluster = self.e.get_cluster()
+        self.ha.cluster = get_cluster_not_initialized_without_leader()
+        self.ha.load_cluster_from_dcs = Mock()
 
-    def test_load_cluster_from_dcs(self):
-        ha = Ha(self.p, self.e)
-        ha.load_cluster_from_dcs()
-        self.e.get_cluster = get_unlocked_cluster
-        ha.load_cluster_from_dcs()
+    def test_update_lock(self):
+        self.p.last_operation = Mock(side_effect=PostgresException(''))
+        self.assertTrue(self.ha.update_lock())
 
-    def test_start_as_slave(self):
+    def test_touch_member(self):
+        self.p.xlog_position = Mock(side_effect=Exception)
+        self.ha.touch_member()
+
+    def test_start_as_replica(self):
         self.p.is_healthy = false
         self.assertEquals(self.ha.run_cycle(), 'started as a secondary')
 
+    def test_recover_replica_failed(self):
+        self.p.controldata = lambda: {'Database cluster state': 'in production'}
+        self.p.is_healthy = false
+        self.p.follow_the_leader = false
+        self.assertEquals(self.ha.run_cycle(), 'failed to start postgres')
+
+    def test_recover_master_failed(self):
+        self.p.follow_the_leader = false
+        self.p.is_healthy = false
+        self.ha.has_lock = true
+        self.assertEquals(self.ha.run_cycle(), 'removed leader key after trying and failing to start postgres')
+
+    @patch.object(Cluster, 'is_unlocked', Mock(return_value=False))
     def test_start_as_readonly(self):
-        self.ha.cluster.is_unlocked = false
         self.p.is_leader = self.p.is_healthy = false
         self.ha.has_lock = true
         self.assertEquals(self.ha.run_cycle(), 'promoted self to leader because i had the session lock')
@@ -102,6 +138,7 @@ class TestHa(unittest.TestCase):
         self.assertEquals(self.ha.run_cycle(), 'acquired session lock as a leader')
 
     def test_promoted_by_acquiring_lock(self):
+        self.ha.is_healthiest_node = true
         self.p.is_leader = false
         self.assertEquals(self.ha.run_cycle(), 'promoted self to leader by acquiring session lock')
 
@@ -110,16 +147,17 @@ class TestHa(unittest.TestCase):
         self.assertEquals(self.ha.run_cycle(), 'demoted self due after trying and failing to obtain lock')
 
     def test_follow_new_leader_after_failing_to_obtain_lock(self):
+        self.ha.is_healthiest_node = true
         self.ha.acquire_lock = false
         self.p.is_leader = false
         self.assertEquals(self.ha.run_cycle(), 'following new leader after trying and failing to obtain lock')
 
     def test_demote_because_not_healthiest(self):
-        self.p.is_healthiest_node = false
+        self.ha.is_healthiest_node = false
         self.assertEquals(self.ha.run_cycle(), 'demoting self because i am not the healthiest node')
 
     def test_follow_new_leader_because_not_healthiest(self):
-        self.p.is_healthiest_node = false
+        self.ha.is_healthiest_node = false
         self.p.is_leader = false
         self.assertEquals(self.ha.run_cycle(), 'following a different leader because i am not the healthiest node')
 
@@ -138,11 +176,124 @@ class TestHa(unittest.TestCase):
         self.ha.cluster.is_unlocked = false
         self.assertEquals(self.ha.run_cycle(), 'demoting self because i do not have the lock and i was a leader')
 
+    def test_demote_because_update_lock_failed(self):
+        self.ha.cluster.is_unlocked = false
+        self.ha.has_lock = true
+        self.ha.update_lock = false
+        self.assertEquals(self.ha.run_cycle(), 'demoting self because i do not have the lock and i was a leader')
+
     def test_follow_the_leader(self):
         self.ha.cluster.is_unlocked = false
         self.p.is_leader = false
         self.assertEquals(self.ha.run_cycle(), 'no action.  i am a secondary and i am following a leader')
 
     def test_no_etcd_connection_master_demote(self):
-        self.ha.load_cluster_from_dcs = dead_etcd
+        self.ha.load_cluster_from_dcs = Mock(side_effect=DCSError('Etcd is not responding properly'))
         self.assertEquals(self.ha.run_cycle(), 'demoted self because DCS is not accessible and i was a leader')
+
+    def test_bootstrap_from_leader(self):
+        self.ha.cluster = get_cluster_initialized_with_leader()
+        self.p.bootstrap = false
+        self.assertEquals(self.ha.bootstrap(), 'trying to bootstrap from leader')
+
+    def test_bootstrap_waiting_for_leader(self):
+        self.ha.cluster = get_cluster_initialized_without_leader()
+        self.assertEquals(self.ha.bootstrap(), 'waiting for leader to bootstrap')
+
+    def test_bootstrap_initialize_lock_failed(self):
+        self.ha.cluster = get_cluster_not_initialized_without_leader()
+        self.assertEquals(self.ha.bootstrap(), 'failed to acquire initialize lock')
+
+    def test_bootstrap_initialized_new_cluster(self):
+        self.ha.cluster = get_cluster_not_initialized_without_leader()
+        self.e.initialize = true
+        self.assertEquals(self.ha.bootstrap(), 'initialized a new cluster')
+
+    def test_bootstrap_release_initialize_key_on_failure(self):
+        self.ha.cluster = get_cluster_not_initialized_without_leader()
+        self.e.initialize = true
+        self.p.bootstrap = Mock(side_effect=PostgresException("Could not bootstrap master PostgreSQL"))
+        self.assertRaises(PostgresException, self.ha.bootstrap)
+
+    def test_reinitialize(self):
+        self.ha.schedule_reinitialize()
+        self.ha.schedule_reinitialize()
+        self.ha.run_cycle()
+        self.assertIsNone(self.ha._async_executor.scheduled_action)
+
+        self.ha.cluster = get_cluster_initialized_with_leader()
+        self.ha.has_lock = true
+        self.ha.schedule_reinitialize()
+        self.ha.run_cycle()
+        self.assertIsNone(self.ha._async_executor.scheduled_action)
+
+        self.ha.has_lock = false
+        self.ha.schedule_reinitialize()
+        self.ha.run_cycle()
+
+    def test_restart(self):
+        self.assertEquals(self.ha.restart(), (True, 'restarted successfully'))
+        self.p.restart = false
+        self.assertEquals(self.ha.restart(), (False, 'restart failed'))
+        self.ha.schedule_reinitialize()
+        self.assertEquals(self.ha.restart(), (False, 'reinitialize already in progress'))
+
+    def test_restart_in_progress(self):
+        self.ha._async_executor.schedule('restart', True)
+        self.assertTrue(self.ha.restart_scheduled())
+        self.assertEquals(self.ha.run_cycle(), 'not healthy enough for leader race')
+
+        self.ha.cluster = get_cluster_initialized_with_leader()
+        self.assertEquals(self.ha.run_cycle(), 'restart in progress')
+
+        self.ha.has_lock = true
+        self.assertEquals(self.ha.run_cycle(), 'updated leader lock during restart')
+
+        self.ha.update_lock = false
+        self.assertEquals(self.ha.run_cycle(), 'failed to update leader lock during restart')
+
+    @patch('requests.get', requests_get)
+    def test_manual_failover_from_leader(self):
+        self.ha.has_lock = true
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, 'blabla', ''))
+        self.assertEquals(self.ha.run_cycle(), 'no action.  i am the leader with the lock')
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, '', MockPostgresql.name))
+        self.assertEquals(self.ha.run_cycle(), 'no action.  i am the leader with the lock')
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, '', 'blabla'))
+        self.assertEquals(self.ha.run_cycle(), 'no action.  i am the leader with the lock')
+        f = Failover(0, MockPostgresql.name, '')
+        self.ha.cluster = get_cluster_initialized_with_leader(f)
+        self.assertEquals(self.ha.run_cycle(), 'manual failover: demoting myself')
+
+    @patch('requests.get', requests_get)
+    def test_manual_failover_process_no_leader(self):
+        self.p.is_leader = false
+        self.ha.cluster = get_cluster_initialized_without_leader(failover=Failover(0, '', MockPostgresql.name))
+        self.assertEquals(self.ha.run_cycle(), 'promoted self to leader by acquiring session lock')
+        self.ha.cluster = get_cluster_initialized_without_leader(failover=Failover(0, '', 'leader'))
+        self.assertEquals(self.ha.run_cycle(), 'promoted self to leader by acquiring session lock')
+        self.ha.fetch_node_status = lambda e: (e, True, True, 0)  # accessible, in_recovery
+        self.assertEquals(self.ha.run_cycle(), 'following a different leader because i am not the healthiest node')
+        self.ha.cluster = get_cluster_initialized_without_leader(failover=Failover(0, MockPostgresql.name, ''))
+        self.assertEquals(self.ha.run_cycle(), 'following a different leader because i am not the healthiest node')
+        self.ha.fetch_node_status = lambda e: (e, False, True, 0)  # accessible, in_recovery
+        self.assertEquals(self.ha.run_cycle(), 'promoted self to leader by acquiring session lock')
+
+    def test__is_healthiest_node(self):
+        self.assertTrue(self.ha._is_healthiest_node(self.ha.old_cluster.members))
+        self.p.is_leader = false
+        self.ha.fetch_node_status = lambda e: (e, True, True, 0)  # accessible, in_recovery
+        self.assertTrue(self.ha._is_healthiest_node(self.ha.old_cluster.members))
+        self.ha.fetch_node_status = lambda e: (e, True, False, 0)  # accessible, not in_recovery
+        self.assertFalse(self.ha._is_healthiest_node(self.ha.old_cluster.members))
+        self.ha.fetch_node_status = lambda e: (e, True, True, 1)  # accessible, in_recovery, xlog location ahead
+        self.assertFalse(self.ha._is_healthiest_node(self.ha.old_cluster.members))
+        self.p.check_replication_lag = false
+        self.assertFalse(self.ha._is_healthiest_node(self.ha.old_cluster.members))
+
+    @patch('requests.get', requests_get)
+    def test_fetch_node_status(self):
+        member = Member(0, 'test', 1, {'api_url': 'http://127.0.0.1:8011/patroni'})
+        self.ha.fetch_node_status(member)
+        member = Member(0, 'test', 1, {'api_url': 'http://localhost:8011/patroni'})
+        self.ha.fetch_node_status(member)

tests/test_patroni.py

@@ -1,144 +1,76 @@
-import datetime
-import helpers.zookeeper
-import psycopg2
-import subprocess
 import sys
 import time
 import unittest
 import yaml
 
-from helpers.api import RestApiServer
-from helpers.dcs import Cluster, Member
-from helpers.etcd import Etcd
-from helpers.zookeeper import ZooKeeper
 from mock import Mock, patch
+from patroni.api import RestApiServer
+from patroni.async_executor import AsyncExecutor
+from patroni.etcd import Etcd
 from patroni import Patroni, main
+from patroni.zookeeper import ZooKeeper
 from six.moves import BaseHTTPServer
-from test_etcd import Client, etcd_read, etcd_write
-from test_ha import true, false
-from test_postgresql import Postgresql, subprocess_call, psycopg2_connect
+from test_etcd import Client, SleepException, etcd_read, etcd_write
+from test_postgresql import Postgresql, psycopg2_connect
 from test_zookeeper import MockKazooClient
 
 
-def nop(*args, **kwargs):
-    pass
-
-
 def time_sleep(*args):
-    raise Exception()
-
-
-class Mock_BaseServer__is_shut_down:
-
-    def set(self):
-        pass
-
-    def clear(self):
-        pass
+    raise SleepException()
 
 
+@patch('time.sleep', Mock())
+@patch('subprocess.call', Mock(return_value=0))
+@patch('psycopg2.connect', psycopg2_connect)
+@patch.object(Postgresql, 'write_pg_hba', Mock())
+@patch.object(Postgresql, 'write_recovery_conf', Mock())
+@patch.object(BaseHTTPServer.HTTPServer, '__init__', Mock())
+@patch.object(AsyncExecutor, 'run', Mock())
 class TestPatroni(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        self.tearDown = self.tear_down
-        super(TestPatroni, self).__init__(method_name)
-
-    def set_up(self):
+    @patch.object(Client, 'machines')
+    def setUp(self, mock_machines):
+        mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
         self.touched = False
-        subprocess.call = subprocess_call
-        psycopg2.connect = psycopg2_connect
-        self.time_sleep = time.sleep
-        time.sleep = nop
-        self.write_pg_hba = Postgresql.write_pg_hba
-        self.write_recovery_conf = Postgresql.write_recovery_conf
-        Postgresql.write_pg_hba = nop
-        Postgresql.write_recovery_conf = nop
-        BaseHTTPServer.HTTPServer.__init__ = nop
-        RestApiServer._BaseServer__is_shut_down = Mock_BaseServer__is_shut_down()
+        self.init_cancelled = False
+        RestApiServer._BaseServer__is_shut_down = Mock()
         RestApiServer._BaseServer__shutdown_request = True
         RestApiServer.socket = 0
         with open('postgres0.yml', 'r') as f:
             config = yaml.load(f)
-            with patch.object(Client, 'machines') as mock_machines:
-                mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
-                self.p = Patroni(config)
-
-    def tear_down(self):
-        time.sleep = self.time_sleep
-        Postgresql.write_pg_hba = self.write_pg_hba
-        Postgresql.write_recovery_conf = self.write_recovery_conf
+            self.p = Patroni(config)
+            self.p.ha.dcs.client.write = etcd_write
+            self.p.ha.dcs.client.read = etcd_read
 
+    @patch('patroni.zookeeper.KazooClient', MockKazooClient())
     def test_get_dcs(self):
-        helpers.zookeeper.KazooClient = MockKazooClient
         self.assertIsInstance(self.p.get_dcs('', {'zookeeper': {'scope': '', 'hosts': ''}}), ZooKeeper)
         self.assertRaises(Exception, self.p.get_dcs, '', {})
 
-    def test_patroni_main(self):
+    @patch('time.sleep', Mock(side_effect=SleepException()))
+    @patch.object(Etcd, 'delete_leader', Mock())
+    @patch.object(Client, 'machines')
+    def test_patroni_main(self, mock_machines):
         main()
         sys.argv = ['patroni.py', 'postgres0.yml']
-        time.sleep = time_sleep
 
-        with patch.object(Client, 'machines') as mock_machines:
-            mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
-            Patroni.initialize = nop
-            touch_member = Patroni.touch_member
-            run = Patroni.run
+        mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
+        with patch.object(Patroni, 'run', Mock(side_effect=SleepException())):
+            self.assertRaises(SleepException, main)
+        with patch.object(Patroni, 'run', Mock(side_effect=KeyboardInterrupt())):
+            main()
 
-            Patroni.touch_member = self.touch_member
-            Patroni.run = time_sleep
+    @patch('time.sleep', Mock(side_effect=SleepException()))
+    def test_run(self):
+        self.p.ha.dcs.watch = time_sleep
+        self.assertRaises(SleepException, self.p.run)
 
-            Etcd.delete_leader = nop
-
-            self.assertRaises(Exception, main)
-
-            Patroni.run = run
-            Patroni.touch_member = touch_member
-
-    def test_patroni_run(self):
-        time.sleep = time_sleep
-        self.p.touch_member = self.touch_member
-        self.p.ha.state_handler.sync_replication_slots = time_sleep
-        self.p.ha.dcs.client.read = etcd_read
-        self.assertRaises(Exception, self.p.run)
-        self.p.ha.state_handler.is_leader = lambda: False
-        self.p.api.start = nop
-        self.assertRaises(Exception, self.p.run)
-
-    def touch_member(self, ttl=None):
-        if not self.touched:
-            self.touched = True
-            return False
-        return True
-
-    def test_touch_member(self):
-        self.p.ha.dcs.client.write = etcd_write
-        self.p.touch_member()
-        now = datetime.datetime.utcnow()
-        member = Member(0, self.p.postgresql.name, 'b', 'c', (now + datetime.timedelta(
-            seconds=self.p.shutdown_member_ttl + 10)).strftime('%Y-%m-%dT%H:%M:%S.%fZ'), None)
-        self.p.ha.cluster = Cluster(True, member, 0, [member])
-        self.p.touch_member()
-
-    def test_patroni_initialize(self):
-        self.p.postgresql.should_use_s3_to_create_replica = false
-        self.p.ha.dcs.client.write = etcd_write
-        self.p.touch_member = self.touch_member
-        self.p.postgresql.data_directory_empty = true
-        self.p.ha.dcs.race = true
-        self.p.initialize()
-
-        self.p.ha.dcs.race = false
-        time.sleep = time_sleep
-        self.p.ha.dcs.client.read = etcd_read
-        self.p.initialize()
-
-        self.p.ha.dcs.current_leader = nop
-        self.assertRaises(Exception, self.p.initialize)
-
-        self.p.postgresql.data_directory_empty = false
-        self.p.initialize()
+        self.p.ha.state_handler.is_leader = Mock(return_value=False)
+        self.p.api.start = Mock()
+        self.assertRaises(SleepException, self.p.run)
 
     def test_schedule_next_run(self):
+        self.p.ha.dcs.watch = Mock(return_value=True)
+        self.p.schedule_next_run()
         self.p.next_run = time.time() - self.p.nap_time - 1
         self.p.schedule_next_run()

tests/test_postgresql.py

@@ -1,51 +1,41 @@
+import mock  # for the mock.call method, importing it without a namespace breaks python3
 import os
 import psycopg2
 import shutil
-import subprocess
 import unittest
 
-from helpers.dcs import Cluster, Member
-from helpers.postgresql import Postgresql
+from sys import version_info
+if version_info.major == 2:
+    import __builtin__ as builtins
+else:
+    import builtins
 
-
-def nop(*args, **kwargs):
-    pass
-
-
-def subprocess_call(cmd, shell=False, env=None):
-    return 0
-
-
-def false(*args, **kwargs):
-    return False
+from mock import Mock, MagicMock, PropertyMock, patch, mock_open
+from patroni.dcs import Cluster, Leader, Member
+from patroni.exceptions import PostgresException, PostgresConnectionException
+from patroni.postgresql import Postgresql
+from patroni.utils import RetryFailedError
+from test_ha import false
+import subprocess
 
 
 class MockCursor:
 
-    def __init__(self):
+    def __init__(self, connection):
+        self.connection = connection
         self.closed = False
-        self.current = 0
         self.results = []
 
     def execute(self, sql, *params):
-        if sql.startswith('blabla'):
+        if sql.startswith('blabla') or sql == 'CHECKPOINT':
             raise psycopg2.OperationalError()
-        elif sql.startswith('InterfaceError'):
-            raise psycopg2.InterfaceError()
+        elif sql.startswith('RetryFailedError'):
+            raise RetryFailedError('retry')
         elif sql.startswith('SELECT slot_name'):
             self.results = [('blabla',), ('foobar',)]
-        elif sql.startswith('SELECT pg_current_xlog_location()'):
+        elif sql.startswith('SELECT pg_xlog_location_diff'):
             self.results = [(0,)]
-        elif sql.startswith('SELECT pg_is_in_recovery(), %s'):
-            if params[0][0] == 1:
-                raise psycopg2.OperationalError()
-            elif params[0][0] == 2:
-                self.results = [(True, -1)]
-            else:
-                self.results = [(False, 0)]
-        elif sql.startswith('SELECT CASE WHEN pg_is_in_recovery()'):
-            self.results = [(0,)]
-        elif sql.startswith('SELECT pg_is_in_recovery()'):
+        elif sql == 'SELECT pg_is_in_recovery()':
             self.results = [(False, )]
         elif sql.startswith('SELECT to_char(pg_postmaster_start_time'):
             self.results = [('', True, '', '', '', False)]
@@ -73,44 +63,108 @@ class MockCursor:
         for i in self.results:
             yield i
 
+    def __enter__(self):
+        return self
 
-class MockConnect:
-
-    def __init__(self):
-        self.autocommit = False
-        self.closed = 0
-
-    def cursor(self):
-        return MockCursor()
-
-    def close(self):
+    def __exit__(self, *args):
         pass
 
 
-def psycopg2_connect(*args, **kwargs):
+class MockConnect(Mock):
 
+    autocommit = False
+    closed = 0
+
+    def cursor(self):
+        return MockCursor(self)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+
+def pg_controldata_string(*args, **kwargs):
+    return """
+pg_control version number:            942
+Catalog version number:               201509161
+Database system identifier:           6200971513092291716
+Database cluster state:               shut down in recovery
+pg_control last modified:             Fri Oct  2 10:57:06 2015
+Latest checkpoint location:           0/30000C8
+Prior checkpoint location:            0/2000060
+Latest checkpoint's REDO location:    0/3000090
+Latest checkpoint's REDO WAL file:    000000020000000000000003
+Latest checkpoint's TimeLineID:       2
+Latest checkpoint's PrevTimeLineID:   2
+Latest checkpoint's full_page_writes: on
+Latest checkpoint's NextXID:          0/943
+Latest checkpoint's NextOID:          24576
+Latest checkpoint's NextMultiXactId:  1
+Latest checkpoint's NextMultiOffset:  0
+Latest checkpoint's oldestXID:        931
+Latest checkpoint's oldestXID's DB:   1
+Latest checkpoint's oldestActiveXID:  943
+Latest checkpoint's oldestMultiXid:   1
+Latest checkpoint's oldestMulti's DB: 1
+Latest checkpoint's oldestCommitTs:   0
+Latest checkpoint's newestCommitTs:   0
+Time of latest checkpoint:            Fri Oct  2 10:56:54 2015
+Fake LSN counter for unlogged rels:   0/1
+Minimum recovery ending location:     0/30241F8
+Min recovery ending loc's timeline:   2
+Backup start location:                0/0
+Backup end location:                  0/0
+End-of-backup record required:        no
+wal_level setting:                    hot_standby
+Current wal_log_hints setting:                on
+Current max_connections setting:              100
+Current max_worker_processes setting:         8
+Current max_prepared_xacts setting:           0
+Current max_locks_per_xact setting:           64
+Current track_commit_timestamp setting:       off
+Maximum data alignment:               8
+Database block size:                  8192
+Blocks per segment of large relation: 131072
+WAL block size:                       8192
+Bytes per WAL segment:                16777216
+Maximum length of identifiers:        64
+Maximum columns in an index:          32
+Maximum size of a TOAST chunk:        1996
+Size of a large-object chunk:         2048
+Date/time type storage:               64-bit integers
+Float4 argument passing:              by value
+Float8 argument passing:              by value
+Data page checksum version:           0
+"""
+
+
+def postmaster_opts_string(*args, **kwargs):
+    return '/usr/local/pgsql/bin/postgres "-D" "data/postgresql0" "--listen_addresses=127.0.0.1" "--port=5432"'\
+        ' "--hot_standby=on" "--wal_keep_segments=8" "--wal_level=hot_standby" "--archive_command=mkdir -p ../wal_archive \n'\
+        '&& cp %p ../wal_archive/%f" "--wal_log_hints=on" "--max_wal_senders=5" "--archive_timeout=1800s" "--archive_mode=on"'\
+        ' "--max_replication_slots=5"\n'
+
+
+def psycopg2_connect(*args, **kwargs):
     return MockConnect()
 
 
-def is_running():
-    return False
-
-
+@patch('subprocess.call', Mock(return_value=0))
+@patch('psycopg2.connect', psycopg2_connect)
+@patch('shutil.copy', Mock())
 class TestPostgresql(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        self.tearDown = self.tear_down
-        super(TestPostgresql, self).__init__(method_name)
-
-    def set_up(self):
-        subprocess.call = subprocess_call
-        shutil.copy = nop
+    @patch('subprocess.call', Mock(return_value=0))
+    @patch('psycopg2.connect', psycopg2_connect)
+    def setUp(self):
         self.p = Postgresql({'name': 'test0', 'scope': 'batman', 'data_dir': 'data/test0',
                              'listen': '127.0.0.1, *:5432', 'connect_address': '127.0.0.2:5432',
                              'pg_hba': ['hostssl all all 0.0.0.0/0 md5', 'host all all 0.0.0.0/0 md5'],
                              'superuser': {'password': ''},
                              'admin': {'username': 'admin', 'password': 'admin'},
+                             'pg_rewind': {'username': 'admin', 'password': 'admin'},
                              'replication': {'username': 'replicator',
                                              'password': 'rep-pass',
                                              'network': '127.0.0.1/32'},
@@ -119,20 +173,17 @@ class TestPostgresql(unittest.TestCase):
                                            'on_restart': 'true', 'on_role_change': 'true',
                                            'on_reload': 'true'
                                            },
-                             'restore': '/usr/bin/true'})
-        psycopg2.connect = psycopg2_connect
+                             'restore': 'true'})
         if not os.path.exists(self.p.data_dir):
             os.makedirs(self.p.data_dir)
-        self.leader = Member(0, 'leader', 'postgres://replicator:[email protected]:5435/postgres', None, None, 28)
-        self.other = Member(0, 'test1', 'postgres://replicator:[email protected]:5433/postgres', None, None, 28)
-        self.me = Member(0, 'test0', 'postgres://replicator:[email protected]:5434/postgres', None, None, 28)
+        self.leadermem = Member(0, 'leader', 28, {'conn_url': 'postgres://replicator:[email protected]:5435/postgres'})
+        self.leader = Leader(-1, 28, self.leadermem)
+        self.other = Member(0, 'test1', 28, {'conn_url': 'postgres://replicator:[email protected]:5433/postgres'})
+        self.me = Member(0, 'test0', 28, {'conn_url': 'postgres://replicator:[email protected]:5434/postgres'})
 
-    def tear_down(self):
+    def tearDown(self):
         shutil.rmtree('data')
 
-    def mock_query(self, p):
-        raise psycopg2.OperationalError("not supported")
-
     def test_data_directory_empty(self):
         self.assertTrue(self.p.data_directory_empty())
 
@@ -140,23 +191,74 @@ class TestPostgresql(unittest.TestCase):
         self.assertTrue(self.p.initialize())
         self.assertTrue(os.path.exists(os.path.join(self.p.data_dir, 'pg_hba.conf')))
 
-    def test_start_stop(self):
-        self.assertFalse(self.p.start())
-        self.p.is_running = is_running
-        with open(os.path.join(self.p.data_dir, 'postmaster.pid'), 'w'):
-            pass
+    def test_start(self):
         self.assertTrue(self.p.start())
+        self.p.is_running = false
+        open(os.path.join(self.p.data_dir, 'postmaster.pid'), 'w').close()
+        self.assertTrue(self.p.start())
+
+    def test_stop(self):
         self.assertTrue(self.p.stop())
+        with patch('subprocess.call', Mock(return_value=1)):
+            self.assertTrue(self.p.stop())
+            self.p.is_running = Mock(return_value=True)
+            self.assertFalse(self.p.stop())
+
+    def test_restart(self):
+        self.p.start = false
+        self.p.is_running = false
+        self.assertFalse(self.p.restart())
+        self.assertEquals(self.p.state, 'restart failed (restarting)')
 
     def test_sync_from_leader(self):
         self.assertTrue(self.p.sync_from_leader(self.leader))
 
-    def test_follow_the_leader(self):
-        self.p.demote(self.leader)
+    @patch('subprocess.call', side_effect=Exception("Test"))
+    def test_pg_rewind(self, mock_call):
+        self.assertTrue(self.p.rewind(self.leader))
+        subprocess.call = mock_call
+        self.assertFalse(self.p.rewind(self.leader))
+
+    @patch('patroni.postgresql.Postgresql.rewind', return_value=False)
+    @patch('patroni.postgresql.Postgresql.remove_data_directory', MagicMock(return_value=True))
+    @patch('patroni.postgresql.Postgresql.single_user_mode', MagicMock(return_value=1))
+    def test_follow_the_leader(self, mock_pg_rewind):
+        self.p.demote()
         self.p.follow_the_leader(None)
-        self.p.demote(self.leader)
+        self.p.demote()
         self.p.follow_the_leader(self.leader)
-        self.p.follow_the_leader(self.other)
+        self.p.follow_the_leader(Leader(-1, 28, self.other))
+        self.p.rewind = mock_pg_rewind
+        self.p.follow_the_leader(self.leader)
+        self.p.require_rewind()
+        with mock.patch('os.path.islink', MagicMock(return_value=True)):
+            with mock.patch('os.unlink', MagicMock(return_value=True)):
+                with mock.patch('patroni.postgresql.Postgresql.can_rewind', new_callable=PropertyMock(return_value=True)):
+                    self.p.follow_the_leader(self.leader, recovery=True)
+        self.p.require_rewind()
+        with mock.patch('patroni.postgresql.Postgresql.can_rewind', new_callable=PropertyMock(return_value=True)):
+            self.p.rewind.return_value = True
+            self.p.follow_the_leader(self.leader, recovery=True)
+            self.p.rewind.return_value = False
+            self.p.follow_the_leader(self.leader, recovery=True)
+
+    def test_can_rewind(self):
+        tmp = self.p.pg_rewind
+        self.p.pg_rewind = None
+        self.assertFalse(self.p.can_rewind)
+        self.p.pg_rewind = tmp
+        with mock.patch('subprocess.call', MagicMock(return_value=1)):
+            self.assertFalse(self.p.can_rewind)
+        with mock.patch('subprocess.call', side_effect=OSError("foo")):
+            self.assertFalse(self.p.can_rewind)
+        tmp = self.p.controldata()
+        self.p.controldata = lambda: {'wal_log_hints setting': 'on'}
+        self.assertTrue(self.p.can_rewind)
+        self.p.controldata = tmp
+
+    def test_create_replica(self):
+        self.p.delete_trigger_file = Mock(side_effect=OSError())
+        self.assertEquals(self.p.create_replica({'host': '', 'port': '', 'user': ''}, ''), 1)
 
     def test_create_connection_users(self):
         cfg = self.p.config
@@ -164,56 +266,166 @@ class TestPostgresql(unittest.TestCase):
         p = Postgresql(cfg)
         p.create_connection_users()
 
-    def test_create_replication_slots(self):
+    def test_sync_replication_slots(self):
         self.p.start()
-        cluster = Cluster(True, self.leader, 0, [self.me, self.other, self.leader])
-        self.p.create_replication_slots(cluster)
+        cluster = Cluster(True, self.leader, 0, [self.me, self.other, self.leadermem], None)
+        self.p.sync_replication_slots(cluster)
+        self.p.query = Mock(side_effect=psycopg2.OperationalError)
+        self.p.schedule_load_slots = True
+        self.p.sync_replication_slots(cluster)
+
+    @patch.object(MockConnect, 'closed', 2)
+    def test__query(self):
+        self.assertRaises(PostgresConnectionException, self.p._query, 'blabla')
+        self.p._state = 'restarting'
+        self.assertRaises(RetryFailedError, self.p._query, 'blabla')
 
     def test_query(self):
         self.p.query('select 1')
-        self.assertRaises(psycopg2.InterfaceError, self.p.query, 'InterfaceError')
+        self.assertRaises(PostgresConnectionException, self.p.query, 'RetryFailedError')
         self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
-        self.p._connection.closed = 2
-        self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
-        self.p._connection.closed = 2
-        self.p.disconnect = false
-        self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla')
-
-    def test_is_healthiest_node(self):
-        cluster = Cluster(True, self.leader, 0, [self.me, self.other, self.leader])
-        self.assertTrue(self.p.is_healthiest_node(cluster))
-        self.p.is_leader = false
-        self.assertFalse(self.p.is_healthiest_node(cluster))
-        self.p.xlog_position = lambda: 1
-        self.assertTrue(self.p.is_healthiest_node(cluster))
-        self.p.xlog_position = lambda: 2
-        self.assertFalse(self.p.is_healthiest_node(cluster))
-        self.p.config['maximum_lag_on_failover'] = -2
-        self.assertFalse(self.p.is_healthiest_node(cluster))
 
     def test_is_leader(self):
-        self.p.is_promoted = True
         self.assertTrue(self.p.is_leader())
-        self.assertFalse(self.p.is_promoted)
 
     def test_reload(self):
         self.assertTrue(self.p.reload())
 
     def test_is_healthy(self):
         self.assertTrue(self.p.is_healthy())
-        self.p.is_running = is_running
+        self.p.is_running = false
         self.assertFalse(self.p.is_healthy())
 
     def test_promote(self):
+        self.p._role = 'replica'
+        self.assertTrue(self.p.promote())
         self.assertTrue(self.p.promote())
 
     def test_last_operation(self):
         self.assertEquals(self.p.last_operation(), '0')
 
+    @patch('subprocess.Popen', Mock(side_effect=OSError()))
+    def test_call_nowait(self):
+        self.assertFalse(self.p.call_nowait('on_start'))
+
     def test_non_existing_callback(self):
         self.assertFalse(self.p.call_nowait('foobar'))
 
     def test_is_leader_exception(self):
         self.p.start()
-        self.p.query = self.mock_query
+        self.p.query = Mock(side_effect=psycopg2.OperationalError("not supported"))
         self.assertTrue(self.p.stop())
+
+    def test_check_replication_lag(self):
+        self.assertTrue(self.p.check_replication_lag(0))
+
+    @patch('os.rename', Mock())
+    @patch('os.path.isdir', Mock(return_value=True))
+    def test_move_data_directory(self):
+        self.p.is_running = false
+        self.p.move_data_directory()
+        with patch('os.rename', Mock(side_effect=OSError())):
+            self.p.move_data_directory()
+
+    def test_bootstrap(self):
+        with patch('subprocess.call', Mock(return_value=1)):
+            self.assertRaises(PostgresException, self.p.bootstrap)
+        self.p.bootstrap()
+        self.p.bootstrap(self.leader)
+
+    def test_remove_data_directory(self):
+        self.p.data_dir = 'data_dir'
+        self.p.remove_data_directory()
+        os.mkdir(self.p.data_dir)
+        self.p.remove_data_directory()
+        open(self.p.data_dir, 'w').close()
+        self.p.remove_data_directory()
+        os.symlink('unexisting', self.p.data_dir)
+        with patch('os.unlink', Mock(side_effect=Exception)):
+            self.p.remove_data_directory()
+        self.p.remove_data_directory()
+
+    @patch('subprocess.check_output', MagicMock(return_value=0, side_effect=pg_controldata_string))
+    @patch('subprocess.check_output', side_effect=subprocess.CalledProcessError)
+    @patch('subprocess.check_output', side_effect=Exception('Failed'))
+    def test_controldata(self, check_output_call_error, check_output_generic_exception):
+        data = self.p.controldata()
+        self.assertEquals(len(data), 50)
+        self.assertEquals(data['Database cluster state'], 'shut down in recovery')
+        self.assertEquals(data['wal_log_hints setting'], 'on')
+        self.assertEquals(int(data['Database block size']), 8192)
+
+        subprocess.check_output = check_output_call_error
+        data = self.p.controldata()
+        self.assertEquals(data, dict())
+
+        subprocess.check_output = check_output_generic_exception
+        self.assertRaises(Exception, self.p.controldata())
+
+    def test_read_postmaster_opts(self):
+        m = mock_open(read_data=postmaster_opts_string())
+        with patch.object(builtins, 'open', m):
+            data = self.p.read_postmaster_opts()
+            self.assertEquals(data['wal_level'], 'hot_standby')
+            self.assertEquals(int(data['max_replication_slots']), 5)
+            self.assertEqual(data.get('D'), None)
+
+            m.side_effect = IOError("foo")
+            data = self.p.read_postmaster_opts()
+            self.assertEqual(data, dict())
+
+            m.side_effect = Exception("foo")
+            self.assertRaises(Exception, self.p.read_postmaster_opts())
+
+    @patch('subprocess.Popen')
+    @patch.object(builtins, 'open', MagicMock(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(options=dict(archive_mode='on', archive_command='false')), 0)
+        subprocess_popen_mock.assert_called_once_with(['postgres', '--single', '-D', self.p.data_dir,
+                                                      '-c', 'archive_command=false', '-c', 'archive_mode=on',
+                                                       'postgres'], stdin=subprocess.PIPE,
+                                                      stdout=42,
+                                                      stderr=subprocess.STDOUT)
+        subprocess_popen_mock.reset_mock()
+        self.assertEquals(self.p.single_user_mode(command="CHECKPOINT"), 0)
+        subprocess_popen_mock.assert_called_once_with(['postgres', '--single', '-D', self.p.data_dir,
+                                                      'postgres'], stdin=subprocess.PIPE,
+                                                      stdout=42,
+                                                      stderr=subprocess.STDOUT)
+        subprocess_popen_mock.return_value = None
+        self.assertEquals(self.p.single_user_mode(), 1)
+
+    def fake_listdir(path):
+        if path.endswith(os.path.join('pg_xlog', 'archive_status')):
+            return ["a", "b", "c"]
+        return []
+
+    @patch('os.listdir', MagicMock(side_effect=fake_listdir))
+    @patch('os.path.isdir', MagicMock(return_value=True))
+    @patch('os.unlink', return_value=True)
+    @patch('os.remove', return_value=True)
+    @patch('os.path.islink', return_value=False)
+    @patch('os.path.isfile', return_value=True)
+    def test_cleanup_archive_status(self, mock_file, mock_link, mock_remove, mock_unlink):
+        ap = os.path.join(self.p.data_dir, 'pg_xlog', 'archive_status/')
+        self.p.cleanup_archive_status()
+        mock_remove.assert_has_calls([mock.call(ap+'a'), mock.call(ap+'b'), mock.call(ap+'c')])
+        mock_unlink.assert_not_called()
+
+        mock_remove.reset_mock()
+
+        mock_file.return_value = False
+        mock_link.return_value = True
+        self.p.cleanup_archive_status()
+        mock_unlink.assert_has_calls([mock.call(ap+'a'), mock.call(ap+'b'), mock.call(ap+'c')])
+        mock_remove.assert_not_called()
+
+        mock_unlink.reset_mock()
+        mock_remove.reset_mock()
+
+        mock_file.side_effect = Exception("foo")
+        mock_link.side_effect = Exception("foo")
+        self.p.cleanup_archive_status()
+        mock_unlink.assert_not_called()
+        mock_remove.assert_not_called()

tests/test_restore.py

@@ -1,7 +1,7 @@
 import unittest
 from mock import MagicMock, patch
 import os
-from scripts.restore import Restore, WALERestore
+from patroni.scripts.restore import Restore, WALERestore
 
 
 def fake_cursor_fetchone(*args, **kwargs):

tests/test_utils.py

@@ -1,16 +1,8 @@
-import os
-import time
 import unittest
 
-from helpers.utils import reap_children, sigchld_handler, sigterm_handler, sleep
-
-
-def nop(*args, **kwargs):
-    pass
-
-
-def os_waitpid(a, b):
-    return (0, 0)
+from mock import Mock, patch
+from patroni.exceptions import PatroniException
+from patroni.utils import Retry, RetryFailedError, reap_children, sigchld_handler, sigterm_handler, sleep
 
 
 def time_sleep(_):
@@ -19,27 +11,63 @@ def time_sleep(_):
 
 class TestUtils(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        self.tearDown = self.tear_down
-        super(TestUtils, self).__init__(method_name)
-
-    def set_up(self):
-        self.time_sleep = time.sleep
-        time.sleep = nop
-
-    def tear_down(self):
-        time.sleep = self.time_sleep
-
     def test_sigterm_handler(self):
         self.assertRaises(SystemExit, sigterm_handler, None, None)
 
+    @patch('time.sleep', Mock())
     def test_reap_children(self):
         reap_children()
-        os.waitpid = os_waitpid
-        sigchld_handler(None, None)
-        reap_children()
+        with patch('os.waitpid', Mock(return_value=(0, 0))):
+            sigchld_handler(None, None)
+            reap_children()
 
+    @patch('time.sleep', time_sleep)
     def test_sleep(self):
-        time.sleep = time_sleep
         sleep(0.01)
+
+
+@patch('time.sleep', Mock())
+class TestRetrySleeper(unittest.TestCase):
+
+    def _fail(self, times=1):
+        scope = dict(times=0)
+
+        def inner():
+            if scope['times'] >= times:
+                pass
+            else:
+                scope['times'] += 1
+                raise PatroniException('Failed!')
+        return inner
+
+    def _makeOne(self, *args, **kwargs):
+        return Retry(*args, **kwargs)
+
+    def test_reset(self):
+        retry = self._makeOne(delay=0, max_tries=2)
+        retry(self._fail())
+        self.assertEquals(retry._attempts, 1)
+        retry.reset()
+        self.assertEquals(retry._attempts, 0)
+
+    def test_too_many_tries(self):
+        retry = self._makeOne(delay=0)
+        self.assertRaises(RetryFailedError, retry, self._fail(times=999))
+        self.assertEquals(retry._attempts, 1)
+
+    def test_maximum_delay(self):
+        retry = self._makeOne(delay=10, max_tries=100)
+        retry(self._fail(times=10))
+        self.assertTrue(retry._cur_delay < 4000, retry._cur_delay)
+        # gevent's sleep function is picky about the type
+        self.assertEquals(type(retry._cur_delay), float)
+
+    def test_deadline(self):
+        retry = self._makeOne(deadline=0.0001)
+        self.assertRaises(RetryFailedError, retry, self._fail(times=100))
+
+    def test_copy(self):
+        _sleep = lambda t: None
+        retry = self._makeOne(sleep_func=_sleep)
+        rcopy = retry.copy()
+        self.assertTrue(rcopy.sleep_func is _sleep)

tests/test_zookeeper.py

@@ -1,125 +1,112 @@
-import helpers.zookeeper
-import requests
+import six
 import unittest
 
-from helpers.zookeeper import ExhibitorEnsembleProvider, ZooKeeper, ZooKeeperError
+from mock import Mock, patch
+from patroni.dcs import Leader
+from patroni.zookeeper import ExhibitorEnsembleProvider, ZooKeeper, ZooKeeperError
 from kazoo.client import KazooState
 from kazoo.exceptions import NoNodeError, NodeExistsError
 from kazoo.protocol.states import ZnodeStat
-from test_etcd import MockPostgresql, requests_get
+from test_etcd import SleepException, requests_get
 
 
-class MockEvent:
+class MockKazooClient(Mock):
 
-    def clear(self):
-        pass
-
-    def set(self):
-        pass
-
-    def wait(self, timeout):
-        pass
-
-    def isSet(self):
-        return True
-
-
-class MockEventHandler:
-
-    def event_object(self):
-        return MockEvent()
-
-
-class MockKazooClient:
-
-    def __init__(self, **kwargs):
-        self.handler = MockEventHandler()
-        self.leader = False
-        self.exists = True
-
-    def start(self, timeout):
-        pass
+    leader = False
+    exists = True
+    handler = Mock()
 
     @property
     def client_id(self):
         return (-1, '')
 
-    def add_listener(self, cb):
-        pass
-
     def retry(self, func, *args, **kwargs):
         func(*args, **kwargs)
 
     def get(self, path, watch=None):
-        if path == '/service/test/no_node':
+        if not isinstance(path, six.string_types):
+            raise TypeError("Invalid type for 'path' (string expected)")
+        if path == '/no_node':
             raise NoNodeError
-        elif path == '/service/test/other_exception':
-            raise Exception()
         elif '/members/' in path:
             return (
-                'postgres://repuser:rep-pass@localhost:5434/postgres?application_name=http://127.0.0.1:8009/patroni',
-                ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
+                b'postgres://repuser:rep-pass@localhost:5434/postgres?application_name=http://127.0.0.1:8009/patroni',
+                ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0 if self.exists else -1, 0, 0, 0)
             )
         elif path.endswith('/optime/leader'):
-            return '1'
+            return (b'1', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))
         elif path.endswith('/leader'):
             if self.leader:
-                return ('foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0))
-            return ('foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))
+                return (b'foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0))
+            return (b'foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))
+        elif path.endswith('/initialize'):
+            return (b'foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))
+        return (b'', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))
 
     def get_children(self, path, watch=None, include_data=False):
+        if not isinstance(path, six.string_types):
+            raise TypeError("Invalid type for 'path' (string expected)")
+        if path.startswith('/no_node'):
+            raise NoNodeError
+        elif path in ['/service/bla/', '/service/test/']:
+            return ['initialize', 'leader', 'members', 'optime', 'failover']
         return ['foo', 'bar', 'buzz']
 
-    def create(self, path, value="", acl=None, ephemeral=False, sequence=False, makepath=False):
+    def create(self, path, value=b"", acl=None, ephemeral=False, sequence=False, makepath=False):
+        if not isinstance(path, six.string_types):
+            raise TypeError("Invalid type for 'path' (string expected)")
+        if not isinstance(value, (six.binary_type,)):
+            raise TypeError("Invalid type for 'value' (must be a byte string)")
         if path.endswith('/initialize') or path == '/service/test/optime/leader':
             raise Exception
-        elif value == 'retry' or (value == 'exists' and self.exists):
+        elif value == b'retry' or (value == b'exists' and self.exists):
             raise NodeExistsError
 
     def set(self, path, value, version=-1):
+        if not isinstance(path, six.string_types):
+            raise TypeError("Invalid type for 'path' (string expected)")
+        if not isinstance(value, (six.binary_type,)):
+            raise TypeError("Invalid type for 'value' (must be a byte string)")
         if path == '/service/bla/optime/leader':
             raise Exception
+        if path == '/service/test/members/bar':
+            if value == b'retry':
+                return
+        if path == '/service/test/failover':
+            if value == b'Exception':
+                raise Exception
+            elif value == b'ok':
+                return
         raise NoNodeError
 
     def delete(self, path, version=-1, recursive=False):
+        if not isinstance(path, six.string_types):
+            raise TypeError("Invalid type for 'path' (string expected)")
         self.exists = False
         if path == '/service/test/leader':
             if self.leader:
                 return
             self.leader = True
             raise Exception
-
-    def set_hosts(self, hosts, randomize_hosts=None):
-        pass
-
-
-def exhibitor_sleep(_):
-    raise Exception
+        elif path == '/service/test/members/buzz':
+            raise Exception
+        elif path.endswith('/initialize') or path == '/service/test/members/bar':
+            raise NoNodeError
 
 
+@patch('requests.get', requests_get)
+@patch('patroni.zookeeper.sleep', Mock(side_effect=SleepException()))
 class TestExhibitorEnsembleProvider(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestExhibitorEnsembleProvider, self).__init__(method_name)
-
-    def set_up(self):
-        requests.get = requests_get
-        helpers.zookeeper.sleep = exhibitor_sleep
-
     def test_init(self):
-        self.assertRaises(Exception, ExhibitorEnsembleProvider, ['localhost'], 8181)
+        self.assertRaises(SleepException, ExhibitorEnsembleProvider, ['localhost'], 8181)
 
 
 class TestZooKeeper(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestZooKeeper, self).__init__(method_name)
-
-    def set_up(self):
-        requests.get = requests_get
-        helpers.zookeeper.KazooClient = MockKazooClient
+    @patch('requests.get', requests_get)
+    @patch('patroni.zookeeper.KazooClient', MockKazooClient)
+    def setUp(self):
         self.zk = ZooKeeper('foo', {'exhibitor': {'hosts': ['localhost', 'exhibitor'], 'port': 8181}, 'scope': 'test'})
 
     def test_session_listener(self):
@@ -127,36 +114,63 @@ class TestZooKeeper(unittest.TestCase):
 
     def test_get_node(self):
         self.assertIsNone(self.zk.get_node('/no_node'))
-        self.assertIsNone(self.zk.get_node('/other_exception'))
+
+    def test_get_children(self):
+        self.assertListEqual(self.zk.get_children('/no_node'), [])
 
     def test__inner_load_cluster(self):
         self.zk._base_path = self.zk._base_path.replace('test', 'bla')
         self.zk._inner_load_cluster()
+        self.zk._base_path = self.zk._base_path = '/no_node'
+        self.zk._inner_load_cluster()
 
     def test_get_cluster(self):
         self.assertRaises(ZooKeeperError, self.zk.get_cluster)
         self.zk.exhibitor.poll = lambda: True
-        self.zk.get_cluster()
+        cluster = self.zk.get_cluster()
+        self.assertIsInstance(cluster.leader, Leader)
         self.zk.touch_member('foo')
         self.zk.delete_leader()
 
-    def test_race(self):
-        self.assertFalse(self.zk.race('/initialize'))
+    def test_set_failover_value(self):
+        self.zk.set_failover_value('')
+        self.zk.set_failover_value('ok')
+        self.zk.set_failover_value('Exception')
+
+    def test_initialize(self):
+        self.assertFalse(self.zk.initialize())
+
+    def test_cancel_initialization(self):
+        self.zk.cancel_initialization()
 
     def test_touch_member(self):
+        self.zk._name = 'buzz'
+        self.zk.get_cluster()
         self.zk.touch_member('new')
+        self.zk._name = 'bar'
+        self.zk.touch_member('new')
+        self.zk._name = 'na'
+        self.zk.client.exists = 1
         self.zk.touch_member('exists')
+        self.zk._name = 'bar'
+        self.zk.touch_member('retry')
+        self.zk.fetch_cluster = True
+        self.zk.get_cluster()
         self.zk.touch_member('retry')
 
     def test_take_leader(self):
         self.zk.take_leader()
 
     def test_update_leader(self):
-        self.zk.last_leader_operation = -1
-        self.assertTrue(self.zk.update_leader(MockPostgresql()))
-        self.zk._base_path = self.zk._base_path.replace('test', 'bla')
-        self.zk.last_leader_operation = -1
-        self.assertTrue(self.zk.update_leader(MockPostgresql()))
+        self.assertTrue(self.zk.update_leader())
 
-    def test_sleep(self):
-        self.zk.sleep(0)
+    def test_write_leader_optime(self):
+        self.zk.last_leader_operation = '0'
+        self.zk.write_leader_optime('1')
+        self.zk._base_path = self.zk._base_path.replace('test', 'bla')
+        self.zk.write_leader_optime('2')
+
+    def test_watch(self):
+        self.zk.watch(0)
+        self.zk.event.isSet = lambda: True
+        self.zk.watch(0)