Bumped version to 0.80

Feature/acceptance tests behave

.gitignore

@@ -1,3 +1,11 @@
 data/*
 *.pyc
-helpers/*.pyc
+*.egg/
+*.egg-info/
+.cache/
+.coverage
+.eggs/
+build/
+coverage.xml
+junit.xml
+pgpass

.travis.yml

@@ -1,14 +1,32 @@
+sudo: required
 language: python
+addons:
+  postgresql: "9.5"
+env:
+  global:
+    - BOTO_CONFIG='' ETCDVERSION=2.2.5
+  matrix:
+    - TEST_SUITE="python setup.py test"
+    - TEST_SUITE="behave"
 python:
   - "2.7"
-  - "3.3"
   - "3.4"
+  - "3.5"
 install:
-  - if [[ $TRAVIS_PYTHON_VERSION == 2* ]]; then pip install -r requirements-py2.txt --use-mirrors; fi
-  - if [[ $TRAVIS_PYTHON_VERSION == 3* ]]; then pip install -r requirements-py3.txt --use-mirrors; fi
-  - pip install coveralls
+  - sudo /etc/init.d/postgresql stop
+  - sudo apt-get -y remove --purge postgresql-9.1 postgresql-9.2 postgresql-9.3 postgresql-9.4
+  - sudo apt-get -y autoremove
+  - sudo apt-key adv --keyserver keys.gnupg.net --recv-keys 7FCC7D46ACCC4CF8
+  - sudo sh -c 'echo "deb http://apt.postgresql.org/pub/repos/apt/ precise-pgdg main 9.5" >> /etc/apt/sources.list.d/postgresql.list'
+  - sudo apt-get update
+  - sudo apt-get -y install postgresql-9.5
+  - sudo /etc/init.d/postgresql stop
+  - pip install -r requirements.txt
+  - curl -L https://github.com/coreos/etcd/releases/download/v${ETCDVERSION}/etcd-v${ETCDVERSION}-linux-amd64.tar.gz | tar xz -C . --strip=1 --wildcards --no-anchored etcd
+  - pip install behave codacy-coverage coverage coveralls 
 script:
-  - python setup.py test
+  - PATH=.:$PATH $TEST_SUITE
   - python setup.py flake8
 after_success:
   - coveralls
+  - if [[ -f coverage.xml ]]; then python-codacy-coverage -r coverage.xml; fi

.zappr.yml

@@ -0,0 +1,12 @@
+approvals:
+  # PR needs at least 4 approvals
+  minimum: 1
+  # approval = comment that matches this regex
+  pattern: "^:?\\+1:?$"
+  from:
+    # commenter must be either one of:
+    # a public zalando org member
+    orgs:
+      - zalando
+    # a collaborator of the repo
+    collaborators: true

Dockerfile

@@ -8,26 +8,29 @@ 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
 
-ENV PGVERSION 9.4
-RUN apt-get install python python-psycopg2 python-yaml python-requests python-boto postgresql-${PGVERSION} python-dnspython python-kazoo python-pip -y
-RUN pip install python-etcd
+ENV PGVERSION 9.5
+RUN apt-get install postgresql-${PGVERSION} postgresql-server-dev-${PGVERSION} -y
+RUN apt-get install python python-dev python-pip -y
+ADD requirements-py2.txt /requirements-py2.txt
+RUN pip install -r /requirements-py2.txt
 
 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 patronictl.py /patronictl.py
+ADD patroni/ /patroni
 
-ENV ETCDVERSION 2.0.13
+RUN ln -s /patroni.py /usr/local/bin/patroni
+RUN ln -s /patronictl.py /usr/local/bin/patronictl
+
+ENV ETCDVERSION 2.2.5
 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

MAINTAINERS

@@ -0,0 +1,3 @@
+Alexander Kukushkin <[email protected]>
+Feike Steenbergen <[email protected]>
+Oleksii Kliukin <[email protected]>

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,202 @@
+|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 customized, high-availability 
+solution using Python and — for maximum accessibility — a distributed 
+configuration store like ZooKeeper or etcd.*
+
+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 its behavior changes. Kill
+some of the 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, review this PDF:
+`postgres-ha.pdf <https://github.com/zalando/patroni/blob/master/postgres-ha.pdf>`__
+
+YAML Configuration
+------------------
+
+For an example file, see ``postgres0.yml``. Regarding settings:
+
+-  *ttl*: the TTL to acquire the leader lock. Think of it as the length of time before initiation of the automatic failover process.
+-  *loop\_wait*: the number of seconds the loop will sleep
+
+-  *restapi*:
+    -  *listen*: IP address + port that Patroni will listen to, to provide health-check information for haproxy.
+    -  *connect\_address*: IP address + port through which restapi is accessible.
+    -  *auth*: (optional) 'username:password' to protect dangerous REST API endpoints.
+    -  *certfile*: (optional) Specifies a file with the certificate in the PEM format. If the certfile is not specified or is left empty, the 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; makes it possible to 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 initiation of the automatic failover process.
+    -  *host*: the host:port for the etcd endpoint.
+
+-  *zookeeper*:
+    -  *scope*: the relative path used on etcd's HTTP API for this deployment; makes it possible to 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 initiation of the automatic failover process.
+    -  *reconnect\_timeout*: how long we should try to reconnect to ZooKeeper after a connection loss. After this timeout, assume that you no longer have a lock and restart in read-only mode.
+    -  *hosts*: list of ZooKeeper cluster members in format: ['host1:port1', 'host2:port2', 'etc...']
+    -  *exhibitor*: if you are running a ZooKeeper cluster under the Exhibitor supervisory, the following section might interest you:
+        -  *poll\_interval*: how often the 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 updates automatically whenever the 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 listens to; must be accessible from other nodes in the cluster, if you're using streaming replication. Multiple comma-separated addresses are permitted, as long as the port component is appended after to the last one with a colon, i.e. ``listen: 127.0.0.1,127.0.0.2:5432``. The first address from this list will be used by Patroni to establish local connections to the PostgreSQL node. 
+    
+    -  *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. You should comment out max_replication_slots before it becomes ineligible for leader status.
+
+    -  *initdb*:  List options to be passed on to initdb
+        -  *encoding*: default encoding for new databases
+        -  *locale*: default locale for new databases
+        -  *data-checksums*  # When pg_rewind is needed on 9.3, this needs to be enabled
+
+    -  *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 the 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 the Postgres user, set during initialization.
+
+    -  *admin*:
+        -  *username*: admin username; user is created during initialization. It will have CREATEDB and CREATEROLE privileges.
+        -  *password*: admin password; user is 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.
+
+    -  *create\_replica\_methods*: an ordered list of the create methods for turning a patroni node into a new replica.
+       "basebackup" is the default method; other methods are assumed to refer to scripts, each of which is configured
+       as its own config item.
+
+    -  *replica\_method* for each create_replica_method other than basebackup, you would add a configuration section
+       of the same name.  At a minimum, this should include "command" with a full path to the actual script to be
+       executed.  Other configuration parameters will be passed along to the script in the form "parameter=value".
+
+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 optimal 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: 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 significantly increases the variability of write 
+performance. If followers become inaccessible from the leader, the
+leader effectively becomes 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 three Postgres data nodes
+to ensure write availability if one host fails.
+
+Choosing your replication schema is dependent on your business
+considerations. 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 an 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 Patroni a "template" because it is far from being a one-size-fits-all
+or plug-and-play replication system. It will have its 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

TODO.md

@@ -1,4 +1,16 @@
 Failover
 ========
-- When determining who should become master, include the minor version of PostgreSQL in the decision
-- Create a way to disable governance of a cluster, something like the existence of a "nogover" or "admin" file in PGDATA will stop governor from changing the cluster state
+- When determining who should become master, include the minor version of PostgreSQL in the decision.
+- Create a way to disable governance of a cluster, something like the existence of a "nogover" or "admin" file in PGDATA will stop patroni from changing the cluster state.
+
+Configuration
+==============
+- Provide a way to change postgresql.conf and pg_hba.conf of a running cluster on the Patroni level, without changing individual nodes.
+- Provide hooks to store and retrieve cluster-wide passwords without exposing them in a plain-text form to unauthorized users.
+- Implement patronictl command to create initial configuration of the cluster with leader and member keys fixed to the user-supplied values in order to simplify migrations.
+- Implement support for consul in addtion to etcd and zookeeper
+- Complete zookeeper support in patronictl
+
+Documentation
+==============
+- Document how to run cascading replication and possibly initialize the cluster without an access to the master node.

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/dev_patroni_cluster.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
 
-DOCKER_IMAGE="os-registry.stups.zalan.do/acid/patroni:1.0-SNAPSHOT"
+DOCKER_IMAGE="registry.opensource.zalan.do/acid/patroni:1.0-SNAPSHOT"
 MEMBERS=3
 
 

docker/entrypoint.sh

@@ -3,25 +3,25 @@
 function usage()
 {
     cat <<__EOF__
-Usage: $0 
+Usage: $0
 
 Options:
 
     --etcd      ETCD    Provide an external etcd to connect to
-    --name      NAME    Give the cluster a specific name 
+    --name      NAME    Give the cluster a specific name
     --etcd-only         Do not run Patroni, run a standalone etcd
 
 Examples:
 
     $0 --etcd=127.17.0.84:4001
     $0 --etcd-only
-    $0 
+    $0
     $0 --name=true_scotsman
 __EOF__
 }
 
 DOCKER_IP=$(hostname --ip-address)
-PATRONI_SCOPE=batman
+PATRONI_SCOPE=${PATRONI_SCOPE:-batman}
 
 optspec=":vh-:"
 while getopts "$optspec" optchar; do
@@ -32,7 +32,7 @@ while getopts "$optspec" optchar; do
                     exec etcd --data-dir /tmp/etcd.data \
                         -advertise-client-urls=http://${DOCKER_IP}:4001 \
                         -listen-client-urls=http://0.0.0.0:4001 \
-                        -listen-peer-urls=http://0.0.0.0:2380 
+                        -listen-peer-urls=http://0.0.0.0:2380
                     exit 0
                     ;;
                 cheat)
@@ -79,20 +79,26 @@ then
     ETCD_CLUSTER="127.0.0.1:4001"
 fi
 
-cat > /patroni/postgres.yml <<__EOF__
+mkdir -p ~postgres/.config/patroni
+cat > ~postgres/.config/patroni/patronictl.yaml <<__EOF__
+{dcs_api: 'etcd://${ETCD_CLUSTER}', namespace: /service/}
+__EOF__
+
+cat > /patroni/postgres.yaml <<__EOF__
 
 ttl: &ttl 30
 loop_wait: &loop_wait 10
-scope: &scope ${PATRONI_SCOPE}
+scope: &scope '${PATRONI_SCOPE}'
+namespace: 'patroni'
 restapi:
-  listen: 127.0.0.1:8008
-  connect_address: 127.0.0.1:8008
+  listen: 0.0.0.0:8008
+  connect_address: ${DOCKER_IP}:8008
 etcd:
   scope: *scope
   ttl: *ttl
   host: ${ETCD_CLUSTER}
 postgresql:
-  name: postgresql_${DOCKER_IP//./_} ## Replication slots do not allow dots in their name
+  name: ${HOSTNAME}
   scope: *scope
   listen: 0.0.0.0:5432
   connect_address: ${DOCKER_IP}:5432
@@ -115,16 +121,18 @@ postgresql:
   parameters:
     archive_mode: "on"
     wal_level: hot_standby
-    archive_command: mkdir -p ../wal_archive && cp %p ../wal_archive/%f
+    archive_command: 'true'
     max_wal_senders: 20
     listen_addresses: 0.0.0.0
-    wal_keep_segments: 8
+    max_wal_size: 1GB
+    min_wal_size: 128MB
+    wal_keep_segments: 64
     archive_timeout: 1800s
     max_replication_slots: 20
     hot_standby: "on"
 __EOF__
 
-cat /patroni/postgres.yml
+cat /patroni/postgres.yaml
 
 if [ ! -z $CHEAT ]
 then
@@ -133,5 +141,5 @@ then
         sleep 60
     done
 else
-    exec /patroni/patroni.py /patroni/postgres.yml
+    exec python /patroni.py /patroni/postgres.yaml
 fi

extras/startup-scripts/README.md

@@ -0,0 +1,13 @@
+# 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.
+
+### patroni.service
+Systemd service file, to be copied to /etc/systemd/system/patroni.service, tested on Centos 7.1 with Patroni installed from pip. 

extras/startup-scripts/patroni.service

@@ -0,0 +1,28 @@
+# This is an example systemd config file for Patroni
+# You can copy it to "/etc/systemd/system/patroni.service",
+
+[Unit]
+Description=Runners to orchestrate a high-availability PostgreSQL
+After=syslog.target network.target
+
+[Service]
+Type=simple
+
+User=postgres
+Group=postgres
+
+# Where to send early-startup messages from the server
+# This is normally controlled by the global default set by systemd
+# StandardOutput=syslog
+
+ExecStart=/bin/patroni /etc/patroni.yml
+
+# Give a reasonable amount of time for the server to start up/shut down
+TimeoutSec=10
+
+# Do not restart the service if it crashes, we want to manually inspect database on failure
+Restart=no
+
+[Install]
+WantedBy=multi-user.target
+

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

features/basic_replication.feature

@@ -0,0 +1,17 @@
+Feature: basic replication
+  We should check that the basic bootstrapping, replication and failover works.
+
+  Scenario: check replication of a single table
+    Given I start postgres0
+    And postgres0 is a leader after 10 seconds
+    And I start postgres1
+    When I add the table foo to postgres0
+    Then table foo is present on postgres1 after 15 seconds
+
+  Scenario: check the basic failover
+    When I kill postgres0
+    Then postgres1 role is the primary after 30 seconds
+    When I start postgres0
+    Then postgres0 role is the secondary after 15 seconds
+    When I add the table bar to postgres1
+    Then table bar is present on postgres0 after 10 seconds

features/cascading_replication.feature

@@ -0,0 +1,13 @@
+Feature: cascading replication
+	We should check that patroni can do base backup and streaming from the replica
+
+Scenario: check a base backup from the replica
+	Given I start postgres0
+	And postgres0 is a leader after 10 seconds
+	And I start postgres1
+	And replication works from postgres0 to postgres1 after 15 seconds
+	And I create label with "postgres0" in postgres0 data directory
+	And I create label with "postgres1" in postgres1 data directory
+	And I configure and start postgres2 with a tag clonefrom postgres1
+	Then replication works from postgres0 to postgres2 after 30 seconds
+	And there is a label with "postgres1" in postgres2 data directory

features/environment.py

@@ -0,0 +1,305 @@
+import os
+import psycopg2
+import requests
+import shutil
+import subprocess
+import tempfile
+import time
+import yaml
+
+
+class PatroniController(object):
+    PATRONI_CONFIG = '{}.yml'
+    """ starts and stops individual patronis"""
+
+    def __init__(self):
+        self._output_dir = None
+        self._patroni_path = None
+        self._connections = {}
+        self._config = {}
+        self._connstring = {}
+        self._cursors = {}
+        self._log = {}
+        self._processes = {}
+
+    @property
+    def patroni_path(self):
+        if self._patroni_path is None:
+            cwd = os.path.realpath(__file__)
+            while True:
+                path, entry = os.path.split(cwd)
+                cwd = path
+                if entry == 'features' or cwd == '/':
+                    break
+            self._patroni_path = cwd
+        return self._patroni_path
+
+    def data_dir(self, pg_name):
+        return os.path.join(self.patroni_path, 'data', pg_name)
+
+    def write_label(self, pg_name, content):
+        with open(os.path.join(self.data_dir(pg_name), 'label'), 'w') as f:
+            f.write(content)
+
+    def read_label(self, pg_name):
+        content = None
+        try:
+            with open(os.path.join(self.data_dir(pg_name), 'label'), 'r') as f:
+                content = f.read()
+        except IOError:
+            return None
+        return content.strip()
+
+    def start(self, pg_name, max_wait_limit=20, tags=None):
+        if not self._is_running(pg_name):
+            if pg_name in self._processes:
+                del self._processes[pg_name]
+            cwd = self.patroni_path
+            self._log[pg_name] = open(os.path.join(self._output_dir, 'patroni_{0}.log'.format(pg_name)), 'a')
+
+            self._config[pg_name] = self._make_patroni_test_config(pg_name, tags=tags)
+
+            p = subprocess.Popen(['coverage', 'run', '--branch', '--source=patroni', '-p', 'patroni.py', self._config[pg_name]],
+                                 stdout=self._log[pg_name], stderr=subprocess.STDOUT, cwd=cwd)
+            if not (p and p.pid and p.poll() is None):
+                assert False, "PostgreSQL {0} is not running after being started".format(pg_name)
+            self._processes[pg_name] = p
+        # wait while patroni is available for queries, but not more than 10 seconds.
+        for _ in range(max_wait_limit):
+            if self.query(pg_name, "SELECT 1", fail_ok=True) is not None:
+                break
+            time.sleep(1)
+        else:
+            assert False,\
+                "Patroni instance is not available for queries after {0} seconds".format(max_wait_limit)
+
+    def stop(self, pg_name, kill=False, timeout=15):
+        start_time = time.time()
+        while self._is_running(pg_name):
+            if not kill:
+                self._processes[pg_name].terminate()
+            else:
+                self._processes[pg_name].kill()
+            time.sleep(1)
+            if not kill and time.time() - start_time > timeout:
+                kill = True
+        if self._log.get('pg_name') and not self._log['pg_name'].closed:
+            self._log[pg_name].close()
+        if pg_name in self._processes:
+            del self._processes[pg_name]
+
+    def query(self, pg_name, query, fail_ok=False):
+        try:
+            cursor = self._cursor(pg_name)
+            cursor.execute(query)
+            return cursor
+        except psycopg2.Error:
+            if fail_ok:
+                return None
+            else:
+                raise
+
+    def check_role_has_changed_to(self, pg_name, new_role, timeout=10):
+        bound_time = time.time() + timeout
+        recovery_status = False if new_role == 'primary' else True
+        role_has_changed = False
+        while not role_has_changed:
+            cur = self.query(pg_name, "SELECT pg_is_in_recovery()", fail_ok=True)
+            if cur:
+                row = cur.fetchone()
+                if row and len(row) > 0 and row[0] == recovery_status:
+                    role_has_changed = True
+            if time.time() > bound_time:
+                break
+            time.sleep(1)
+        return role_has_changed
+
+    def stop_all(self):
+        for patroni in self._processes.copy():
+            self.stop(patroni)
+
+    def create_and_set_output_directory(self, feature_name):
+        feature_dir = os.path.join(self.patroni_path, "features", "output",
+                                   feature_name.replace(' ', '_'))
+        if os.path.exists(feature_dir):
+            shutil.rmtree(feature_dir)
+        os.makedirs(feature_dir)
+        self._output_dir = feature_dir
+
+    def _is_running(self, pg_name):
+        return pg_name in self._processes and self._processes[pg_name].pid and (self._processes[pg_name].poll() is None)
+
+    def _make_patroni_test_config(self, pg_name, tags=None):
+        patroni_config_name = PatroniController.PATRONI_CONFIG.format(pg_name)
+        patroni_config_path = os.path.join(self._output_dir, patroni_config_name)
+
+        with open(patroni_config_name) as f:
+            config = yaml.load(f)
+        postgresql = config['postgresql']
+        postgresql['name'] = pg_name
+        postgresql['data_dir'] = 'data/{0}'.format(pg_name)
+        postgresql_params = postgresql['parameters']
+        postgresql_params['logging_collector'] = 'on'
+        postgresql_params['log_destination'] = 'csvlog'
+        postgresql_params['log_directory'] = self._output_dir
+        postgresql_params['log_filename'] = '{0}.log'.format(pg_name)
+        postgresql_params['log_statement'] = 'all'
+        postgresql_params['log_min_messages'] = 'debug1'
+        postgresql_params['unix_socket_directories'] = '.'
+
+        if tags:
+            config['tags'] = tags
+
+        with open(patroni_config_path, 'w') as f:
+            yaml.dump(config, f, default_flow_style=False)
+
+        return patroni_config_path
+
+    def _make_connstring(self, pg_name):
+        if pg_name in self._connstring:
+            return self._connstring[pg_name]
+        try:
+            patroni_path = self.patroni_path
+            with open(os.path.join(patroni_path, PatroniController.PATRONI_CONFIG.format(pg_name)), 'r') as f:
+                config = yaml.load(f)
+        except IOError:
+            return None
+        connstring = config['postgresql']['connect_address']
+        if ':' in connstring:
+            address, port = connstring.split(':')
+        else:
+            address = connstring
+            port = '5432'
+        user = "postgres"
+        dbname = "postgres"
+        self._connstring[pg_name] = "host={0} port={1} dbname={2} user={3}".format(address, port, dbname, user)
+        return self._connstring[pg_name]
+
+    def _connection(self, pg_name):
+        if pg_name not in self._connections or self._connections[pg_name].closed:
+            conn = psycopg2.connect(self._make_connstring(pg_name))
+            conn.autocommit = True
+            self._connections[pg_name] = conn
+        return self._connections[pg_name]
+
+    def _cursor(self, pg_name):
+        if pg_name not in self._cursors or self._cursors[pg_name].closed:
+            cursor = self._connection(pg_name).cursor()
+            self._cursors[pg_name] = cursor
+        return self._cursors[pg_name]
+
+
+class EtcdController(object):
+
+    """ handles all etcd related tasks, used for the tests setup and cleanup """
+    ETCD_VERSION_URL = 'http://127.0.0.1:2379/version'
+    ETCD_CLEANUP_URL = 'http://127.0.0.1:2379/v2/keys/service/batman?recursive=true'
+
+    def __init__(self, log_directory):
+        self.handle = None
+        self.work_directory = None
+        self.log_directory = log_directory
+        self.log_file = None
+        self.pid = None
+        self.start_timeout = 5
+
+    def start(self):
+        """ start etcd if it's not already running """
+        if self._is_running():
+            return True
+        self.work_directory = tempfile.mkdtemp()
+        # etcd is running throughout the tests, no need to append to the log
+        output_dir = os.path.join(self.log_directory, "features", "output")
+        if not os.path.exists(output_dir):
+            os.makedirs(output_dir)
+        self.log_file = open(os.path.join(output_dir, 'etcd.log'), 'w')
+        self.handle =\
+            subprocess.Popen(["etcd", "--debug", "--data-dir", self.work_directory],
+                             stdout=self.log_file, stderr=subprocess.STDOUT)
+        start_time = time.time()
+        while (not self._is_running()):
+            if time.time() - start_time > self.start_timeout:
+                assert False, "Failed to start etcd"
+            time.sleep(1)
+        return True
+
+    @staticmethod
+    def query(key):
+        """ query etcd for a value of a given key """
+        r = requests.get("http://127.0.0.1:2379/v2/keys/service/batman/{0}".format(key))
+        if r.ok:
+            content = r.json()
+            if content:
+                return content.get('node', {}).get('value')
+        return None
+
+    def stop_and_remove_work_directory(self, timeout=15):
+        """ terminate etcd and wipe out the temp work directory, but only if we actually started it"""
+        kill = False
+        start_time = time.time()
+        while self._is_running() and self.handle:
+            if not kill:
+                self.handle.terminate()
+            else:
+                self.handle.kill()
+            time.sleep(1)
+            if not kill and time.time() - start_time > timeout:
+                kill = True
+        self.handle = None
+        if self.log_file and not self.log_file.closed:
+            self.log_file.close()
+        if self.work_directory:
+            shutil.rmtree(self.work_directory)
+            self.work_directory = None
+
+    @staticmethod
+    def cleanup_service_tree():
+        """ clean all contents stored in the tree used for the tests """
+        r = None
+        try:
+            r = requests.delete(EtcdController.ETCD_CLEANUP_URL)
+            if r and not r.ok:
+                assert False,\
+                    "request to cleanup the etcd contents was not successfull: status code {0}".format(r.status_code)
+        except requests.exceptions.RequestException as e:
+            assert False, "exception when cleaning up etcd contents: {0}".format(e)
+
+    @staticmethod
+    def _is_running():
+        # if etcd is running, but we didn't start it
+        try:
+            r = requests.get(EtcdController.ETCD_VERSION_URL)
+            running = (r and r.ok and b'etcdserver' in r.content)
+        except requests.ConnectionError:
+            running = False
+        return running
+
+
+# actions to execute on start/stop of the tests and before running invidual features
+def before_all(context):
+    context.pctl = PatroniController()
+    context.etcd_ctl = EtcdController(context.pctl.patroni_path)
+    context.etcd_ctl.start()
+    try:
+        context.etcd_ctl.cleanup_service_tree()
+    except AssertionError:  # after.all handlers won't be executed in before.all
+        context.etcd_ctl.stop_and_remove_work_directory()
+        raise
+
+
+def after_all(context):
+    context.etcd_ctl.stop_and_remove_work_directory()
+    subprocess.call(['coverage', 'combine'])
+    subprocess.call(['coverage', 'report'])
+
+
+def before_feature(context, feature):
+    """ create per-feature output directory to collect Patroni and PostgreSQL logs """
+    context.pctl.create_and_set_output_directory(feature.name)
+
+
+def after_feature(context, feature):
+    """ stop all Patronis, remove their data directory and cleanup the keys in etcd """
+    context.pctl.stop_all()
+    shutil.rmtree(os.path.join(context.pctl.patroni_path, 'data'))
+    context.etcd_ctl.cleanup_service_tree()

features/patroni_api.feature

@@ -0,0 +1,48 @@
+Feature: patroni api
+	We should check that patroni correctly responds to valid and not-valid API requests.
+
+Scenario: check API requests on a stand-alone server
+	Given I start postgres0
+	And postgres0 is a leader after 10 seconds
+	When I issue a GET request to http://127.0.0.1:8008/
+	Then I receive a response code 200
+	And I receive a response state running
+	And I receive a response role master
+	When I issue a GET request to http://127.0.0.1:8008/replica
+	Then I receive a response code 503
+	When I issue an empty POST request to http://127.0.0.1:8008/reinitialize
+	Then I receive a response code 503
+	And I receive a response text "I am the leader, can not reinitialize"
+	When I issue a POST request to http://127.0.0.1:8008/failover with leader=postgres0
+	Then I receive a response code 500
+	And I receive a response text "failover is not possible: cluster does not have members except leader"
+	When I issue an empty POST request to http://127.0.0.1:8008/failover
+	Then I receive a response code 400
+	And I receive a response text "No values given for required parameters leader and member"
+
+Scenario: check API requests for the primary-replica pair
+	Given I start postgres1
+	And replication works from postgres0 to postgres1 after 15 seconds
+	When I issue a GET request to http://127.0.0.1:8009/replica
+	Then I receive a response code 200
+	And I receive a response state running
+	And I receive a response role replica
+	When I issue an empty POST request to http://127.0.0.1:8009/reinitialize
+	Then I receive a response code 200
+	Given replication works from postgres0 to postgres1 after 10 seconds
+	When I issue an empty POST request to http://127.0.0.1:8008/restart
+	Then I receive a response code 200
+	And postgres0 is a leader after 5 seconds
+
+Scenario: check the failover via the API
+	Given I issue a POST request to http://127.0.0.1:8008/failover with leader=postgres0,candidate=postgres1
+	Then I receive a response code 200
+	And postgres1 is a leader after 5 seconds
+	And replication works from postgres1 to postgres0 after 15 seconds
+
+Scenario: check the scheduled failover
+	Given I issue a scheduled failover at http://127.0.0.1:8009 from postgres1 to postgresq0 in 10 seconds
+	Then I receive a response code 200
+	And postgres0 is a leader after 15 seconds
+	And replication works from postgres0 to postgres1 after 25 seconds
+

features/steps/basic_replication.py

@@ -0,0 +1,55 @@
+import psycopg2 as pg
+
+from behave import step, then
+from time import sleep, time
+
+
+@step('I start {name:w}')
+def start_patroni(context, name):
+    return context.pctl.start(name)
+
+
+@step('I shut down {name:w}')
+def stop_patroni(context, name):
+    return context.pctl.stop(name)
+
+
+@step('I kill {name:w}')
+def kill_patroni(context, name):
+    return context.pctl.stop(name, kill=True)
+
+
+@step('I add the table {table_name:w} to {pg_name:w}')
+def add_table(context, table_name, pg_name):
+    # parse the configuration file and get the port
+    try:
+        context.pctl.query(pg_name, "CREATE TABLE {0}()".format(table_name))
+    except pg.Error as e:
+        assert False, "Error creating table {0} on {1}: {2}".format(table_name, pg_name, e)
+
+
+@then('Table {table_name:w} is present on {pg_name:w} after {max_replication_delay:d} seconds')
+def table_is_present_on(context, table_name, pg_name, max_replication_delay):
+    for _ in range(int(max_replication_delay)):
+        if context.pctl.query(pg_name, "SELECT 1 FROM {0}".format(table_name), fail_ok=True) is not None:
+            break
+        sleep(1)
+    else:
+        assert False,\
+            "Table {0} is not present on {1} after {2} seconds".format(table_name, pg_name, max_replication_delay)
+
+
+@then('{pg_name:w} role is the {pg_role:w} after {max_promotion_timeout:d} seconds')
+def check_role(context, pg_name, pg_role, max_promotion_timeout):
+    if not context.pctl.check_role_has_changed_to(pg_name, pg_role, timeout=int(max_promotion_timeout)):
+        assert False,\
+            "{0} role didn't change to {1} after {2} seconds".format(pg_name, pg_role, max_promotion_timeout)
+
+
+@step('replication works from {master:w} to {replica:w} after {time_limit:d} seconds')
+@then('replication works from {master:w} to {replica:w} after {time_limit:d} seconds')
+def replication_works(context, master, replica, time_limit):
+    context.execute_steps(u"""
+        When I add the table test_{0} to {1}
+        Then table test_{0} is present on {2} after {3} seconds
+    """.format(int(time()), master, replica, time_limit))

features/steps/cascading_replication.py

@@ -0,0 +1,17 @@
+from behave import step, then
+
+
+@step('I configure and start {name:w} with a tag {tag_name:w} {tag_value:w}')
+def start_patroni_with_a_name_value_tag(context, name, tag_name, tag_value):
+    return context.pctl.start(name, tags={tag_name: tag_value})
+
+
+@then('There is a label with "{content:w}" in {name:w} data directory')
+def check_label(context, content, name):
+    label = context.pctl.read_label(name)
+    assert label == content, "{0} is not equal to {1}".format(label, content)
+
+
+@step('I create label with "{content:w}" in {name:w} data directory')
+def write_label(context, content, name):
+    context.pctl.write_label(name, content)

features/steps/patroni_api.py

@@ -0,0 +1,101 @@
+import parse
+import pytz
+import requests
+import time
+
+from behave import register_type, step, then
+from datetime import datetime, timedelta
+
+
+@parse.with_pattern(r'https?://(?:\w|\.|:|/)+')
+def parse_url(text):
+    return text
+
+
+@parse.with_pattern(r'(?:\w+=(?:\w|\.|:|-|\+|\s)+,?)+')
+def parse_data(text):
+    return text
+
+
+register_type(url=parse_url, data=parse_data)
+
+
+# there is no way we can find out if the node has already
+# started as a leader without checking the DCS. We cannot
+# just rely on the database availability, since there is
+# a short gap between the time PostgreSQL becomes available
+# and Patroni assuming the leader role.
+@step('{name:w} is a leader after {time_limit:d} seconds')
+@then('{name:w} is a leader after {time_limit:d} seconds')
+def is_a_leader(context, name, time_limit):
+    max_time = time.time() + int(time_limit)
+    while (context.etcd_ctl.query("leader") != name):
+        time.sleep(1)
+        if time.time() > max_time:
+            assert False, "{0} is not a leader in etcd after {1} seconds".format(name, time_limit)
+
+
+@step('I sleep for {value:d} seconds')
+def sleep_for_n_seconds(context, value):
+    time.sleep(int(value))
+
+
+@step('I issue a GET request to {url:url}')
+def do_get(context, url):
+    try:
+        r = requests.get(url)
+    except requests.exceptions.RequestException:
+        context.status_code = None
+        context.response = None
+    else:
+        context.status_code = r.status_code
+        try:
+            context.response = r.json()
+        except ValueError:
+            context.response = r.content.decode('utf-8')
+
+
+@step('I issue an empty POST request to {url:url}')
+def do_post_empty(context, url):
+    do_post(context, url, None)
+
+
+@step('I issue a POST request to {url:url} with {data:data}')
+def do_post(context, url, data):
+    post_data = {}
+    if data:
+        post_components = data.split(',')
+        for pc in post_components:
+            if '=' in pc:
+                k, v = pc.split('=', 2)
+                post_data[k.strip()] = v.strip()
+    try:
+        r = requests.post(url, json=post_data)
+    except requests.exceptions.RequestException:
+        context.status_code = None
+        context.response = None
+    else:
+        context.status_code = r.status_code
+        try:
+            context.response = r.json()
+        except ValueError:
+            context.response = r.content.decode('utf-8')
+
+
+@then('I receive a response {component:w} {data}')
+def check_response(context, component, data):
+    if component == 'code':
+        assert context.status_code == int(data),\
+            "status code {0} != {1}, response: {2}".format(context.status_code, int(data), context.response)
+    elif component == 'text':
+        assert context.response == data.strip('"'), "response {0} does not contain {1}".format(context.response, data)
+    else:
+        assert component in context.response, "{0} is not part of the response".format(component)
+        assert context.response[component] == data, "{0} does not contain {1}".format(component, data)
+
+
+@step('I issue a scheduled failover at {at_url:url} from {from_host:w} to {to_host:w} in {in_seconds:d} seconds')
+def scheduled_failover(context, at_url, from_host, to_host, in_seconds):
+    context.execute_steps(u"""
+        Given I issue a POST request to {0}/failover with leader={1},candidate={2},scheduled_at={3}
+    """.format(at_url, from_host, to_host, datetime.now(pytz.utc) + timedelta(seconds=int(in_seconds))))

haproxy.cfg

@@ -15,7 +15,7 @@ frontend ft_postgresql
 	default_backend bk_db
 
 backend bk_db
-	option httpchk GET
+	option httpchk
 
-  server postgresql_127.0.0.1_5432 127.0.0.1:5432 maxconn 100 check port 8008
-  server postgresql_127.0.0.1_5433 127.0.0.1:5433 maxconn 100 check port 8009
+    server postgresql_127.0.0.1_5432 127.0.0.1:5432 maxconn 100 check port 8008
+    server postgresql_127.0.0.1_5433 127.0.0.1:5433 maxconn 100 check port 8009

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/etcd.py

@@ -1,214 +0,0 @@
-from __future__ import absolute_import
-import etcd
-import logging
-import os
-import random
-import requests
-import socket
-
-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 requests.exceptions import RequestException
-
-logger = logging.getLogger(__name__)
-
-
-class EtcdError(DCSError):
-    pass
-
-
-class Client(etcd.Client):
-
-    def __init__(self, config):
-        super(Client, self).__init__(read_timeout=5)
-        self._config = config
-        self._load_machines_cache()
-        self._allow_reconnect = True
-
-    @property
-    def machines(self):
-        """Original `machines` method(property) of `etcd.Client` class raise exception
-        when it failed to get list of etcd cluster members. This method is being called
-        only when request failed on one of the etcd members during `api_execute` call.
-        For us it's more important to execute original request rather then get new
-        topology of etcd cluster. So we will catch this exception and return valid list
-        of machines with setting flag `self._update_machines_cache` to `!True`.
-        Later, during next `api_execute` call we will forcefully update machines_cache"""
-        try:
-            ret = super(Client, self).machines
-            random.shuffle(ret)
-            return ret
-        except etcd.EtcdException:
-            if self._update_machines_cache:  # We are updating machines_cache
-                raise  # This exception is fatal, we should re-raise it.
-            self._update_machines_cache = True
-            return [self._base_uri]
-
-    def api_execute(self, path, method, **kwargs):
-        # Update machines_cache if previous attempt of update has failed
-        self._update_machines_cache and self._load_machines_cache()
-        return super(Client, self).api_execute(path, method, **kwargs)
-
-    @staticmethod
-    def get_srv_record(host):
-        try:
-            return [(str(r.target).rstrip('.'), r.port) for r in resolver.query('_etcd-server._tcp.' + host, 'SRV')]
-        except DNSException:
-            logger.exception('Can not resolve SRV for %s', host)
-        return []
-
-    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
-        'GET http://{host}:{port}/members' request (peer protocol)"""
-
-        ret = []
-        for host, port in self.get_srv_record(discovery_srv):
-            url = '{}://{}:{}/members'.format(self._protocol, host, port)
-            try:
-                response = requests.get(url)
-                if response.ok:
-                    for member in response.json():
-                        ret.extend(member['clientURLs'])
-                    break
-            except RequestException:
-                logger.exception('GET %s', url)
-        return list(set(ret))
-
-    def _get_machines_cache_from_dns(self, addr):
-        """One host might be resolved into multiple ip addresses. We will make list out of it"""
-
-        ret = []
-        host, port = addr.split(':')
-        try:
-            for r in set(socket.getaddrinfo(host, port, socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP)):
-                ret.append('{}://{}:{}'.format(self._protocol, r[4][0], r[4][1]))
-        except socket.error:
-            logger.exception('Can not resolve %s', host)
-        return list(set(ret)) if ret else ['{}://{}:{}'.format(self._protocol, host, port)]
-
-    def _load_machines_cache(self):
-        """This method should fill up `_machines_cache` from scratch.
-        It could happen only in two cases:
-        1. During class initialization
-        2. When all etcd members failed"""
-
-        self._update_machines_cache = True
-
-        if 'discovery_srv' not in self._config and 'host' not in self._config:
-            raise Exception('Neither discovery_srv nor host are defined in etcd section of config')
-
-        self._machines_cache = []
-
-        if 'discovery_srv' in self._config:
-            self._machines_cache = self._get_machines_cache_from_srv(self._config['discovery_srv'])
-
-        if not self._machines_cache and 'host' in self._config:
-            self._machines_cache = self._get_machines_cache_from_dns(self._config['host'])
-
-        # Can not bootstrap list of etcd-cluster members, giving up
-        if not self._machines_cache:
-            raise etcd.EtcdException
-
-        # After filling up initial list of machines_cache we should ask etcd-cluster about actual list
-        self._base_uri = self._machines_cache.pop(0)
-        self._machines_cache = self.machines
-        self._base_uri in self._machines_cache and self._machines_cache.remove(self._base_uri)
-
-        self._update_machines_cache = False
-
-
-def catch_etcd_errors(func):
-    def wrapper(*args, **kwargs):
-        try:
-            return not func(*args, **kwargs) is None
-        except etcd.EtcdException:
-            return False
-    return wrapper
-
-
-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.client = self.get_etcd_client(config)
-
-    def get_etcd_client(self, config):
-        client = None
-        while not client:
-            try:
-                client = Client(config)
-            except etcd.EtcdException:
-                logger.info('waiting on etcd')
-                sleep(5)
-        return client
-
-    @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)
-
-    def get_cluster(self):
-        try:
-            result = 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))
-
-            # get last leader operation
-            last_leader_operation = nodes.get('optime/leader', 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]
-
-            # get leader
-            leader = nodes.get('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]
-
-            return Cluster(initialize, leader, last_leader_operation, members)
-        except etcd.EtcdKeyNotFound:
-            return Cluster(False, None, None, [])
-        except:
-            logger.exception('get_cluster')
-
-        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)
-
-    @catch_etcd_errors
-    def take_leader(self):
-        return self.client.set(self.client_path('/leader'), 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
-
-    @catch_etcd_errors
-    def write_leader_optime(self, state_handler):
-        return self.client.set(self.client_path('/optime/leader'), state_handler.last_operation())
-
-    @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
-
-    @catch_etcd_errors
-    def race(self, path):
-        return self.client.write(self.client_path(path), self._name, prevExist=False)
-
-    @catch_etcd_errors
-    def delete_leader(self):
-        return self.client.delete(self.client_path('/leader'), prevValue=self._name)

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

helpers/zookeeper.py

@@ -1,229 +0,0 @@
-import logging
-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 requests.exceptions import RequestException
-
-logger = logging.getLogger(__name__)
-
-
-class ZooKeeperError(DCSError):
-    pass
-
-
-class ExhibitorEnsembleProvider:
-
-    TIMEOUT = 3.1
-
-    def __init__(self, hosts, port, uri_path='/exhibitor/v1/cluster/list', poll_interval=300):
-        self._exhibitor_port = port
-        self._uri_path = uri_path
-        self._poll_interval = poll_interval
-        self._exhibitors = hosts
-        self._master_exhibitors = hosts
-        self._zookeeper_hosts = ''
-        self._next_poll = None
-        while not self.poll():
-            logger.info('waiting on exhibitor')
-            sleep(5)
-
-    def poll(self):
-        if self._next_poll and self._next_poll > time.time():
-            return False
-
-        json = self._query_exhibitors(self._exhibitors)
-        if not json:
-            json = self._query_exhibitors(self._master_exhibitors)
-
-        if isinstance(json, dict) and 'servers' in json and 'port' in json:
-            self._next_poll = time.time() + self._poll_interval
-            zookeeper_hosts = ','.join([h + ':' + str(json['port']) for h in sorted(json['servers'])])
-            if self._zookeeper_hosts != zookeeper_hosts:
-                logger.info('ZooKeeper connection string has changed: %s => %s', self._zookeeper_hosts, zookeeper_hosts)
-                self._zookeeper_hosts = zookeeper_hosts
-                self._exhibitors = json['servers']
-                return True
-        return False
-
-    def _query_exhibitors(self, exhibitors):
-        random.shuffle(exhibitors)
-        for host in exhibitors:
-            uri = 'http://{}:{}{}'.format(host, self._exhibitor_port, self._uri_path)
-            try:
-                response = requests.get(uri, timeout=self.TIMEOUT)
-                return response.json()
-            except RequestException:
-                pass
-        return None
-
-    @property
-    def zookeeper_hosts(self):
-        return self._zookeeper_hosts
-
-
-class ZooKeeper(AbstractDCS):
-
-    def __init__(self, name, config):
-        super(ZooKeeper, self).__init__(name, config)
-
-        hosts = config.get('hosts', [])
-        if isinstance(hosts, list):
-            hosts = ','.join(hosts)
-
-        self.exhibitor = None
-        if 'exhibitor' in config:
-            exhibitor = config['exhibitor']
-            interval = exhibitor.get('poll_interval', 300)
-            self.exhibitor = ExhibitorEnsembleProvider(exhibitor['hosts'], exhibitor['port'], poll_interval=interval)
-            hosts = self.exhibitor.zookeeper_hosts
-
-        self.client = KazooClient(hosts=hosts,
-                                  timeout=(config.get('session_timeout', None) or 30),
-                                  command_retry={
-                                      'deadline': (config.get('reconnect_timeout', None) or 10),
-                                      'max_delay': 1,
-                                      '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.fetch_cluster = True
-        self.members = []
-        self.leader = None
-        self.last_leader_operation = 0
-
-        self.client.start(None)
-
-    def session_listener(self, state):
-        if state in [KazooState.SUSPENDED, KazooState.LOST]:
-            self.cluster_watcher(None)
-
-    def cluster_watcher(self, event):
-        self.fetch_cluster = True
-        self.cluster_event.set()
-
-    def get_node(self, name, watch=None):
-        try:
-            return self.client.get(self.client_path(name), watch)
-        except NoNodeError:
-            pass
-        except:
-            logger.exception('get_node')
-        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)
-
-    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)
-            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()
-        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
-
-            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])
-
-    def get_cluster(self):
-        if self.exhibitor and self.exhibitor.poll():
-            self.client.set_hosts(self.exhibitor.zookeeper_hosts)
-
-        if self.fetch_cluster:
-            try:
-                self.client.retry(self._inner_load_cluster)
-            except:
-                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)
-            return True
-        except:
-            return False
-
-    def attempt_to_acquire_leader(self):
-        ret = self._create('/leader', 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)
-        try:
-            self.client.retry(self.client.create, path, connection_string, makepath=True, ephemeral=True)
-            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)
-                return True
-            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()
-        if last_operation != self.last_leader_operation:
-            self.last_leader_operation = last_operation
-            path = self.client_path('/optime/leader')
-            try:
-                self.client.retry(self.client.set, path, last_operation)
-            except NoNodeError:
-                try:
-                    self.client.retry(self.client.create, path, last_operation, makepath=True)
-                except:
-                    logger.exception('Failed to create %s', path)
-            except:
-                logger.exception('Failed to update %s', path)
-        return True
-
-    def delete_leader(self):
-        if isinstance(self.leader, Member) and self.leader.name == self._name:
-            self.client.delete(self.client_path('/leader'))
-
-    def sleep(self, timeout):
-        self.cluster_event.wait(timeout)
-        if self.cluster_event.isSet():
-            self.fetch_cluster = True

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,89 @@
+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
+from .version import __version__
+
+logger = logging.getLogger(__name__)
+
+
+class Patroni(object):
+
+    def __init__(self, config):
+        self.nap_time = config['loop_wait']
+        self.tags = config.get('tags', dict())
+        self.postgresql = Postgresql(config['postgresql'])
+        self.dcs = self.get_dcs(self.postgresql.name, config)
+        self.version = __version__
+        self.api = RestApiServer(self, config['restapi'])
+        self.ha = Ha(self)
+        self.next_run = time.time()
+
+    @property
+    def nofailover(self):
+        return self.tags.get('nofailover', False)
+
+    @property
+    def replicatefrom(self):
+        return self.tags.get('replicatefrom')
+
+    @property
+    def clonefrom(self):
+        return self.tags.get('clonefrom')
+
+    @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 suitable 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: {0} 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,342 @@
+import base64
+import fcntl
+import json
+import logging
+import psycopg2
+import socket
+import time
+import dateutil
+import datetime
+import pytz
+
+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 finish(self, *args, **kwargs):
+        try:
+            if not self.wfile.closed:
+                self.wfile.flush()
+                self.wfile.close()
+        except socket.error:
+            pass
+        self.rfile.close()
+
+    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_OPTIONS(self):
+        self.do_GET(options=True)
+
+    def do_GET(self, options=False):
+        """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()
+        response.update(self.get_tags())
+
+        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)
+        if not options:
+            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)
+        response.update(self.get_tags())
+        response['patroni'] = {'version': self.server.patroni.version, 'scope': self.server.patroni.postgresql.scope}
+
+        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 = 500
+        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 poll_failover_result(self, leader, member):
+        for _ in range(0, 15):
+            time.sleep(1)
+            try:
+                cluster = self.server.patroni.dcs.get_cluster()
+                if cluster.leader and cluster.leader.name != leader:
+                    return 200, ('Successfully failed over to ' + cluster.leader.name).encode('utf-8')
+                if not cluster.failover:
+                    return 503, b'Failover failed'
+            except:
+                pass
+        return 503, b'Failover status unknown'
+
+    def is_failover_possible(self, cluster, leader, member):
+        if leader and not cluster.leader or cluster.leader.name != leader:
+            return b'leader name does not match'
+        if member:
+            members = [m for m in cluster.members if m.name == member]
+            if not members:
+                return b'member does not exists'
+        else:
+            members = [m for m in cluster.members if m.name != cluster.leader.name and m.api_url]
+            if not members:
+                return b'failover is not possible: cluster does not have members except leader'
+        for member, reachable, _, xlog_location, tags in self.server.patroni.ha.fetch_nodes_statuses(members):
+            if reachable and not tags.get('nofailover', False):
+                return None
+        return b'failover is not possible: no good candidates have been found'
+
+    @check_auth
+    def do_POST_failover(self):
+        content_length = int(self.headers.get('content-length', 0))
+        try:
+            request = json.loads(self.rfile.read(content_length).decode('utf-8'))
+        except ValueError:
+            request = {}
+        leader = request.get('leader')
+        member = request.get('member')
+        cluster = self.server.patroni.ha.dcs.get_cluster()
+        status_code = 500
+
+        logger.info("received failover request with leader {0} member {1} scheduled_at {2}".
+                    format(leader, member, request.get("scheduled_at")))
+
+        data = b''
+        if leader or member:
+            if request.get('scheduled_at'):
+                try:
+                    scheduled_at = dateutil.parser.parse(request['scheduled_at'])
+                    if scheduled_at.tzinfo is None:
+                        data = b'Timezone information is mandatory for scheduled_at'
+                        status_code = 400
+                    elif scheduled_at < datetime.datetime.now(pytz.utc):
+                        data = b'Cannot schedule failover in the past'
+                        status_code = 422
+                    elif self.server.patroni.dcs.manual_failover(leader, member, scheduled_at):
+                        data = b'Failover scheduled'
+                        status_code = 200
+                except (ValueError, TypeError):
+                    logger.exception('Invalid scheduled failover time: {}'.format(request['scheduled_at']))
+                    data = b'Unable to parse scheduled timestamp. It should be in an unambiguous format, e.g. ISO 8601'
+                    status_code = 422
+            else:
+                data = self.is_failover_possible(cluster, leader, member)
+                if not data:
+                    if not self.server.patroni.dcs.manual_failover(leader, member):
+                        data = b'failed to write failover key into DCS'
+                        status_code = 503
+                    else:
+                        self.server.patroni.dcs.event.set()
+                        status_code, data = self.poll_failover_result(cluster.leader and cluster.leader.name, member)
+        else:
+            status_code = 400
+            data = b'No values given for required parameters leader and member'
+
+        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 handle_one_request(self):
+        try:
+            BaseHTTPRequestHandler.handle_one_request(self)
+        except socket.error:
+            pass
+
+    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,
+                                       to_char(pg_last_xact_replay_timestamp(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'),
+                                       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',
+                'server_version': self.server.patroni.postgresql.server_version,
+                'xlog': ({
+                    'received_location': row[3],
+                    'replayed_location': row[4],
+                    'replayed_timestamp': row[5],
+                    'paused': row[6]} if row[1] else {
+                    'location': row[2]
+                })
+            }
+        except (psycopg2.Error, RetryFailedError, PostgresConnectionException):
+            state = self.server.patroni.postgresql.state
+            if state == 'running':
+                logger.exception('get_postgresql_status')
+                state = 'unknown'
+            return {'state': state}
+
+    def get_tags(self):
+        return {'tags': self.server.patroni.tags}
+
+    def log_message(self, fmt, *args):
+        logger.debug("API thread: %s - - [%s] %s", self.client_address[0], self.log_date_time_string(), fmt % args)
+
+
+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'):
+            import ssl
+            self.socket = ssl.wrap_socket(self.socket, server_side=True, **options)
+            protocol = 'https'
+
+        self.connection_string = '{0}://{1}/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,54 @@
+import logging
+from threading import Lock, Thread
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncExecutor(object):
+
+    def __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, *args):
+        self._thread_lock.release()

patroni/ctl.py

@@ -0,0 +1,653 @@
+'''
+Patroni Control
+'''
+
+import click
+import os
+import yaml
+import json
+import time
+import psycopg2
+import random
+import requests
+import datetime
+from prettytable import PrettyTable
+from six.moves.urllib_parse import urlparse
+import logging
+import dateutil
+import tzlocal
+
+from .etcd import Etcd
+from .zookeeper import ZooKeeper
+from .exceptions import PatroniCtlException
+from .postgresql import parseurl
+
+CONFIG_DIR_PATH = click.get_app_dir('patroni')
+CONFIG_FILE_PATH = os.path.join(CONFIG_DIR_PATH, 'patronictl.yaml')
+LOGLEVEL = 'WARNING'
+
+
+def parse_dcs(dcs):
+    """
+    Break up the provided dcs string
+    >>> parse_dcs('localhost') == {'scheme': 'etcd', 'hostname': 'localhost', 'port': 4001}
+    True
+    >>> parse_dcs('localhost:8500') == {'scheme': 'consul', 'hostname': 'localhost', 'port': 8500}
+    True
+    >>> parse_dcs('zookeeper://localhost') == {'scheme': 'zookeeper', 'hostname': 'localhost', 'port': 2181}
+    True
+    """
+
+    if not dcs:
+        return {}
+
+    parsed = urlparse(dcs)
+    scheme = parsed.scheme
+    if scheme == '' and parsed.netloc == '':
+        parsed = urlparse('//' + dcs)
+
+    if scheme == '':
+        default_schemes = {'2181': 'zookeeper', '8181': 'exhibitor', '8500': 'consul'}
+        scheme = default_schemes.get(str(parsed.port), 'etcd')
+
+    port = parsed.port
+    if port is None:
+        default_ports = {'consul': 8500, 'zookeeper': 2181, 'exhibitor': 8181}
+        port = default_ports.get(str(scheme), 4001)
+
+    return {'scheme': str(scheme), 'hostname': str(parsed.hostname), 'port': int(port)}
+
+
+def load_config(path, dcs):
+    logging.debug('Loading configuration from file %s', path)
+    config = dict()
+    try:
+        with open(path, 'rb') as fd:
+            config = yaml.safe_load(fd)
+    except (IOError, yaml.YAMLError):
+        logging.exception('Could not load configuration file')
+
+    if dcs:
+        config['dcs'] = parse_dcs(dcs)
+    else:
+        config['dcs'] = parse_dcs(config.get('dcs_api'))
+
+    return config
+
+
+def store_config(config, path):
+    dir_path = os.path.dirname(path)
+    if dir_path and not os.path.isdir(dir_path):
+        os.makedirs(dir_path)
+    with open(path, 'w') as fd:
+        yaml.dump(config, fd)
+
+
+option_config_file = click.option('--config-file', '-c', help='Configuration file', default=CONFIG_FILE_PATH)
+option_format = click.option('--format', '-f', 'fmt', help='Output format (pretty, json)', default='pretty')
+option_dcs = click.option('--dcs', '-d', help='Use this DCS', envvar='DCS')
+option_watchrefresh = click.option('-w', '--watch', type=float, help='Auto update the screen every X seconds')
+option_watch = click.option('-W', is_flag=True, help='Auto update the screen every 2 seconds')
+option_force = click.option('--force', is_flag=True, help='Do not ask for confirmation at any point')
+
+
+@click.group()
+@click.pass_context
+def ctl(ctx):
+    global LOGLEVEL
+    LOGLEVEL = os.environ.get('LOGLEVEL', LOGLEVEL)
+
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s', level=LOGLEVEL)
+
+
+def get_dcs(config, scope):
+    scheme, hostname, port = map(config.get('dcs', {}).get, ('scheme', 'hostname', 'port'))
+
+    if scheme == 'etcd':
+        return Etcd(name=scope, config={'scope': scope, 'host': '{0}:{1}'.format(hostname, port)})
+
+    if scheme == 'zookeeper':
+        return ZooKeeper(name=scope, config={'scope': scope, 'hosts': [hostname], 'port': port})
+
+    if scheme == 'exhibitor':
+        return ZooKeeper(name=scope, config={'scope': scope, 'exhibitor': {'hosts': [hostname], 'port': port}})
+
+    raise PatroniCtlException('Can not find suitable configuration of distributed configuration store')
+
+
+def post_patroni(member, endpoint, content, headers=None):
+    url = urlparse(member.api_url)
+    logging.debug(url)
+    return requests.post('{0}://{1}/{2}'.format(url.scheme, url.netloc, endpoint),
+                         headers=headers or {'Content-Type': 'application/json'},
+                         data=json.dumps(content), timeout=60)
+
+
+def print_output(columns, rows=None, alignment=None, fmt='pretty', header=True, delimiter='\t'):
+    rows = rows or []
+    if fmt == 'pretty':
+        t = PrettyTable(columns)
+        for k, v in (alignment or {}).items():
+            t.align[k] = v
+        for r in rows:
+            t.add_row(r)
+        click.echo(t)
+        return
+
+    if fmt == 'json':
+        elements = list()
+        for r in rows:
+            elements.append(dict(zip(columns, r)))
+
+        click.echo(json.dumps(elements))
+
+    if fmt == 'tsv':
+        if columns is not None and header:
+            click.echo(delimiter.join(columns) + '\n')
+
+        for r in rows:
+            c = [str(c) for c in r]
+            click.echo(delimiter.join(c))
+
+
+def watching(w, watch, max_count=None, clear=True):
+    """
+    >>> len(list(watching(True, 1, 0)))
+    1
+    >>> len(list(watching(True, 1, 1)))
+    2
+    >>> len(list(watching(True, None, 0)))
+    1
+    """
+
+    if w and not watch:
+        watch = 2
+    if watch and clear:
+        click.clear()
+    yield 0
+
+    if max_count is not None and max_count < 1:
+        return
+
+    counter = 1
+    while watch and counter <= (max_count or counter):
+        time.sleep(watch)
+        counter += 1
+        if clear:
+            click.clear()
+        yield 0
+
+
+def build_connect_parameters(conn_url, connect_parameters=None):
+    params = (connect_parameters or {}).copy()
+    parsed = parseurl(conn_url)
+    params['host'] = parsed['host']
+    params['port'] = parsed['port']
+    params['fallback_application_name'] = 'Patroni ctl'
+    params['connect_timeout'] = '5'
+
+    return params
+
+
+def get_all_members(cluster, role='master'):
+    if role == 'master':
+        if cluster.leader is not None:
+            yield cluster.leader
+        return
+
+    leader_name = (cluster.leader.member.name if cluster.leader else None)
+    for m in cluster.members:
+        if role == 'any' or role == 'replica' and m.name != leader_name:
+            yield m
+
+
+def get_any_member(cluster, role='master', member=None):
+    members = get_all_members(cluster=cluster, role=role)
+    for m in members:
+        if member is None or m.name == member:
+            return m
+
+    return None
+
+
+def get_cursor(cluster, role='master', member=None, connect_parameters=None):
+    member = get_any_member(cluster=cluster, role=role, member=member)
+    if member is None:
+        return None
+
+    params = build_connect_parameters(member.conn_url, connect_parameters)
+
+    conn = psycopg2.connect(**params)
+    conn.autocommit = True
+    cursor = conn.cursor()
+    if role == 'any':
+        return cursor
+
+    cursor.execute('SELECT pg_is_in_recovery()')
+    in_recovery = cursor.fetchone()[0]
+
+    if in_recovery and role == 'replica' or not in_recovery and role == 'master':
+        return cursor
+
+    conn.close()
+
+    return None
+
+
+@ctl.command('dsn', help='Generate a dsn for the provided member, defaults to a dsn of the master')
+@click.option('--role', '-r', help='Give a dsn of any member with this role', type=click.Choice(['master', 'replica',
+              'any']), default=None)
+@click.option('--member', '-m', help='Generate a dsn for this member', type=str)
+@option_dcs
+@option_config_file
+@click.argument('cluster_name')
+def dsn(cluster_name, config_file, dcs, role, member):
+    if role is not None and member is not None:
+        raise PatroniCtlException('--role and --member are mutually exclusive options')
+    if member is None and role is None:
+        role = 'master'
+
+    _, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+    m = get_any_member(cluster=cluster, role=role, member=member)
+    if m is None:
+        raise PatroniCtlException('Can not find a suitable member')
+
+    params = build_connect_parameters(m.conn_url)
+    click.echo('host={host} port={port}'.format(**params))
+
+
+@ctl.command('query', help='Query a Patroni PostgreSQL member')
+@click.argument('cluster_name')
+@option_config_file
+@option_format
+@click.option('--format', 'fmt', help='Output format (pretty, json)', default='tsv')
+@click.option('--file', '-f', 'p_file', help='Execute the SQL commands from this file', type=click.File('rb'))
+@click.option('--password', help='force password prompt', is_flag=True)
+@click.option('-U', '--username', help='database user name', type=str)
+@option_dcs
+@option_watch
+@option_watchrefresh
+@click.option('--role', '-r', help='The role of the query', type=click.Choice(['master', 'replica', 'any']),
+              default=None)
+@click.option('--member', '-m', help='Query a specific member', type=str)
+@click.option('--delimiter', help='The column delimiter', default='\t')
+@click.option('--command', '-c', help='The SQL commands to execute')
+@click.option('-d', '--dbname', help='database name to connect to', type=str)
+def query(
+    cluster_name,
+    config_file,
+    dcs,
+    role,
+    member,
+    w,
+    watch,
+    delimiter,
+    command,
+    p_file,
+    password,
+    username,
+    dbname,
+    fmt='tsv',
+):
+    if role is not None and member is not None:
+        raise PatroniCtlException('--role and --member are mutually exclusive options')
+    if member is None and role is None:
+        role = 'master'
+
+    if p_file is not None and command is not None:
+        raise PatroniCtlException('--file and --command are mutually exclusive options')
+
+    if p_file is None and command is None:
+        raise PatroniCtlException('You need to specify either --command or --file')
+
+    connect_parameters = dict()
+    if username:
+        connect_parameters['user'] = username
+    if password:
+        connect_parameters['password'] = click.prompt('Password', hide_input=True, type=str)
+    if dbname:
+        connect_parameters['database'] = dbname
+
+    if p_file is not None:
+        command = p_file.read()
+
+    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+
+    cursor = None
+    for _ in watching(w, watch, clear=False):
+
+        output, cursor = query_member(cluster=cluster, cursor=cursor, member=member, role=role, command=command,
+                                      connect_parameters=connect_parameters)
+        print_output(None, output, fmt=fmt, delimiter=delimiter)
+
+        if cursor is None:
+            cluster = dcs.get_cluster()
+
+
+def query_member(cluster, cursor, member, role, command, connect_parameters=None):
+    try:
+        if cursor is None:
+            cursor = get_cursor(cluster, role=role, member=member, connect_parameters=connect_parameters)
+
+        if cursor is None:
+            if role is None:
+                message = 'No connection to member {0} is available'.format(member)
+            else:
+                message = 'No connection to role={0} is available'.format(role)
+            logging.debug(message)
+            return [[timestamp(0), message]], None
+
+        cursor.execute('SELECT pg_is_in_recovery()')
+        in_recovery = cursor.fetchone()[0]
+
+        if in_recovery and role == 'master' or not in_recovery and role == 'replica':
+            cursor.connection.close()
+            return None, None
+
+        cursor.execute(command)
+        return cursor.fetchall(), cursor
+    except (psycopg2.OperationalError, psycopg2.DatabaseError) as oe:
+        logging.debug(oe)
+        if cursor is not None and not cursor.connection.closed:
+            cursor.connection.close()
+        message = oe.pgcode or oe.pgerror or str(oe)
+        message = message.replace('\n', ' ')
+        return [[timestamp(0), 'ERROR, SQLSTATE: {0}'.format(message)]], None
+
+
+@ctl.command('remove', help='Remove cluster from DCS')
+@click.argument('cluster_name')
+@option_config_file
+@option_format
+@option_dcs
+def remove(config_file, cluster_name, fmt, dcs):
+    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+
+    if not isinstance(dcs, Etcd):
+        raise PatroniCtlException('We have not implemented this for DCS of type {0}'.format(type(dcs)))
+
+    output_members(cluster, fmt=fmt)
+
+    confirm = click.prompt('Please confirm the cluster name to remove', type=str)
+    if confirm != cluster_name:
+        raise PatroniCtlException('Cluster names specified do not match')
+
+    message = 'Yes I am aware'
+    confirm = \
+        click.prompt('You are about to remove all information in DCS for {0}, please type: "{1}"'.format(cluster_name,
+                     message), type=str)
+    if message != confirm:
+        raise PatroniCtlException('You did not exactly type "{0}"'.format(message))
+
+    if cluster.leader:
+        confirm = click.prompt('This cluster currently is healthy. Please specify the master name to continue')
+        if confirm != cluster.leader.name:
+            raise PatroniCtlException('You did not specify the current master of the cluster')
+
+    dcs.client.delete(dcs.client_path(''), recursive=True)
+
+
+def wait_for_leader(dcs, timeout=30):
+    t_stop = time.time() + timeout
+    timeout /= 2
+
+    while time.time() < t_stop:
+        dcs.watch(timeout)
+        cluster = dcs.get_cluster()
+
+        if cluster.leader:
+            return cluster
+
+    raise PatroniCtlException('Timeout occured')
+
+
+def empty_post_to_members(cluster, member_names, force, endpoint):
+    candidates = dict()
+    for m in cluster.members:
+        candidates[m.name] = m
+
+    if not member_names:
+        member_names = [click.prompt('Which member do you want to {0} [{1}]?'.format(endpoint,
+                        ', '.join(candidates.keys())), type=str, default='')]
+
+    for mn in member_names:
+        if mn not in candidates.keys():
+            raise PatroniCtlException('{0} is not a member of cluster'.format(mn))
+
+    if not force:
+        confirm = click.confirm('Are you sure you want to {0} members {1}?'.format(endpoint, ', '.join(member_names)))
+        if not confirm:
+            raise PatroniCtlException('Aborted {0}'.format(endpoint))
+
+    for mn in member_names:
+        r = post_patroni(candidates[mn], endpoint, '')
+        if r.status_code != 200:
+            click.echo('{0} failed for member {1}, status code={2}, ({3})'.format(endpoint, mn, r.status_code, r.text))
+        else:
+            click.echo('Succesful {0} on member {1}'.format(endpoint, mn))
+
+
+def ctl_load_config(cluster_name, config_file, dcs):
+    config = load_config(config_file, dcs)
+    dcs = get_dcs(config, cluster_name)
+    cluster = dcs.get_cluster()
+
+    return config, dcs, cluster
+
+
+@ctl.command('restart', help='Restart cluster member')
+@click.argument('cluster_name')
+@click.argument('member_names', nargs=-1)
+@click.option('--role', '-r', help='Restart only members with this role', default='any',
+              type=click.Choice(['master', 'replica', 'any']))
+@click.option('--any', 'p_any', help='Restart a single member only', is_flag=True)
+@option_config_file
+@option_force
+@option_dcs
+def restart(cluster_name, member_names, config_file, dcs, force, role, p_any):
+    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+
+    role_names = [m.name for m in get_all_members(cluster=cluster, role=role)]
+
+    if member_names:
+        member_names = list(set(member_names) & set(role_names))
+    else:
+        member_names = role_names
+
+    if p_any:
+        random.shuffle(member_names)
+        member_names = member_names[:1]
+
+    output_members(cluster)
+    empty_post_to_members(cluster, member_names, force, 'restart')
+
+
+@ctl.command('reinit', help='Reinitialize cluster member')
+@click.argument('cluster_name')
+@click.argument('member_names', nargs=-1)
+@option_config_file
+@option_force
+@option_dcs
+def reinit(cluster_name, member_names, config_file, dcs, force):
+    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+    empty_post_to_members(cluster, member_names, force, 'reinitialize')
+
+
+@ctl.command('failover', help='Failover to a replica')
+@click.argument('cluster_name')
+@click.option('--master', help='The name of the current master', default=None)
+@click.option('--candidate', help='The name of the candidate', default=None)
+@click.option('--scheduled', help='Timestamp of a scheduled failover in unambiguous format (e.g. ISO 8601)',
+              default=None)
+@click.option('--force', is_flag=True)
+@option_config_file
+@option_dcs
+def failover(config_file, cluster_name, master, candidate, force, dcs, scheduled):
+    """
+        We want to trigger a failover for the specified cluster name.
+
+        We verify that the cluster name, master name and candidate name are correct.
+        If so, we trigger a failover and keep the client up to date.
+    """
+
+    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+
+    if cluster.leader is None:
+        raise PatroniCtlException('This cluster has no master')
+
+    if master is None:
+        if force:
+            master = cluster.leader.member.name
+        else:
+            master = click.prompt('Master', type=str, default=cluster.leader.member.name)
+
+    if cluster.leader.member.name != master:
+        raise PatroniCtlException('Member {0} is not the leader of cluster {1}'.format(master, cluster_name))
+
+    candidate_names = [str(m.name) for m in cluster.members if m.name != master]
+    # We sort the names for consistent output to the client
+    candidate_names.sort()
+
+    if not candidate_names:
+        raise PatroniCtlException('No candidates found to failover to')
+
+    if candidate is None and not force:
+        candidate = click.prompt('Candidate ' + str(candidate_names), type=str, default='')
+
+    if candidate == master:
+        raise PatroniCtlException('Failover target and source are the same.')
+
+    if candidate and candidate not in candidate_names:
+        raise PatroniCtlException('Member {0} does not exist in cluster {1}'.format(candidate, cluster_name))
+
+    if scheduled is None and not force:
+        scheduled = click.prompt('When should the failover take place (e.g. 2015-10-01T14:30) ', type=str,
+                                 default='now')
+
+    if (scheduled or 'now') == 'now':
+        scheduled_at = None
+    else:
+        try:
+            scheduled_at = dateutil.parser.parse(scheduled)
+            if scheduled_at.tzinfo is None:
+                scheduled_at = tzlocal.get_localzone().localize(scheduled_at)
+        except (ValueError, TypeError):
+            message = 'Unable to parse scheduled timestamp ({}). It should be in an unambiguous format (e.g. ISO 8601)'
+            raise PatroniCtlException(message.format(scheduled))
+        scheduled_at = scheduled_at.isoformat()
+
+    failover_value = {'leader': master, 'member': candidate, 'scheduled_at': scheduled_at}
+    logging.debug(failover_value)
+
+    # By now we have established that the leader exists and the candidate exists
+    click.echo('Current cluster topology')
+    output_members(dcs.get_cluster(), name=cluster_name)
+
+    if not force:
+        a = \
+            click.confirm('Are you sure you want to failover cluster {0}, demoting current master {1}?'.format(
+                cluster_name, master))
+        if not a:
+            raise PatroniCtlException('Aborting failover')
+
+    r = None
+    try:
+        r = post_patroni(cluster.leader.member, 'failover', failover_value)
+        if r.status_code == 200:
+            logging.debug(r)
+            cluster = dcs.get_cluster()
+            logging.debug(cluster)
+            click.echo('{0} {1}'.format(timestamp(), r.text))
+        else:
+            click.echo('Failover failed, details: {0}, {1}'.format(r.status_code, r.text))
+            return
+    except:
+        logging.exception(r)
+        logging.warning('Failing over to DCS')
+        click.echo(timestamp() + ' Could not failover using Patroni api, falling back to DCS')
+        click.echo(timestamp() + ' Initializing failover from master {0}'.format(master))
+        dcs.manual_failover(leader=master, member=candidate, scheduled_at=failover_value)
+
+    output_members(cluster, name=cluster_name)
+
+
+def output_members(cluster, name=None, fmt='pretty'):
+    rows = []
+    logging.debug(cluster)
+    leader_name = None
+    if cluster.leader:
+        leader_name = cluster.leader.member.name
+
+    xlog_location_cluster = cluster.last_leader_operation or 0
+
+    # Mainly for consistent pretty printing and watching we sort the output
+    cluster.members.sort(key=lambda x: x.name)
+    for m in cluster.members:
+        logging.debug(m)
+
+        leader = ''
+        if m.name == leader_name:
+            leader = '*'
+
+        host = build_connect_parameters(m.conn_url)['host']
+
+        xlog_location = m.data.get('xlog_location') or 0
+        lag = ''
+        if (xlog_location_cluster >= xlog_location):
+            lag = round((xlog_location_cluster - xlog_location)/1024/1024)
+
+        rows.append([
+            name,
+            m.name,
+            host,
+            leader,
+            m.data.get('state', ''),
+            lag
+        ])
+
+    columns = [
+        'Cluster',
+        'Member',
+        'Host',
+        'Leader',
+        'State',
+        'Lag in MB',
+    ]
+    alignment = {'Cluster': 'l', 'Member': 'l', 'Host': 'l', 'Lag in MB': 'r'}
+
+    print_output(columns, rows, alignment, fmt)
+
+
+@ctl.command('list', help='List the Patroni members for a given Patroni')
+@click.argument('cluster_names', nargs=-1)
+@option_config_file
+@option_format
+@option_watch
+@option_watchrefresh
+@option_dcs
+def members(config_file, cluster_names, fmt, watch, w, dcs):
+    if not cluster_names:
+        logging.warning('Listing members: No cluster names were provided')
+        return
+
+    config = load_config(config_file, dcs)
+    for cn in cluster_names:
+        dcs = get_dcs(config, cn)
+
+        for _ in watching(w, watch):
+            output_members(dcs.get_cluster(), name=cn, fmt=fmt)
+
+
+def timestamp(precision=6):
+    return datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:precision - 7]
+
+
+@ctl.command('configure', help='Create configuration file')
+@click.option('--config-file', '-c', help='Configuration file', prompt='Configuration file', default=CONFIG_FILE_PATH)
+@click.option('--dcs', '-d', help='The DCS connect url', prompt='DCS connect url', default='etcd://localhost:4001')
+@click.option('--namespace', '-n', help='The namespace', prompt='Namespace', default='/service/')
+def configure(config_file, dcs, namespace):
+    config = dict()
+    config['dcs_api'] = str(dcs)
+    config['namespace'] = str(namespace)
+    store_config(config, config_file)

patroni/dcs.py

@@ -0,0 +1,319 @@
+import abc
+import json
+import dateutil
+
+from collections import namedtuple
+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 (TypeError, ValueError):
+                data = {}
+        return Member(index, name, session, data)
+
+    @property
+    def conn_url(self):
+        return self.data.get('conn_url')
+
+    @property
+    def api_url(self):
+        return self.data.get('api_url')
+
+    @property
+    def nofailover(self):
+        return self.data.get('tags', {}).get('nofailover', False)
+
+    @property
+    def replicatefrom(self):
+        return self.data.get('tags', {}).get('replicatefrom')
+
+
+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,scheduled_at')):
+
+    """
+    >>> 'Failover' in str(Failover.from_node(1, '{"leader": "cluster_leader"}'))
+    True
+    >>> 'Failover' in str(Failover.from_node(1, '{"leader": "cluster_leader", "member": "cluster:member"}'))
+    True
+    >>> Failover.from_node(1, 'null') is None
+    True
+    >>> n = '{"leader": "cluster_leader", "member": "cluster:member", "scheduled_at": "2016-01-14T10:09:57.1394Z"}'
+    >>> 'tzinfo=' in str(Failover.from_node(1, n))
+    True
+    >>> Failover.from_node(1, None) is None
+    True
+    >>> Failover.from_node(1, '{}') is None
+    True
+    >>> 'abc' in Failover.from_node(1, 'abc:def')
+    True
+    """
+    @staticmethod
+    def from_node(index, value):
+        if not value:
+            return None
+
+        try:
+            data = json.loads(value)
+            if not data:
+                return None
+        except ValueError:
+            t = [a.strip() for a in value.split(':')]
+            leader = t[0]
+            candidate = t[1] if len(t) > 1 else None
+            return Failover(index, leader, candidate, None) if leader or candidate else None
+
+        if data.get('scheduled_at'):
+            data['scheduled_at'] = dateutil.parser.parse(data['scheduled_at'])
+
+        return Failover(index, data.get('leader'), data.get('member'), data.get('scheduled_at'))
+
+
+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)
+
+    def has_member(self, member_name):
+        return any(m for m in self.members if m.name == member_name)
+
+    def get_member(self, member_name):
+        return ([m for m in self.members if m.name == member_name] or [None])[0]
+
+
+class AbstractDCS(object):
+
+    __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._namespace = '/{0}'.format(config.get('namespace', '/service/').strip('/'))
+        self._base_path = '/'.join([self._namespace, config['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, scheduled_at=None, index=None):
+        failover_value = dict()
+        if leader:
+            failover_value['leader'] = leader
+
+        if member:
+            failover_value['member'] = member
+
+        if scheduled_at:
+            failover_value['scheduled_at'] = scheduled_at.isoformat()
+
+        return self.set_failover_value(json.dumps(failover_value), index)
+
+    @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, create_new=True, sysid=""):
+        """Race for cluster initialization.
+
+        :param create_new: False if the key should already exist (in the case we are setting the system_id)
+        :param sysid: PostgreSQL cluster system identifier, if specified, is written to the key
+        :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

@@ -0,0 +1,322 @@
+from __future__ import absolute_import
+import etcd
+import logging
+import os
+import random
+import requests
+import socket
+import time
+import urllib3
+
+from dns.exception import DNSException
+from dns import resolver
+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
+from six.moves.http_client import HTTPException
+
+logger = logging.getLogger(__name__)
+
+
+class EtcdError(DCSError):
+    pass
+
+
+class Client(etcd.Client):
+
+    def __init__(self, config):
+        super(Client, self).__init__(read_timeout=5)
+        self._config = config
+        self._load_machines_cache()
+        self._allow_reconnect = True
+
+    @property
+    def machines(self):
+        """Original `machines` method(property) of `etcd.Client` class raise exception
+        when it failed to get list of etcd cluster members. This method is being called
+        only when request failed on one of the etcd members during `api_execute` call.
+        For us it's more important to execute original request rather then get new
+        topology of etcd cluster. So we will catch this exception and return valid list
+        of machines with setting flag `self._update_machines_cache` to `!True`.
+        Later, during next `api_execute` call we will forcefully update machines_cache"""
+        try:
+            ret = super(Client, self).machines
+            random.shuffle(ret)
+            return ret
+        except etcd.EtcdException:
+            if self._update_machines_cache:  # We are updating machines_cache
+                raise  # This exception is fatal, we should re-raise it.
+            self._update_machines_cache = True
+            return [self._base_uri]
+
+    def _do_http_request(self, request_executor, method, url, fields=None, **kwargs):
+        try:
+            response = request_executor(method, url, fields=fields, **kwargs)
+            response.data.decode('utf-8')
+            self._check_cluster_id(response)
+        except (urllib3.exceptions.HTTPError, HTTPException, socket.error) as e:
+            if (isinstance(fields, dict) and fields.get("wait") == "true" and
+                    isinstance(e, urllib3.exceptions.ReadTimeoutError)):
+                logger.debug("Watch timed out.")
+                raise etcd.EtcdWatchTimedOut("Watch timed out: {0}".format(e), cause=e)
+            logger.error("Request to server %s failed: %r", self._base_uri, e)
+            logger.info("Reconnection allowed, looking for another server.")
+            self._base_uri = self._next_server(cause=e)
+            response = False
+        return response
+
+    def api_execute(self, path, method, params=None, timeout=None):
+        if not path.startswith('/'):
+            raise ValueError('Path does not start with /')
+
+        if timeout is None:
+            timeout = self.read_timeout
+
+        if timeout == 0:
+            timeout = None
+
+        kwargs = {'timeout': timeout, 'fields': params, 'redirect': self.allow_redirect,
+                  'headers': self._get_headers(), 'preload_content': False}
+
+        if method in [self._MGET, self._MDELETE]:
+            request_executor = self.http.request
+        elif method in [self._MPUT, self._MPOST]:
+            request_executor = self.http.request_encode_body
+            kwargs['encode_multipart'] = False
+        else:
+            raise etcd.EtcdException('HTTP method {0} not supported'.format(method))
+
+        # Update machines_cache if previous attempt of update has failed
+        if self._update_machines_cache:
+            self._load_machines_cache()
+
+        response = False
+
+        try:
+            while not response:
+                response = self._do_http_request(request_executor, method, self._base_uri + path, **kwargs)
+
+                if response is False and not self._use_proxies:
+                    self._machines_cache = self.machines
+                    self._machines_cache.remove(self._base_uri)
+            return self._handle_server_response(response)
+        except etcd.EtcdConnectionFailed:
+            self._update_machines_cache = True
+            raise
+
+    @staticmethod
+    def get_srv_record(host):
+        try:
+            return [(str(r.target).rstrip('.'), r.port) for r in resolver.query('_etcd-server._tcp.' + host, 'SRV')]
+        except DNSException:
+            logger.exception('Can not resolve SRV for %s', host)
+        return []
+
+    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
+        'GET http://{host}:{port}/members' request (peer protocol)"""
+
+        ret = []
+        for host, port in self.get_srv_record(discovery_srv):
+            url = '{0}://{1}:{2}/members'.format(self._protocol, host, port)
+            try:
+                response = requests.get(url, timeout=5)
+                if response.ok:
+                    for member in response.json():
+                        ret.extend(member['clientURLs'])
+                    break
+            except RequestException:
+                logger.exception('GET %s', url)
+        return list(set(ret))
+
+    def _get_machines_cache_from_dns(self, addr):
+        """One host might be resolved into multiple ip addresses. We will make list out of it"""
+
+        ret = []
+        host, port = addr.split(':')
+        try:
+            for r in set(socket.getaddrinfo(host, port, socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP)):
+                ret.append('{0}://{1}:{2}'.format(self._protocol, r[4][0], r[4][1]))
+        except socket.error:
+            logger.exception('Can not resolve %s', host)
+        return list(set(ret)) if ret else ['{0}://{1}:{2}'.format(self._protocol, host, port)]
+
+    def _load_machines_cache(self):
+        """This method should fill up `_machines_cache` from scratch.
+        It could happen only in two cases:
+        1. During class initialization
+        2. When all etcd members failed"""
+
+        self._update_machines_cache = True
+
+        if 'discovery_srv' not in self._config and 'host' not in self._config:
+            raise Exception('Neither discovery_srv nor host are defined in etcd section of config')
+
+        self._machines_cache = []
+
+        if 'discovery_srv' in self._config:
+            self._machines_cache = self._get_machines_cache_from_srv(self._config['discovery_srv'])
+
+        if not self._machines_cache and 'host' in self._config:
+            self._machines_cache = self._get_machines_cache_from_dns(self._config['host'])
+
+        # Can not bootstrap list of etcd-cluster members, giving up
+        if not self._machines_cache:
+            raise etcd.EtcdException
+
+        # After filling up initial list of machines_cache we should ask etcd-cluster about actual list
+        self._base_uri = self._machines_cache.pop(0)
+        self._machines_cache = self.machines
+
+        if self._base_uri in self._machines_cache:
+            self._machines_cache.remove(self._base_uri)
+
+        self._update_machines_cache = False
+
+
+def catch_etcd_errors(func):
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs) is not None
+        except (RetryFailedError, etcd.EtcdException):
+            return False
+        except:
+            logger.exception("")
+            raise EtcdError("unexpected error")
+
+    return wrapper
+
+
+class Etcd(AbstractDCS):
+
+    def __init__(self, name, config):
+        super(Etcd, self).__init__(name, config)
+        self.ttl = config.get('ttl', 30)
+        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)
+
+    @staticmethod
+    def get_etcd_client(config):
+        client = None
+        while not client:
+            try:
+                client = Client(config)
+            except etcd.EtcdException:
+                logger.info('waiting on etcd')
+                sleep(5)
+        return client
+
+    @staticmethod
+    def member(node):
+        return Member.from_node(node.modifiedIndex, os.path.basename(node.key), node.ttl, node.value)
+
+    def _load_cluster(self):
+        try:
+            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 = nodes.get(self._INITIALIZE)
+            initialize = initialize and initialize.value
+
+            # get last leader operation
+            last_leader_operation = nodes.get(self._LEADER_OPTIME)
+            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(self._MEMBERS) and k.count('/') == 1]
+
+            # get leader
+            leader = nodes.get(self._LEADER)
+            if leader:
+                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)
+
+            # failover key
+            failover = nodes.get(self._FAILOVER)
+            if failover:
+                failover = Failover.from_node(failover.modifiedIndex, failover.value)
+
+            self._cluster = Cluster(initialize, leader, last_leader_operation, members, failover)
+        except etcd.EtcdKeyNotFound:
+            self._cluster = Cluster(False, None, None, [], None)
+        except:
+            logger.exception('get_cluster')
+            raise EtcdError('Etcd is not responding properly')
+
+    @catch_etcd_errors
+    def touch_member(self, connection_string, ttl=None):
+        return self.retry(self.client.set, self.member_path, connection_string, ttl or self.ttl)
+
+    @catch_etcd_errors
+    def take_leader(self):
+        return self.retry(self.client.set, self.leader_path, self._name, self.ttl)
+
+    def attempt_to_acquire_leader(self):
+        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 set_failover_value(self, value, index=None):
+        return self.client.write(self.failover_path, value, prevIndex=index or 0)
+
+    @catch_etcd_errors
+    def write_leader_optime(self, last_operation):
+        return self.client.set(self.leader_optime_path, last_operation)
+
+    @catch_etcd_errors
+    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, create_new=True, sysid=""):
+        return self.retry(self.client.write, self.initialize_path, sysid, prevExist=(not create_new))
+
+    @catch_etcd_errors
+    def delete_leader(self):
+        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)
+
+    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 etcd.EtcdWatchTimedOut:
+                    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,32 @@
+from click import ClickException
+
+
+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 PatroniCtlException(ClickException):
+    pass
+
+
+class PostgresException(PatroniException):
+    pass
+
+
+class DCSError(PatroniException):
+    pass
+
+
+class PostgresConnectionException(PostgresException):
+    pass

patroni/ha.py

@@ -0,0 +1,497 @@
+import json
+import logging
+import psycopg2
+import requests
+import sys
+import datetime
+import pytz
+
+from multiprocessing.pool import ThreadPool
+from patroni.async_executor import AsyncExecutor
+from patroni.exceptions import DCSError, PostgresConnectionException
+from patroni.utils import sleep
+
+logger = logging.getLogger(__name__)
+
+
+class Ha(object):
+
+    def __init__(self, patroni):
+        self.patroni = patroni
+        self.state_handler = patroni.postgresql
+        self.dcs = patroni.dcs
+        self.cluster = None
+        self.old_cluster = None
+        self.recovering = False
+        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 healthy
+        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,
+            'tags': self.patroni.tags
+        }
+        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 clone(self, clone_member, clone_member_name="leader"):
+        if self.state_handler.bootstrap(cluster_initialized=True, clone_member=clone_member):
+            logger.info('bootstrapped from {0}'.format(clone_member_name)
+                        if clone_member else 'bootstrapped without leader')
+        else:
+            self.state_handler.stop('immediate')
+            self.state_handler.remove_data_directory()
+            logger.error('failed to bootstrap from {0}'.format(clone_member_name)
+                         if clone_member else 'failed to bootstrap (without leader)')
+
+    def bootstrap(self):
+        if not self.cluster.is_unlocked():  # cluster already has leader
+            clonefrom = self.patroni.clonefrom
+            clone_member = self.cluster.get_member(clonefrom)\
+                if self.cluster.has_member(clonefrom) else self.cluster.leader
+            clone_member_name = 'leader' if clone_member == self.cluster.leader else 'replica \'{0}\''.format(clonefrom)
+            self._async_executor.schedule('bootstrap from {0}'.format(clone_member_name))
+            self._async_executor.run_async(self.clone, args=(clone_member, clone_member_name))
+            return 'trying to bootstrap from {0}'.format(clone_member_name)
+        elif not self.cluster.initialize and not self.patroni.nofailover:  # no initialize key
+            if self.dcs.initialize(create_new=True):  # race for initialization
+                try:
+                    self.state_handler.bootstrap()
+                    self.dcs.initialize(create_new=False, sysid=self.state_handler.sysid)
+                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()
+                self.load_cluster_from_dcs()
+                return 'initialized a new cluster'
+            else:
+                return 'failed to acquire initialize lock'
+        else:
+            if self.state_handler.can_create_replica_without_replication_connection():
+                self._async_executor.run_async(self.clone, args=(None, ))
+                return "trying to bootstrap without leader"
+            return 'waiting for leader to bootstrap'
+
+    def recover(self):
+        # 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 (self.state_handler.role == 'master') and pg_controldata and\
+                pg_controldata.get('Database cluster state', '') == 'in production':  # crashed master
+            self.state_handler.require_rewind()
+        self.recovering = True
+        return self.follow("started as readonly because i had the session lock",
+                           "started as a secondary",
+                           refresh=True, recovery=True)
+
+    def follow(self, demote_reason, follow_reason, refresh=True, recovery=False):
+        if refresh:
+            self.load_cluster_from_dcs()
+
+        if not recovery and self.state_handler.is_leader() or recovery and self.state_handler.role == 'master':
+            ret = demote_reason
+        else:
+            ret = follow_reason
+
+        # determine the node to follow. If replicatefrom tag is set,
+        # try to follow the node mentioned there, otherwise, follow the leader.
+        if self.patroni.replicatefrom:
+            node_to_follow = [m for m in self.cluster.members if m.name == self.patroni.replicatefrom]
+            node_to_follow = node_to_follow[0] if node_to_follow else self.cluster.leader
+        else:
+            node_to_follow = self.cluster.leader
+        node_to_follow = None if node_to_follow and node_to_follow.name == self.state_handler.name else node_to_follow
+        if not self.state_handler.check_recovery_conf(node_to_follow) or recovery:
+            self._async_executor.schedule('changing primary_conninfo and restarting')
+            self._async_executor.run_async(self.state_handler.follow, (node_to_follow, recovery))
+        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()
+            self.touch_member()
+            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.
+        tags - dictionary with values of different tags (i.e. nofailover)
+        """
+
+        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']
+            tags = json.get('tags', dict())
+            return (member, True, not is_master, xlog_location, tags)
+        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 self.patroni.nofailover is True:
+            return False
+
+        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 not m.nofailover and m.api_url]
+
+        if members:
+            my_xlog_location = self.state_handler.xlog_position()
+            for member, reachable, in_recovery, xlog_location, tags in self.fetch_nodes_statuses(members):
+                if reachable and not tags.get('nofailover', False):  # 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 not m.nofailover and m.api_url]
+        if members:
+            for member, reachable, _, _, tags in self.fetch_nodes_statuses(members):
+                if reachable and not tags.get('nofailover', False):
+                    ret = True  # TODO: check xlog_location
+                elif not reachable:
+                    logger.info('Member %s is not reachable', member.name)
+                elif tags.get('nofailover', False):
+                    logger.info('Member %s is not allowed to promote', 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, _, _, tags = self.fetch_node_status(members[0])
+                if reachable and not tags.get('nofailover', False):  # 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
+                if not reachable:
+                    logger.warning('manual failover: member %s is unhealthy', member.name)
+                elif tags.get('nofailover', False):
+                    logger.warning('manual failover: member %s is not allowed to promote', 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 not in (failover.member, failover.leader)]
+                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.state_handler.is_leader():  # leader is always the healthiest
+            return True
+
+        if self.patroni.nofailover:  # nofailover tag makes node always unhealthy
+            return False
+
+        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.touch_member()
+            self.dcs.reset_cluster()
+        self.state_handler.follow(None)
+
+    def process_manual_failover_from_leader(self):
+        failover = self.cluster.failover
+
+        if failover.scheduled_at:
+            # If the failover is in the far future, we shouldn't do anything and just return.
+            # If the failover is in the past, we consider the value to be stale and we remove
+            # the value.
+            # If the value is close to now, we initiate the failover
+            now = datetime.datetime.now(pytz.utc)
+            try:
+                delta = (failover.scheduled_at - now).total_seconds()
+
+                if delta > self.patroni.nap_time:
+                    logging.info('Awaiting failover at %s (in %.0f seconds)', failover.scheduled_at.isoformat(), delta)
+                    return
+                elif delta < - int(self.patroni.nap_time * 1.5):
+                    logger.warning('Found a stale failover value, cleaning up: %s', failover.scheduled_at)
+                    self.dcs.manual_failover('', '', self.cluster.failover.index)
+                    return
+
+                # The value is very close to now
+                sleep(max(delta, 0))
+                logger.info('Manual scheduled failover at {}'.format(failover.scheduled_at.isoformat()))
+            except TypeError:
+                logger.warning('Incorrect value in of scheduled_at: %s', failover.scheduled_at)
+
+        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('Cleaning up failover key after acquiring leader lock...')
+                    self.dcs.manual_failover('', '')
+                self.dcs.get_cluster()
+                return self.enforce_master_role('acquired session lock as a leader',
+                                                'promoted self to leader by acquiring session lock')
+            else:
+                return self.follow('demoted self after trying and failing to obtain lock',
+                                   'following new leader after trying and failing to obtain lock')
+        else:
+            if self.patroni.nofailover:
+                return self.follow('demoting self because I am not allowed to become master',
+                                   'following a different leader because I am not allowed to promote')
+            return self.follow('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('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.clone(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'
+
+    @staticmethod
+    def sysid_valid(sysid):
+        # sysid does tv_sec << 32, where tv_sec is the number of seconds sine 1970,
+        # so even 1 << 32 would have 10 digits.
+        return str(sysid) and len(str(sysid)) >= 10 and str(sysid).isdigit()
+
+    def post_recover(self):
+        if not self.state_handler.is_running():
+            if self.has_lock():
+                self.dcs.delete_leader()
+                self.dcs.reset_cluster()
+                return 'removed leader key after trying and failing to start postgres'
+            return 'failed to start postgres'
+        return None
+
+    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.sysid_valid(self.cluster.initialize) and self.has_lock():
+                self.dcs.initialize(create_new=(self.cluster.initialize is None), sysid=self.state_handler.sysid)
+
+            if self._async_executor.busy:
+                return self.handle_long_action_in_progress()
+
+            # we've got here, so any async action has finished. Check if we tried to recover and failed
+            if self.recovering:
+                self.recovering = False
+                msg = self.post_recover()
+                if msg is not None:
+                    return msg
+
+            # 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.sysid_valid(self.cluster.initialize) and self.cluster.is_unlocked():
+                self.dcs.initialize(create_new=(self.cluster.initialize is None), sysid=self.state_handler.sysid)
+            else:
+                # check if we are allowed to join
+                if self.sysid_valid(self.cluster.initialize) and self.cluster.initialize != self.state_handler.sysid:
+                    logger.fatal("system ID mismatch, node %s belongs to a different cluster: %s != %s",
+                                 self.state_handler.name, self.cluster.initialize, self.state_handler.sysid)
+                    sys.exit(1)
+
+            # 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:
+                # we might not have a valid PostgreSQL connection here if another thread
+                # stops PostgreSQL, therefore, we only reload replication slots if no
+                # asynchronous processes are running (should be always the case for the master)
+                if not self._async_executor.busy:
+                    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,786 @@
+import logging
+import os
+import psycopg2
+import shlex
+import shutil
+import subprocess
+import tempfile
+import time
+
+from patroni.exceptions import PostgresConnectionException, PostgresException
+from patroni.utils import Retry, RetryFailedError
+from six import string_types
+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(object):
+
+    def __init__(self, config):
+        self.config = config
+        self.name = config['name']
+        self.server_parameters = config.get('parameters', {})
+        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.initdb_options = config.get('initdb', [])
+        self.pgpass = config.get('pgpass') or os.path.join(os.path.expanduser('~'), 'pgpass')
+        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') 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') 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._sysid = None
+        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()
+        return data.get('wal_log_hints setting', 'off') == 'on' or data.get('Data page checksum version', '0') != '0'
+
+    @property
+    def sysid(self):
+        if not self._sysid:
+            data = self.controldata()
+            self._sysid = data.get('Database system identifier', "")
+        return self._sysid
+
+    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
+
+    @property
+    def _connect_kwargs(self):
+        r = parseurl('postgres://{0}/postgres'.format(self.local_address))
+        if 'username' in self.superuser:
+            r['user'] = self.superuser['username']
+        if 'password' in self.superuser:
+            r['password'] = self.superuser['password']
+        return r
+
+    def connection(self):
+        if not self._connection or self._connection.closed != 0:
+            self._connection = psycopg2.connect(**self._connect_kwargs)
+            self._connection.autocommit = True
+            self.server_version = self._connection.server_version
+        return self._connection
+
+    def _cursor(self):
+        if not self._cursor_holder or self._cursor_holder.closed or self._cursor_holder.connection.closed != 0:
+            logger.info("establishing a new patroni connection to the postgres cluster")
+            self._cursor_holder = self.connection().cursor()
+        return self._cursor_holder
+
+    def close_connection(self):
+        if self._cursor_holder and self._cursor_holder.connection and self._cursor_holder.connection.closed == 0:
+            self._cursor_holder.connection.close()
+            logger.info("closed patroni connection to the postgresql cluster")
+
+    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) == []
+
+    @staticmethod
+    def initdb_allowed_option(name):
+        if name in ['pgdata', 'nosync', 'pwfile', 'sync-only']:
+            raise Exception('{0} option for initdb is not allowed'.format(name))
+        return True
+
+    def get_initdb_options(self):
+        options = []
+        for o in self.initdb_options:
+            if isinstance(o, string_types) and self.initdb_allowed_option(o):
+                options.append('--{0}'.format(o))
+            elif isinstance(o, dict):
+                keys = list(o.keys())
+                if len(keys) != 1 or not isinstance(keys[0], string_types) or not self.initdb_allowed_option(keys[0]):
+                    raise Exception('Invalid option: {0}'.format(o))
+                options.append('--{0}={1}'.format(keys[0], o[keys[0]]))
+            else:
+                raise Exception('Unknown type of initdb option: {0}'.format(o))
+        return options
+
+    def initialize(self):
+        self.set_state('initalizing new cluster')
+        options = self.get_initdb_options()
+        pwfile = None
+
+        if self.superuser:
+            if 'username' in self.superuser:
+                options.append('--username={0}'.format(self.superuser['username']))
+            if 'password' in self.superuser:
+                (fd, pwfile) = tempfile.mkstemp()
+                os.write(fd, self.superuser['password'].encode('utf-8'))
+                os.close(fd)
+                options.append('--pwfile={0}'.format(pwfile))
+
+        ret = subprocess.call(self._pg_ctl + ['initdb'] + (['-o', ' '.join(options)] if options else [])) == 0
+        if pwfile:
+            os.remove(pwfile)
+        if ret:
+            self.write_pg_hba()
+        else:
+            self.set_state('initdb failed')
+        return ret
+
+    def delete_trigger_file(self):
+        if os.path.exists(self.trigger_file):
+            os.unlink(self.trigger_file)
+
+    def write_pgpass(self, record):
+        with open(self.pgpass, 'w') as f:
+            os.fchmod(f.fileno(), 0o600)
+            f.write('{host}:{port}:*:{user}:{password}\n'.format(**record))
+
+        env = os.environ.copy()
+        env['PGPASSFILE'] = self.pgpass
+        return env
+
+    def sync_replica(self, clone_member):
+        # add the credentials to connect to the replica origin to pgpass.
+        env = self.write_pgpass(parseurl(clone_member.conn_url)) if clone_member else os.environ.copy()
+        if self.create_replica(clone_member, env) == 0:
+            self.delete_trigger_file()
+            return True
+        return False
+
+    @staticmethod
+    def build_connstring(conn):
+        """
+        >>> Postgresql.build_connstring({'host': '127.0.0.1', 'port': '5432'}) == 'host=127.0.0.1 port=5432'
+        True
+        """
+        return ' '.join('{0}={1}'.format(param, val) for param, val in sorted(conn.items()))
+
+    def replica_method_can_work_without_replication_connection(self, method):
+        return method != 'basebackup' and self.config and self.config.get(method, {}).get('no_master')
+
+    def can_create_replica_without_replication_connection(self):
+        """ go through the replication methods to see if there are ones
+            that does not require a working replication connection.
+        """
+        replica_methods = self.config.get('create_replica_method', [])
+        return any(self.replica_method_can_work_without_replication_connection(replica_method)
+                   for replica_method in replica_methods)
+
+    def create_replica(self, clone_member, env):
+        # create the replica according to the replica_method
+        # defined by the user.  this is a list, so we need to
+        # loop through all methods the user supplies
+        connstring = clone_member.conn_url if clone_member else ""
+        # get list of replica methods from config.
+        # If there is no configuration key, or no value is specified, use basebackup
+        replica_methods = self.config.get('create_replica_method') or ['basebackup']
+        # if we don't have any source, leave only replica methods that work without it
+        replica_methods = \
+            [r for r in replica_methods if self.replica_method_can_work_without_replication_connection(r)]\
+            if not clone_member else replica_methods
+        # go through them in priority order
+        ret = 1
+        for replica_method in replica_methods:
+            # if the method is basebackup, then use the built-in
+            if replica_method == "basebackup":
+                ret = self.basebackup(clone_member, env)
+                if ret == 0:
+                    logger.info("replica has been created using basebackup")
+                    # if basebackup succeeds, exit with success
+                    break
+            else:
+                cmd = replica_method
+                method_config = {}
+                # user-defined method; check for configuration
+                # not required, actually
+                if replica_method in self.config:
+                    method_config = self.config[replica_method].copy()
+                    # look to see if the user has supplied a full command path
+                    # if not, use the method name as the command
+                    cmd = method_config.pop('command', cmd)
+                    # add the default parameters
+                try:
+                    method_config.update({"scope": self.scope,
+                                          "role": "replica",
+                                          "datadir": self.data_dir,
+                                          "connstring": connstring})
+                    params = ["--{0}={1}".format(arg, val) for arg, val in method_config.items()]
+                    # call script with the full set of parameters
+                    ret = subprocess.call(shlex.split(cmd) + params, env=env)
+                    # if we succeeded, stop
+                    if ret == 0:
+                        logger.info("replica has been created using {0}".format(replica_method))
+                        break
+                except Exception as e:
+                    logger.exception('Error creating replica using method {0}: {1}'.format(replica_method, str(e)))
+                    ret = 1
+
+        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 OSError:
+            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')
+
+        env = os.environ.copy()
+        if 'username' in self.superuser:
+            env['PGUSER'] = self.superuser['username']
+        ret = subprocess.call(self._pg_ctl + ['start', '-o', self.server_options()], env=env, preexec_fn=os.setsid) == 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
+        if ret and not block_callbacks:
+            self.call_nowait(ACTION_ON_START)
+        return ret
+
+    def checkpoint(self, connect_kwargs=None):
+        connect_kwargs = connect_kwargs or self._connect_kwargs
+        for p in ['connect_timeout', 'options']:
+            connect_kwargs.pop(p, None)
+        try:
+            with psycopg2.connect(**connect_kwargs) as conn:
+                conn.autocommit = True
+                with conn.cursor() as cur:
+                    cur.execute("SET statement_timeout = 0")
+                    cur.execute('CHECKPOINT')
+        except psycopg2.Error:
+            logging.exception('Exception during CHECKPOINT')
+
+    def stop(self, mode='fast', block_callbacks=False):
+        # make sure we close all connections established against
+        # the former node, otherwise, we might get a stalled one
+        # after kill -9, which would report incorrect data to
+        # patroni.
+
+        self.close_connection()
+        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
+        if ret:
+            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 ({0})'.format(self.state))
+        return ret
+
+    def server_options(self):
+        options = "--listen_addresses='{0}' --port={1}".format(self.listen_addresses, self.port)
+        for setting, value in self.server_parameters.items():
+            options += " --{0}='{1}'".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 or 0) - 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.server_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, bootstrap=False):
+        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 = '{0}'\n""".format(self.primary_conninfo(leader.conn_url)))
+                if self.use_slots:
+                    f.write("""primary_slot_name = '{0}'\n""".format(self.name))
+            if (leader and leader.conn_url) or bootstrap:
+                for name, value in self.config.get('recovery_conf', {}).items():
+                    f.write("{0} = '{1}'\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.pop('username')
+        env = self.write_pgpass(r)
+        pc = "user={user} host={host} port={port} dbname=postgres sslmode=prefer sslcompression=1".format(**r)
+        # first run a checkpoint on a promoted master in order
+        # to make it store the new timeline ([email protected])
+        self.checkpoint(r)
+        logger.info("running pg_rewind from %s", pc)
+        pg_rewind = ['pg_rewind', '-D', self.data_dir, '--source-server', pc]
+        try:
+            ret = subprocess.call(pg_rewind, env=env) == 0
+        except OSError:
+            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.decode('utf-8').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")
+        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')
+        return result
+
+    def single_user_mode(self, command=None, options=None):
+        """ 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, val in sorted((options or {}).items()):
+            cmd.extend(['-c', '{0}={1}'.format(opt, val)])
+        # 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:
+            if command:
+                p.communicate('{0}\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')
+        try:
+            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 OSError:
+                    logger.exception("Unable to remove %s", path)
+        except OSError:
+            logger.exception("Unable to list %s", status_dir)
+
+    def follow(self, leader, recovery=False):
+        if self.check_recovery_conf(leader) and not recovery:
+            return True
+
+        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 leader and self._need_rewind:  # 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.update({'archive_mode': 'on', '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
+        else:  # do not rewind until the leader becomes available
+            ret = self.restart()
+        if change_role and ret:
+            self.call_nowait(ACTION_ON_ROLE_CHANGE)
+        return ret
+
+    def save_configuration_files(self):
+        """
+            copy postgresql.conf to postgresql.conf.backup to be able to retrive configuration files
+            - originally stored as symlinks, those are normally skipped by pg_basebackup
+            - in case of WAL-E basebackup (see http://comments.gmane.org/gmane.comp.db.postgresql.wal-e/239)
+        """
+        try:
+            for f in self.configuration_to_save:
+                if os.path.isfile(f):
+                    shutil.copy(f, f + '.backup')
+        except IOError:
+            logger.exception('unable to create backup copies of configuration files')
+
+    def restore_configuration_files(self):
+        """ restore a previously saved postgresql.conf """
+        try:
+            for f in self.configuration_to_save:
+                if not os.path.isfile(f) and os.path.isfile(f + '.backup'):
+                    shutil.copy(f + '.backup', f)
+        except IOError:
+            logger.exception('unable to restore configuration files from 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 create_or_update_role(self, name, password, options):
+        self.query("""DO $$
+BEGIN
+    SET local synchronous_commit = 'local';
+    PERFORM * FROM pg_authid WHERE rolname = %s;
+    IF FOUND THEN
+        ALTER ROLE "{0}" WITH LOGIN {1} PASSWORD %s;
+    ELSE
+        CREATE ROLE "{0}" WITH LOGIN {1} PASSWORD %s;
+    END IF;
+END;
+$$""".format(name, options), name, password, password)
+
+    def create_replication_user(self):
+        self.create_or_update_role(self.replication['username'], self.replication['password'], 'REPLICATION')
+
+    def create_connection_user(self):
+        if self.admin:
+            self.create_or_update_role(self.admin['username'], self.admin['password'], 'CREATEDB CREATEROLE')
+
+    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()
+                # if the replicatefrom tag is set on the member - we should not create the replication slot for it on
+                # the current master, because that member would replicate from elsewhere. We still create the slot if
+                # the replicatefrom destination member is currently not a member of the cluster (fallback to the
+                # master), or if replicatefrom destination member happens to be the current master
+                if self.role == 'master':
+                    slots = [m.name for m in cluster.members if m.name != self.name and
+                             (m.replicatefrom is None or m.replicatefrom == self.name or
+                              not cluster.has_member(m.replicatefrom))]
+                else:
+                    # only manage slots for replicas that replicate from this one, except for the leader among them
+                    slots = [m.name for m in cluster.members if m.replicatefrom == self.name and
+                             m.name != cluster.leader.name]
+                # 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 psycopg2.Error:
+                logger.exception('Exception when changing replication slots')
+
+    def last_operation(self):
+        return str(self.xlog_position())
+
+    def bootstrap(self, cluster_initialized=False, clone_member=None):
+        """
+            Populate PostgreSQL data directory by doing one of the following:
+             - create with initdb if there is no master.
+             - initialize the replica from an existing member (master or replica)
+             - initialize the replica using the replica creation method that
+               works without the replication connection (i.e. restore from on-disk
+               base backup)
+
+            The choice between the last 2 is triggered by the initialize flag.
+            We should never try to initdb an already initialized cluster, nor
+            try to bootstrap the cluster that lacks the initialize key using the
+            master-less replica creation method (in the latter case, there is
+            no clear inidicator of the moment we should abandon our attempts and
+            swich to initdb).
+
+            Failure during initdb always leads to an exception, since there is
+            no point in continuing if initdb fails. For the rest of the cases,
+            the function returns False in order to inidicate a failed attempt
+            that should be retried in the future.
+        """
+        ret = False
+        if not (cluster_initialized or clone_member):
+            ret = self.initialize() and self.start()
+            if ret:
+                self.create_replication_user()
+                self.create_connection_user()
+            else:
+                raise PostgresException("Could not bootstrap master PostgreSQL")
+        else:
+            if self.sync_replica(clone_member):
+                self.restore_configuration_files()
+                self.write_recovery_conf(clone_member, True)
+                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 OSError:
+                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 (IOError, OSError):
+            logger.exception('Could not remove data directory %s', self.data_dir)
+            self.move_data_directory()
+
+    def basebackup(self, clone_member, env):
+        # creates a replica data dir using pg_basebackup.
+        # this is the default, built-in create_replica_method
+        # tries twice, then returns failure (as 1)
+        # uses "stream" as the xlog-method to avoid sync issues
+        master_connection = clone_member.conn_url
+        maxfailures = 2
+        ret = 1
+        for bbfailures in range(0, maxfailures):
+            try:
+                ret = subprocess.call(['pg_basebackup', '--pgdata=' + self.data_dir,
+                                       '--xlog-method=stream', "--dbname=" + master_connection], env=env)
+                if ret == 0:
+                    break
+
+            except Exception as e:
+                logger.error('Error when fetching backup with pg_basebackup: {0}'.format(e))
+
+            if bbfailures < maxfailures - 1:
+                logger.error('Trying again in 5 seconds')
+                time.sleep(5)
+
+        return ret

patroni/scripts/aws.py

@@ -1,4 +1,4 @@
-#!/usr/bin/python
+#!/usr/bin/env python
 
 import logging
 import requests
@@ -9,7 +9,7 @@ import boto.ec2
 logger = logging.getLogger(__name__)
 
 
-class AWSConnection:
+class AWSConnection(object):
     def __init__(self, cluster_name):
         self.available = False
         self.cluster_name = cluster_name if cluster_name is not None else 'unknown'
@@ -56,7 +56,7 @@ class AWSConnection:
             conn = boto.ec2.connect_to_region(self.region)
             conn.create_tags([self.instance_id], tags)
         except Exception as e:
-            logger.info("could not set tags for EC2 instance {}: {}".format(self.instance_id, e))
+            logger.info("could not set tags for EC2 instance %s: %s", self.instance_id, e)
             return False
         return True
 
@@ -65,8 +65,12 @@ class AWSConnection:
         return self._tag_ebs(new_role) and ret
 
 
-if __name__ == '__main__':
+def main():
+    logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO)
     if len(sys.argv) == 4 and sys.argv[1] in ('on_start', 'on_stop', 'on_role_change'):
         AWSConnection(cluster_name=sys.argv[3]).on_role_change(sys.argv[2])
     else:
         sys.exit("Usage: {0} action role name".format(sys.argv[0]))
+
+if __name__ == '__main__':
+    main()

patroni/scripts/wale_restore.py

@@ -0,0 +1,169 @@
+#!/usr/bin/env python
+
+# sample script to clone new replicas using WAL-E restore
+# falls back to pg_basebackup if WAL-E restore fails, or if
+# WAL-E backup is too far behind
+# note that pg_basebackup still expects to use restore from
+# WAL-E for transaction logs
+
+# theoretically should work with SWIFT, but not tested on it
+
+# arguments are:
+#   - cluster scope
+#   - cluster role
+#   - master connection string
+#   - number of retries
+#   - envdir for the WALE env
+# - WALE_BACKUP_THRESHOLD_MEGABYTES if WAL amount is above that - use pg_basebackup
+# - WALE_BACKUP_THRESHOLD_PERCENTAGE if WAL size exceeds a certain percentage of the
+
+# this script depends on an envdir defining the S3 bucket (or SWIFT dir),and login
+# credentials per WALE Documentation.
+
+# currently also requires that you configure the restore_command to use wal_e, example:
+#       recovery_conf:
+#               restore_command: envdir /etc/wal-e.d/env wal-e wal-fetch "%f" "%p" -p 1
+
+from collections import namedtuple
+import logging
+import os
+import psycopg2
+import subprocess
+import sys
+import argparse
+
+
+if sys.hexversion >= 0x03000000:
+    long = int
+
+logger = logging.getLogger(__name__)
+
+
+class WALERestore(object):
+
+    def __init__(self, scope, datadir, connstring, env_dir, threshold_mb, threshold_pct, use_iam, no_master):
+        self.scope = scope
+        self.master_connection = connstring
+        self.data_dir = datadir
+        self.wal_e = namedtuple('wale', 'dir,threshold_mb,threshold_pct,iam_string,cmd')
+        self.wal_e.dir = env_dir
+        self.wal_e.threshold_mb = threshold_mb
+        self.wal_e.threshold_pct = threshold_pct
+        self.wal_e.iam_string = ' --aws-instance-profile ' if use_iam == 1 else ''
+        self.no_master = no_master
+        self.wal_e.cmd = 'envdir {0} wal-e {1} '.format(self.wal_e.dir, self.wal_e.iam_string)
+        self.init_error = (not os.path.exists(self.wal_e.dir))
+
+    def run(self):
+        """ creates a new replica using WAL-E """
+        if not self.init_error and self.should_use_s3_to_create_replica():
+            return self.create_replica_with_s3()
+        return 2
+
+    def should_use_s3_to_create_replica(self):
+        """ determine whether it makes sense to use S3 and not pg_basebackup """
+
+        threshold_megabytes = self.wal_e.threshold_mb
+        threshold_backup_size_percentage = self.wal_e.threshold_pct
+
+        try:
+            latest_backup = subprocess.check_output(self.wal_e.cmd.split() + ['backup-list', '--detail', 'LATEST'])
+            # name    last_modified   expanded_size_bytes wal_segment_backup_start    wal_segment_offset_backup_start
+            #                                                   wal_segment_backup_stop wal_segment_offset_backup_stop
+            # base_00000001000000000000007F_00000040  2015-05-18T10:13:25.000Z
+            # 20310671    00000001000000000000007F    00000040
+            # 00000001000000000000007F    00000240
+            backup_strings = latest_backup.splitlines() if latest_backup else ()
+            if len(backup_strings) != 2:
+                return False
+
+            names = backup_strings[0].split()
+            vals = backup_strings[1].split()
+            if (len(names) != len(vals)) or (len(names) != 7):
+                return False
+
+            backup_info = dict(zip(names, vals))
+        except subprocess.CalledProcessError as e:
+            logger.error("could not query wal-e latest backup: {}".format(e))
+            return False
+
+        try:
+            backup_size = backup_info['expanded_size_bytes']
+            backup_start_segment = backup_info['wal_segment_backup_start']
+            backup_start_offset = backup_info['wal_segment_offset_backup_start']
+        except Exception as e:
+            logger.error("unable to get some of WALE backup parameters: {}".format(e))
+            return False
+
+        # WAL filename is XXXXXXXXYYYYYYYY000000ZZ, where X - timeline, Y - LSN logical log file,
+        # ZZ - 2 high digits of LSN offset. The rest of the offset is the provided decimal offset,
+        # that we have to convert to hex and 'prepend' to the high offset digits.
+
+        lsn_segment = backup_start_segment[8:16]
+        # first 2 characters of the result are 0x and the last one is L
+        lsn_offset = hex((long(backup_start_segment[16:32], 16) << 24) + long(backup_start_offset))[2:-1]
+
+        # construct the LSN from the segment and offset
+        backup_start_lsn = '{0}/{1}'.format(lsn_segment, lsn_offset)
+
+        diff_in_bytes = long(backup_size)
+        if not self.no_master:
+            try:
+                # get the difference in bytes between the current WAL location and the backup start offset
+                with psycopg2.connect(self.master_connection) as con:
+                    con.autocommit = True
+                    with con.cursor() as cur:
+                        cur.execute("SELECT pg_xlog_location_diff(pg_current_xlog_location(), %s)", (backup_start_lsn,))
+                        diff_in_bytes = long(cur.fetchone()[0])
+            except psycopg2.Error as e:
+                logger.error('could not determine difference with the master location: %s', e)
+                return False
+        else:
+            # always try to use WAL-E if base backup is available
+            diff_in_bytes = 0
+
+        # if the size of the accumulated WAL segments is more than a certan percentage of the backup size
+        # or exceeds the pre-determined size - pg_basebackup is chosen instead.
+        return (diff_in_bytes < long(threshold_megabytes) * 1048576) and\
+            (diff_in_bytes < long(backup_size) * float(threshold_backup_size_percentage) / 100)
+
+    def create_replica_with_s3(self):
+        # if we're set up, restore the replica using fetch latest
+        try:
+            ret = subprocess.call(self.wal_e.cmd.split() + ['backup-fetch', '{}'.format(self.data_dir), 'LATEST'])
+        except Exception as e:
+            logger.error('Error when fetching backup with WAL-E: {0}'.format(e))
+            return 1
+
+        return ret
+
+
+def main():
+    logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO)
+    parser = argparse.ArgumentParser(description='Script to image replicas using WAL-E')
+    parser.add_argument('--scope', required=True)
+    parser.add_argument('--role', required=False)
+    parser.add_argument('--datadir', required=True)
+    parser.add_argument('--connstring', required=True)
+    parser.add_argument('--retries', type=int, default=1)
+    parser.add_argument('--envdir', required=True)
+    parser.add_argument('--threshold_megabytes', type=int, default=10240)
+    parser.add_argument('--threshold_backup_size_percentage', type=int, default=30)
+    parser.add_argument('--use_iam', type=int, default=0)
+    parser.add_argument('--no_master', type=int, default=0)
+    args = parser.parse_args()
+
+    # retry cloning in a loop
+    for _ in range(0, args.retries + 1):
+        restore = WALERestore(scope=args.scope, datadir=args.datadir, connstring=args.connstring,
+                              env_dir=args.envdir, threshold_mb=args.threshold_megabytes,
+                              threshold_pct=args.threshold_backup_size_percentage, use_iam=args.use_iam,
+                              no_master=args.no_master)
+        ret = restore.run()
+        if ret == 0:
+            break
+
+    sys.exit(ret)
+
+if __name__ == '__main__':
+    main()

patroni/utils.py

@@ -0,0 +1,150 @@
+import datetime
+import os
+import random
+import signal
+import sys
+import time
+import pytz
+import dateutil.parser
+
+from patroni.exceptions import PatroniException
+
+__ignore_sigterm = False
+__interrupted_sleep = False
+__reap_children = False
+
+
+def calculate_ttl(expiration):
+    """
+    >>> calculate_ttl(None)
+    >>> calculate_ttl('2015-06-10 12:56:30.552539016Z') < 0
+    True
+    >>> calculate_ttl('2015-06-10T12:56:30.552539016Z') < 0
+    True
+    >>> calculate_ttl('fail-06-10T12:56:30.552539016Z')
+    """
+    if not expiration:
+        return None
+    try:
+        expiration = dateutil.parser.parse(expiration)
+    except (ValueError, TypeError):
+        return None
+    now = datetime.datetime.now(pytz.utc)
+    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(object):
+
+    """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.80'

patroni/zookeeper.py

@@ -0,0 +1,287 @@
+import logging
+import random
+import requests
+import time
+
+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__)
+
+
+class ZooKeeperError(DCSError):
+    pass
+
+
+class ExhibitorEnsembleProvider(object):
+
+    TIMEOUT = 3.1
+
+    def __init__(self, hosts, port, uri_path='/exhibitor/v1/cluster/list', poll_interval=300):
+        self._exhibitor_port = port
+        self._uri_path = uri_path
+        self._poll_interval = poll_interval
+        self._exhibitors = hosts
+        self._master_exhibitors = hosts
+        self._zookeeper_hosts = ''
+        self._next_poll = None
+        while not self.poll():
+            logger.info('waiting on exhibitor')
+            sleep(5)
+
+    def poll(self):
+        if self._next_poll and self._next_poll > time.time():
+            return False
+
+        json = self._query_exhibitors(self._exhibitors)
+        if not json:
+            json = self._query_exhibitors(self._master_exhibitors)
+
+        if isinstance(json, dict) and 'servers' in json and 'port' in json:
+            self._next_poll = time.time() + self._poll_interval
+            zookeeper_hosts = ','.join([h + ':' + str(json['port']) for h in sorted(json['servers'])])
+            if self._zookeeper_hosts != zookeeper_hosts:
+                logger.info('ZooKeeper connection string has changed: %s => %s', self._zookeeper_hosts, zookeeper_hosts)
+                self._zookeeper_hosts = zookeeper_hosts
+                self._exhibitors = json['servers']
+                return True
+        return False
+
+    def _query_exhibitors(self, exhibitors):
+        random.shuffle(exhibitors)
+        for host in exhibitors:
+            uri = 'http://{0}:{1}{2}'.format(host, self._exhibitor_port, self._uri_path)
+            try:
+                response = requests.get(uri, timeout=self.TIMEOUT)
+                return response.json()
+            except RequestException:
+                pass
+        return None
+
+    @property
+    def zookeeper_hosts(self):
+        return self._zookeeper_hosts
+
+
+class ZooKeeper(AbstractDCS):
+
+    def __init__(self, name, config):
+        super(ZooKeeper, self).__init__(name, config)
+
+        hosts = config.get('hosts', [])
+        if isinstance(hosts, list):
+            hosts = ','.join(hosts)
+
+        self.exhibitor = None
+        if 'exhibitor' in config:
+            exhibitor = config['exhibitor']
+            interval = exhibitor.get('poll_interval', 300)
+            self.exhibitor = ExhibitorEnsembleProvider(exhibitor['hosts'], exhibitor['port'], poll_interval=interval)
+            hosts = self.exhibitor.zookeeper_hosts
+
+        self.client = KazooClient(hosts=hosts,
+                                  timeout=(config.get('session_timeout') or 30),
+                                  command_retry={
+                                      'deadline': (config.get('reconnect_timeout') or 10),
+                                      'max_delay': 1,
+                                      'max_tries': -1},
+                                  connection_retry={'max_delay': 1, 'max_tries': -1})
+        self.client.add_listener(self.session_listener)
+
+        self._my_member_data = None
+        self.fetch_cluster = True
+        self.last_leader_operation = 0
+
+        self.client.start(None)
+
+    def session_listener(self, state):
+        if state in [KazooState.SUSPENDED, KazooState.LOST]:
+            self.cluster_watcher(None)
+
+    def cluster_watcher(self, event):
+        self.fetch_cluster = True
+        self.event.set()
+
+    def get_node(self, key, watch=None):
+        try:
+            ret = self.client.get(key, watch)
+            return (ret[0].decode('utf-8'), ret[1])
+        except NoNodeError:
+            return None
+
+    @staticmethod
+    def member(name, value, znode):
+        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.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.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.get_node(self.initialize_path) or [None])[0] if self._INITIALIZE in nodes else None
+
+        # 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:
+            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:
+                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
+
+        # 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)
+
+        if self.fetch_cluster:
+            try:
+                self.client.retry(self._inner_load_cluster)
+            except:
+                logger.exception('get_cluster')
+                self.session_listener(KazooState.LOST)
+                raise ZooKeeperError('ZooKeeper in not responding properly')
+
+    def _create(self, path, value, **kwargs):
+        try:
+            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(self.leader_path, self._name, makepath=True, ephemeral=True)
+        if ret:
+            logger.info('Could not take out TTL lock')
+        return ret
+
+    def set_failover_value(self, value, index=None):
+        try:
+            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))
+        except:
+            logging.exception('set_failover_value')
+            return False
+
+    def initialize(self, create_new=True, sysid=""):
+        return self._create(self.initialize_path, sysid, makepath=True) if create_new \
+            else self.client.retry(self.client.set, self.initialize_path,  sysid.encode("utf-8"))
+
+    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.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 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.leader_optime_path
+            try:
+                self.client.retry(self.client.set, path, last_operation)
+            except NoNodeError:
+                try:
+                    self.client.retry(self.client.create, path, last_operation, makepath=True)
+                except:
+                    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):
+        self.client.restart()
+        self._my_member_data = None
+        return True
+
+    def _cancel_initialization(self):
+        node = self.get_node(self.initialize_path)
+        if node:
+            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

patronictl.py

@@ -0,0 +1,5 @@
+#!/usr/bin/env python
+from patroni.ctl import ctl
+
+if __name__ == '__main__':
+    ctl(None)

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,28 @@ postgresql:
   connect_address: 127.0.0.1:5432
   data_dir: data/postgresql0
   maximum_lag_on_failover: 1048576 # 1 megabyte in bytes
+  use_slots: True
+  pgpass: /tmp/pgpass0
+  initdb:  ## We allow the following options to be passed on to initdb
+  # - auth: authmethod
+  # - auth-host: authmethod
+  # - auth-local: authmethod
+  - encoding: UTF8
+  # - data-checksums  # When pg_rewind is needed on 9.3, this needs to be enabled
+  # - locale: locale
+  # - lc-collate: locale
+  # - lc-ctype: locale
+  # - lc-messages: locale
+  # - lc-monetary: locale
+  # - lc-numeric: locale
+  # - lc-time: locale
+  # - text-search-config: CFG
+  # - xlogdir: directory
+  # - debug
+  # - noclean
+  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,23 +63,37 @@ postgresql:
     password: rep-pass
     network:  127.0.0.1/32
   superuser:
+    username: postgres
     password: zalando
   admin:
     username: admin
     password: admin
-  wal_e:
-    env_dir: /home/postgres/etc/wal-e.d/env
-    threshold_megabytes: 10240
-    threshold_backup_size_percentage: 30
-  restore: scripts/restore.py
+  create_replica_method:
+    - basebackup
+#   - wal_e
+# commented-out example for wal-e provisioning
+  #wal_e:
+    #command: /patroni/scripts/wale_restore.py
+    #env_dir: /etc/wal-e.d/env
+    #threshold_megabytes: 10240
+    #threshold_backup_size_percentage: 30
+    #retries: 2
+    #use_iam: 1
   #recovery_conf:
-    #restore_command: cp ../wal_archive/%f %p
+    #restore_command: envdir /etc/wal-e.d/env wal-e wal-fetch "%f" "%p" -p 1
+  recovery_conf:
+    restore_command: cp ../wal_archive/%f %p
   parameters:
     archive_mode: "on"
     wal_level: hot_standby
-    archive_command: mkdir -p ../wal_archive && cp %p ../wal_archive/%f
-    max_wal_senders: 5
+    archive_command: mkdir -p ../wal_archive && test ! -f ../wal_archive/%f && cp %p ../wal_archive/%f
+    max_wal_senders: 10
     wal_keep_segments: 8
     archive_timeout: 1800s
-    max_replication_slots: 5
+    max_replication_slots: 10
     hot_standby: "on"
+    wal_log_hints: "on"
+tags:
+    nofailover: False
+    noloadbalance: False
+    clonefrom: False

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,28 @@ postgresql:
   connect_address: 127.0.0.1:5433
   data_dir: data/postgresql1
   maximum_lag_on_failover: 1048576 # 1 megabyte in bytes
+  use_slots: True
+  pgpass: /tmp/pgpass1
+  initdb:  ## We allow the following options to be passed on to initdb
+  # - auth: authmethod
+  # - auth-host: authmethod
+  # - auth-local: authmethod
+  - encoding: UTF8
+  # - data-checksums  # When pg_rewind is needed on 9.3, this needs to be enabled
+  # - locale: locale
+  # - lc-collate: locale
+  # - lc-ctype: locale
+  # - lc-messages: locale
+  # - lc-monetary: locale
+  # - lc-numeric: locale
+  # - lc-time: locale
+  # - text-search-config: CFG
+  # - xlogdir: directory
+  # - debug
+  # - noclean
+  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,23 +63,38 @@ postgresql:
     password: rep-pass
     network: 127.0.0.1/32
   superuser:
+    username: postgres
     password: zalando
   admin:
     username: admin
     password: admin
+# commented-out example for wal-e provisioning
+  create_replica_method:
+    - basebackup
+#   - wal_e
+# commented-out example for wal-e provisioning
+  #wal_e:
+    #command: /patroni/scripts/wale_restore.py
+    #env_dir: /home/postgres/etc/wal-e.d/env
+    #threshold_megabytes: 10240
+    #threshold_backup_size_percentage: 30
+    #retries: 2
+    #use_iam: 1
   #recovery_conf:
-    #restore_command: cp ../wal_archive/%f %p
-  wal_e:
-    env_dir: /home/postgres/etc/wal-e.d/env
-    threshold_megabytes: 10240
-    threshold_backup_size_percentage: 30
-  restore: scripts/restore.py
+    #restore_command: envdir /etc/wal-e.d/env wal-e wal-fetch "%f" "%p" -p 1
+  recovery_conf:
+    restore_command: cp ../wal_archive/%f %p
   parameters:
     archive_mode: "on"
     wal_level: hot_standby
-    archive_command: mkdir -p ../wal_archive && cp %p ../wal_archive/%f
-    max_wal_senders: 5
+    archive_command: mkdir -p ../wal_archive && test ! -f ../wal_archive/%f && cp %p ../wal_archive/%f
+    max_wal_senders: 10
     wal_keep_segments: 8
     archive_timeout: 1800s
-    max_replication_slots: 5
+    max_replication_slots: 10
     hot_standby: "on"
+    wal_log_hints: "on"
+tags:
+    nofailover: False
+    noloadbalance: False
+    clonefrom: False

postgres2.yml

@@ -0,0 +1,101 @@
+ttl: &ttl 30
+loop_wait: &loop_wait 10
+scope: &scope batman
+restapi:
+  listen: 127.0.0.1:8010
+  connect_address: 127.0.0.1:8010
+  auth: 'username:password'
+#  certfile: /etc/ssl/certs/ssl-cert-snakeoil.pem
+#  keyfile: /etc/ssl/private/ssl-cert-snakeoil.key
+etcd:
+  scope: *scope
+  ttl: *ttl
+  host: 127.0.0.1:4001
+  #discovery_srv: my-etcd.domain
+#zookeeper:
+#  scope: *scope
+#  session_timeout: *ttl
+#  reconnect_timeout: *loop_wait
+#  hosts:
+#    - 127.0.0.1:2181
+#    - 127.0.0.2:2181
+#  exhibitor:
+#    poll_interval: 300
+#    port: 8181
+#    hosts:
+#      - host1
+#      - host2
+#      - host3
+postgresql:
+  name: postgresql2
+  scope: *scope
+  listen: 127.0.0.1:5434
+  connect_address: 127.0.0.1:5434
+  data_dir: data/postgresql2
+  maximum_lag_on_failover: 1048576 # 1 megabyte in bytes
+  use_slots: True
+  pgpass: /tmp/pgpass2
+  initdb:  ## We allow the following options to be passed on to initdb
+  # - auth: authmethod
+  # - auth-host: authmethod
+  # - auth-local: authmethod
+  - encoding: UTF8
+  # - data-checksums  # When pg_rewind is needed on 9.3, this needs to be enabled
+  # - locale: locale
+  # - lc-collate: locale
+  # - lc-ctype: locale
+  # - lc-messages: locale
+  # - lc-monetary: locale
+  # - lc-numeric: locale
+  # - lc-time: locale
+  # - text-search-config: CFG
+  # - xlogdir: directory
+  # - debug
+  # - noclean
+  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
+  replication:
+    username: replicator
+    password: rep-pass
+    network: 127.0.0.1/32
+  superuser:
+    username: postgres
+    password: zalando
+  admin:
+    username: admin
+    password: admin
+# commented-out example for wal-e provisioning
+  create_replica_method:
+    - basebackup
+#   - wal_e
+# commented-out example for wal-e provisioning
+  #wal_e:
+    #command: /patroni/scripts/wale_restore.py
+    #env_dir: /home/postgres/etc/wal-e.d/env
+    #threshold_megabytes: 10240
+    #threshold_backup_size_percentage: 30
+    #retries: 2
+    #use_iam: 1
+  #recovery_conf:
+    #restore_command: envdir /etc/wal-e.d/env wal-e wal-fetch "%f" "%p" -p 1
+  recovery_conf:
+    restore_command: cp ../wal_archive/%f %p
+  parameters:
+    archive_mode: "on"
+    wal_level: hot_standby
+    archive_command: mkdir -p ../wal_archive && test ! -f ../wal_archive/%f && cp %p ../wal_archive/%f
+    max_wal_senders: 10
+    wal_keep_segments: 8
+    archive_timeout: 1800s
+    max_replication_slots: 10
+    hot_standby: "on"
+    wal_log_hints: "on"
+tags:
+    nofailover: False
+    noloadbalance: False
+    clonefrom: False
+    replicatefrom: postgresql1

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

requirements-py2.txt

@@ -1,9 +0,0 @@
-boto
-dnspython
-mock
-psycopg2
-PyYAML
-requests
-six >= 1.7
-kazoo>=2.2.1
-python-etcd>=0.4.1

requirements-py3.txt

@@ -1,9 +0,0 @@
-boto
-mock
-dnspython3
-psycopg2
-PyYAML
-requests
-six
-kazoo>=2.2.1
-python-etcd>=0.4.1

requirements.txt

@@ -0,0 +1,11 @@
+boto
+psycopg2>=2.6.1
+PyYAML
+requests
+six >= 1.7
+kazoo>=2.2.1
+python-etcd==0.4.3
+click>=4.1
+prettytable>=0.7
+tzlocal
+python-dateutil

scripts/restore.py

@@ -1,216 +0,0 @@
-#!/usr/bin/python
-# arguments are:
-#   - cluster scope
-#   - cluster role
-#   - master connection string
-
-# for the AWS, the folliowing environment variables should be defined:
-# - WALE_ENV_DIR: directory where WAL-E environment is kept
-# - WAL_S3_BUCKET: a name of the S3 bucket for WAL-E
-# - WALE_BACKUP_THRESHOLD_MEGABYTES if WAL amount is above that - use pg_basebackup
-# - WALE_BACKUP_THRESHOLD_PERCENTAGE if WAL size exceeds a certain percentage of the
-#   latest backup size
-from collections import namedtuple
-import logging
-import os
-import psycopg2
-import subprocess
-import sys
-
-
-if sys.hexversion >= 0x03000000:
-    long = int
-
-logger = logging.getLogger(__name__)
-
-
-class Restore(object):
-
-    def __init__(self, scope, role, datadir, connstring, env=None):
-        self.scope = scope
-        self.role = role
-        self.master_connection = Restore.parse_connstring(connstring)
-        self.data_dir = datadir
-        self.env = os.environ.copy() if not env else env
-
-    @staticmethod
-    def parse_connstring(connstring):
-        # the connection string is in the form host= port= user=
-        # return the dictionary with all components as separare keys
-        result = {}
-        if connstring:
-            for x in connstring.split():
-                if x and '=' in x:
-                    key, val = x.split('=')
-                result[key.strip()] = val.strip()
-        return result
-
-    def setup(self):
-        pass
-
-    def replica_method(self):
-        return self.create_replica_with_pg_basebackup
-
-    def replica_fallback_method(self):
-        return None
-
-    def run(self):
-        """ creates a new replica using either pg_basebackup or WAL-E """
-        method_fn = self.replica_method()
-        ret = method_fn() if method_fn else 1
-        if ret != 0 and self.replica_fallback_method() is not None:
-            ret = (self.replica_fallback_method())()
-        return ret
-
-    def create_replica_with_pg_basebackup(self):
-        try:
-            ret = subprocess.call(['pg_basebackup', '-R', '-D',
-                                   self.data_dir, '--host=' + self.master_connection['host'],
-                                   '--port=' + str(self.master_connection['port']),
-                                   '-U', self.master_connection['user']],
-                                  env=self.env)
-        except Exception as e:
-            logger.error('Error when fetching backup with pg_basebackup: {0}'.format(e))
-            return 1
-        return ret
-
-
-class WALERestore(Restore):
-
-    def __init__(self, scope, role, datadir, connstring, env=None):
-        super(WALERestore, self).__init__(scope, role, datadir, connstring, env)
-        # check the environment variables
-        self.init_error = False
-
-    def setup(self):
-        if (self.env.get('WAL_S3_BUCKET') and
-            self.env.get('WALE_BACKUP_THRESHOLD_PERCENTAGE') and
-           self.env.get('WALE_BACKUP_THRESHOLD_MEGABYTES')) is None:
-                self.init_error = True
-        else:
-            self.wal_e = namedtuple('WALE',
-                                    'threshold_megabytes threshold_backup_size_percentage s3_bucket cmd dir env_file')
-
-            self.wal_e.dir = self.env.get('WALE_ENV_DIR', '/home/postgres/etc/wal-e.d/env')
-            self.wal_e.env_file = os.path.join(self.wal_e.dir, 'WALE_S3_PREFIX')
-
-            self.wal_e.cmd = 'envdir {} wal-e --aws-instance-profile '.\
-                format(self.wal_e.dir)
-            self.wal_e.s3_bucket = self.env['WAL_S3_BUCKET']
-            self.wal_e.threshold_megabytes = self.env['WALE_BACKUP_THRESHOLD_MEGABYTES']
-            self.wal_e.threshold_backup_size_percentage = self.env['WALE_BACKUP_THRESHOLD_PERCENTAGE']
-
-            # check that the env file exists, create it otherwise
-            try:
-                if not os.path.exists(self.wal_e.dir):
-                    os.makedirs(self.wal_e.dir)
-                # if this is a directory - make sure we have full access there
-                elif not (os.path.isdir(self.wal_e.dir) and os.access(self.wal_e.dir, os.R_OK | os.W_OK | os.X_OK)):
-                    logger.error("Unable to access {} or not a directory".format(self.wal_e.dir))
-                    self.init_error = True
-                # if WAL_S3_PREFIX is not there - create it and write the full path to bucket
-                if not self.init_error and not os.path.exists(self.wal_e.env_file):
-                    with open(self.wal_e.env_file, 'w') as f:
-                        f.write("s3://{0}/spilo/{1}/wal/\n".format(self.wal_e.s3_bucket, self.scope))
-
-            except (os.error, IOError) as e:
-                logger.error("{0}: WAL-e archiving is disabled".format(e))
-                self.init_error = True
-
-    def replica_method(self):
-        if self.should_use_s3_to_create_replica():
-            return self.create_replica_with_s3
-        return None
-
-    def replica_fallback_method(self):
-        return self.create_replica_with_pg_basebackup
-
-    def should_use_s3_to_create_replica(self):
-        """ determine whether it makes sense to use S3 and not pg_basebackup """
-        if self.init_error:
-            return False
-
-        threshold_megabytes = self.wal_e.threshold_megabytes
-        threshold_backup_size_percentage = self.wal_e.threshold_backup_size_percentage
-
-        try:
-            latest_backup = subprocess.check_output(self.wal_e.cmd.split() + ['backup-list', '--detail', 'LATEST'],
-                                                    env=self.env)
-            # name    last_modified   expanded_size_bytes wal_segment_backup_start    wal_segment_offset_backup_start
-            #                                                   wal_segment_backup_stop wal_segment_offset_backup_stop
-            # base_00000001000000000000007F_00000040  2015-05-18T10:13:25.000Z
-            # 20310671    00000001000000000000007F    00000040
-            # 00000001000000000000007F    00000240
-            backup_strings = latest_backup.splitlines() if latest_backup else ()
-            if len(backup_strings) != 2:
-                return False
-
-            names = backup_strings[0].split()
-            vals = backup_strings[1].split()
-            if (len(names) != len(vals)) or (len(names) != 7):
-                return False
-
-            backup_info = dict(zip(names, vals))
-        except subprocess.CalledProcessError as e:
-            logger.error("could not query wal-e latest backup: {}".format(e))
-            return False
-
-        try:
-            backup_size = backup_info['expanded_size_bytes']
-            backup_start_segment = backup_info['wal_segment_backup_start']
-            backup_start_offset = backup_info['wal_segment_offset_backup_start']
-        except Exception as e:
-            logger.error("unable to get some of S3 backup parameters: {}".format(e))
-            return False
-
-        # WAL filename is XXXXXXXXYYYYYYYY000000ZZ, where X - timeline, Y - LSN logical log file,
-        # ZZ - 2 high digits of LSN offset. The rest of the offset is the provided decimal offset,
-        # that we have to convert to hex and 'prepend' to the high offset digits.
-
-        lsn_segment = backup_start_segment[8:16]
-        # first 2 characters of the result are 0x and the last one is L
-        lsn_offset = hex((long(backup_start_segment[16:32], 16) << 24) + long(backup_start_offset))[2:-1]
-
-        # construct the LSN from the segment and offset
-        backup_start_lsn = '{}/{}'.format(lsn_segment, lsn_offset)
-
-        conn = None
-        cursor = None
-        diff_in_bytes = long(backup_size)
-        try:
-            # get the difference in bytes between the current WAL location and the backup start offset
-            conn = psycopg2.connect(**(self.master_connection))
-            conn.autocommit = True
-            cursor = conn.cursor()
-            cursor.execute("SELECT pg_xlog_location_diff(pg_current_xlog_location(), %s)", (backup_start_lsn,))
-            diff_in_bytes = long(cursor.fetchone()[0])
-        except psycopg2.Error as e:
-            logger.error('could not determine difference with the master location: {}'.format(e))
-            return False
-        finally:
-            cursor and cursor.close()
-            conn and conn.close()
-
-        # if the size of the accumulated WAL segments is more than a certan percentage of the backup size
-        # or exceeds the pre-determined size - pg_basebackup is chosen instead.
-        return (diff_in_bytes < long(threshold_megabytes) * 1048576) and\
-               (diff_in_bytes < long(backup_size) * float(threshold_backup_size_percentage) / 100)
-
-    def create_replica_with_s3(self):
-        if self.init_error:
-            return 1
-        try:
-            ret = subprocess.call(self.wal_e.cmd + ' backup-fetch {} LATEST'.format(self.data_dir), env=self.env)
-        except Exception as e:
-            logger.error('Error when fetching backup with WAL-E: {0}'.format(e))
-            return 1
-        return ret
-
-
-if __name__ == '__main__':
-    if len(sys.argv) == 5:
-        # scope, role, datadir, connstring
-        restore = WALERestore(*(sys.argv[1:]))
-        restore.setup()
-        sys.exit(restore.run())
-    sys.exit("Usage: {0} scope role datadir connstring".format(sys.argv[0]))

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, Oleksii Kliukin, 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,15 +47,20 @@ 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',
-    'Programming Language :: Python :: 3.3',
     'Programming Language :: Python :: 3.4',
+    'Programming Language :: Python :: 3.5',
     'Programming Language :: Python :: Implementation :: CPython',
 ]
 
+CONSOLE_SCRIPTS = ['patroni = patroni:main',
+                   'patronictl = patroni.ctl:ctl',
+                   "patroni_wale_restore = patroni.scripts.wale_restore:main",
+                   "patroni_aws = patroni.scripts.aws:main"]
+
 
 class PyTest(TestCommand):
 
@@ -62,8 +76,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-report', 'term-missing']
             if self.cov_xml:
                 self.cov.extend(['--cov-report', 'xml'])
             if self.cov_html:
@@ -79,10 +92,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)
 
@@ -104,8 +116,7 @@ def setup_package():
     # Some helper variables
     version = os.getenv('GO_PIPELINE_LABEL', VERSION)
 
-    requirements = 'requirements-py2.txt' if sys.version_info[0] == 2 else 'requirements-py3.txt'
-    install_reqs = get_install_requirements(requirements)
+    install_reqs = get_install_requirements('requirements.txt')
 
     command_options = {'test': {'test_suite': ('setup.py', 'tests')}}
     if JUNIT_XML:
@@ -118,19 +129,23 @@ 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.*']),
         package_data={MAIN_PACKAGE: ["*.json"]},
         install_requires=install_reqs,
-        setup_requires=['six', 'flake8'],
+        setup_requires=['flake8'],
         cmdclass=cmdclass,
-        tests_require=['pytest-cov', 'pytest'],
+        tests_require=['mock', 'pytest-cov', 'pytest'],
         command_options=command_options,
+        entry_points={'console_scripts': CONSOLE_SCRIPTS},
     )
 
 

tests/test_api.py

@@ -1,31 +1,69 @@
 import psycopg2
 import unittest
 
-from helpers.api import RestApiHandler, RestApiServer
+from mock import Mock, patch
+from patroni.api import RestApiHandler, RestApiServer
+from patroni.dcs import Member
 from six import BytesIO as IO
-from test_postgresql import psycopg2_connect
+from six.moves import BaseHTTPServer
+from six.moves.BaseHTTPServer import BaseHTTPRequestHandler
+import socket
+from test_postgresql import psycopg2_connect, MockCursor
 
 
-def throws(*args, **kwargs):
-    raise psycopg2.OperationalError()
+class MockPostgresql(Mock):
 
+    name = 'test'
+    state = 'running'
+    role = 'master'
+    server_version = '999999'
+    scope = 'dummy'
 
-class MockPostgresql:
-
-    def connection(self):
+    @staticmethod
+    def connection():
         return psycopg2_connect()
 
-    def is_running(self):
+    @staticmethod
+    def is_running():
         return True
 
 
-class MockPatroni:
+class MockHa(Mock):
 
-    def __init__(self):
-        self.postgresql = MockPostgresql()
+    dcs = Mock()
+    state_handler = MockPostgresql()
+
+    @staticmethod
+    def schedule_restart():
+        return 'restart'
+
+    @staticmethod
+    def schedule_reinitialize():
+        return 'reinitialize'
+
+    @staticmethod
+    def restart():
+        return (True, '')
+
+    @staticmethod
+    def restart_scheduled():
+        return False
+
+    @staticmethod
+    def fetch_nodes_statuses(members):
+        return [[None, True, None, None, {}]]
 
 
-class MockRequest:
+class MockPatroni(Mock):
+
+    postgresql = MockPostgresql()
+    ha = MockHa()
+    dcs = Mock()
+    tags = {}
+    version = '0.00'
+
+
+class MockRequest(object):
 
     def __init__(self, path):
         self.path = path
@@ -36,18 +74,129 @@ 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')
+        self.assertIsNotNone(MockRestApiServer(RestApiHandler, b'GET /master'))
+
+    def test_do_OPTIONS(self):
+        self.assertIsNotNone(MockRestApiServer(RestApiHandler, b'OPTIONS / HTTP/1.0'))
+
+        with patch.object(BaseHTTPRequestHandler, 'handle_one_request') as mock_handle_request:
+            mock_handle_request.side_effect = socket.error("foo")
+            MockRestApiServer(RestApiHandler, b'OPTIONS / HTTP/1.0')
+
+        # make sure socket.error gets propagated via wfile object in finalize()
+        with patch.object(MockRequest, 'makefile') as makefile:
+            makefile.return_value.closed = False
+            makefile.return_value.readline.side_effect = lambda x: b"foo"
+            makefile.return_value.flush = Mock(side_effect=socket.error("foo"))
+            MockRestApiServer(RestApiHandler, b'OPTIONS / HTTP/1.0')
+
+    def test_do_GET_patroni(self):
+        self.assertIsNotNone(MockRestApiServer(RestApiHandler, b'GET /patroni'))
+
+    def test_basicauth(self):
+        self.assertIsNotNone(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'
+        self.assertIsNotNone(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'
+        self.assertIsNotNone(MockRestApiServer(RestApiHandler, request))
+
+    @patch('time.sleep', Mock())
+    def test_RestApiServer_query(self):
+        with patch.object(MockCursor, 'execute', Mock(side_effect=psycopg2.OperationalError)):
+            self.assertIsNotNone(MockRestApiServer(RestApiHandler, b'GET /patroni'))
+        with patch.object(MockPostgresql, 'connection', Mock(side_effect=psycopg2.OperationalError)):
+            self.assertIsNotNone(MockRestApiServer(RestApiHandler, b'GET /patroni'))
+
+    @patch('time.sleep', Mock())
+    @patch.object(MockHa, 'dcs')
+    def test_do_POST_failover(self, dcs):
+        cluster = dcs.get_cluster.return_value
+
+        request = b'POST /failover HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0\n' +\
+                  b'Content-Length: 0\n\n'
+        MockRestApiServer(RestApiHandler, request)
+        request = b'POST /failover HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0\n' +\
+                  b'Content-Length: 25\n\n{"leader": "postgresql1"}'
+        MockRestApiServer(RestApiHandler, request)
+        cluster.leader.name = 'postgresql1'
+        MockRestApiServer(RestApiHandler, request)
+        cluster.members = [Member(0, 'postgresql0', 30, {'api_url': 'http'})]
+        MockRestApiServer(RestApiHandler, request)
+        with patch.object(MockPatroni, 'dcs') as d:
+            cluster = d.get_cluster.return_value
+            cluster.leader.name = 'postgresql0'
+            MockRestApiServer(RestApiHandler, request)
+            cluster.leader.name = 'postgresql1'
+            cluster.failover = None
+            MockRestApiServer(RestApiHandler, request)
+            d.get_cluster = Mock(side_effect=Exception())
+            MockRestApiServer(RestApiHandler, request)
+            d.manual_failover.return_value = False
+            MockRestApiServer(RestApiHandler, request)
+        with patch.object(MockHa, 'fetch_nodes_statuses', Mock(return_value=[])):
+            MockRestApiServer(RestApiHandler, request)
+        request = b'POST /failover HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0\n' +\
+                  b'Content-Length: 50\n\n{"leader": "postgresql1", "member": "postgresql2"}'
+        MockRestApiServer(RestApiHandler, request)
+
+        # Valid future date
+        request = b'POST /failover HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0\nContent-Length: 103\n\n{"leader": ' +\
+                  b'"postgresql1", "member": "postgresql2", "scheduled_at": "6016-02-15T18:13:30.568224+01:00"}'
+        MockRestApiServer(RestApiHandler, request)
+
+        # Exception: No timezone specified
+        request = b'POST /failover HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0\nContent-Length: 97\n\n{"leader": ' +\
+                  b'"postgresql1", "member": "postgresql2", "scheduled_at": "6016-02-15T18:13:30.568224"}'
+        MockRestApiServer(RestApiHandler, request)
+
+        # Exception: Scheduled in the past
+        request = b'POST /failover HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0\nContent-Length: 103\n\n{"leader": ' +\
+                  b'"postgresql1", "member": "postgresql2", "scheduled_at": "1016-02-15T18:13:30.568224+01:00"}'
+        MockRestApiServer(RestApiHandler, request)
+
+        # Invalid date
+        request = b'POST /failover HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0\nContent-Length: 103\n\n{"leader": ' +\
+                  b'"postgresql1", "member": "postgresql2", "scheduled_at": "2010-02-29T18:13:30.568224+01:00"}'
+        self.assertIsNotNone(MockRestApiServer(RestApiHandler, request))

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

@@ -1,12 +1,15 @@
-import unittest
-import requests
 import boto.ec2
+import requests
+import sys
+import unittest
+
+from mock import Mock, patch
 from collections import namedtuple
-from scripts.aws import AWSConnection
+from patroni.scripts.aws import AWSConnection, main as _main
 from requests.exceptions import RequestException
 
 
-class MockEc2Connection:
+class MockEc2Connection(object):
 
     def __init__(self, error=False):
         self.error = error
@@ -23,7 +26,7 @@ class MockEc2Connection:
         return True
 
 
-class MockResponse:
+class MockResponse(object):
 
     def __init__(self, content):
         self.content = content
@@ -35,15 +38,6 @@ class MockResponse:
 
 class TestAWSConnection(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        super(TestAWSConnection, self).__init__(method_name)
-
-    def set_error(self):
-        self.error = True
-
-    def set_json_error(self):
-        self.json_error = True
-
     def boto_ec2_connect_to_region(self, region):
         return MockEc2Connection(self.error)
 
@@ -74,21 +68,27 @@ class TestAWSConnection(unittest.TestCase):
         self.assertTrue(self.conn.on_role_change('master'))
 
     def test_non_aws(self):
-        self.set_error()
+        self.error = True
         conn = AWSConnection('test')
         self.assertFalse(conn.aws_available())
         self.assertFalse(conn._tag_ebs('master'))
         self.assertFalse(conn._tag_ec2('master'))
 
     def test_aws_bizare_response(self):
-        self.set_json_error()
+        self.json_error = True
         conn = AWSConnection('test')
         self.assertFalse(conn.aws_available())
 
     def test_aws_tag_ebs_error(self):
-        self.set_error()
+        self.error = True
         self.assertFalse(self.conn._tag_ebs("master"))
 
     def test_aws_tag_ec2_error(self):
-        self.set_error()
+        self.error = True
         self.assertFalse(self.conn._tag_ec2("master"))
+
+    @patch('sys.exit', Mock())
+    def test_main(self):
+        self.assertIsNone(_main())
+        sys.argv = ['aws.py', 'on_start', 'replica', 'foo']
+        self.assertIsNone(_main())

tests/test_ctl.py

@@ -0,0 +1,390 @@
+import os
+import pytest
+import requests.exceptions
+import unittest
+
+from click.testing import CliRunner
+from etcd import EtcdException
+from mock import patch, Mock, MagicMock
+from patroni.ctl import ctl, members, store_config, load_config, output_members, post_patroni, get_dcs, \
+    wait_for_leader, get_all_members, get_any_member, get_cursor, query_member, configure
+from patroni.etcd import Etcd, Client
+from patroni.exceptions import PatroniCtlException
+from psycopg2 import OperationalError
+from test_etcd import etcd_read, etcd_write, requests_get, socket_getaddrinfo, MockResponse
+from test_zookeeper import MockKazooClient
+from test_ha import get_cluster_initialized_without_leader, get_cluster_initialized_with_leader, \
+    get_cluster_initialized_with_only_leader
+from test_postgresql import MockConnect, psycopg2_connect
+
+CONFIG_FILE_PATH = './test-ctl.yaml'
+
+
+def test_rw_config():
+    runner = CliRunner()
+    config = {'a': 'b'}
+    with runner.isolated_filesystem():
+        store_config(config, CONFIG_FILE_PATH + '/dummy')
+        os.remove(CONFIG_FILE_PATH + '/dummy')
+        os.rmdir(CONFIG_FILE_PATH)
+
+        with pytest.raises(Exception):
+            result = load_config(CONFIG_FILE_PATH, None)
+            assert 'Could not load configuration file' in result.output
+
+        os.mkdir(CONFIG_FILE_PATH)
+        with pytest.raises(Exception):
+            store_config(config, CONFIG_FILE_PATH)
+
+        os.rmdir(CONFIG_FILE_PATH)
+
+    store_config(config, CONFIG_FILE_PATH)
+    load_config(CONFIG_FILE_PATH, None)
+    load_config(CONFIG_FILE_PATH, '0.0.0.0')
+
+
+@patch('patroni.ctl.load_config', Mock(return_value={'dcs': {'scheme': 'etcd', 'hostname': 'localhost', 'port': 4001}}))
+class TestCtl(unittest.TestCase):
+
+    @patch('socket.getaddrinfo', socket_getaddrinfo)
+    def setUp(self):
+        self.runner = CliRunner()
+        with patch.object(Client, 'machines') as mock_machines:
+            mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
+            self.e = Etcd('foo', {'ttl': 30, 'host': 'ok:2379', 'scope': 'test'})
+            self.e.client.read = etcd_read
+            self.e.client.write = etcd_write
+            self.e.client.delete = Mock(side_effect=EtcdException)
+
+    @patch('psycopg2.connect', psycopg2_connect)
+    def test_get_cursor(self):
+        self.assertIsNone(get_cursor(get_cluster_initialized_without_leader(), role='master'))
+
+        self.assertIsNotNone(get_cursor(get_cluster_initialized_with_leader(), role='master'))
+
+        # MockCursor returns pg_is_in_recovery as false
+        self.assertIsNone(get_cursor(get_cluster_initialized_with_leader(), role='replica'))
+
+        self.assertIsNotNone(get_cursor(get_cluster_initialized_with_leader(), role='any'))
+
+    def test_output_members(self):
+        cluster = get_cluster_initialized_with_leader()
+        self.assertIsNone(output_members(cluster, name='abc', fmt='pretty'))
+        self.assertIsNone(output_members(cluster, name='abc', fmt='json'))
+        self.assertIsNone(output_members(cluster, name='abc', fmt='tsv'))
+
+    @patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_leader()))
+    @patch('patroni.etcd.Etcd.get_etcd_client', Mock(return_value=None))
+    @patch('patroni.etcd.Etcd.set_failover_value', Mock(return_value=None))
+    @patch('patroni.ctl.wait_for_leader', Mock(return_value=get_cluster_initialized_with_leader()))
+    @patch('requests.get', requests_get)
+    @patch('requests.post', requests_get)
+    @patch('patroni.ctl.post_patroni', Mock(return_value=MockResponse()))
+    def test_failover(self):
+        with patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_leader())):
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+
+y''')
+            assert 'leader' in result.output
+
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+2100-01-01T12:23:00
+y''')
+            assert result.exit_code == 0
+
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+2030-01-01T12:23:00
+y''')
+            assert result.exit_code == 0
+
+            # Aborting failover,as we anser NO to the confirmation
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+
+N''')
+            assert result.exit_code == 1
+
+            # Target and source are equal
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+leader
+
+y''')
+            assert result.exit_code == 1
+
+            # Reality is not part of this cluster
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+Reality
+
+y''')
+            assert result.exit_code == 1
+
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--force'])
+            assert 'Member' in result.output
+
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--force',
+                                              '--scheduled', '2015-01-01T12:00:00+01:00'])
+            assert result.exit_code == 0
+
+            # Invalid timestamp
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--force', '--scheduled', 'invalid'])
+            assert result.exit_code != 0
+
+            # Invalid timestamp
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--force',
+                                              '--scheduled', '2115-02-30T12:00:00+01:00'])
+            assert result.exit_code != 0
+
+            # Specifying wrong leader
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='dummy')
+            assert result.exit_code == 1
+
+        with patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_only_leader())):
+            # No members available
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+
+y''')
+            assert result.exit_code == 1
+
+        with patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_without_leader())):
+            # No master available
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+
+y''')
+            assert result.exit_code == 1
+
+        with patch('patroni.ctl.post_patroni', Mock(side_effect=Exception())):
+            # Non-responding patroni
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+
+y''')
+            assert 'falling back to DCS' in result.output
+
+        mocked = Mock()
+        mocked.return_value.status_code = 500
+        with patch('patroni.ctl.post_patroni', Mock(return_value=mocked)):
+            result = self.runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+
+y''')
+            assert 'Failover failed' in result.output
+
+    @patch('patroni.zookeeper.KazooClient', MockKazooClient)
+    @patch('requests.get', requests_get)
+    def test_get_dcs(self):
+        self.assertIsNotNone(get_dcs({'dcs': {'scheme': 'zookeeper', 'hostname': 'foo', 'port': 2181}}, 'dummy'))
+        self.assertIsNotNone(get_dcs({'dcs': {'scheme': 'exhibitor', 'hostname': 'exhibitor', 'port': 8181}}, 'dummy'))
+        self.assertRaises(PatroniCtlException, get_dcs, {'scheme': 'dummy'}, 'dummy')
+
+    @patch('psycopg2.connect', psycopg2_connect)
+    @patch('patroni.ctl.query_member', Mock(return_value=([['mock column']], None)))
+    def test_query(self):
+        with patch('patroni.ctl.get_dcs', Mock(return_value=self.e)):
+            # Mutually exclusive
+            result = self.runner.invoke(ctl, [
+                'query',
+                'alpha',
+                '--member',
+                'abc',
+                '--role',
+                'master',
+            ])
+            assert result.exit_code == 1
+
+            with self.runner.isolated_filesystem():
+                with open('dummy', 'w') as dummy_file:
+                    dummy_file.write('SELECT 1')
+
+                # Mutually exclusive
+                result = self.runner.invoke(ctl, [
+                    'query',
+                    'alpha',
+                    '--file',
+                    'dummy',
+                    '--command',
+                    'dummy',
+                ])
+                assert result.exit_code == 1
+
+                result = self.runner.invoke(ctl, ['query', 'alpha', '--file', 'dummy'])
+
+                os.remove('dummy')
+
+            result = self.runner.invoke(ctl, ['query', 'alpha', '--command', 'SELECT 1'])
+            assert 'mock column' in result.output
+
+            # --command or --file is mandatory
+            result = self.runner.invoke(ctl, ['query', 'alpha'])
+            assert result.exit_code == 1
+
+            result = self.runner.invoke(ctl, ['query', 'alpha', '--command', 'SELECT 1', '--username', 'root',
+                                              '--password', '--dbname', 'postgres'], input='ab\nab')
+            assert 'mock column' in result.output
+
+    @patch('patroni.ctl.get_cursor', Mock(return_value=MockConnect().cursor()))
+    def test_query_member(self):
+        rows = query_member(None, None, None, 'master', 'SELECT pg_is_in_recovery()')
+        self.assertTrue('False' in str(rows))
+
+        rows = query_member(None, None, None, 'replica', 'SELECT pg_is_in_recovery()')
+        self.assertEquals(rows, (None, None))
+
+        with patch('patroni.ctl.get_cursor', Mock(return_value=None)):
+            rows = query_member(None, None, None, None, 'SELECT pg_is_in_recovery()')
+            self.assertTrue('No connection to' in str(rows))
+
+            rows = query_member(None, None, None, 'replica', 'SELECT pg_is_in_recovery()')
+            self.assertTrue('No connection to' in str(rows))
+
+        with patch('patroni.ctl.get_cursor', Mock(side_effect=OperationalError('bla'))):
+            rows = query_member(None, None, None, 'replica', 'SELECT pg_is_in_recovery()')
+
+        with patch('test_postgresql.MockCursor.execute', Mock(side_effect=OperationalError('bla'))):
+            rows = query_member(None, None, None, 'replica', 'SELECT pg_is_in_recovery()')
+
+    @patch('patroni.dcs.AbstractDCS.get_cluster', Mock(return_value=get_cluster_initialized_with_leader()))
+    def test_dsn(self):
+        with patch('patroni.ctl.get_dcs', Mock(return_value=self.e)):
+            result = self.runner.invoke(ctl, ['dsn', 'alpha', '--dcs', '8.8.8.8'])
+            assert 'host=127.0.0.1 port=5435' in result.output
+
+            # Mutually exclusive options
+            result = self.runner.invoke(ctl, [
+                'dsn',
+                'alpha',
+                '--role',
+                'master',
+                '--member',
+                'dummy',
+            ])
+            assert result.exit_code == 1
+
+            # Non-existing member
+            result = self.runner.invoke(ctl, ['dsn', 'alpha', '--member', 'dummy'])
+            assert result.exit_code == 1
+
+    @patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_leader()))
+    @patch('patroni.etcd.Etcd.get_etcd_client', Mock(return_value=None))
+    @patch('requests.get', requests_get)
+    @patch('requests.post', requests_get)
+    def test_restart_reinit(self):
+        result = self.runner.invoke(ctl, ['restart', 'alpha', '--dcs', '8.8.8.8'], input='y')
+        assert result.exit_code == 0
+
+        result = self.runner.invoke(ctl, ['reinit', 'alpha', '--dcs', '8.8.8.8'], input='y')
+        assert result.exit_code == 1
+
+        # Aborted restart
+        result = self.runner.invoke(ctl, ['restart', 'alpha', '--dcs', '8.8.8.8'], input='N')
+        assert result.exit_code == 1
+
+        # Not a member
+        result = self.runner.invoke(ctl, [
+            'restart',
+            'alpha',
+            '--dcs',
+            '8.8.8.8',
+            'dummy',
+            '--any',
+        ], input='y')
+        assert result.exit_code == 1
+
+        with patch('requests.post', Mock(return_value=MockResponse())):
+            result = self.runner.invoke(ctl, ['restart', 'alpha', '--dcs', '8.8.8.8'], input='y')
+
+    @patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_leader()))
+    @patch('patroni.etcd.Etcd.get_etcd_client', Mock(return_value=None))
+    def test_remove(self):
+        result = self.runner.invoke(ctl, ['remove', 'alpha', '--dcs', '8.8.8.8'], input='alpha\nslave')
+        assert 'Please confirm' in result.output
+        assert 'You are about to remove all' in result.output
+        # Not typing an exact confirmation
+        assert result.exit_code == 1
+
+        # master specified does not match master of cluster
+        result = self.runner.invoke(ctl, ['remove', 'alpha', '--dcs', '8.8.8.8'], input='''alpha
+Yes I am aware
+slave''')
+        assert result.exit_code == 1
+
+        # cluster specified on cmdline does not match verification prompt
+        result = self.runner.invoke(ctl, ['remove', 'alpha', '--dcs', '8.8.8.8'], input='beta\nleader')
+        assert result.exit_code == 1
+
+        with patch('patroni.etcd.Etcd.get_cluster', get_cluster_initialized_with_leader):
+            result = self.runner.invoke(ctl, ['remove', 'alpha', '--dcs', '8.8.8.8'],
+                                        input='''alpha
+Yes I am aware
+leader''')
+            assert 'object has no attribute' in str(result.exception)
+
+        with patch('patroni.ctl.get_dcs', Mock(return_value=Mock())):
+            # Not implemented DCS
+            result = self.runner.invoke(ctl, ['remove', 'alpha', '--dcs', '8.8.8.8'], input='''alpha
+Yes I am aware
+leader''')
+            assert result.exit_code == 1
+
+    @patch('patroni.etcd.Etcd.watch', Mock(return_value=None))
+    @patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_leader()))
+    def test_wait_for_leader(self):
+        dcs = self.e
+        self.assertRaises(PatroniCtlException, wait_for_leader, dcs, 0)
+
+        cluster = wait_for_leader(dcs=dcs, timeout=2)
+        assert cluster.leader.member.name == 'leader'
+
+    def test_post_patroni(self):
+        with patch('requests.post', MagicMock(side_effect=requests.exceptions.ConnectionError('foo'))):
+            member = get_cluster_initialized_with_leader().leader.member
+            self.assertRaises(requests.exceptions.ConnectionError, post_patroni, member, 'dummy', {})
+
+    def test_ctl(self):
+        self.runner.invoke(ctl, ['list'])
+
+        result = self.runner.invoke(ctl, ['--help'])
+        assert 'Usage:' in result.output
+
+    def test_get_any_member(self):
+        self.assertIsNone(get_any_member(get_cluster_initialized_without_leader(), role='master'))
+
+        m = get_any_member(get_cluster_initialized_with_leader(), role='master')
+        self.assertEquals(m.name, 'leader')
+
+    def test_get_all_members(self):
+        self.assertEquals(list(get_all_members(get_cluster_initialized_without_leader(), role='master')), [])
+
+        r = list(get_all_members(get_cluster_initialized_with_leader(), role='master'))
+        self.assertEquals(len(r), 1)
+        self.assertEquals(r[0].name, 'leader')
+
+        r = list(get_all_members(get_cluster_initialized_with_leader(), role='replica'))
+        self.assertEquals(len(r), 1)
+        self.assertEquals(r[0].name, 'other')
+
+        self.assertEquals(len(list(get_all_members(get_cluster_initialized_without_leader(), role='replica'))), 2)
+
+    @patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_leader()))
+    @patch('patroni.etcd.Etcd.get_etcd_client', Mock(return_value=None))
+    @patch('requests.get', requests_get)
+    @patch('requests.post', requests_get)
+    def test_members(self):
+        result = self.runner.invoke(members, ['alpha'])
+        assert result.exit_code == 0
+
+    def test_configure(self):
+        result = self.runner.invoke(configure, [
+            '--dcs',
+            'abc',
+            '-c',
+            'dummy',
+            '-n',
+            'bla',
+        ])
+
+        assert result.exit_code == 0

tests/test_etcd.py

@@ -1,45 +1,49 @@
-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
+from patroni.etcd import Client, Etcd, EtcdError
+from patroni.exceptions import DCSError
 
 
-class MockResponse:
+class MockResponse(object):
 
     def __init__(self):
         self.status_code = 200
         self.content = '{}'
         self.ok = True
+        self.text = ''
 
     def json(self):
         return json.loads(self.content)
 
     @property
     def data(self):
-        return self.content
+        return self.content.encode('utf-8')
 
     @property
     def status(self):
         return self.status_code
 
+    @staticmethod
     def getheader(*args):
         return ''
 
 
-class MockPostgresql:
-    name = ''
+class MockPostgresql(Mock):
 
-    def last_operation(self):
-        return 0
+    server_version = '999999'
+    scope = 'dummy'
+
+    @staticmethod
+    def last_operation():
+        return '0'
 
 
 def requests_get(url, **kwargs):
@@ -48,11 +52,10 @@ 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}, "tags": {}}'
     elif url.endswith('/members'):
-        if url.startswith('http://error'):
-            response.content = '[{}]'
-        else:
-            response.content = members
+        response.content = '[{}]' if url.startswith('http://error') else members
     elif url.startswith('http://exhibitor'):
         response.content = '{"servers":["127.0.0.1","127.0.0.2","127.0.0.3"],"port":2181}'
     else:
@@ -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 etcd.EtcdWatchTimedOut
+    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/test/leader':
-        if kwargs.get('prevValue', None) == 'foo' or not kwargs.get('prevExist', True):
-            return True
-    raise etcd.EtcdException
-
-
-def etcd_delete(key, **kwargs):
+    if key == '/service/exists/leader':
+        raise etcd.EtcdAlreadyExist
+    if key in ['/service/test/leader', '/patroni/test/leader'] and \
+            (kwargs.get('prevValue') == 'foo' or not kwargs.get('prevExist', True)):
+        return True
     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",
@@ -90,33 +106,29 @@ def etcd_read(key, **kwargs):
                  "modifiedIndex": 20437, "createdIndex": 20437},
                 {"key": "/service/batman5/members", "dir": True, "nodes": [
                     {"key": "/service/batman5/members/postgresql1",
-                     "value": "postgres://replicator:[email protected]:5434/postgres"
-                        + "?application_name=http://127.0.0.1:8009/patroni",
+                     "value": "postgres://replicator:[email protected]:5434/postgres" +
+                        "?application_name=http://127.0.0.1:8009/patroni",
                      "expiration": "2015-05-15T09:10:59.949384522Z", "ttl": 21,
                      "modifiedIndex": 20727, "createdIndex": 20727},
                     {"key": "/service/batman5/members/postgresql0",
-                     "value": "postgres://replicator:[email protected]:5433/postgres"
-                        + "?application_name=http://127.0.0.1:8008/patroni",
+                     "value": "postgres://replicator:[email protected]:5433/postgres" +
+                        "?application_name=http://127.0.0.1:8008/patroni",
                      "expiration": "2015-05-15T09:11:09.611860899Z", "ttl": 30,
                      "modifiedIndex": 20730, "createdIndex": 20730}],
                  "modifiedIndex": 1581, "createdIndex": 1581}], "modifiedIndex": 1581, "createdIndex": 1581}}
     return etcd.EtcdResult(**response)
 
 
-def time_sleep(_):
+class SleepException(Exception):
     pass
 
 
-def time_sleep_exception(_):
-    raise Exception()
-
-
-class MockSRV:
+class MockSRV(object):
     port = 2380
     target = '127.0.0.1'
 
 
-def dns_query(name, type):
+def dns_query(name, _):
     if name == '_etcd-server._tcp.blabla':
         return []
     elif name == '_etcd-server._tcp.exception':
@@ -127,57 +139,56 @@ 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):
+    if url == 'http://localhost:2379/timeout':
+        raise urllib3.exceptions.ReadTimeoutError(None, None, None)
     if url == 'http://localhost:2379/':
         return MockResponse()
     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'})
             self.client.http.request = http_request
+            self.client.http.request_encode_body = http_request
 
     def test_api_execute(self):
         self.client._base_uri = 'http://localhost:4001'
         self.client._machines_cache = ['http://localhost:2379']
-        self.client.api_execute('/', 'GET')
+        self.assertRaises(etcd.EtcdWatchTimedOut, self.client.api_execute, '/timeout', 'POST', params={'wait': 'true'})
+        self.client._update_machines_cache = False
+        self.client.api_execute('/', 'POST', timeout=0)
+        self.client._update_machines_cache = False
+        self.client._base_uri = 'http://localhost:4001'
+        self.client._machines_cache = []
+        self.assertRaises(etcd.EtcdConnectionFailed, self.client.api_execute, '/', 'GET')
+        self.assertTrue(self.client._update_machines_cache)
+        self.assertRaises(etcd.EtcdException, self.client.api_execute, '/', 'GET')
+        self.assertRaises(etcd.EtcdException, self.client.api_execute, '/', '')
+        self.assertRaises(ValueError, self.client.api_execute, '', '')
 
     def test_get_srv_record(self):
         self.assertEquals(self.client.get_srv_record('blabla'), [])
         self.assertEquals(self.client.get_srv_record('exception'), [])
 
     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 +197,26 @@ 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 = Etcd('foo', {'namespace': '/patroni/', '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())
 
+    def test_base_path(self):
+        self.assertEquals(self.etcd._base_path, '/patroni/test')
+
+    @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)
@@ -212,11 +224,8 @@ class TestEtcd(unittest.TestCase):
         cluster = self.etcd.get_cluster()
         self.assertIsInstance(cluster, Cluster)
         self.assertIsNone(cluster.leader)
-
-    def test_current_leader(self):
-        self.assertIsInstance(self.etcd.current_leader(), Member)
         self.etcd._base_path = '/service/noleader'
-        self.assertIsNone(self.etcd.current_leader())
+        self.assertRaises(EtcdError, self.etcd.get_cluster)
 
     def test_touch_member(self):
         self.assertFalse(self.etcd.touch_member('', ''))
@@ -224,12 +233,36 @@ 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)
+
+    @patch('patroni.etcd.Etcd.retry', Mock(side_effect=AttributeError("foo")))
+    def test_other_exceptions(self):
+        self.assertRaises(EtcdError, self.etcd.cancel_initialization)

tests/test_ha.py

@@ -1,10 +1,15 @@
 import unittest
+import datetime
+import pytz
 
-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 etcd import EtcdException
+from mock import Mock, MagicMock, patch
+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 patroni.postgresql import Postgresql
+from test_etcd import socket_getaddrinfo, etcd_read, etcd_write, requests_get
 
 
 def true(*args, **kwargs):
@@ -15,86 +20,125 @@ def false(*args, **kwargs):
     return False
 
 
-class MockPostgresql:
-
-    def __init__(self):
-        self.name = 'postgresql0'
-        self.is_promoted = False
-
-    def is_healthy(self):
-        return True
-
-    def write_recovery_conf(self, _):
-        return True
-
-    def start(self):
-        return True
-
-    def is_healthiest_node(self, members):
-        return True
-
-    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 last_operation(self):
-        return 0
+def get_cluster(initialize, leader, members, failover):
+    return Cluster(initialize, leader, 10, members, failover)
 
 
-def nop(*args, **kwargs):
-    pass
+def get_cluster_not_initialized_without_leader():
+    return get_cluster(None, None, [], None)
 
 
-def dead_etcd():
-    raise DCSError('Etcd is not responding properly')
+def get_cluster_initialized_without_leader(leader=False, failover=None):
+    m1 = Member(0, 'leader', 28, {'conn_url': 'postgres://replicator:[email protected]:5435/postgres',
+                                  'api_url': 'http://127.0.0.1:8008/patroni', 'xlog_location': 4})
+    l = Leader(0, 0, m1) if leader else None
+    m2 = 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, [m1, m2], failover)
 
 
-def get_unlocked_cluster():
-    return Cluster(False, None, None, [])
+def get_cluster_initialized_with_leader(failover=None):
+    return get_cluster_initialized_without_leader(leader=True, failover=failover)
 
 
+def get_cluster_initialized_with_only_leader(failover=None):
+    l = get_cluster_initialized_without_leader(leader=True, failover=failover).leader
+    return get_cluster(True, l, [l], failover)
+
+
+class MockPatroni(object):
+
+    def __init__(self, p, d):
+        self.postgresql = p
+        self.dcs = d
+        self.api = Mock()
+        self.tags = {}
+        self.nofailover = None
+        self.nap_time = 10
+        self.replicatefrom = None
+        self.api.connection_string = 'http://127.0.0.1:8008'
+        self.clonefrom = None
+
+
+def run_async(func, args=()):
+    return func(*args) if args else func()
+
+
+@patch.object(Postgresql, 'is_running', Mock(return_value=True))
+@patch.object(Postgresql, 'is_leader', Mock(return_value=True))
+@patch.object(Postgresql, 'xlog_position', Mock(return_value=0))
+@patch.object(Postgresql, 'call_nowait', Mock(return_value=True))
+@patch.object(Postgresql, 'data_directory_empty', Mock(return_value=False))
+@patch.object(Postgresql, 'controldata', Mock(return_value={'Database system identifier': '1234567890'}))
+@patch.object(Postgresql, 'sync_replication_slots', Mock())
+@patch.object(Postgresql, 'write_pg_hba', Mock())
+@patch.object(Postgresql, 'write_pgpass', Mock())
+@patch.object(Postgresql, 'write_recovery_conf', Mock())
+@patch.object(Postgresql, 'query', Mock())
+@patch.object(Postgresql, 'checkpoint', Mock())
+@patch('subprocess.call', Mock(return_value=0))
 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):
-        self.p = MockPostgresql()
+    @patch('socket.getaddrinfo', socket_getaddrinfo)
+    def setUp(self):
         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.p = Postgresql({'name': 'postgresql0', 'scope': 'dummy', 'listen': '127.0.0.1:5432',
+                                 'data_dir': 'data/postgresql0', 'superuser': {}, 'admin': {},
+                                 'replication': {'username': '', 'password': '', 'network': ''}})
+            self.p.set_state('running')
+            self.p.check_replication_lag = true
+            self.p.can_create_replica_without_replication_connection = MagicMock(return_value=False)
+            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(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.client.delete = Mock(side_effect=EtcdException())
+            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.is_running = false
+        self.p.follow = false
+        self.assertEquals(self.ha.run_cycle(), 'started as a secondary')
+        self.assertEquals(self.ha.run_cycle(), 'failed to start postgres')
+
+    def test_recover_master_failed(self):
+        self.p.follow = false
+        self.p.is_healthy = false
+        self.p.is_running = false
+        self.ha.has_lock = true
+        self.p.set_role('master')
+        self.p.controldata = lambda: {'Database cluster state': 'in production'}
+        self.assertEquals(self.ha.run_cycle(), 'started as readonly because i had the session lock')
+        self.assertEquals(self.ha.run_cycle(), 'removed leader key after trying and failing to start postgres')
+
+    @patch('sys.exit', return_value=1)
+    @patch('patroni.ha.Ha.sysid_valid', MagicMock(return_value=True))
+    def test_sysid_no_match(self, exit_mock):
+        self.ha.run_cycle()
+        exit_mock.assert_called_once_with(1)
+
+    @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.p.is_leader = false
+        self.p.is_healthy = true
         self.ha.has_lock = true
         self.assertEquals(self.ha.run_cycle(), 'promoted self to leader because i had the session lock')
 
@@ -102,24 +146,26 @@ 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')
 
     def test_demote_after_failing_to_obtain_lock(self):
         self.ha.acquire_lock = false
-        self.assertEquals(self.ha.run_cycle(), 'demoted self due after trying and failing to obtain lock')
+        self.assertEquals(self.ha.run_cycle(), 'demoted self 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 +184,191 @@ 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_follow_the_leader(self):
+    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(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')
+        self.ha.patroni.replicatefrom = "foo"
+        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_from_another_member(self):
+        self.ha.cluster = get_cluster_initialized_with_leader()
+        self.ha.patroni.clonefrom = 'other'
+        self.assertEquals(self.ha.bootstrap(), 'trying to bootstrap from replica \'other\'')
+
+    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_without_leader(self):
+        self.ha.cluster = get_cluster_initialized_without_leader()
+        self.p.can_create_replica_without_replication_connection = MagicMock(return_value=True)
+        self.assertEquals(self.ha.bootstrap(), "trying to bootstrap without leader")
+
+    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', '', None))
+        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, '', self.p.name, None))
+        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', None))
+        self.assertEquals(self.ha.run_cycle(), 'no action.  i am the leader with the lock')
+        f = Failover(0, self.p.name, '', None)
+        self.ha.cluster = get_cluster_initialized_with_leader(f)
+        self.assertEquals(self.ha.run_cycle(), 'manual failover: demoting myself')
+        self.ha.fetch_node_status = lambda e: (e, True, True, 0, {'nofailover': 'True'})
+        self.assertEquals(self.ha.run_cycle(), 'no action.  i am the leader with the lock')
+        # manual failover from the previous leader to us won't happen if we hold the nofailover flag
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, 'blabla', self.p.name, None))
+        self.assertEquals(self.ha.run_cycle(), 'no action.  i am the leader with the lock')
+
+        # Failover scheduled time must include timezone
+        scheduled = datetime.datetime.now()
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, 'blabla', self.p.name, scheduled))
+        self.ha.run_cycle()
+
+        scheduled = datetime.datetime.utcnow().replace(tzinfo=pytz.UTC)
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, 'blabla', self.p.name, scheduled))
+        self.assertEquals('no action.  i am the leader with the lock', self.ha.run_cycle())
+
+        scheduled = scheduled + datetime.timedelta(seconds=30)
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, 'blabla', self.p.name, scheduled))
+        self.assertEquals('no action.  i am the leader with the lock', self.ha.run_cycle())
+
+        scheduled = scheduled + datetime.timedelta(seconds=-600)
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, 'blabla', self.p.name, scheduled))
+        self.assertEquals('no action.  i am the leader with the lock', self.ha.run_cycle())
+
+        scheduled = None
+        self.ha.cluster = get_cluster_initialized_with_leader(Failover(0, 'blabla', self.p.name, scheduled))
+        self.assertEquals('no action.  i am the leader with the lock', self.ha.run_cycle())
+
+    @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, '', self.p.name, None))
+        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', None))
+        self.p.set_role('replica')
+        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, self.p.name, '', None))
+        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, {})  # inaccessible, in_recovery
+        self.p.set_role('replica')
+        self.assertEquals(self.ha.run_cycle(), 'promoted self to leader by acquiring session lock')
+        # set failover flag to True for all members of the cluster
+        # this should elect the current member, as we are not going to call the API for it.
+        self.ha.cluster = get_cluster_initialized_without_leader(failover=Failover(0, '', 'other', None))
+        self.ha.fetch_node_status = lambda e: (e, True, True, 0, {'nofailover': 'True'})  # accessible, in_recovery
+        self.p.set_role('replica')
+        self.assertEquals(self.ha.run_cycle(), 'promoted self to leader by acquiring session lock')
+        # same as previous, but set the current member to nofailover. In no case it should be elected as a leader
+        self.ha.patroni.nofailover = True
+        self.assertEquals(self.ha.run_cycle(), 'following a different leader because I am not allowed to promote')
+
+    def test_is_healthiest_node(self):
+        self.ha.state_handler.is_leader = false
+        self.ha.patroni.nofailover = False
+        self.ha.fetch_node_status = lambda e: (e, True, True, 0, {})
+        self.assertTrue(self.ha.is_healthiest_node())
+
+    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))
+        self.ha.patroni.nofailover = True
+        self.assertFalse(self.ha._is_healthiest_node(self.ha.old_cluster.members))
+        self.ha.patroni.nofailover = False
+
+    @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)
+
+    def test_post_recover(self):
+        self.p.is_running = false
+        self.ha.has_lock = true
+        self.assertEqual(self.ha.post_recover(), 'removed leader key after trying and failing to start postgres')
+        self.ha.has_lock = false
+        self.assertEqual(self.ha.post_recover(), 'failed to start postgres')
+        self.p.is_running = true
+        self.assertIsNone(self.ha.post_recover())

tests/test_patroni.py

@@ -1,144 +1,83 @@
-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 import Patroni, main
+from patroni.api import RestApiServer
+from patroni.async_executor import AsyncExecutor
+from patroni.etcd import Etcd
+from patroni import Patroni, main as _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
-
-
+@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):
-        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()
-        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'])
+    def setUp(self):
+        with patch.object(Client, 'machines') as mock_machines:
+            mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
+            self.touched = False
+            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)
                 self.p = Patroni(config)
+                self.p.ha.dcs.client.write = etcd_write
+                self.p.ha.dcs.client.read = etcd_read
 
-    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
-
+    @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):
-        main()
+    @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 = Mock(side_effect=SleepException())
+        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()
+
+    def test_nofailover(self):
+        self.p.tags['nofailover'] = True
+        self.assertTrue(self.p.nofailover)
+        self.p.tags['nofailover'] = None
+        self.assertFalse(self.p.nofailover)
+
+    def test_replicatefrom(self):
+        self.assertIsNone(self.p.replicatefrom)
+        self.p.tags['replicatefrom'] = 'foo'
+        self.assertEqual(self.p.replicatefrom, 'foo')

tests/test_postgresql.py

@@ -1,54 +1,39 @@
+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 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 six.moves import builtins
+from test_ha import false
 
 
-def nop(*args, **kwargs):
-    pass
+class MockCursor(object):
 
-
-def subprocess_call(cmd, shell=False, env=None):
-    return 0
-
-
-def false(*args, **kwargs):
-    return False
-
-
-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)]
+            self.results = [('', True, '', '', '', '', False)]
         else:
             self.results = [(
                 None,
@@ -66,51 +51,122 @@ class MockCursor:
     def fetchone(self):
         return self.results[0]
 
-    def close(self):
+    def fetchall(self):
+        return self.results
+
+    @staticmethod
+    def close():
         pass
 
     def __iter__(self):
         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 b"""
+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 && 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
+def fake_listdir(path):
+    return ["a", "b", "c"] if path.endswith('pg_xlog/archive_status') else []
 
 
+@patch('subprocess.call', Mock(return_value=0))
+@patch('psycopg2.connect', psycopg2_connect)
 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': ''},
+                             'superuser': {'username': 'test', 'password': 'test'},
                              'admin': {'username': 'admin', 'password': 'admin'},
+                             'pg_rewind': {'username': 'admin', 'password': 'admin'},
                              'replication': {'username': 'replicator',
                                              'password': 'rep-pass',
                                              'network': '127.0.0.1/32'},
@@ -119,101 +175,313 @@ 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',
+                            'tags': {'replicatefrom': 'leader'}})
+        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())
 
+    def test_get_initdb_options(self):
+        self.p.initdb_options = [{'encoding': 'UTF8'}, 'data-checksums']
+        self.assertEquals(self.p.get_initdb_options(), ['--encoding=UTF8', '--data-checksums'])
+        self.p.initdb_options = [{'pgdata': 'bar'}]
+        self.assertRaises(Exception, self.p.get_initdb_options)
+        self.p.initdb_options = [{'foo': 'bar', 1: 2}]
+        self.assertRaises(Exception, self.p.get_initdb_options)
+        self.p.initdb_options = [1]
+        self.assertRaises(Exception, self.p.get_initdb_options)
+
     def test_initialize(self):
         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
+    @patch('os.path.exists', Mock(return_value=True))
+    @patch('os.unlink', Mock())
+    def test_delete_trigger_file(self):
+        self.p.delete_trigger_file()
+
+    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_sync_from_leader(self):
-        self.assertTrue(self.p.sync_from_leader(self.leader))
+    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_follow_the_leader(self):
-        self.p.demote(self.leader)
-        self.p.follow_the_leader(None)
-        self.p.demote(self.leader)
-        self.p.follow_the_leader(self.leader)
-        self.p.follow_the_leader(self.other)
+    @patch.object(builtins, 'open', MagicMock())
+    def test_write_pgpass(self):
+        self.p.write_pgpass({'host': 'localhost', 'port': '5432', 'user': 'foo', 'password': 'bar'})
 
-    def test_create_connection_users(self):
-        cfg = self.p.config
-        cfg['superuser']['username'] = 'test'
-        p = Postgresql(cfg)
-        p.create_connection_users()
+    @patch('patroni.postgresql.Postgresql.write_pgpass', MagicMock(return_value=dict()))
+    def test_sync_replica(self):
+        self.assertTrue(self.p.sync_replica(self.leader))
+        self.p.create_replica = Mock(return_value=1)
+        self.assertFalse(self.p.sync_replica(self.leader))
 
-    def test_create_replication_slots(self):
+    @patch('subprocess.call', side_effect=OSError)
+    @patch('patroni.postgresql.Postgresql.write_pgpass', MagicMock(return_value=dict()))
+    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))
+    @patch('patroni.postgresql.Postgresql.write_pgpass', MagicMock(return_value=dict()))
+    @patch('subprocess.check_output', Mock(return_value=0, side_effect=pg_controldata_string))
+    def test_follow(self, mock_pg_rewind):
+        self.p.follow(None)
+        self.p.follow(self.leader)
+        self.p.follow(Leader(-1, 28, self.other))
+        self.p.rewind = mock_pg_rewind
+        self.p.follow(self.leader)
+        self.p.require_rewind()
+        with mock.patch('os.path.islink', MagicMock(return_value=True)):
+            with mock.patch('patroni.postgresql.Postgresql.can_rewind', new_callable=PropertyMock(return_value=True)):
+                with mock.patch('os.unlink', MagicMock(return_value=True)):
+                    self.p.follow(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(self.leader, recovery=True)
+            self.p.rewind.return_value = False
+            self.p.follow(self.leader, recovery=True)
+        with mock.patch('patroni.postgresql.Postgresql.check_recovery_conf', MagicMock(return_value=True)):
+            self.assertTrue(self.p.follow(None))
+
+    @patch('subprocess.check_output', Mock(return_value=0, side_effect=pg_controldata_string))
+    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):
+            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
+
+    @patch('time.sleep', Mock())
+    def test_create_replica(self):
+        self.p.delete_trigger_file = Mock(side_effect=OSError)
+        with patch('subprocess.call', Mock(side_effect=[1, 0])):
+            self.assertEquals(self.p.create_replica(self.leader, ''), 0)
+        with patch('subprocess.call', Mock(side_effect=[Exception(), 0])):
+            self.assertEquals(self.p.create_replica(self.leader, ''), 0)
+
+        self.p.config['create_replica_method'] = ['wale', 'basebackup']
+        self.p.config['wale'] = {'command': 'foo'}
+        with patch('subprocess.call', Mock(return_value=0)):
+            self.assertEquals(self.p.create_replica(self.leader, ''), 0)
+            del self.p.config['wale']
+            self.assertEquals(self.p.create_replica(self.leader, ''), 0)
+
+        with patch('subprocess.call', Mock(side_effect=Exception("foo"))):
+            self.assertEquals(self.p.create_replica(self.leader, ''), 1)
+
+    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)
+        self.p.schedule_load_slots = False
+        with mock.patch('patroni.postgresql.Postgresql.role', new_callable=PropertyMock(return_value='replica')):
+            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()
+
+    @patch('patroni.postgresql.Postgresql.write_pgpass', MagicMock(return_value=dict()))
+    def test_bootstrap(self):
+        with patch('subprocess.call', Mock(return_value=1)):
+            self.assertRaises(PostgresException, self.p.bootstrap)
+        self.p.bootstrap()
+        with patch('patroni.postgresql.Postgresql.sync_replica', MagicMock(return_value=True)):
+            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=OSError)):
+            self.p.remove_data_directory()
+        self.p.remove_data_directory()
+
+    def test_controldata(self):
+        with patch('subprocess.check_output', Mock(return_value=0, side_effect=pg_controldata_string)):
+            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)
+
+        with patch('subprocess.check_output', Mock(side_effect=subprocess.CalledProcessError(1, ''))):
+            self.assertEquals(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
+            data = self.p.read_postmaster_opts()
+            self.assertEqual(data, dict())
+
+    @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)
+
+    @patch('os.listdir', MagicMock(side_effect=fake_listdir))
+    @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 = OSError
+        mock_link.side_effect = OSError
+        self.p.cleanup_archive_status()
+        mock_unlink.assert_not_called()
+        mock_remove.assert_not_called()
+
+    @patch('subprocess.check_output', MagicMock(return_value=0, side_effect=pg_controldata_string))
+    def test_sysid(self):
+        self.assertEqual(self.p.sysid, "6200971513092291716")
+
+    @patch('os.path.isfile', Mock(return_value=True))
+    @patch('shutil.copy', Mock(side_effect=IOError))
+    def test_save_configuration_files(self):
+        self.p.save_configuration_files()
+
+    @patch('os.path.isfile', Mock(side_effect=[False, True]))
+    @patch('shutil.copy', Mock(side_effect=IOError))
+    def test_restore_configuration_files(self):
+        self.p.restore_configuration_files()
+
+    def test_can_create_replica_without_replication_connection(self):
+        self.p.config['create_replica_method'] = []
+        self.assertFalse(self.p.can_create_replica_without_replication_connection())
+        self.p.config['create_replica_method'] = ['wale', 'basebackup']
+        self.p.config['wale'] = {'command': 'foo', 'no_master': 1}
+        self.assertTrue(self.p.can_create_replica_without_replication_connection())
+
+    def test_replica_method_can_work_without_replication_connection(self):
+        self.assertFalse(self.p.replica_method_can_work_without_replication_connection('basebackup'))
+        self.assertFalse(self.p.replica_method_can_work_without_replication_connection('foobar'))
+        self.p.config['foo'] = {'command': 'bar', 'no_master': 1}
+        self.assertTrue(self.p.replica_method_can_work_without_replication_connection('foo'))
+        self.p.config['foo'] = {'command': 'bar'}
+        self.assertFalse(self.p.replica_method_can_work_without_replication_connection('foo'))

tests/test_restore.py

@@ -1,111 +0,0 @@
-import unittest
-from mock import MagicMock, patch
-import os
-from scripts.restore import Restore, WALERestore
-
-
-def fake_cursor_fetchone(*args, **kwargs):
-    return ('16777216',)
-
-
-def fake_call_fail_for_wal_e(*args, **kwargs):
-    if len(args) > 0 and 'backup-fetch' in args[0]:
-        return 1
-    return 0
-
-
-def fake_call_fail_for_base_backup(*args, **kwargs):
-    if len(args) > 0 and 'backup-fetch' in args[0]:
-        return 0
-    return 1
-
-
-def fake_backup_data(self, *args, **kwargs):
-    """ return the fake result of WAL-E backup-list"""
-    return """name    last_modified   expanded_size_bytes wal_segment_backup_start    wal_segment_offset_backup_start wal_segment_backup_stop wal_segment_offset_backup_stop
-base_00000001000000000000007F_00000040  2015-05-18T10:13:25.000Z 167772160   00000001000000000000007F    00000040 00000001000000000000007F    00000240
-"""
-
-
-class TestRestore(unittest.TestCase):
-
-    def setUp(self):
-        self.restore = Restore("batman", "master", "/data", "host=batman port=5432 user=batman")
-        pass
-
-    def tearDown(self):
-        pass
-
-    def test_parse_connstring(self):
-        self.assertDictEqual(self.restore.master_connection, {'host': 'batman', 'port': '5432', 'user': 'batman'})
-
-    @patch('subprocess.call', MagicMock(return_value=0))
-    def test_run(self):
-        ret = self.restore.run()
-        self.assertEqual(ret, 0)
-
-    @patch('subprocess.call', MagicMock(return_value=1))
-    def test_run_fail(self):
-        ret = self.restore.run()
-        self.assertEqual(ret, 1)
-
-
-@patch('os.access', MagicMock(return_value=True))
-@patch('os.makedirs', MagicMock(return_value=True))
-@patch('os.path.exists', MagicMock(return_value=True))
-@patch('os.path.isdir', MagicMock(return_value=True))
-@patch('psycopg2.extensions.cursor.fetchone', MagicMock(side_effect=fake_cursor_fetchone))
-@patch('psycopg2.extensions.cursor', MagicMock(autospec=True))
-@patch('psycopg2.extensions.connection', MagicMock(autospec=True))
-@patch('psycopg2.connect', MagicMock(autospec=True))
-@patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data))
-class TestWALERestore(unittest.TestCase):
-
-    def setUp(self):
-        env = {}
-        env['WAL_S3_BUCKET'] = 'batman'
-        env['WALE_BACKUP_THRESHOLD_PERCENTAGE'] = 100
-        env['WALE_BACKUP_THRESHOLD_MEGABYTES'] = 100
-        self.wale_restore = WALERestore("batman", "master", "/data", "host=batman port=5432 user=batman", env=env)
-
-    def tearDown(self):
-        pass
-
-    def test_setup(self):
-        self.wale_restore.setup()
-        self.assertFalse(self.wale_restore.init_error)
-
-    # have to redefine the class-level os.access mock inside the function
-    # since the class-level mock will be applied after the function level one.
-    @patch('os.access', return_value=False)
-    def test_setup_fail(self, mock_no_access):
-        os.access = mock_no_access
-        self.wale_restore.setup()
-        self.assertTrue(self.wale_restore.init_error)
-
-    # The 3 tests above only differ with the mock function instead of a subprocess call
-    # in the first one, subprocess call should return success only for wal-e command,
-    # checking the primary use-case of restoring from WAL-E backup.
-    # In the second one, we test fallbacks by failing at WAL-E, but succeeding at
-    # pg_basebackup.
-    # Finally, the last use case is when all subprocess.call fails. resulting in a
-    # failure to restore from replica
-    @patch('subprocess.call',
-           MagicMock(side_effect=lambda *args, **kwargs: 0 if 'wal-e' in args[0] else 1))
-    def test_run(self):
-        self.wale_restore.setup()
-        ret = self.wale_restore.run()
-        self.assertEqual(ret, 0)
-
-    @patch('subprocess.call',
-           MagicMock(side_effect=lambda *args, **kwargs: 0 if 'pg_basebackup' in args[0] else 1))
-    def test_run_fallback(self):
-        self.wale_restore.setup()
-        ret = self.wale_restore.run()
-        self.assertEqual(ret, 0)
-
-    @patch('subprocess.call', MagicMock(return_value=1))
-    def test_run_all_fail(self):
-        self.wale_restore.setup()
-        ret = self.wale_restore.run()
-        self.assertEqual(ret, 1)

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()
+        self.assertIsNone(reap_children())
+        with patch('os.waitpid', Mock(return_value=(0, 0))):
+            sigchld_handler(None, None)
+            self.assertIsNone(reap_children())
 
+    @patch('time.sleep', time_sleep)
     def test_sleep(self):
-        time.sleep = time_sleep
-        sleep(0.01)
+        self.assertIsNone(sleep(0.01))
+
+
+@patch('time.sleep', Mock())
+class TestRetrySleeper(unittest.TestCase):
+
+    @staticmethod
+    def _fail(times=1):
+        scope = dict(times=0)
+
+        def inner():
+            if scope['times'] >= times:
+                pass
+            else:
+                scope['times'] += 1
+                raise PatroniException('Failed!')
+        return inner
+
+    def test_reset(self):
+        retry = Retry(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 = Retry(delay=0)
+        self.assertRaises(RetryFailedError, retry, self._fail(times=999))
+        self.assertEquals(retry._attempts, 1)
+
+    def test_maximum_delay(self):
+        retry = Retry(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 = Retry(deadline=0.0001)
+        self.assertRaises(RetryFailedError, retry, self._fail(times=100))
+
+    def test_copy(self):
+        def _sleep(t):
+            pass
+
+        retry = Retry(sleep_func=_sleep)
+        rcopy = retry.copy()
+        self.assertTrue(rcopy.sleep_func is _sleep)

tests/test_wale_restore.py

@@ -0,0 +1,101 @@
+import psycopg2
+import subprocess
+import unittest
+
+from mock import MagicMock, patch, PropertyMock
+from patroni.scripts.wale_restore import WALERestore, main as _main
+
+
+def fake_cursor_fetchone(*args, **kwargs):
+    return ('16777216',)
+
+
+def fake_call_fail_for_wal_e(*args, **kwargs):
+    if len(args) > 0 and 'backup-fetch' in args[0]:
+        return 1
+    return 0
+
+
+def fake_call_fail_for_base_backup(*args, **kwargs):
+    if len(args) > 0 and 'backup-fetch' in args[0]:
+        return 0
+    return 1
+
+
+def fake_backup_data(self, *args, **kwargs):
+    """ return the fake result of WAL-E backup-list"""
+    return """name    last_modified   expanded_size_bytes wal_segment_backup_start    wal_segment_offset_backup_start wal_segment_backup_stop wal_segment_offset_backup_stop
+base_00000001000000000000007F_00000040  2015-05-18T10:13:25.000Z 167772160   00000001000000000000007F    00000040 00000001000000000000007F    00000240
+"""
+
+
+def fake_backup_data_2(self, *args, **kwargs):
+    """ return the fake result of WAL-E backup-list"""
+    return """name    last_modified   expanded_size_bytes wal_segment_backup_start    wal_segment_offset_backup_start wal_segment_backup_stop wal_segment_offset_backup_stop """
+
+
+def fake_backup_data_3(self, *args, **kwargs):
+    """ return the fake result of WAL-E backup-list"""
+    return """name    last_modified   expanded_size_bytes wal_segment_backup_start    wal_segment_offset_backup_start wal_segment_backup_stop
+base_00000001000000000000007F_00000040  2015-05-18T10:13:25.000Z 167772160   00000001000000000000007F    00000040 00000001000000000000007F    00000240
+"""
+
+
+def fake_backup_data_4(self, *args, **kwargs):
+    """ return the fake result of WAL-E backup-list"""
+    return """name    last_modified   expanded_size_foo wal_segment_backup_start    wal_segment_offset_backup_start wal_segment_backup_stop wal_segment_offset_backup_stop
+base_00000001000000000000007F_00000040  2015-05-18T10:13:25.000Z 167772160   00000001000000000000007F    00000040 00000001000000000000007F    00000240
+"""
+
+
+@patch('os.access', MagicMock(return_value=True))
+@patch('os.makedirs', MagicMock(return_value=True))
+@patch('os.path.exists', MagicMock(return_value=True))
+@patch('os.path.isdir', MagicMock(return_value=True))
+@patch('psycopg2.extensions.cursor.fetchone', MagicMock(side_effect=fake_cursor_fetchone))
+@patch('psycopg2.extensions.cursor', MagicMock(autospec=True))
+@patch('psycopg2.extensions.connection', MagicMock(autospec=True))
+@patch('psycopg2.connect', MagicMock(autospec=True))
+@patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data))
+class TestWALERestore(unittest.TestCase):
+
+    def setUp(self):
+        self.wale_restore = WALERestore("batman", "/data",
+                                        "host=batman port=5432 user=batman", "/etc", 100, 100, 1, 0)
+
+    def tearDown(self):
+        pass
+
+    def test_should_use_s3_to_create_replica(self):
+        with patch('psycopg2.connect', MagicMock(side_effect=psycopg2.Error("foo"))):
+            self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
+        with patch('subprocess.check_output', MagicMock(side_effect=subprocess.CalledProcessError(1, "cmd", "foo"))):
+            self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
+        with patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data_2)):
+            self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
+        with patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data_3)):
+            self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
+        with patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data_4)):
+            self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
+
+        self.wale_restore.should_use_s3_to_create_replica()
+        self.wale_restore.no_master = 1
+        self.assertTrue(self.wale_restore.should_use_s3_to_create_replica())
+
+    def test_create_replica_with_s3(self):
+        with patch('subprocess.call', MagicMock(return_value=0)):
+            self.assertEqual(self.wale_restore.create_replica_with_s3(), 0)
+        with patch('subprocess.call', MagicMock(side_effect=Exception("foo"))):
+            self.assertEqual(self.wale_restore.create_replica_with_s3(), 1)
+
+    def test_run(self):
+        with patch.object(self.wale_restore, 'init_error', PropertyMock(return_value=True)):
+            self.assertEqual(self.wale_restore.run(), 2)
+        with patch.object(self.wale_restore, 'should_use_s3_to_create_replica', MagicMock(return_value=True)):
+            with patch.object(self.wale_restore, 'create_replica_with_s3', MagicMock(return_value=0)):
+                self.assertEqual(self.wale_restore.run(), 0)
+
+    @patch('sys.exit', MagicMock())
+    @patch.object(WALERestore, 'run', MagicMock(return_value=0))
+    def test_main(self):
+        self.assertEqual(_main(), None)

tests/test_zookeeper.py

@@ -1,125 +1,114 @@
-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):
+    @staticmethod
+    def retry(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):
+    @staticmethod
+    def get_children(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):
+    @staticmethod
+    def set(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' and 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 +116,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)