Bumped version to 0.90

urllib3.exceptions.HTTPError fixes for python 3.5.1

.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,82 @@
+sudo: false
 language: python
 python:
-  - "2.7"
+- "3.5"
-  - "3.3"
+addons:
-  - "3.4"
+  apt:
+    packages:
+    - postgresql-contrib-9.5
+  postgresql: "9.5"
+env:
+  global:
+  - ETCDVERSION=2.3.2 ZKVERSION=3.4.6 CONSULVERSION=0.6.4
+  matrix:
+  - TEST_SUITE="python setup.py"
+  - DCS="etcd" TEST_SUITE="behave"
+  - DCS="exhibitor" TEST_SUITE="behave"
+  - DCS="consul" TEST_SUITE="behave"
+cache:
+  directories:
+  - $HOME/virtualenv/python2.7.9
+  - $HOME/virtualenv/python3.4.2
+  - $HOME/virtualenv/python3.5.0
 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
+    set -e
-  - pip install coveralls
+
+    if [[ $TEST_SUITE == "behave" ]]; then
+        if [[ $DCS == "consul" ]]; then
+            curl -L https://releases.hashicorp.com/consul/${CONSULVERSION}/consul_${CONSULVERSION}_linux_amd64.zip \
+                | gunzip > consul
+            chmod +x consul
+        fi
+
+        if [[ $DCS == "etcd" ]]; then
+            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
+        fi
+
+        if [[ $DCS == "exhibitor" ]]; then
+            curl -L http://www.apache.org/dist/zookeeper/zookeeper-${ZKVERSION}/zookeeper-${ZKVERSION}.tar.gz | tar xz
+            mv zookeeper-${ZKVERSION}/conf/zoo_sample.cfg zookeeper-${ZKVERSION}/conf/zoo.cfg
+            zookeeper-${ZKVERSION}/bin/zkServer.sh start
+            # following lines are 'emulating' exhibitor REST API
+            while true; do
+                echo -e 'HTTP/1.0 200 OK\nContent-Type: application/json\n\n{"servers":["127.0.0.1"],"port":2181}' \
+                    | nc -l 8181 &> /dev/null
+            done&
+        fi
+    fi
+
+    for pv in "2.7" "3.4" "3.5"; do
+        source ~/virtualenv/python${pv}/bin/activate
+        # explicitly install all needed python modules to cache them
+        for p in '-r requirements.txt' 'behave codacy-coverage coverage coveralls flake8 mock pytest-cov pytest'; do
+            pip install $p
+        done
+    done
 script:
-  - python setup.py test
+  - |
-  - python setup.py flake8
+    for pv in "2.7" "3.4" "3.5"; do
+        source ~/virtualenv/python${pv}/bin/activate
+
+        if [[ $TEST_SUITE == "behave" ]]; then
+            if [[ $pv != "3.4" ]]; then
+                echo Running acceptance tests using python${pv}
+                if ! PATH=.:$PATH $TEST_SUITE; then
+                    # output all log files when tests are failing
+                    grep . features/output/*/*postgres?.*
+                    exit 1
+                fi
+            fi
+        else
+            echo Running unit tests using python${pv}
+            $TEST_SUITE test
+            $TEST_SUITE flake8
+        fi
+    done
+
+    set +e
 after_success:
   - coveralls
+  - if [[ $TEST_SUITE != "behave" ]]; 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

@@ -12,16 +12,22 @@ 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
+ENV PGVERSION 9.5
-RUN apt-get install python python-psycopg2 python-yaml python-requests python-boto postgresql-${PGVERSION} python-dnspython python-kazoo python-pip -y
+RUN apt-get install postgresql-${PGVERSION} postgresql-server-dev-${PGVERSION} -y
-RUN pip install python-etcd
+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
 
 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

MAINTAINERS

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

README.rst

@@ -1,14 +1,14 @@
 |Build Status| |Coverage Status|
 
-Patroni: A Template for PostgreSQL HA with ZooKeeper or etcd
+Patroni: A Template for PostgreSQL HA with ZooKeeper, etcd or Consul
 ------------------------------------------------------------
 
 Patroni was previously known as Governor.
 
-*There are many ways to run high availability with PostgreSQL; here we
+*There are many ways to run high availability with PostgreSQL. Here, we
-present a template for you to create your own custom fit high
+present a template for you to create your own customized, high-availability 
-availability solution using python and distributed configuration store
+solution using Python and — for maximum accessibility — a distributed 
-(like ZooKeeper or etcd) for maximum accessibility.*
+configuration store like ZooKeeper, etcd or Consul.*
 
 Getting Started
 ---------------
@@ -22,8 +22,8 @@ To get started, do the following from different terminals:
     > ./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
+different settings in the YAML files to see how its behavior changes. Kill
-some of the different components to see how the system behaves.
+some of the components to see how the system behaves.
 
 Add more ``postgres*.yml`` files to create an even larger cluster.
 
@@ -39,109 +39,95 @@ run:
 
     > psql --host 127.0.0.1 --port 5000 postgres
 
-How Patroni works
+How Patroni Works
 -----------------
 
-For a diagram of the high availability decision loop, see the included a
+For a diagram of the high availability decision loop, review this PDF:
-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
+For an example file, see ``postgres0.yml``. Regarding settings:
-settings:
 
--  *ttl*: the TTL to acquire the leader lock. Think of it as the length
+-  *ttl*: the TTL to acquire the leader lock. Think of it as the length of time before initiation of the automatic failover process.
-   of time before automatic failover process is initiated.
 -  *loop\_wait*: the number of seconds the loop will sleep
 
--  *restapi*
+-  *restapi*:
--  *listen*: ip address + port that Patroni will listen to provide
+    -  *listen*: IP address + port that Patroni will listen to, to provide health-check information for haproxy.
-   health-check information for haproxy.
+    -  *connect\_address*: IP address + port through which restapi is accessible.
--  *connect\_address*: ip address + port through which restapi is
+    -  *auth*: (optional) 'username:password' to protect dangerous REST API endpoints.
-   accessible.
+    -  *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*
+-  *etcd*:
--  *scope*: the relative path used on etcd's http api for this
+    -  *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 cluster.
-   deployment, thus you can run multiple HA deployments from a single
+    -  *ttl*: the TTL to acquire the leader lock. Think of it as the length of time before initiation of the automatic failover process.
-   etcd
+    -  *host*: the host:port for the etcd endpoint.
--  *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*
+-  *consul*:
--  *scope*: the relative path used on etcd's http api for this
+    -  *scope*: the relative path used on Consul's HTTP API for this deployment; makes it possible to run multiple HA deployments from a single Consul cluster.
-   deployment, thus you can run multiple HA deployments from a single
+    -  *ttl*: the TTL to acquire the leader lock. Think of it as the length of time before initiation of the automatic failover process.
-   etcd
+    -  *host*: the host:port for the Consul endpoint.
--  *session\_timeout*: the TTL to acquire the leader lock. Think of it
-   as the length of time before automatic failover process is initiated.
--  *reconnect\_timeout*: how long we should try to reconnect to
-   ZooKeeper after connection loss. After this timeout we assume that we
-   don't have lock anymore and will restart in read-only mode.
--  *hosts*: list of ZooKeeper cluster members in format: [
-   'host1:port1', 'host2:port2', 'etc...']
--  *exhibitor*: if you are running ZooKeeper cluster under Exhibitor
-   supervisory the following section could be interesting for you
 
-   -  *poll\_interval*: how often list of ZooKeeper and Exhibitor nodes
+-  *zookeeper*:
-      should be updated from Exhibitor
+    -  *scope*: the relative path used on ZooKeeper for this deployment; makes it possible to run multiple HA deployments from a single ZooKeeper cluster.
-   -  *port*: Exhibitor port
+    -  *session\_timeout*: the TTL to acquire the leader lock. Think of it as the length of time before initiation of the automatic failover process.
-   -  *hosts*: initial list of Exhibitor (ZooKeeper) nodes in format: [
+    -  *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.
-      'host1', 'host2', 'etc...' ]. This list would be updated
+    -  *hosts*: list of ZooKeeper cluster members in format: ['host1:port1', 'host2:port2', 'etc...']
-      automatically when Exhibitor (ZooKeeper) cluster topology changes.
+    -  *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*
+-  *postgresql*:
--  *name*: the name of the Postgres host, must be unique for the cluster
+    -  *name*: the name of the Postgres host. Must be unique for the cluster.
--  *listen*: ip address + port that Postgres listening. Must be
+    -  *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. 
-   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.
--  *connect\_address*: ip address + port through which Postgres is
+    -  *data\_dir*: file path to initialize and store Postgres data files.
-   accessible from other nodes and applications.
+    -  *maximum\_lag\_on\_failover*: the maximum bytes a follower may lag.
--  *data\_dir*: file path to initialize and store Postgres data files
+    -  *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.
--  *maximum\_lag\_on\_failover*: the maximum bytes a follower may lag
--  *use\_slots*: whether or not to use replication_slots.  Must be False for PostgreSQL 9.3, and you should comment out max_replication_slots.
-   before it is not eligible become leader
--  *pg\_hba*: list of lines which should be added to pg\_hba.conf
 
-   -  *- host all all 0.0.0.0/0 md5*
+    -  *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
 
--  *replication*
+    -  *pg\_hba*: list of lines which should be added to pg\_hba.conf.
+        -  *- host all all 0.0.0.0/0 md5*.
+        -  *- host replication replicator 127.0.0.1/32 md5* # A line like this is required for replication
 
-   -  *username*: replication username, user will be created during
+    -  *replication*:
-      initialization
+        -  *username*: replication username; user will be created during initialization.
-   -  *password*: replication password, user will be created during
+        -  *password*: replication password; user will be created during initialization.
-      initialization
-   -  *network*: network setting for replication in pg\_hba.conf
 
--  *callbacks* callback scripts to run on certain actions. Patroni will
+    -  *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.
-   pass current action, role and cluster name. See scripts/aws.py as an
+        -  *on\_start*: a script to run when the cluster starts.
-   example on how to write them.
+        -  *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.
 
-   -  *on\_start*: a script to run when the cluster starts
+    -  *superuser*:
-   -  *on\_stop*: a script to run when the cluster stops
+        -  *password*: password for the Postgres user, set during initialization.
-   -  *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*
+    -  *admin*:
+        -  *username*: admin username; user is created during initialization. It will have CREATEDB and CREATEROLE privileges.
+        -  *password*: admin password; user is created during initialization.
 
-   -  *password*: password for postgres user. It would be set during
+    -  *recovery\_conf*: additional configuration settings written to recovery.conf when configuring follower.
-      initialization
+        -  *parameters*: list of configuration settings for Postgres.  Many of these are required for replication to work.
 
--  *admin*:
+    -  *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.
 
-   -  *username*: admin username, user will be created during
+    -  *replica\_method* for each create_replica_method other than basebackup, you would add a configuration section
-      initialization. It would have CREATEDB and CREATEROLE privileges
+       of the same name.  At a minimum, this should include "command" with a full path to the actual script to be
-   -  *password*: admin password, user will be created during
+       executed.  Other configuration parameters will be passed along to the script in the form "parameter=value".
-      initialization.
 
--  *recovery\_conf*: additional configuration settings written to recovery.conf when configuring follower
+Replication Choices
--  *parameters*: list of configuration settings for Postgres.  Many of these are required for replication to work.
-
-Replication choices
 -------------------
 
 Patroni uses Postgres' streaming replication. By default, this
@@ -155,43 +141,43 @@ 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
+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 will be reduced throughput on writes. This throughput will
+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 increases the variability of write performance
+synchrous replication significantly increases the variability of write 
-significantly. If followers become inaccessible from the leader, the
+performance. If followers become inaccessible from the leader, the
-leader will becomes effectively readonly.
+leader effectively becomes readonly.
 
 To enable a simple synchronous replication test, add the follow lines to
-the ``parameters`` section of your YAML configuration files.
+the ``parameters`` section of your YAML configuration files:
 
 .. code:: YAML
 
         synchronous_commit: "on"
         synchronous_standby_names: "*"
 
-When using synchronous replication, use at least a 3-Postgres data nodes
+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 the many business
+Choosing your replication schema is dependent on your business
-decisions. Investigate both async and sync replication, as well as other
+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
+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
+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
+the Primary because the connection pool is full, behavior will be
 undesireable.
 
 Requirements on a Mac
@@ -207,14 +193,13 @@ Run the following on a Mac to install requirements:
 Notice
 ------
 
-There are many different ways to do HA with PostgreSQL, see `the
+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
+We call Patroni a "template" because it is far from being a one-size-fits-all
-all, or a plug-and-play replication system. It will have it's own
+or plug-and-play replication system. It will have its own caveats. Use wisely.
-caveats. Use wisely.
 
 .. |Build Status| image:: https://travis-ci.org/zalando/patroni.svg?branch=master
    :target: https://travis-ci.org/zalando/patroni

TODO.md

@@ -1,4 +1,16 @@
 Failover
 ========
-- When determining who should become master, include the minor version of PostgreSQL in the decision
+- 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
+- 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,7 +1,7 @@
 # Patroni Dockerfile
 You can run Patroni in a docker container using this Dockerfile, or by using one of the Docker image at
 
-    https://os-registry.stups.zalan.do/v1/repositories/acid/patroni/tags
+    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
@@ -10,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
 
@@ -36,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
+    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        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
+    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        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
+    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
+  listen: 0.0.0.0:8008
-  connect_address: 127.0.0.1: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 python /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,31 @@
+# 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
+
+# only kill the patroni process, not it's children, so it will gracefully stop postgres
+KillMode=process
+
+# 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 20 seconds
+
+  Scenario: check the basic failover
+    When I kill postgres0
+    Then postgres1 role is the primary after 32 seconds
+    When I start postgres0
+    Then postgres0 role is the secondary after 20 seconds
+    When I add the table bar to postgres1
+    Then table bar is present on postgres0 after 20 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 and streaming replication from a replica
+	Given I start postgres0
+	And postgres0 is a leader after 10 seconds
+	And I configure and start postgres1 with a tag clonefrom true
+	And replication works from postgres0 to postgres1 after 20 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 replicatefrom 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,385 @@
+import abc
+import consul
+import etcd
+import kazoo.client
+import kazoo.exceptions
+import os
+import psycopg2
+import shutil
+import six
+import subprocess
+import tempfile
+import time
+import yaml
+
+
+@six.add_metaclass(abc.ABCMeta)
+class AbstractController(object):
+
+    def __init__(self, name, work_directory, output_dir):
+        self._name = name
+        self._work_directory = work_directory
+        self._output_dir = output_dir
+        self._handle = None
+        self._log = None
+
+    def _has_started(self):
+        return self._handle and self._handle.pid and self._handle.poll() is None
+
+    def _is_running(self):
+        return self._has_started()
+
+    @abc.abstractmethod
+    def _is_accessible(self):
+        """process is accessible for queries"""
+
+    @abc.abstractmethod
+    def _start(self):
+        """start process"""
+
+    def start(self, max_wait_limit=5):
+        if self._is_running():
+            return True
+
+        self._log = open(os.path.join(self._output_dir, self._name + '.log'), 'a')
+        self._handle = self._start()
+
+        assert self._has_started(), "Process {0} is not running after being started".format(self._name)
+
+        for _ in range(max_wait_limit):
+            if self._is_accessible():
+                break
+            time.sleep(1)
+        else:
+            assert False,\
+                "{0} instance is not available for queries after {1} seconds".format(self._name, max_wait_limit)
+
+    def stop(self, kill=False, timeout=15):
+        term = False
+        start_time = time.time()
+
+        while self._handle and self._is_running():
+            if kill:
+                self._handle.kill()
+            elif not term:
+                self._handle.terminate()
+                term = True
+            time.sleep(1)
+            if not kill and time.time() - start_time > timeout:
+                kill = True
+
+        if self._log:
+            self._log.close()
+
+
+class PatroniController(AbstractController):
+    __PORT = 5440
+    PATRONI_CONFIG = '{}.yml'
+    """ starts and stops individual patronis"""
+
+    def __init__(self, dcs, name, work_directory, output_dir, tags=None):
+        super(PatroniController, self).__init__('patroni_' + name, work_directory, output_dir)
+        PatroniController.__PORT += 1
+        self._data_dir = os.path.join(work_directory, 'data', name)
+        self._connstring = None
+        self._config = self._make_patroni_test_config(name, dcs, tags)
+
+        self._conn = None
+        self._curs = None
+
+    def write_label(self, content):
+        with open(os.path.join(self._data_dir, 'label'), 'w') as f:
+            f.write(content)
+
+    def read_label(self):
+        try:
+            with open(os.path.join(self._data_dir, 'label'), 'r') as f:
+                return f.read().strip()
+        except IOError:
+            return None
+
+    def _start(self):
+        return subprocess.Popen(['coverage', 'run', '--source=patroni', '-p', 'patroni.py', self._config],
+                                stdout=self._log, stderr=subprocess.STDOUT, cwd=self._work_directory)
+
+    def _is_accessible(self):
+        return self.query("SELECT 1", fail_ok=True) is not None
+
+    def _make_patroni_test_config(self, name, dcs, tags):
+        patroni_config_name = self.PATRONI_CONFIG.format(name)
+        patroni_config_path = os.path.join(self._output_dir, patroni_config_name)
+
+        with open(patroni_config_name) as f:
+            config = yaml.load(f)
+
+        host = config['postgresql']['listen'].split(':')[0]
+
+        config['postgresql']['listen'] = config['postgresql']['connect_address'] = '{0}:{1}'.format(host, self.__PORT)
+
+        user = config['postgresql'].get('superuser', {})
+        self._connkwargs = {k: user[n] for n, k in [('username', 'user'), ('password', 'password')] if n in user}
+        self._connkwargs.update({'host': host, 'port': self.__PORT, 'database': 'postgres'})
+
+        config['postgresql'].update({'name': name, 'data_dir': self._data_dir})
+        config['postgresql']['parameters'].update({
+            'logging_collector': 'on', 'log_destination': 'csvlog', 'log_directory': self._output_dir,
+            'log_filename': name + '.log', 'log_statement': 'all', 'log_min_messages': 'debug1'})
+
+        if tags:
+            config['tags'] = tags
+
+        if dcs != 'etcd':
+            dcs_config = config.pop('etcd')
+            dcs_config.pop('host')
+
+            if dcs == 'consul':
+                config[dcs] = dcs_config
+            else:
+                dcs_config.update({'session_timeout': dcs_config.pop('ttl'), 'reconnect_timeout': config['loop_wait']})
+                if dcs == 'exhibitor':
+                    dcs_config['exhibitor'] = {'hosts': ['127.0.0.1'], 'port': 8181}
+                else:
+                    dcs_config['hosts'] = ['127.0.0.1:2181']
+                config['zookeeper'] = dcs_config
+
+        with open(patroni_config_path, 'w') as f:
+            yaml.dump(config, f, default_flow_style=False)
+
+        return patroni_config_path
+
+    def _connection(self):
+        if not self._conn or self._conn.closed != 0:
+            self._conn = psycopg2.connect(**self._connkwargs)
+            self._conn.autocommit = True
+        return self._conn
+
+    def _cursor(self):
+        if not self._curs or self._curs.closed or self._curs.connection.closed != 0:
+            self._curs = self._connection().cursor()
+        return self._curs
+
+    def query(self, query, fail_ok=False):
+        try:
+            cursor = self._cursor()
+            cursor.execute(query)
+            return cursor
+        except psycopg2.Error:
+            if not fail_ok:
+                raise
+
+    def check_role_has_changed_to(self, new_role, timeout=10):
+        bound_time = time.time() + timeout
+        recovery_status = new_role != 'primary'
+        while time.time() < bound_time:
+            cur = self.query("SELECT pg_is_in_recovery()", fail_ok=True)
+            if cur:
+                row = cur.fetchone()
+                if row and row[0] == recovery_status:
+                    return True
+            time.sleep(1)
+        return False
+
+
+class AbstractDcsController(AbstractController):
+
+    _CLUSTER_NODE = 'service/batman'
+
+    def _is_accessible(self):
+        return self._is_running()
+
+    def stop_and_remove_work_directory(self, timeout=15):
+        """ terminate process and wipe out the temp work directory, but only if we actually started it"""
+        self.stop(timeout=timeout)
+        if self._work_directory:
+            shutil.rmtree(self._work_directory)
+
+    @abc.abstractmethod
+    def query(self, key):
+        """ query for a value of a given key """
+
+    @abc.abstractmethod
+    def cleanup_service_tree(self):
+        """ clean all contents stored in the tree used for the tests """
+
+
+class ConsulController(AbstractDcsController):
+
+    def __init__(self, output_dir):
+        super(ConsulController, self).__init__('consul', tempfile.mkdtemp(), output_dir)
+        self._client = consul.Consul()
+
+    def _start(self):
+        return subprocess.Popen(['consul', 'agent', '-server', '-bootstrap', '-advertise=127.0.0.1',
+                                 '-data-dir', self._work_directory], stdout=self._log, stderr=subprocess.STDOUT)
+
+    def _is_running(self):
+        try:
+            return bool(self._client.status.leader())
+        except Exception:
+            return False
+
+    def query(self, key):
+        _, value = self._client.kv.get('{0}/{1}'.format(self._CLUSTER_NODE, key))
+        return value and value['Value'].decode('utf-8')
+
+    def cleanup_service_tree(self):
+        self._client.kv.delete(self._CLUSTER_NODE, recurse=True)
+
+
+class EtcdController(AbstractDcsController):
+
+    """ handles all etcd related tasks, used for the tests setup and cleanup """
+
+    def __init__(self, output_dir):
+        super(EtcdController, self).__init__('etcd', tempfile.mkdtemp(), output_dir)
+        self._client = etcd.Client()
+
+    def _start(self):
+        return subprocess.Popen(["etcd", "--debug", "--data-dir", self._work_directory],
+                                stdout=self._log, stderr=subprocess.STDOUT)
+
+    def query(self, key):
+        try:
+            return self._client.get('/{0}/{1}'.format(self._CLUSTER_NODE, key)).value
+        except etcd.EtcdKeyNotFound:
+            return None
+
+    def cleanup_service_tree(self):
+        try:
+            self._client.delete('/' + self._CLUSTER_NODE, recursive=True)
+        except (etcd.EtcdKeyNotFound, etcd.EtcdConnectionFailed):
+            return
+        except Exception as e:
+            assert False, "exception when cleaning up etcd contents: {0}".format(e)
+
+    def _is_running(self):
+        # if etcd is running, but we didn't start it
+        try:
+            return bool(self._client.machines)
+        except Exception:
+            return False
+
+
+class ZooKeeperController(AbstractDcsController):
+
+    """ handles all zookeeper related tasks, used for the tests setup and cleanup """
+
+    def __init__(self, output_dir):
+        super(ZooKeeperController, self).__init__('zookeeper', None, output_dir)
+        self._client = kazoo.client.KazooClient()
+
+    def _start(self):
+        pass  # TODO: implement later
+
+    def query(self, key):
+        try:
+            return self._client.get('/{0}/{1}'.format(self._CLUSTER_NODE, key))[0].decode('utf-8')
+        except kazoo.exceptions.NoNodeError:
+            return None
+
+    def cleanup_service_tree(self):
+        try:
+            self._client.delete('/' + self._CLUSTER_NODE, recursive=True)
+        except (kazoo.exceptions.NoNodeError):
+            return
+        except Exception as e:
+            assert False, "exception when cleaning up zookeeper contents: {0}".format(e)
+
+    def _is_running(self):
+        # if zookeeper is running, but we didn't start it
+        if self._client.connected:
+            return True
+        try:
+            return self._client.start(1) or True
+        except Exception:
+            return False
+
+
+class PatroniPoolController(object):
+
+    KNOWN_DCS = {'consul': ConsulController, 'etcd': EtcdController,
+                 'zookeeper': ZooKeeperController, 'exhibitor': ZooKeeperController}
+
+    def __init__(self):
+        self._dcs = None
+        self._output_dir = None
+        self._patroni_path = None
+        self._processes = {}
+        self.create_and_set_output_directory('')
+
+    @property
+    def patroni_path(self):
+        if self._patroni_path is None:
+            cwd = os.path.realpath(__file__)
+            while True:
+                cwd, entry = os.path.split(cwd)
+                if entry == 'features' or cwd == '/':
+                    break
+            self._patroni_path = cwd
+        return self._patroni_path
+
+    @property
+    def output_dir(self):
+        return self._output_dir
+
+    def start(self, pg_name, max_wait_limit=20, tags=None):
+        if pg_name not in self._processes:
+            self._processes[pg_name] = PatroniController(self.dcs, pg_name, self.patroni_path, self._output_dir, tags)
+        self._processes[pg_name].start(max_wait_limit)
+
+    def __getattr__(self, func):
+        if func not in ['stop', 'query', 'write_label', 'read_label', 'check_role_has_changed_to']:
+            raise AttributeError("PatroniPoolController instance has no attribute '{0}'".format(func))
+
+        def wrapper(pg_name, *args, **kwargs):
+            return getattr(self._processes[pg_name], func)(*args, **kwargs)
+        return wrapper
+
+    def stop_all(self):
+        for ctl in self._processes.values():
+            ctl.stop()
+        self._processes.clear()
+
+    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
+
+    @property
+    def dcs(self):
+        if self._dcs is None:
+            self._dcs = os.environ.get('DCS', 'etcd')
+            assert self._dcs in self.KNOWN_DCS, 'Unsupported dcs: ' + self.dcs
+        return self._dcs
+
+
+# actions to execute on start/stop of the tests and before running invidual features
+def before_all(context):
+    context.pctl = PatroniPoolController()
+    context.dcs_ctl = context.pctl.KNOWN_DCS[context.pctl.dcs](context.pctl.output_dir)
+    context.dcs_ctl.start()
+    try:
+        context.dcs_ctl.cleanup_service_tree()
+    except AssertionError:  # after_all handlers won't be executed in before_all
+        context.dcs_ctl.stop_and_remove_work_directory()
+        raise
+
+
+def after_all(context):
+    context.dcs_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.dcs_ctl.cleanup_service_tree()

features/patroni_api.feature

@@ -0,0 +1,53 @@
+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 candidate"
+
+Scenario: check API requests for the primary-replica pair
+	Given I start postgres1
+	And replication works from postgres0 to postgres1 after 20 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
+	When I issue an empty POST request to http://127.0.0.1:8008/restart
+	Then I receive a response code 200
+        And postgres0 role is the primary after 5 seconds
+        When I sleep for 10 seconds
+        Then postgres1 role is the secondary after 15 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 postgres1 role is the primary after 5 seconds
+        And postgres0 role is the secondary after 10 seconds
+	And replication works from postgres1 to postgres0 after 20 seconds
+
+Scenario: check the scheduled failover
+	Given I issue a scheduled failover at http://127.0.0.1:8009 from postgres1 to postgres0 in 1 seconds
+	Then I receive a response code 200
+	And postgres0 is a leader after 20 seconds
+        And postgres0 role is the primary after 5 seconds
+        And postgres1 role is the secondary after 10 seconds
+	And replication works from postgres0 to postgres1 after 25 seconds
+

features/steps/basic_replication.py

@@ -0,0 +1,54 @@
+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):
+    assert context.pctl.check_role_has_changed_to(pg_name, pg_role, timeout=int(max_promotion_timeout)),\
+        "{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,100 @@
+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.dcs_ctl.query("leader") != name):
+        time.sleep(1)
+        assert time.time() < max_time, "{0} is not a leader in dcs 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_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_5433 127.0.0.1:5433 maxconn 100 check port 8009

patroni/__init__.py

@@ -5,94 +5,69 @@ import time
 import yaml
 
 from patroni.api import RestApiServer
-from patroni.etcd import Etcd
+from patroni.exceptions import PatroniException
 from patroni.ha import Ha
 from patroni.postgresql import Postgresql
-from patroni.utils import setup_signal_handlers, sleep, reap_children
+from patroni.utils import reap_children, set_ignore_sigterm, setup_signal_handlers
-from patroni.zookeeper import ZooKeeper
+from patroni.version import __version__
 
 logger = logging.getLogger(__name__)
 
 
-class Patroni:
+class Patroni(object):
+    PATRONI_CONFIG_VARIABLE = 'PATRONI_CONFIGURATION'
 
     def __init__(self, config):
         self.nap_time = config['loop_wait']
+        self.tags = {tag: value for tag, value in config.get('tags', {}).items()
+                     if tag not in ('clonefrom', 'nofailover', 'noloadbalance') or value}
         self.postgresql = Postgresql(config['postgresql'])
-        self.ha = Ha(self.postgresql, self.get_dcs(self.postgresql.name, config))
+        self.dcs = self.get_dcs(self.postgresql.name, config)
-        host, port = config['restapi']['listen'].split(':')
+        self.version = __version__
         self.api = RestApiServer(self, config['restapi'])
+        self.ha = Ha(self)
         self.next_run = time.time()
-        self.shutdown_member_ttl = 300
+
+    @property
+    def noloadbalance(self):
+        return self.tags.get('noloadbalance', False)
+
+    @property
+    def nofailover(self):
+        return self.tags.get('nofailover', False)
+
+    @property
+    def replicatefrom(self):
+        return self.tags.get('replicatefrom')
 
     @staticmethod
     def get_dcs(name, config):
         if 'etcd' in config:
+            from patroni.etcd import Etcd
             return Etcd(name, config['etcd'])
         if 'zookeeper' in config:
+            from patroni.zookeeper import ZooKeeper
             return ZooKeeper(name, config['zookeeper'])
-        raise Exception('Can not find sutable configuration of distributed configuration store')
+        if 'consul' in config:
-
+            from patroni.consul import Consul
-    def touch_member(self, ttl=None):
+            return Consul(name, config['consul'])
-        connection_string = self.postgresql.connection_string + '?application_name=' + self.api.connection_string
+        raise PatroniException('Can not find suitable configuration of distributed configuration store')
-        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):
-        if self.postgresql.is_promoted:
-            self.next_run = time.time()
         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:
+        elif self.dcs.watch(nap_time):
-            self.ha.dcs.watch(nap_time)
+            self.next_run = time.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()
 
@@ -102,20 +77,30 @@ def main():
     logging.getLogger('requests').setLevel(logging.WARNING)
     setup_signal_handlers()
 
-    if len(sys.argv) < 2 or not os.path.isfile(sys.argv[1]):
+    # Patroni reads the configuration from the command-line argument if it exists, and from the environment otherwise.
-        print('Usage: {} config.yml'.format(sys.argv[0]))
+    use_env = False
-        return
+    use_file = (len(sys.argv) >= 2 and os.path.isfile(sys.argv[1]))
+    if not use_file:
+        config_env = os.environ.get(Patroni.PATRONI_CONFIG_VARIABLE)
+        use_env = config_env is not None
+        if not use_env:
+            print('Usage: {0} config.yml'.format(sys.argv[0]))
+            print('\tPatroni may also read the configuration from the {} environment variable'.
+                  format(Patroni.PATRONI_CONFIG_VARIABLE))
+            return
 
-    with open(sys.argv[1], 'r') as f:
+    if use_file:
-        config = yaml.load(f)
+        with open(sys.argv[1], 'r') as f:
+            config = yaml.load(f)
+    elif use_env:
+        config = yaml.load(config_env)
 
     patroni = Patroni(config)
     try:
-        patroni.initialize()
         patroni.run()
     except KeyboardInterrupt:
-        pass
+        set_ignore_sigterm()
     finally:
-        patroni.touch_member(patroni.shutdown_member_ttl)  # schedule member removal
+        patroni.api.shutdown()
-        patroni.postgresql.stop()
+        patroni.postgresql.stop(checkpoint=False)
-        patroni.ha.dcs.delete_leader()
+        patroni.dcs.delete_leader()

patroni/api.py

@@ -1,8 +1,16 @@
+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
@@ -10,64 +18,324 @@ 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 do_GET(self):
+    def _write_response(self, status_code, body, headers=None):
-        response = self.get_postgresql_status()
+        self.send_response(status_code)
+        if body is not None:
+            headers = headers or {}
+            if 'Content-Type' not in headers:
+                headers['Content-Type'] = 'text/html'
+            for name, value in (headers or {}).items():
+                self.send_header(name, value)
+            self.end_headers()
+            self.wfile.write(body.encode('utf-8'))
+
+    def send_auth_request(self, body):
+        self._write_response(401, body, {'WWW-Authenticate': 'Basic realm=\"Patroni\"'})
+
+    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 _write_status_response(self, status_code, response, options=False):
+        if options:
+            body = None
+        else:
+            patroni = self.server.patroni
+            response.update({'tags': patroni.tags} if patroni.tags else {})
+            if patroni.postgresql.sysid:
+                response['database_system_identifier'] = patroni.postgresql.sysid
+            response['patroni'] = {'version': patroni.version, 'scope': patroni.postgresql.scope}
+            body = json.dumps(response)
+        self._write_response(status_code, body, {'Content-Type': 'application/json'})
+
+    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
-        status_code = 200 if response['running'] and 'role' in response and response['role'] in path else 503
+        response = self.get_postgresql_status()
 
-        self.send_response(status_code)
+        patroni = self.server.patroni
-        self.send_header('Content-Type', 'application/json')
+        cluster = patroni.dcs.cluster
-        self.end_headers()
+        if cluster:  # dcs available
-        self.wfile.write(json.dumps(response).encode('utf-8'))
+            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:  # response['role'] != 'master'
+                status_code = 503 if patroni.noloadbalance else 200
+            else:
+                status_code = 503
+        elif 'role' in response and response['role'] in path:
+            status_code = 503 if response['role'] != 'master' and patroni.noloadbalance else 200
+        elif patroni.ha.restart_scheduled() and patroni.postgresql.role == 'master' and 'master' in path:
+            # exceptional case for master node when the postgres is being restarted via API
+            status_code = 200
+        else:
+            status_code = 503
+        self._write_status_response(status_code, response, options)
 
-    def get_postgresql_status(self):
+    def do_OPTIONS(self):
+        self.do_GET(options=True)
+
+    def do_GET_patroni(self):
+        response = self.get_postgresql_status(True)
+        self._write_status_response(200, response)
+
+    @check_auth
+    def do_POST_restart(self):
+        status_code = 500
+        data = 'restart failed'
         try:
-            row = self.server.query("""SELECT to_char(pg_postmaster_start_time(), 'YYYY-MM-DD HH24:MI:SS.MS TZ'),
+            status, data = self.server.patroni.ha.restart()
-                                              pg_is_in_recovery(),
+            status_code = 200 if status else 503
-                                              CASE WHEN pg_is_in_recovery()
+        except Exception:
-                                                   THEN null
+            logger.exception('Exception during restart')
-                                                   ELSE pg_current_xlog_location() END,
+        self._write_response(status_code, data)
-                                              pg_last_xlog_receive_location(),
+
-                                              pg_last_xlog_replay_location(),
+    @check_auth
-                                              pg_is_in_recovery() AND pg_is_xlog_replay_paused()""")[0]
+    def do_POST_reinitialize(self):
+        ha = self.server.patroni.ha
+        cluster = ha.dcs.get_cluster()
+        if cluster.is_unlocked():
+            status_code = 503
+            data = 'Cluster has no leader, can not reinitialize'
+        elif cluster.leader.name == ha.state_handler.name:
+            status_code = 503
+            data = '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'
+            else:
+                status_code = 200
+                data = 'reinitialize scheduled'
+        self._write_response(status_code, data)
+
+    def poll_failover_result(self, leader, candidate):
+        for _ in range(0, 15):
+            time.sleep(1)
+            try:
+                cluster = self.server.patroni.dcs.get_cluster()
+                if cluster.leader and cluster.leader.name != leader:
+                    if not candidate or candidate == cluster.leader.name:
+                        return 200, 'Successfully failed over to "{0}"'.format(cluster.leader.name)
+                    else:
+                        return 200, 'Failed over to "{0}" instead of "{1}"'.format(cluster.leader.name, candidate)
+                if not cluster.failover:
+                    return 503, 'Failover failed'
+            except Exception as e:
+                logger.debug('Exception occured during polling failover result: %s', e)
+        return 503, 'Failover status unknown'
+
+    def is_failover_possible(self, cluster, leader, candidate):
+        if leader and not cluster.leader or cluster.leader.name != leader:
+            return 'leader name does not match'
+        if candidate:
+            members = [m for m in cluster.members if m.name == candidate]
+            if not members:
+                return 'candidate 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 'failover is not possible: cluster does not have members except leader'
+        for _, reachable, _, _, tags in self.server.patroni.ha.fetch_nodes_statuses(members):
+            if reachable and not tags.get('nofailover', False):
+                return None
+        return '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')
+        candidate = request.get('candidate') or request.get('member')
+        scheduled_at = request.get('scheduled_at')
+        cluster = self.server.patroni.ha.dcs.get_cluster()
+        status_code = 500
+
+        logger.info("received failover request with leader=%s candidate=%s scheduled_at=%s",
+                    leader, candidate, scheduled_at)
+
+        data = ''
+        if leader or candidate:
+            if scheduled_at:
+                try:
+                    scheduled_at = dateutil.parser.parse(scheduled_at)
+                    if scheduled_at.tzinfo is None:
+                        data = 'Timezone information is mandatory for scheduled_at'
+                        status_code = 400
+                    elif scheduled_at < datetime.datetime.now(pytz.utc):
+                        data = 'Cannot schedule failover in the past'
+                        status_code = 422
+                    elif self.server.patroni.dcs.manual_failover(leader, candidate, scheduled_at=scheduled_at):
+                        self.server.patroni.dcs.event.set()
+                        data = 'Failover scheduled'
+                        status_code = 200
+                    else:
+                        data = 'failed to write failover key into DCS'
+                        status_code = 503
+                except (ValueError, TypeError):
+                    logger.exception('Invalid scheduled failover time: %s', request['scheduled_at'])
+                    data = 'Unable to parse scheduled timestamp. It should be in an unambiguous format, e.g. ISO 8601'
+                    status_code = 422
+            else:
+                data = self.is_failover_possible(cluster, leader, candidate)
+                if not data:
+                    if self.server.patroni.dcs.manual_failover(leader, candidate):
+                        self.server.patroni.dcs.event.set()
+                        status_code, data = self.poll_failover_result(cluster.leader and cluster.leader.name, candidate)
+                    else:
+                        data = 'failed to write failover key into DCS'
+                        status_code = 503
+        else:
+            status_code = 400
+            data = 'No values given for required parameters leader and candidate'
+        self._write_response(status_code, 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 {
-                'running': True,
+                'state': self.server.patroni.postgresql.state,
                 'postmaster_start_time': row[0],
-                'role': 'slave' if row[1] else 'master',
+                'role': 'replica' if row[1] else 'master',
+                'server_version': self.server.patroni.postgresql.server_version,
                 'xlog': ({
                     'received_location': row[3],
                     'replayed_location': row[4],
-                    'paused': row[5]} if row[1] else {
+                    'replayed_timestamp': row[5],
+                    'paused': row[6]} if row[1] else {
                     'location': row[2]
                 })
             }
-        except (psycopg2.OperationalError, psycopg2.InterfaceError):
+        except (psycopg2.Error, RetryFailedError, PostgresConnectionException):
-            logger.exception('get_postgresql_status')
+            state = self.server.patroni.postgresql.state
-            return {'running': self.server.patroni.postgresql.is_running()}
+            if state == 'running':
+                logger.exception('get_postgresql_status')
+                state = 'unknown'
+            return {'state': state}
+
+    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.connection_string = 'http://{}/patroni'.format(config.get('connect_address', None) or config['listen'])
+        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 = self.patroni.postgresql.connection().cursor()
+        cursor = None
-        cursor.execute(sql, params)
+        try:
-        ret = [r for r in cursor]
+            with self.patroni.postgresql.connection().cursor() as cursor:
-        cursor.close()
+                cursor.execute(sql, params)
-        return ret
+                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/consul.py

@@ -0,0 +1,244 @@
+from __future__ import absolute_import
+import logging
+import os
+import time
+import six
+
+from consul import ConsulException, NotFound, base, std
+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 ConsulError(DCSError):
+    pass
+
+
+class HTTPClient(std.HTTPClient):
+
+    def __init__(self, *args, **kwargs):
+        super(HTTPClient, self).__init__(*args, **kwargs)
+        self._patch_default_timeout()
+
+    def _patch_default_timeout(self):
+        # Set a default timeout for the `request.session.request` method, that is used
+        # internally  by the methods request.session.get, request.session.post and
+        # others. We monkey-patch here to avoid reimplementing each individual method from
+        # `std.HTTPClient`. By default, the timeout is not set. It means that a new
+        # session may hang almost indefinitely waiting for the server to respond,
+        # which is not what we want in Patroni.
+
+        request_func = getattr(self.session.request, '__func__' if six.PY3 else 'im_func')
+        defaults_attr_name = '__defaults__' if six.PY3 else 'func_defaults'
+        defaults = list(getattr(request_func, defaults_attr_name))
+        code = request_func.__code__ if six.PY3 else request_func.func_code
+        defaults[code.co_varnames[code.co_argcount - len(defaults):code.co_argcount].index('timeout')] = 5
+        setattr(request_func, defaults_attr_name, tuple(defaults))  # monkeypatching
+
+    def get(self, callback, path, params=None):
+        # The get function is overridden to handle a special case of it being called
+        # with an index and wait parameters. That form indicates that a user needs to
+        # wait for the given key to change its value, with a wait timeout supplied. We
+        # don't want our monkey-patched timeout to be less than the value of the wait
+        # parameter, therefore, we set it to either the value of wait or a default of 5 minutes.
+
+        if isinstance(params, dict) and 'index' in params:
+            timeout = (float(params['wait'][:-1]) if 'wait' in params else 300) + 1
+        else:
+            timeout = None
+        return callback(self.response(self.session.get(self.uri(path, params), verify=self.verify, timeout=timeout)))
+
+
+class ConsulClient(base.Consul):
+
+    @staticmethod
+    def connect(host, port, scheme, verify=True):
+        return HTTPClient(host, port, scheme, verify)
+
+
+def catch_consul_errors(func):
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except (ConsulException, RequestException):
+            return False
+    return wrapper
+
+
+class Consul(AbstractDCS):
+
+    def __init__(self, name, config):
+        super(Consul, self).__init__(name, config)
+        self.ttl = int((config.get('ttl') or 30)/2)  # My experiments have shown that session expires after 2*ttl time
+        host, port = config.get('host', '127.0.0.1:8500').split(':')
+        self._client = ConsulClient(host=host, port=port)
+        self._scope = config['scope']
+        self._session = None
+        self._my_member_data = None
+        self.create_or_restore_session()
+
+    def create_or_restore_session(self):
+        while not self._session:
+            try:
+                _, member = self._client.kv.get(self.member_path)
+                self._session = (member or {}).get('Session')
+                if self.refresh_session():
+                    self._client.kv.delete(self.member_path)
+            except (ConsulException, RequestException):
+                logger.info('waiting on consul')
+                sleep(5)
+
+    def refresh_session(self):
+        """:returns: `!True` if it had to create new session"""
+        if self._session:
+            try:
+                return self._client.session.renew(self._session) is None
+            except NotFound:
+                self._session = None
+        if not self._session:
+            name = self._scope + '-' + self._name
+            try:
+                self._session = self._client.session.create(name=name, lock_delay=0, behavior='delete', ttl=self.ttl)
+            except (ConsulException, RequestException):
+                logger.exception('session.create')
+        if not self._session:
+            raise ConsulError('Failed to renew/create session')
+        return True
+
+    def client_path(self, path):
+        return super(Consul, self).client_path(path)[1:]
+
+    @staticmethod
+    def member(node):
+        return Member.from_node(node['ModifyIndex'], os.path.basename(node['Key']), node.get('Session'), node['Value'])
+
+    def _load_cluster(self):
+        try:
+            path = self.client_path('/')
+            _, results = self._client.kv.get(path, recurse=True)
+
+            if results is None:
+                raise NotFound
+
+            nodes = {}
+            for node in results:
+                node['Value'] = (node['Value'] or b'').decode('utf-8')
+                nodes[os.path.relpath(node['Key'], path)] = node
+
+            # 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 and leader['Value'] == self._name and self._session != leader.get('Session', 'x'):
+                logger.info('I am leader but not owner of the session. Removing leader node')
+                self._client.kv.delete(self.leader_path, cas=leader['ModifyIndex'])
+                leader = None
+
+            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['ModifyIndex'], leader.get('Session'), member)
+
+            # failover key
+            failover = nodes.get(self._FAILOVER)
+            if failover:
+                failover = Failover.from_node(failover['ModifyIndex'], failover['Value'])
+
+            self._cluster = Cluster(initialize, leader, last_leader_operation, members, failover)
+        except NotFound:
+            self._cluster = Cluster(False, None, None, [], None)
+        except:
+            logger.exception('get_cluster')
+            raise ConsulError('Consul is not responding properly')
+
+    def touch_member(self, data, **kwargs):
+        create_member = self.refresh_session()
+        cluster = self.cluster
+        member = cluster and ([m for m in cluster.members if m.name == self._name] or [None])[0]
+        if create_member and member:
+            try:
+                self._client.kv.delete(self.member_path)
+            except Exception:
+                return False
+
+        if not create_member and member and data == self._my_member_data:
+            return True
+
+        try:
+            self._client.kv.put(self.member_path, data, acquire=self._session)
+            self._my_member_data = data
+            return True
+        except Exception:
+            logger.exception('touch_member')
+        return False
+
+    @catch_consul_errors
+    def attempt_to_acquire_leader(self):
+        ret = self._client.kv.put(self.leader_path, self._name, acquire=self._session)
+        if not ret:
+            logger.info('Could not take out TTL lock')
+        return ret
+
+    def take_leader(self):
+        return self.attempt_to_acquire_leader()
+
+    @catch_consul_errors
+    def set_failover_value(self, value, index=None):
+        return self._client.kv.put(self.failover_path, value, cas=index)
+
+    @catch_consul_errors
+    def write_leader_optime(self, last_operation):
+        return self._client.kv.put(self.leader_optime_path, last_operation)
+
+    @staticmethod
+    def update_leader():
+        return True
+
+    @catch_consul_errors
+    def initialize(self, create_new=True, sysid=''):
+        kwargs = {'cas': 0} if create_new else {}
+        return self._client.kv.put(self.initialize_path, sysid, **kwargs)
+
+    @catch_consul_errors
+    def cancel_initialization(self):
+        return self._client.kv.delete(self.initialize_path)
+
+    @catch_consul_errors
+    def delete_cluster(self):
+        return self._client.kv.delete(self.client_path(''), recurse=True)
+
+    @catch_consul_errors
+    def delete_leader(self):
+        cluster = self.cluster
+        if cluster and isinstance(cluster.leader, Leader) and cluster.leader.name == self._name:
+            return self._client.kv.delete(self.leader_path, cas=cluster.leader.index)
+
+    def watch(self, timeout):
+        cluster = self.cluster
+        if cluster and cluster.leader and cluster.leader.name != self._name and cluster.leader.index:
+            end_time = time.time() + timeout
+            while timeout >= 1:
+                try:
+                    idx, _ = self._client.kv.get(self.leader_path, index=cluster.leader.index, wait=str(timeout) + 's')
+                    return str(idx) != str(cluster.leader.index)
+                except (ConsulException, RequestException):
+                    logging.exception('watch')
+
+                timeout = end_time - time.time()
+
+        try:
+            return super(Consul, self).watch(timeout)
+        finally:
+            self.event.clear()

patroni/ctl.py

@@ -0,0 +1,633 @@
+'''
+Patroni Control
+'''
+
+import click
+import datetime
+import dateutil
+import json
+import logging
+import os
+import psycopg2
+import random
+import requests
+import time
+import tzlocal
+import yaml
+
+from click import ClickException
+from patroni import Patroni, PatroniException
+from patroni.postgresql import parseurl
+from prettytable import PrettyTable
+from six.moves.urllib_parse import urlparse
+
+CONFIG_DIR_PATH = click.get_app_dir('patroni')
+CONFIG_FILE_PATH = os.path.join(CONFIG_DIR_PATH, 'patronictl.yaml')
+LOGLEVEL = 'WARNING'
+DCS_DEFAULTS = {'zookeeper': {'port': 2181, 'template': "zookeeper:\n hosts: ['{host}:{port}']"},
+                'exhibitor': {'port': 8181, 'template': "zookeeper:\n exhibitor:\n  hosts: [{host}]\n  port: {port}"},
+                'consul': {'port': 8500, 'template': "consul:\n host: '{host}:{port}'"},
+                'etcd': {'port': 4001, 'template': "etcd:\n host: '{host}:{port}'"}}
+
+
+class PatroniCtlException(ClickException):
+    pass
+
+
+def parse_dcs(dcs):
+    if dcs is None:
+        return None
+
+    parsed = urlparse(dcs)
+    scheme = parsed.scheme
+    if scheme == '' and parsed.netloc == '':
+        parsed = urlparse('//' + dcs)
+    port = int(parsed.port) if parsed.port else None
+
+    if scheme == '':
+        scheme = ([k for k, v in DCS_DEFAULTS.items() if v['port'] == port] or ['etcd'])[0]
+    elif scheme not in DCS_DEFAULTS:
+        raise PatroniCtlException('Unknown dcs scheme: {}'.format(scheme))
+
+    dcs_info = DCS_DEFAULTS[scheme]
+    return yaml.load(dcs_info['template'].format(host=parsed.hostname or 'localhost', port=port or dcs_info['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')
+
+    config.update(parse_dcs(dcs) or parse_dcs(config.get('dcs_api')) or {})
+
+    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):
+    for k in set(DCS_DEFAULTS.keys()) & set(config.keys()):
+        config[k].setdefault('scope', scope)
+    try:
+        return Patroni.get_dcs(scope, config)
+    except PatroniException as e:
+        raise PatroniCtlException(str(e))
+
+
+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, role)
+    for m in members:
+        if member is None or m.name == member:
+            return m
+
+
+def get_cursor(cluster, role='master', member=None, connect_parameters=None):
+    member = get_any_member(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, 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()
+
+    _, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+
+    cursor = None
+    for _ in watching(w, watch, clear=False):
+
+        output, cursor = query_member(cluster, cursor, member, role, command, 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):
+    _, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+
+    output_members(cluster, cluster_name, 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.delete_cluster()
+
+
+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):
+    _, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+
+    role_names = [m.name for m in get_all_members(cluster, 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, cluster_name)
+    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):
+    _, 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.
+    """
+
+    _, 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 ({0}). 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, 'candidate': 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(), 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 Exception:
+        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(master, candidate, scheduled_at=failover_value)
+
+    output_members(cluster, cluster_name)
+
+
+def output_members(cluster, name, 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 cluster_name in cluster_names:
+        dcs = get_dcs(config, cluster_name)
+
+        for _ in watching(w, watch):
+            output_members(dcs.get_cluster(), cluster_name, 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

@@ -1,9 +1,12 @@
 import abc
+import dateutil
+import json
+import six
 
 from collections import namedtuple
-from patroni.exceptions import DCSError
+from random import randint
-from patroni.utils import calculate_ttl, sleep
 from six.moves.urllib_parse import urlparse, urlunparse, parse_qsl
+from threading import Event, Lock
 
 
 def parse_connection_string(value):
@@ -23,28 +26,68 @@ def parse_connection_string(value):
     return conn_url, api_url
 
 
-class Member(namedtuple('Member', 'index,name,conn_url,api_url,expiration,ttl')):
+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 conn_url: connection string containing host, user and password which could be used to access this member.
+    :param session: either session id or just ttl in seconds
-    :param api_url: REST API url of patroni instance
+    :param data: arbitrary data i.e. conn_url, api_url, xlog location, state, role, tags, etc...
-    :param expiration: expiration time of given member key
-    :param ttl: ttl of given member key in seconds"""
 
-    def real_ttl(self):
+    There are two mandatory keys in a data:
-        return calculate_ttl(self.expiration) or -1
+    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 tags(self):
+        return self.data.get('tags', {})
+
+    @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', False)
 
 
-class Leader(namedtuple('Leader', 'index,expiration,ttl,member')):
+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 expiration: expiration time of the leader key
+    :param session: either session id or just ttl in seconds
-    :param ttl: ttl of the leader key
     :param member: reference to a `Member` object which represents current leader (see `Cluster.members`)"""
 
     @property
@@ -56,7 +99,47 @@ class Leader(namedtuple('Leader', 'index,expiration,ttl,member')):
         return self.member.conn_url
 
 
-class Cluster(namedtuple('Cluster', 'initialize,leader,last_leader_operation,members')):
+class Failover(namedtuple('Failover', 'index,leader,candidate,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_candidate"}'))
+    True
+    >>> Failover.from_node(1, 'null') is None
+    True
+    >>> n = '{"leader": "cluster_leader", "member": "cluster_candidate", "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:
@@ -64,15 +147,32 @@ class Cluster(namedtuple('Cluster', 'initialize,leader,last_leader_operation,mem
     :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 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)
 
-class AbstractDCS:
+    def get_member(self, member_name, fallback_to_leader=True):
+        return ([m for m in self.members if m.name == member_name] or [self.leader if fallback_to_leader else None])[0]
 
-    __metaclass__ = abc.ABCMeta
+    def get_clone_member(self):
+        candidates = [m for m in self.members if m.clonefrom and (not self.leader or m.name != self.leader.name)]
+        return candidates[randint(0, len(candidates) - 1)] if candidates else self.leader
+
+
+@six.add_metaclass(abc.ABCMeta)
+class AbstractDCS(object):
+
+    _INITIALIZE = 'initialize'
+    _LEADER = 'leader'
+    _FAILOVER = 'failover'
+    _MEMBERS = 'members/'
+    _OPTIME = 'optime'
+    _LEADER_OPTIME = _OPTIME + '/' + _LEADER
 
     def __init__(self, name, config):
         """
@@ -81,27 +181,80 @@ class AbstractDCS:
             i.e.: `zookeeper` for zookeeper, `etcd` for etcd, etc...
         """
         self._name = name
-        self._scope = config['scope']
+        self._namespace = '/{0}'.format(config.get('namespace', '/service/').strip('/'))
-        self._base_path = '/service/' + self._scope
+        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 self._base_path + 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):
-        """:returns: `Cluster` object which represent current state and topology of the cluster
+        with self._cluster_thread_lock:
-        raise `~DCSError` in case of communication or other problems with DCS. If current instance was
+            try:
-            running as a master and exception raised instance would be demoted."""
+                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 update_leader(self, state_handler):
+    def write_leader_optime(self, last_operation):
-        """Update leader key (or session) ttl and `/optime/leader` key in DCS.
+        """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
 
-        :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."""
 
@@ -114,12 +267,22 @@ class AbstractDCS:
         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):
+    @abc.abstractmethod
-        try:
+    def set_failover_value(self, value, index=None):
-            cluster = self.get_cluster()
+        """Create or update `/failover` key"""
-            return None if cluster.is_unlocked() else cluster.leader
+
-        except DCSError:
+    def manual_failover(self, leader, candidate, scheduled_at=None, index=None):
-            return None
+        failover_value = {}
+        if leader:
+            failover_value['leader'] = leader
+
+        if candidate:
+            failover_value['member'] = candidate
+
+        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):
@@ -139,12 +302,14 @@ class AbstractDCS:
         overwriting the key if necessary."""
 
     @abc.abstractmethod
-    def race(self, path):
+    def initialize(self, create_new=True, sysid=""):
         """Race for cluster initialization.
-        :param path: usually this is just '/initialize'
+
+        :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 `path` key and return `!True`
+        this method should create atomically initialize key and return `!True`
         otherwise it should return `!False`"""
 
     @abc.abstractmethod
@@ -152,5 +317,20 @@ class AbstractDCS:
         """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 """
+
+    @abc.abstractmethod
+    def delete_cluster(self):
+        """Delete cluster from DCS"""
+
     def watch(self, timeout):
-        sleep(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

@@ -6,13 +6,15 @@ 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, DCSError, Leader, Member, parse_connection_string
+from patroni.dcs import AbstractDCS, Cluster, Failover, Leader, Member
+from patroni.exceptions import DCSError
 from patroni.utils import Retry, RetryFailedError, sleep
+from urllib3.exceptions import HTTPError, ReadTimeoutError
 from requests.exceptions import RequestException
+from six.moves.http_client import HTTPException
 
 logger = logging.getLogger(__name__)
 
@@ -48,10 +50,60 @@ class Client(etcd.Client):
             self._update_machines_cache = True
             return [self._base_uri]
 
-    def api_execute(self, path, method, **kwargs):
+    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 (HTTPError, HTTPException, socket.error, socket.timeout) as e:
+            if (isinstance(fields, dict) and fields.get("wait") == "true" and
+                    isinstance(e, 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
-        self._update_machines_cache and self._load_machines_cache()
+        if self._update_machines_cache:
-        return super(Client, self).api_execute(path, method, **kwargs)
+            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):
@@ -61,16 +113,6 @@ class Client(etcd.Client):
             logger.exception('Can not resolve SRV for %s', host)
         return []
 
-    # try to workarond bug in python-etcd: https://github.com/jplana/python-etcd/issues/81
-    def _result_from_response(self, response):
-        try:
-            response.data.decode('utf-8')
-        except urllib3.exceptions.TimeoutError:
-            raise
-        except Exception as e:
-            raise etcd.EtcdException('Unable to decode server response: %s' % e)
-        return super(Client, self)._result_from_response(response)
-
     def _get_machines_cache_from_srv(self, discovery_srv):
         """Fetch list of etcd-cluster member by resolving _etcd-server._tcp. SRV record.
         This record should contain list of host and peer ports which could be used to run
@@ -78,9 +120,9 @@ class Client(etcd.Client):
 
         ret = []
         for host, port in self.get_srv_record(discovery_srv):
-            url = '{}://{}:{}/members'.format(self._protocol, host, port)
+            url = '{0}://{1}:{2}/members'.format(self._protocol, host, port)
             try:
-                response = requests.get(url)
+                response = requests.get(url, timeout=5)
                 if response.ok:
                     for member in response.json():
                         ret.extend(member['clientURLs'])
@@ -96,10 +138,10 @@ class Client(etcd.Client):
         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]))
+                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 ['{}://{}:{}'.format(self._protocol, host, port)]
+        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.
@@ -127,7 +169,9 @@ class Client(etcd.Client):
         # 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)
+
+        if self._base_uri in self._machines_cache:
+            self._machines_cache.remove(self._base_uri)
 
         self._update_machines_cache = False
 
@@ -135,9 +179,13 @@ class Client(etcd.Client):
 def catch_etcd_errors(func):
     def wrapper(*args, **kwargs):
         try:
-            return not func(*args, **kwargs) is None
+            return func(*args, **kwargs) is not None
         except (RetryFailedError, etcd.EtcdException):
             return False
+        except:
+            logger.exception("")
+            raise EtcdError("unexpected error")
+
     return wrapper
 
 
@@ -145,20 +193,19 @@ class Etcd(AbstractDCS):
 
     def __init__(self, name, config):
         super(Etcd, self).__init__(name, config)
-        self.ttl = config['ttl']
+        self.ttl = config.get('ttl', 30)
-        self.member_ttl = config.get('member_ttl', 3600)
         self._retry = Retry(deadline=10, max_delay=1, max_tries=-1,
                             retry_exceptions=(etcd.EtcdConnectionFailed,
                                               etcd.EtcdLeaderElectionInProgress,
                                               etcd.EtcdWatcherCleared,
                                               etcd.EtcdEventIndexCleared))
-        self.client = self.get_etcd_client(config)
+        self._client = self.get_etcd_client(config)
-        self.cluster = None
 
     def retry(self, *args, **kwargs):
         return self._retry.copy()(*args, **kwargs)
 
-    def get_etcd_client(self, config):
+    @staticmethod
+    def get_etcd_client(config):
         client = None
         while not client:
             try:
@@ -170,53 +217,54 @@ class Etcd(AbstractDCS):
 
     @staticmethod
     def member(node):
-        conn_url, api_url = parse_connection_string(node.value)
+        return Member.from_node(node.modifiedIndex, os.path.basename(node.key), node.ttl, node.value)
-        return Member(node.modifiedIndex, os.path.basename(node.key), conn_url, api_url, node.expiration, node.ttl)
 
-    def get_cluster(self):
+    def _load_cluster(self):
         try:
-            result = self.retry(self.client.read, self.client_path(''), recursive=True)
+            result = self.retry(self._client.read, self.client_path(''), recursive=True)
             nodes = {os.path.relpath(node.key, result.key): node for node in result.leaves}
 
             # get initialize flag
-            initialize = bool(nodes.get('initialize', False))
+            initialize = nodes.get(self._INITIALIZE)
+            initialize = initialize and initialize.value
 
             # get last leader operation
-            last_leader_operation = nodes.get('optime/leader', None)
+            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('members/') and len(k.split('/')) == 2]
+            members = [self.member(n) for k, n in nodes.items() if k.startswith(self._MEMBERS) and k.count('/') == 1]
 
             # get leader
-            leader = nodes.get('leader', None)
+            leader = nodes.get(self._LEADER)
             if leader:
-                member = Member(-1, leader.value, None, None, None, None)
+                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.expiration, leader.ttl, member)
+                leader = Leader(leader.modifiedIndex, leader.ttl, member)
 
-            self.cluster = Cluster(initialize, leader, last_leader_operation, members)
+            # 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, [])
+            self._cluster = Cluster(False, None, None, [], None)
         except:
-            self.cluster = None
             logger.exception('get_cluster')
             raise EtcdError('Etcd is not responding properly')
-        return self.cluster
 
     @catch_etcd_errors
     def touch_member(self, connection_string, ttl=None):
-        return self.retry(self.client.set, self.client_path('/members/' + self._name),
+        return self.retry(self._client.set, self.member_path, connection_string, ttl or self.ttl)
-                          connection_string, ttl or self.member_ttl)
 
     @catch_etcd_errors
     def take_leader(self):
-        return self.retry(self.client.set, self.client_path('/leader'), self._name, self.ttl)
+        return self.retry(self._client.set, self.leader_path, self._name, self.ttl)
 
     def attempt_to_acquire_leader(self):
         try:
-            return not self.retry(self.client.write, self.client_path('/leader'),
+            return bool(self.retry(self._client.write, self.leader_path, self._name, ttl=self.ttl, prevExist=False))
-                                  self._name, ttl=self.ttl, prevExist=False) is None
         except etcd.EtcdAlreadyExist:
             logger.info('Could not take out TTL lock')
         except (RetryFailedError, etcd.EtcdException):
@@ -224,41 +272,54 @@ class Etcd(AbstractDCS):
         return False
 
     @catch_etcd_errors
-    def write_leader_optime(self, state_handler):
+    def set_failover_value(self, value, index=None):
-        return self.client.set(self.client_path('/optime/leader'), state_handler.last_operation())
+        return self._client.write(self.failover_path, value, prevIndex=index or 0)
 
     @catch_etcd_errors
-    def update_leader(self, state_handler):
+    def write_leader_optime(self, last_operation):
-        ret = self.retry(self.client.test_and_set, self.client_path('/leader'), self._name, self._name, self.ttl)
+        return self._client.set(self.leader_optime_path, last_operation)
-        ret and self.write_leader_optime(state_handler)
-        return ret
 
     @catch_etcd_errors
-    def race(self, path):
+    def update_leader(self):
-        return self.retry(self.client.write, self.client_path(path), self._name, prevExist=False)
+        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.client_path('/leader'), prevValue=self._name)
+        return self._client.delete(self.leader_path, prevValue=self._name)
+
+    @catch_etcd_errors
+    def cancel_initialization(self):
+        return self.retry(self._client.delete, self.initialize_path)
+
+    @catch_etcd_errors
+    def delete_cluster(self):
+        return self.retry(self._client.delete, self.client_path(''), recursive=True)
 
     def watch(self, timeout):
+        cluster = self.cluster
         # watch on leader key changes if it is defined and current node is not lock owner
-        if self.cluster and self.cluster.leader and self.cluster.leader.name != self._name:
+        if cluster and cluster.leader and cluster.leader.name != self._name and cluster.leader.index:
             end_time = time.time() + timeout
-            index = self.cluster.leader.index
 
-            while index and timeout >= 1:  # when timeout is too small urllib3 doesn't have enough time to connect
+            while timeout >= 1:  # when timeout is too small urllib3 doesn't have enough time to connect
                 try:
-                    res = self.client.watch(self.client_path('/leader'), index=index + 1, timeout=timeout)
+                    self._client.watch(self.leader_path, index=cluster.leader.index + 1, timeout=timeout + 0.5)
-                    if res.action not in ['set', 'compareAndSwap'] or res.value != self.cluster.leader.name:
+                    # Synchronous work of all cluster members with etcd is less expensive
-                        return
+                    # than reestablishing http connection every time from every replica.
-                    index = res.modifiedIndex
+                    return True
-                except urllib3.exceptions.TimeoutError:
+                except etcd.EtcdWatchTimedOut:
-                    self.client.http.clear()
+                    self._client.http.clear()
-                    return
+                    return False
                 except etcd.EtcdException:
-                    index = None
+                    logging.exception('watch')
 
                 timeout = end_time - time.time()
 
-        timeout > 0 and super(Etcd, self).watch(timeout)
+        try:
+            return super(Etcd, self).watch(timeout)
+        finally:
+            self.event.clear()

patroni/exceptions.py

@@ -7,11 +7,19 @@ class PatroniException(Exception):
 
     def __str__(self):
         """
-        >>> str(DCSError('foo'))
+        >>> str(PatroniException('foo'))
         "'foo'"
         """
         return repr(self.value)
 
 
+class PostgresException(PatroniException):
+    pass
+
+
 class DCSError(PatroniException):
     pass
+
+
+class PostgresConnectionException(PostgresException):
+    pass

patroni/ha.py

@@ -1,26 +1,35 @@
+import json
 import logging
+import psycopg2
+import requests
+import sys
+import datetime
+import pytz
 
-from patroni.dcs import DCSError
+from multiprocessing.pool import ThreadPool
-from psycopg2 import InterfaceError, OperationalError
+from patroni.async_executor import AsyncExecutor
+from patroni.exceptions import DCSError, PostgresConnectionException
+from patroni.utils import sleep
 
 logger = logging.getLogger(__name__)
 
 
-class Ha:
+class Ha(object):
 
-    def __init__(self, state_handler, etcd):
+    def __init__(self, patroni):
-        self.state_handler = state_handler
+        self.patroni = patroni
-        self.dcs = etcd
+        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 healhy
+        # We want to keep the state of cluster when it was healthy
-        if cluster.is_unlocked() and self.cluster and not self.cluster.is_unlocked():
+        if not cluster.is_unlocked() or not self.old_cluster:
-            self.old_cluster = self.cluster
-        if not self.old_cluster:
             self.old_cluster = cluster
         self.cluster = cluster
 
@@ -28,74 +37,456 @@ class Ha:
         return self.dcs.attempt_to_acquire_leader()
 
     def update_lock(self):
-        return self.dcs.update_leader(self.state_handler)
+        ret = self.dcs.update_leader()
+        if ret and not self._async_executor.busy:
+            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 demote(self):
+    def touch_member(self):
-        return self.state_handler.demote(self.cluster.leader)
+        data = {
+            'conn_url': self.state_handler.connection_string,
+            'api_url': self.patroni.api.connection_string,
+            'state': self.state_handler.state,
+            'role': self.state_handler.role
+        }
+        if self.patroni.tags:
+            data['tags'] = self.patroni.tags
+        if not self._async_executor.busy and 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 follow_the_leader(self):
+    def clone(self, clone_member=None, msg='(without leader)'):
-        return self.state_handler.follow_the_leader(self.cluster.leader)
+        if self.state_handler.bootstrap(cluster_initialized=True, clone_member=clone_member):
+            logger.info('bootstrapped %s', msg)
+        else:
+            logger.error('failed to bootstrap %s', msg)
+            self.state_handler.stop('immediate')
+            self.state_handler.remove_data_directory()
 
-    def run_cycle(self):
+    def bootstrap(self):
+        if not self.cluster.is_unlocked():  # cluster already has leader
+            clone_member = self.cluster.get_clone_member()
+            member_role = 'leader' if clone_member == self.cluster.leader else 'replica'
+            msg = "from {0} '{1}'".format(member_role, clone_member.name)
+            self._async_executor.schedule('bootstrap {0}'.format(msg))
+            self._async_executor.run_async(self.clone, args=(clone_member, msg))
+            return 'trying to bootstrap {0}'.format(msg)
+        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)
+                return "trying to bootstrap (without leader)"
+            return 'waiting for leader to bootstrap'
+
+    def recover(self):
+        self.recovering = True
+        return self.follow("starting as readonly because i had the session lock", "starting as a secondary", True, True)
+
+    def follow(self, demote_reason, follow_reason, refresh=True, recovery=False):
+        if refresh:
+            self.load_cluster_from_dcs()
+
+        ret = demote_reason if not recovery and self.state_handler.is_leader() else 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 = self.cluster.get_member(self.patroni.replicatefrom, fallback_to_leader=True)
+        else:
+            node_to_follow = self.cluster.leader
+        if node_to_follow and node_to_follow.name == self.state_handler.name:
+            ret = demote_reason
+            node_to_follow = None
+        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.candidate:  # manual failover to specific member
+            if failover.candidate == self.state_handler.name:  # manual failover to me
+                return True
+
+            # find specific node and check that it is healthy
+            member = self.cluster.get_member(failover.candidate, fallback_to_leader=False)
+            if member:
+                member, reachable, _, _, tags = self.fetch_node_status(member)
+                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.candidate 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.candidate can be None)
+                members = [m for m in self.cluster.members if m.name not in (failover.candidate, 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()
+            sleep(2)  # Give a time to somebody to promote
+            self.recover()
+        else:
+            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('', '', index=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.candidate or failover.candidate != self.state_handler.name:
+                members = [m for m in self.cluster.members if not failover.candidate or m.name == failover.candidate]
+                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('', '', index=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.load_cluster_from_dcs()
+                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()
+
+        clone_member = cluster.get_clone_member()
+        member_role = 'leader' if clone_member == cluster.leader else 'replica'
+        self.clone(clone_member, "from {0} '{1}'".format(member_role, clone_member.name))
+
+    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()
-            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():
+            self.touch_member()
-                if self.state_handler.is_healthiest_node(self.old_cluster):
+
-                    if self.acquire_lock():
+            # cluster has leader key but not initialize key
-                        if self.state_handler.is_leader() or self.state_handler.is_promoted:
+            if not self.cluster.is_unlocked() and not self.sysid_valid(self.cluster.initialize) and self.has_lock():
-                            return 'acquired session lock as a leader'
+                self.dcs.initialize(create_new=(self.cluster.initialize is None), sysid=self.state_handler.sysid)
-                        else:
+
-                            self.state_handler.promote()
+            if self._async_executor.busy:
-                            return 'promoted self to leader by acquiring session lock'
+                return self.handle_long_action_in_progress()
-                    else:
+
-                        self.load_cluster_from_dcs()
+            # we've got here, so any async action has finished. Check if we tried to recover and failed
-                        if self.state_handler.is_leader():
+            if self.recovering:
-                            self.demote()
+                self.recovering = False
-                            return 'demoted self due after trying and failing to obtain lock'
+                msg = self.post_recover()
-                        else:
+                if msg is not None:
-                            self.follow_the_leader()
+                    return msg
-                            return 'following new leader after trying and failing to obtain lock'
+
-                else:
+            # currently it can trigger only reinitialize
-                    self.load_cluster_from_dcs()
+            msg = self.process_scheduled_action()
-                    if self.state_handler.is_leader():
+            if msg is not None:
-                        self.demote()
+                return msg
-                        return 'demoting self because i am not the healthiest node'
+
-                    else:
+            # is data directory empty?
-                        self.follow_the_leader()
+            if self.state_handler.data_directory_empty():
-                        return 'following a different leader because i am not the healthiest node'
+                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:
-                if self.has_lock() and self.update_lock():
+                # check if we are allowed to join
-                    if self.state_handler.is_leader() or self.state_handler.is_promoted:
+                if self.sysid_valid(self.cluster.initialize) and self.cluster.initialize != self.state_handler.sysid:
-                        return 'no action.  i am the leader with the lock'
+                    logger.fatal("system ID mismatch, node %s belongs to a different cluster: %s != %s",
-                    else:
+                                 self.state_handler.name, self.cluster.initialize, self.state_handler.sysid)
-                        self.state_handler.promote()
+                    sys.exit(1)
-                        return 'promoted self to leader because i had the session lock'
+
+            # 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:
-                    logger.info('does not have lock')
+                    return self.process_healthy_cluster()
-                    if self.state_handler.is_leader():
+            finally:
-                        self.demote()
+                # we might not have a valid PostgreSQL connection here if another thread
-                        return 'demoting self because i do not have the lock and i was a leader'
+                # stops PostgreSQL, therefore, we only reload replication slots if no
-                    else:
+                # asynchronous processes are running (should be always the case for the master)
-                        self.follow_the_leader()
+                if not self._async_executor.busy:
-                        return 'no action.  i am a secondary and i am following a leader'
+                    self.state_handler.sync_replication_slots(self.cluster)
         except DCSError:
             logger.error('Error communicating with DCS')
-            if self.state_handler.is_leader():
+            if self.state_handler.is_running() and self.state_handler.is_leader():
-                self.state_handler.demote(None)
+                self.demote(delete_leader=False)
                 return 'demoted self because DCS is not accessible and i was a leader'
-        except (InterfaceError, OperationalError):
+        except (psycopg2.Error, PostgresConnectionException):
-            logger.error('Error communicating with Postgresql.  Will try again')
+            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

@@ -4,9 +4,14 @@ import psycopg2
 import shlex
 import shutil
 import subprocess
+import tempfile
+import time
 
-from patroni.utils import sleep
+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__)
 
@@ -34,37 +39,81 @@ def parseurl(url):
     return ret
 
 
-class Postgresql:
+class Postgresql(object):
 
     def __init__(self, config):
         self.config = config
         self.name = config['name']
+        self.server_parameters = config.get('parameters') or {}
         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.superuser = config.get('superuser') or {}
-        self.admin = config['admin']
+        self.admin = config.get('admin') or {}
-        self.callback = config.get('callbacks', {})
+
+        self.initdb_options = config.get('initdb') or []
+        self.pgpass = config.get('pgpass') or os.path.join(os.path.expanduser('~'), 'pgpass')
+        self.pg_rewind = config.get('pg_rewind') or {}
+        self.callback = config.get('callbacks') or {}
         self.use_slots = config.get('use_slots', True)
+        self.schedule_load_slots = self.use_slots
         self.recovery_conf = os.path.join(self.data_dir, 'recovery.conf')
         self.configuration_to_save = (os.path.join(self.data_dir, 'pg_hba.conf'),
                                       os.path.join(self.data_dir, 'postgresql.conf'))
         self.postmaster_pid = os.path.join(self.data_dir, 'postmaster.pid')
-        self.trigger_file = config.get('recovery_conf', {}).get('trigger_file', None) or 'promote'
+        self.trigger_file = 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.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
+        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.members = []  # list of already existing replication slots
+        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 = self.get_postgres_role_from_data_directory()
+        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 get_local_address(self):
         listen_addresses = self.listen_addresses.split(',')
@@ -76,159 +125,316 @@ class Postgresql:
                 break
         return local_address + ':' + self.port
 
+    def get_postgres_role_from_data_directory(self):
+        return 'replica' if os.path.exists(self.recovery_conf) else 'master'
+
+    @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:
-            r = parseurl('postgres://{}/postgres'.format(self.local_address))
+            self._connection = psycopg2.connect(**self._connect_kwargs)
-            self._connection = psycopg2.connect(**r)
             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:
+        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 disconnect(self):
+    def close_connection(self):
-        self._connection and self._connection.close()
+        if self._cursor_holder and self._cursor_holder.connection and self._cursor_holder.connection.closed == 0:
-        self._connection = self._cursor_holder = None
+            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):
-        max_attempts = 0
+        try:
-        while True:
+            return self.retry(self._query, sql, *params)
-            ex = None
+        except RetryFailedError as e:
-            try:
+            raise PostgresConnectionException(str(e))
-                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) == []
 
+    @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):
-        ret = subprocess.call(self._pg_ctl + ['initdb', '-o', '--encoding=UTF8']) == 0
+        self.set_state('initalizing new cluster')
-        ret and self.write_pg_hba()
+        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):
-        os.path.exists(self.trigger_file) and os.unlink(self.trigger_file)
+        if os.path.exists(self.trigger_file):
+            os.unlink(self.trigger_file)
 
-    def sync_from_leader(self, leader):
+    def write_pgpass(self, record):
-        r = parseurl(leader.conn_url)
+        with open(self.pgpass, 'w') as f:
-
-        pgpass = 'pgpass'
-        with open(pgpass, 'w') as f:
             os.fchmod(f.fileno(), 0o600)
-            f.write('{host}:{port}:*:{user}:{password}\n'.format(**r))
+            f.write('{host}:{port}:*:{user}:{password}\n'.format(**record))
 
         env = os.environ.copy()
-        env['PGPASSFILE'] = pgpass
+        env['PGPASSFILE'] = self.pgpass
-        return self.create_replica(r, env) == 0
+        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):
-        return "host={host} port={port} user={user}".format(**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 = replica_methods if clone_member else \
+            [r for r in replica_methods if self.replica_method_can_work_without_replication_connection(r)]
+        # 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
 
-    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):
+    def is_leader(self):
-        ret = not self.query('SELECT pg_is_in_recovery()').fetchone()[0]
+        return 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
+        return subprocess.call(' '.join(self._pg_ctl) + ' status > /dev/null 2>&1', shell=True) == 0
 
-    def call_nowait(self, cb_name, is_leader=None):
+    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]
-        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, self.role, self.scope])
-            subprocess.Popen(shlex.split(cmd) + [cb_name, role, self.scope])
+        except OSError:
-        except:
+            logger.exception('callback %s %s %s %s failed', cmd, cb_name, self.role, self.scope)
-            logger.exception('callback %s %s %s %s failed', cmd, cb_name, role, self.scope)
             return False
         return True
 
-    def start(self):
+    @property
-        if self.is_running():
+    def role(self):
-            self.load_replication_slots()
+        with self._role_lock:
-            logger.error('Cannot start PostgreSQL because one is already running.')
+            return self._role
-            return False
 
+    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(self.get_postgres_role_from_data_directory())
         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
+        if not block_callbacks:
-        ret and self.load_replication_slots()
+            self.set_state('starting')
+
+        env = {'PATH': os.environ.get('PATH')}
+        # pg_ctl will write a FATAL if the username is incorrect. exporting PGUSER if necessary
+        if 'username' in self.superuser and self.superuser['username'] != os.environ.get('USER'):
+            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()
-        if ret and ACTION_ON_START in self.callback:
+        # 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 stop(self):
+    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:
-            is_leader = self.is_leader(check_only=True)
+            with psycopg2.connect(**connect_kwargs) as conn:
-        except:
+                conn.autocommit = True
-            is_leader = None
+                with conn.cursor() as cur:
-            pass
+                    cur.execute("SET statement_timeout = 0")
-        ret = subprocess.call(self._pg_ctl + ['stop', '-m', 'fast'])
+                    cur.execute('CHECKPOINT')
-        if ret == 0 and ACTION_ON_STOP in self.callback:
+        except psycopg2.Error:
-            self.call_nowait(ACTION_ON_STOP, is_leader=is_leader)
+            logging.exception('Exception during CHECKPOINT')
-        return ret == 0
+
+    def stop(self, mode='fast', block_callbacks=False, checkpoint=True):
+        # 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 checkpoint:
+            self.checkpoint()
+
+        if not block_callbacks:
+            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'])
+        ret = subprocess.call(self._pg_ctl + ['reload']) == 0
-        if ret == 0 and ACTION_ON_RELOAD in self.callback:
+        if ret:
             self.call_nowait(ACTION_ON_RELOAD)
-        return ret == 0
+        return ret
 
     def restart(self):
-        try:
+        self.set_state('restarting')
-            is_leader = self.is_leader(check_only=True)
+        ret = self.stop(block_callbacks=True) and self.start(block_callbacks=True)
-        except:
+        if ret:
-            is_leader = None
+            self.call_nowait(ACTION_ON_RESTART)
-            pass
+        else:
-        ret = subprocess.call(self._pg_ctl + ['restart', '-m', 'fast'])
+            self.set_state('restart failed ({0})'.format(self.state))
-        if ret == 0 and ACTION_ON_RESTART in self.callback:
+        return ret
-            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)
+        options = "--listen_addresses='{0}' --port={1}".format(self.listen_addresses, self.port)
-        for setting, value in self.config['parameters'].items():
+        for setting, value in self.server_parameters.items():
-            options += " --{}='{}'".format(setting, value)
+            options += " --{0}='{1}'".format(setting, value)
         return options
 
     def is_healthy(self):
@@ -237,44 +443,12 @@ class Postgresql:
             return False
         return True
 
-    def is_healthiest_node(self, cluster):
+    def check_replication_lag(self, last_leader_operation):
-        if self.is_leader():
+        return (last_leader_operation or 0) - self.xlog_position() <= self.config.get('maximum_lag_on_failover', 0)
-            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_xlog_location_diff(pg_last_xlog_replay_location(), '0/0')",
-                    (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))
+            f.write('\n{}\n'.format('\n'.join(self.config.get('pg_hba', []))))
-            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):
@@ -290,104 +464,315 @@ class Postgresql:
         with open(self.recovery_conf, 'r') as f:
             for line in f:
                 if line.startswith('primary_conninfo'):
-                    if not pattern:
+                    return pattern and (pattern in line)
-                        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'
+            f.write("standby_mode = 'on'\nrecovery_target_timeline = 'latest'\n")
-recovery_target_timeline = 'latest'
-""")
             if leader and leader.conn_url:
-                f.write("""primary_conninfo = '{}'\n""".format(self.primary_conninfo(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 = '{}'\n""".format(self.name))
+                    f.write("primary_slot_name = '{0}'\n".format(self.name))
-                for name, value in self.config.get('recovery_conf', {}).items():
+            for name, value in self.config.get('recovery_conf', {}).items():
-                    f.write("{} = '{}'\n".format(name, value))
+                if name not in ('standby_mode', 'recovery_target_timeline', 'primary_conninfo', 'primary_slot_name'):
+                    f.write("{0} = '{1}'\n".format(name, value))
 
-    def follow_the_leader(self, leader):
+    def rewind(self, leader):
-        if not self.check_recovery_conf(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)
-            self.restart()
+        return ret
-            if ACTION_ON_ROLE_CHANGE in self.callback:
+
-                self.call_nowait(ACTION_ON_ROLE_CHANGE)
+    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'
+        need_rewind = change_role and self.can_rewind
+        if need_rewind:
+            logger.info("set the rewind flag after demote")
+        self.write_recovery_conf(leader)
+        if leader and 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
+        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 preserve it in the WAL-e backup.
+            copy postgresql.conf to postgresql.conf.backup to be able to retrive configuration files
-            see http://comments.gmane.org/gmane.comp.db.postgresql.wal-e/239
+            - 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)
         """
-        for f in self.configuration_to_save:
+        try:
-            shutil.copy(f, f + '.backup')
+            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:
-                shutil.copy(f + '.backup', f)
+                if not os.path.isfile(f) and os.path.isfile(f + '.backup'):
-        except:
+                    shutil.copy(f + '.backup', f)
-            logger.exception('unable to restore configuration from WAL-E backup')
+        except IOError:
+            logger.exception('unable to restore configuration files from backup')
 
     def promote(self):
-        self.is_promoted = subprocess.call(self._pg_ctl + ['promote']) == 0
+        if self.role == 'master':
-        if self.is_promoted and ACTION_ON_ROLE_CHANGE in self.callback:
+            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.call_nowait(ACTION_ON_ROLE_CHANGE)
-        return self.is_promoted
+        return ret
 
-    def demote(self, leader):
+    def create_or_update_role(self, name, password, options):
-        self.follow_the_leader(leader)
+        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.query('CREATE USER "{}" WITH REPLICATION ENCRYPTED PASSWORD %s'.format(
+        self.create_or_update_role(self.replication['username'], self.replication['password'], 'REPLICATION')
-            self.replication['username']), self.replication['password'])
 
-    def create_connection_users(self):
+    def create_connection_user(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.create_or_update_role(self.admin['username'], self.admin['password'], 'CREATEDB CREATEROLE')
-                self.admin['username']), self.admin['password'])
 
     def xlog_position(self):
         return self.query("""SELECT pg_xlog_location_diff(CASE WHEN pg_is_in_recovery()
                                                                THEN pg_last_xlog_replay_location()
                                                                ELSE pg_current_xlog_location()
-                                                          END, '0/0')""").fetchone()[0]
+                                                          END, '0/0')::bigint""").fetchone()[0]
 
     def load_replication_slots(self):
-        if self.use_slots:
+        if self.use_slots and self.schedule_load_slots:
             cursor = self.query("SELECT slot_name FROM pg_replication_slots WHERE slot_type='physical'")
-            self.members = [r[0] for r in cursor]
+            self.replication_slots = [r[0] for r in cursor]
+            self.schedule_load_slots = False
 
-    def sync_replication_slots(self, members):
+    def sync_replication_slots(self, cluster):
         if self.use_slots:
-            # drop unused slots
+            try:
-            for slot in set(self.members) - set(members):
+                self.load_replication_slots()
-                self.query("""SELECT pg_drop_replication_slot(%s)
+                # if the replicatefrom tag is set on the member - we should not create the replication slot for it on
-                               WHERE EXISTS(SELECT 1 FROM pg_replication_slots
+                # the current master, because that member would replicate from elsewhere. We still create the slot if
-                               WHERE slot_name = %s)""", slot, slot)
+                # 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
+                # create new slots
-            for slot in set(members) - set(self.members):
+                for slot in set(slots) - set(self.replication_slots):
-                self.query("""SELECT pg_create_physical_replication_slot(%s)
+                    self.query("""SELECT pg_create_physical_replication_slot(%s)
-                               WHERE NOT EXISTS (SELECT 1 FROM pg_replication_slots
+                                   WHERE NOT EXISTS (SELECT 1 FROM pg_replication_slots
-                               WHERE slot_name = %s)""", slot, slot)
+                                   WHERE slot_name = %s)""", slot, slot)
 
-        self.members = members
+                self.replication_slots = slots
-
+            except psycopg2.Error:
-    def create_replication_slots(self, cluster):
+                logger.exception('Exception when changing replication slots')
-        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())
+
+    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)
+                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/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]))

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

@@ -1,70 +1,65 @@
 import datetime
 import os
 import random
-import re
 import signal
 import sys
 import time
+import pytz
+import dateutil.parser
 
-from patroni.exceptions import DCSError
+from patroni.exceptions import PatroniException
 
-interrupted_sleep = False
+__ignore_sigterm = False
-reap_children = False
+__interrupted_sleep = False
-
+__reap_children = False
-_DATE_TIME_RE = re.compile(r'''^
-(?P<year>\d{4})\-(?P<month>\d{2})\-(?P<day>\d{2})  # date
-T
-(?P<hour>\d{2}):(?P<minute>\d{2}):(?P<second>\d{2})\.(?P<microsecond>\d{6})  # time
-\d*Z$''', re.X)
-
-
-def parse_datetime(time_str):
-    """
-    >>> parse_datetime('2015-06-10T12:56:30.552539016Z')
-    datetime.datetime(2015, 6, 10, 12, 56, 30, 552539)
-    >>> parse_datetime('2015-06-10 12:56:30.552539016Z')
-    """
-    m = _DATE_TIME_RE.match(time_str)
-    if not m:
-        return None
-    p = dict((n, int(m.group(n))) for n in 'year month day hour minute second microsecond'.split(' '))
-    return datetime.datetime(**p)
 
 
 def calculate_ttl(expiration):
     """
     >>> calculate_ttl(None)
-    >>> calculate_ttl('2015-06-10 12:56:30.552539016Z')
+    >>> calculate_ttl('2015-06-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
-    expiration = parse_datetime(expiration)
+    try:
-    if not expiration:
+        expiration = dateutil.parser.parse(expiration)
+    except (ValueError, TypeError):
         return None
-    now = datetime.datetime.utcnow()
+    now = datetime.datetime.now(pytz.utc)
     return int((expiration - now).total_seconds())
 
 
+def set_ignore_sigterm(value=True):
+    global __ignore_sigterm
+    __ignore_sigterm = value
+
+
 def sigterm_handler(signo, stack_frame):
-    sys.exit()
+    if not __ignore_sigterm:
+        set_ignore_sigterm()
+        sys.exit()
 
 
 def sigchld_handler(signo, stack_frame):
-    global interrupted_sleep, reap_children
+    global __interrupted_sleep, __reap_children
-    reap_children = interrupted_sleep = True
+    __reap_children = __interrupted_sleep = True
 
 
 def sleep(interval):
-    global interrupted_sleep
+    global __interrupted_sleep
     current_time = time.time()
     end_time = current_time + interval
     while current_time < end_time:
-        interrupted_sleep = False
+        __interrupted_sleep = False
         time.sleep(end_time - current_time)
-        if not interrupted_sleep:  # we will ignore only sigchld
+        if not __interrupted_sleep:  # we will ignore only sigchld
             break
         current_time = time.time()
-    interrupted_sleep = False
+    __interrupted_sleep = False
 
 
 def setup_signal_handlers():
@@ -73,8 +68,8 @@ def setup_signal_handlers():
 
 
 def reap_children():
-    global reap_children
+    global __reap_children
-    if reap_children:
+    if __reap_children:
         try:
             while True:
                 ret = os.waitpid(-1, os.WNOHANG)
@@ -83,20 +78,20 @@ def reap_children():
         except OSError:
             pass
         finally:
-            reap_children = False
+            __reap_children = False
 
 
-class RetryFailedError(DCSError):
+class RetryFailedError(PatroniException):
 
     """Raised when retrying an operation ultimately failed, after retrying the maximum number of attempts."""
 
 
-class Retry:
+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=time.sleep, deadline=None, retry_exceptions=DCSError):
+                 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.
@@ -150,13 +145,10 @@ class Retry:
                 if self._attempts == self.max_tries:
                     raise RetryFailedError("Too many retry attempts")
                 self._attempts += 1
-                sleeptime = self._cur_delay + (
+                sleeptime = self._cur_delay + (random.randint(0, self.max_jitter) / 100.0)
-                    random.randint(0, self.max_jitter) / 100.0)
 
-                if self._cur_stoptime is not None and \
+                if self._cur_stoptime is not None and time.time() + sleeptime >= self._cur_stoptime:
-                   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._cur_delay = min(self._cur_delay * self.backoff, self.max_delay)
-                                      self.max_delay)

patroni/version.py

@@ -1 +1 @@
-__version__ = '0.2'
+__version__ = '0.90'

patroni/zookeeper.py

@@ -5,7 +5,8 @@ import time
 
 from kazoo.client import KazooClient, KazooState
 from kazoo.exceptions import NoNodeError, NodeExistsError
-from patroni.dcs import AbstractDCS, Cluster, DCSError, Leader, Member, parse_connection_string
+from patroni.dcs import AbstractDCS, Cluster, Failover, Leader, Member
+from patroni.exceptions import DCSError
 from patroni.utils import sleep
 from requests.exceptions import RequestException
 
@@ -16,7 +17,7 @@ class ZooKeeperError(DCSError):
     pass
 
 
-class ExhibitorEnsembleProvider:
+class ExhibitorEnsembleProvider(object):
 
     TIMEOUT = 3.1
 
@@ -53,7 +54,7 @@ class ExhibitorEnsembleProvider:
     def _query_exhibitors(self, exhibitors):
         random.shuffle(exhibitors)
         for host in exhibitors:
-            uri = 'http://{}:{}{}'.format(host, self._exhibitor_port, self._uri_path)
+            uri = 'http://{0}:{1}{2}'.format(host, self._exhibitor_port, self._uri_path)
             try:
                 response = requests.get(uri, timeout=self.TIMEOUT)
                 return response.json()
@@ -82,145 +83,207 @@ class ZooKeeper(AbstractDCS):
             self.exhibitor = ExhibitorEnsembleProvider(exhibitor['hosts'], exhibitor['port'], poll_interval=interval)
             hosts = self.exhibitor.zookeeper_hosts
 
-        self.client = KazooClient(hosts=hosts,
+        self._client = KazooClient(hosts=hosts, timeout=(config.get('session_timeout') or 30),
-                                  timeout=(config.get('session_timeout', None) or 30),
+                                   command_retry={'deadline': (config.get('reconnect_timeout') or 10),
-                                  command_retry={
+                                                  'max_delay': 1, 'max_tries': -1},
-                                      'deadline': (config.get('reconnect_timeout', None) or 10),
+                                   connection_retry={'max_delay': 1, 'max_tries': -1})
-                                      'max_delay': 1,
+        self._client.add_listener(self.session_listener)
-                                      '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._my_member_data = None
-        self.members = []
+        self._fetch_cluster = True
-        self.leader = None
+        self._last_leader_operation = 0
-        self.last_leader_operation = 0
 
-        self.client.start(None)
+        self._client.start()
 
     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._fetch_cluster = True
-        self.cluster_event.set()
+        self.event.set()
 
-    def get_node(self, name, watch=None):
+    def get_node(self, key, watch=None):
         try:
-            return self.client.get(self.client_path(name), watch)
+            ret = self._client.get(key, watch)
+            return (ret[0].decode('utf-8'), ret[1])
         except NoNodeError:
-            pass
+            return None
-        except:
-            logger.exception('get_node')
-        return None
 
     @staticmethod
     def member(name, value, znode):
-        conn_url, api_url = parse_connection_string(value)
+        return Member.from_node(znode.version, name, znode.ephemeralOwner, value)
-        return Member(znode.mzxid, name, conn_url, api_url, None, None)
+
+    def get_children(self, key, watch=None):
+        try:
+            return self._client.get_children(key, watch)
+        except NoNodeError:
+            return []
 
     def load_members(self):
         members = []
-        for member in self.client.get_children(self.client_path('/members'), self.cluster_watcher):
+        for member in self.get_children(self.members_path, self.cluster_watcher):
-            data = self.get_node('/members/' + member)
+            data = self.get_node(self.members_path + member)
             if data is not None:
                 members.append(self.member(member, *data))
         return members
 
     def _inner_load_cluster(self):
-        self.cluster_event.clear()
+        self._fetch_cluster = False
-        leader = self.get_node('/leader', self.cluster_watcher)
+        self.event.clear()
-        self.members = self.load_members()
+        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
+            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.client_path('/leader'))
+                self._client.delete(self.leader_path)
                 leader = None
 
             if leader:
-                member = Member(-1, leader[0], None, None, None, None)
+                member = Member(-1, leader[0], None, {})
-                member = ([m for m in self.members if m.name == leader[0]] or [member])[0]
+                member = ([m for m in members if m.name == leader[0]] or [member])[0]
-                leader = Leader(leader[1].mzxid, None, None, member)
+                leader = Leader(leader[1].version, leader[1].ephemeralOwner, member)
-                self.fetch_cluster = member.index == -1
+                self._fetch_cluster = member.index == -1
 
-        self.leader = leader
+        # failover key
-        if self.fetch_cluster:
+        failover = self.get_node(self.failover_path, watch=self.cluster_watcher) if self._FAILOVER in nodes else None
-            last_leader_operation = self.get_node('/optime/leader')
+        if failover:
-            if last_leader_operation:
+            failover = Failover.from_node(failover[1].version, failover[0])
-                self.last_leader_operation = int(last_leader_operation[0])
 
-    def get_cluster(self):
+        # 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)
+            self._client.set_hosts(self.exhibitor.zookeeper_hosts)
 
-        if self.fetch_cluster:
+        if self._fetch_cluster or self._cluster is None:
             try:
-                self.client.retry(self._inner_load_cluster)
+                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)
+            self._client.retry(self._client.create, path, value.encode('utf-8'), **kwargs)
             return True
         except:
             return False
 
     def attempt_to_acquire_leader(self):
-        ret = self._create('/leader', self._name, makepath=True, ephemeral=True)
+        ret = self._create(self.leader_path, self._name, makepath=True, ephemeral=True)
-        ret or logger.info('Could not take out TTL lock')
+        if not ret:
+            logger.info('Could not take out TTL lock')
         return ret
 
-    def race(self, path):
+    def set_failover_value(self, value, index=None):
-        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)
+            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
+        member = 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')
+        if member and self._client.client_id is not None and member.session != self._client.client_id[0]:
+            try:
+                self._client.retry(self._client.delete, path)
+            except NoNodeError:
+                pass
+            except:
+                return False
+            member = None
+
+        if member and data == self._my_member_data:
+            return True
+
+        try:
+            if member:
+                self._client.retry(self._client.set, path, data)
+            else:
+                self._client.retry(self._client.create, path, data, makepath=True, ephemeral=True)
+            self._my_member_data = data
             return True
         except NodeExistsError:
             try:
-                self.client.retry(self.client.delete, path)
+                self._client.retry(self._client.set, path, data)
-                self.client.retry(self.client.create, path, connection_string, makepath=True, ephemeral=True)
+                self._my_member_data = data
                 return True
             except:
                 logger.exception('touch_member')
+        except:
+            logger.exception('touch_member')
         return False
 
     def take_leader(self):
         return self.attempt_to_acquire_leader()
 
-    def update_leader(self, state_handler):
+    def write_leader_optime(self, last_operation):
-        last_operation = state_handler.last_operation()
+        last_operation = last_operation.encode('utf-8')
-        if last_operation != self.last_leader_operation:
+        if last_operation != self._last_leader_operation:
-            self.last_leader_operation = last_operation
+            self._last_leader_operation = last_operation
-            path = self.client_path('/optime/leader')
+            path = self.leader_optime_path
             try:
-                self.client.retry(self.client.set, path, last_operation)
+                self._client.retry(self._client.set, path, last_operation)
             except NoNodeError:
                 try:
-                    self.client.retry(self.client.create, path, last_operation, makepath=True)
+                    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):
-        if isinstance(self.leader, Leader) and self.leader.name == self._name:
+        self._client.restart()
-            self.client.delete(self.client_path('/leader'))
+        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 delete_cluster(self):
+        try:
+            return self._client.retry(self._client.delete, self.client_path(''), recursive=True)
+        except NoNodeError:
+            return True
 
     def watch(self, timeout):
-        self.cluster_event.wait(timeout)
+        if super(ZooKeeper, self).watch(timeout):
-        if self.cluster_event.isSet():
+            self._fetch_cluster = True
-            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,11 +4,18 @@ 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
   host: 127.0.0.1:4001
   #discovery_srv: my-etcd.domain
+#consul:
+#  scope: *scope
+#  ttl: *ttl
+#  host: 127.0.0.1:8500
 #zookeeper:
 #  scope: *scope
 #  session_timeout: *ttl
@@ -31,31 +38,67 @@ postgresql:
   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 replication replicator 127.0.0.1/32 md5
   - host all all 0.0.0.0/0 md5
-  - hostssl 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
-  wal_e:
+  create_replica_method:
-    env_dir: /home/postgres/etc/wal-e.d/env
+    - basebackup
-    threshold_megabytes: 10240
+#   - wal_e
-    threshold_backup_size_percentage: 30
+# commented-out example for wal-e provisioning
-  restore: scripts/restore.py
+  #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
+    archive_command: mkdir -p ../wal_archive && test ! -f ../wal_archive/%f && cp %p ../wal_archive/%f
-    max_wal_senders: 5
+    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"
+    unix_socket_directories: '.'
+tags:
+    nofailover: False
+    noloadbalance: False
+    clonefrom: False

postgres1.yml

@@ -4,11 +4,18 @@ 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
   host: 127.0.0.1:4001
   #discovery_srv: my-etcd.domain
+#consul:
+#  scope: *scope
+#  ttl: *ttl
+#  host: 127.0.0.1:8500
 #zookeeper:
 #  scope: *scope
 #  session_timeout: *ttl
@@ -31,31 +38,68 @@ postgresql:
   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 replication replicator 127.0.0.1/32 md5
   - host all all 0.0.0.0/0 md5
-  - hostssl 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: cp ../wal_archive/%f %p
+    #restore_command: envdir /etc/wal-e.d/env wal-e wal-fetch "%f" "%p" -p 1
-  wal_e:
+  recovery_conf:
-    env_dir: /home/postgres/etc/wal-e.d/env
+    restore_command: cp ../wal_archive/%f %p
-    threshold_megabytes: 10240
-    threshold_backup_size_percentage: 30
-  restore: scripts/restore.py
   parameters:
     archive_mode: "on"
     wal_level: hot_standby
-    archive_command: mkdir -p ../wal_archive && cp %p ../wal_archive/%f
+    archive_command: mkdir -p ../wal_archive && test ! -f ../wal_archive/%f && cp %p ../wal_archive/%f
-    max_wal_senders: 5
+    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"
+    unix_socket_directories: '.'
+tags:
+    nofailover: False
+    noloadbalance: False
+    clonefrom: False

postgres2.yml

@@ -0,0 +1,106 @@
+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
+#consul:
+#  scope: *scope
+#  ttl: *ttl
+#  host: 127.0.0.1:8500
+#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 replication replicator 127.0.0.1/32 md5
+  - host all all 0.0.0.0/0 md5
+  # - hostssl all all 0.0.0.0/0 md5
+  replication:
+    username: replicator
+    password: rep-pass
+  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"
+    unix_socket_directories: '.'
+tags:
+    nofailover: False
+    noloadbalance: False
+    clonefrom: False
+    replicatefrom: postgresql1

release.sh

@@ -27,5 +27,5 @@ git push
 
 python3 setup.py sdist bdist_wheel upload
 
-git tag ${version}
+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,12 @@
+boto
+psycopg2>=2.6.1
+PyYAML
+requests
+six >= 1.7
+kazoo>=2.2.1
+python-etcd==0.4.3
+python-consul==0.6.0
+click>=4.1
+prettytable>=0.7
+tzlocal
+python-dateutil

setup.py

@@ -4,13 +4,12 @@
     Setup file for patroni
 """
 
-import sys
-import os
 import inspect
+import os
+import sys
 
-import setuptools
 from setuptools.command.test import test as TestCommand
-from setuptools import setup
+from setuptools import find_packages, setup
 
 if sys.version_info < (2, 7, 0):
     sys.stderr.write('FATAL: patroni needs to be run with Python 2.7+\n')
@@ -32,7 +31,7 @@ VERSION = read_version(MAIN_PACKAGE)
 DESCRIPTION = 'PostgreSQL High-Available orchestrator and CLI'
 LICENSE = 'The MIT License'
 URL = 'https://github.com/zalando/patroni'
-AUTHOR = 'Alexander Kukushkin, Alexey Klyukin, Feike Steenbergen'
+AUTHOR = 'Alexander Kukushkin, Oleksii Kliukin, Feike Steenbergen'
 AUTHOR_EMAIL = '[email protected], [email protected], [email protected]'
 KEYWORDS = 'etcd governor patroni postgresql postgres ha zookeeper streaming replication'
 
@@ -51,12 +50,15 @@ CLASSIFIERS = [
     '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']
+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):
@@ -73,7 +75,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', MAIN_PACKAGE, '--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:
@@ -89,10 +91,14 @@ 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', MAIN_PACKAGE, '-s', '-vv']
+        params['args'] += ['--doctest-modules', MAIN_PACKAGE, '-vv']
+
+        import logging
+        silence = logging.WARNING
+        logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=os.getenv('LOGLEVEL', silence))
+        params['args'] += ['-s' if logging.getLogger().getEffectiveLevel() < silence else '--capture=fd']
         errno = pytest.main(**params)
         sys.exit(errno)
 
@@ -114,8 +120,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.txt')
-    install_reqs = get_install_requirements(requirements)
 
     command_options = {'test': {'test_suite': ('setup.py', 'tests')}}
     if JUNIT_XML:
@@ -137,12 +142,12 @@ def setup_package():
         long_description=read('README.rst'),
         classifiers=CLASSIFIERS,
         test_suite='tests',
-        packages=setuptools.find_packages(exclude=['tests', 'tests.*']),
+        packages=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,63 @@
 import psycopg2
 import unittest
 
+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):
+class MockPostgresql(object):
-    raise psycopg2.OperationalError()
 
+    name = 'test'
+    state = 'running'
+    role = 'master'
+    server_version = '999999'
+    sysid = 'dummysysid'
+    scope = 'dummy'
 
-class MockPostgresql:
+    @staticmethod
-
+    def connection():
-    def connection(self):
         return psycopg2_connect()
 
-    def is_running(self):
+
-        return True
+class MockHa(object):
+
+    dcs = Mock()
+    state_handler = MockPostgresql()
+
+    @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 MockPatroni:
+class MockPatroni(object):
 
-    def __init__(self):
+    postgresql = MockPostgresql()
-        self.postgresql = MockPostgresql()
+    ha = MockHa()
+    dcs = Mock()
+    tags = {}
+    version = '0.00'
+    noloadbalance = Mock(return_value=False)
 
 
-class MockRequest:
+class MockRequest(object):
 
     def __init__(self, path):
         self.path = path
@@ -36,18 +68,143 @@ class MockRequest:
 
 class MockRestApiServer(RestApiServer):
 
-    def __init__(self, Handler, path, *args):
+    def __init__(self, Handler, path):
-        self.patroni = MockPatroni()
+        self.socket = 0
-        if len(args) > 0:
+        BaseHTTPServer.HTTPServer.__init__ = Mock()
-            self.query = args[0]
+        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 /replica')
-        MockRestApiServer(RestApiHandler, b'GET /', throws)
+        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 = Mock(return_value=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)
+
+        cluster.leader.name = 'postgresql1'
+        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 = 'postgresql2'
+        request = b'POST /failover HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0\n' +\
+                  b'Content-Length: 53\n\n{"leader": "postgresql1", "candidate": "postgresql2"}'
+        MockRestApiServer(RestApiHandler, request)
+
+        cluster.leader.name = 'postgresql1'
+        MockRestApiServer(RestApiHandler, request)
+
+        cluster.members = [Member(0, 'postgresql0', 30, {'api_url': 'http'}),
+                           Member(0, 'postgresql2', 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 = 'postgresql2'
+            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)
+
+        # 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)
+        with patch.object(MockPatroni, 'dcs') as d:
+            d.manual_failover.return_value = False
+            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 patroni.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_consul.py

@@ -0,0 +1,128 @@
+import consul
+import unittest
+
+from patroni.dcs import AbstractDCS
+from mock import Mock, patch
+from patroni.consul import Cluster, Consul, ConsulError, ConsulException, HTTPClient, NotFound
+from test_etcd import SleepException
+
+
+def kv_get(self, key, **kwargs):
+    if key == 'service/test/members/postgresql1':
+        return '1', {'Session': 'fd4f44fe-2cac-bba5-a60b-304b51ff39b7'}
+    if key == 'service/test/':
+        return None, None
+    if key == 'service/good/leader':
+        return '1', None
+    if key == 'service/good/':
+        return ('6429',
+                [{'CreateIndex': 1334, 'Flags': 0, 'Key': key + 'failover', 'LockIndex': 0,
+                  'ModifyIndex': 1334, 'Value': b''},
+                 {'CreateIndex': 1334, 'Flags': 0, 'Key': key + 'initialize', 'LockIndex': 0,
+                  'ModifyIndex': 1334, 'Value': b'postgresql0'},
+                 {'CreateIndex': 2621, 'Flags': 0, 'Key': key + 'leader', 'LockIndex': 1,
+                  'ModifyIndex': 2621, 'Session': 'fd4f44fe-2cac-bba5-a60b-304b51ff39b7', 'Value': b'postgresql1'},
+                 {'CreateIndex': 6156, 'Flags': 0, 'Key': key + 'members/postgresql0', 'LockIndex': 1,
+                  'ModifyIndex': 6156, 'Session': '782e6da4-ed02-3aef-7963-99a90ed94b53',
+                  'Value': ('postgres://replicator:[email protected]:5432/postgres' +
+                            '?application_name=http://127.0.0.1:8008/patroni').encode('utf-8')},
+                 {'CreateIndex': 2630, 'Flags': 0, 'Key': key + 'members/postgresql1', 'LockIndex': 1,
+                  'ModifyIndex': 2630, 'Session': 'fd4f44fe-2cac-bba5-a60b-304b51ff39b7',
+                  'Value': ('postgres://replicator:[email protected]:5433/postgres' +
+                            '?application_name=http://127.0.0.1:8009/patroni').encode('utf-8')},
+                 {'CreateIndex': 1085, 'Flags': 0, 'Key': key + 'optime/leader', 'LockIndex': 0,
+                  'ModifyIndex': 6429, 'Value': b'4496294792'}])
+    raise ConsulException
+
+
+class TestHTTPClient(unittest.TestCase):
+
+    def test_get(self):
+        self.client = HTTPClient('127.0.0.1', '8500', 'http', False)
+        self.client.session.get = Mock()
+        self.client.get(Mock(), '')
+        self.client.get(Mock(), '', {'wait': '1s', 'index': 1})
+
+
+@patch.object(consul.Consul.KV, 'get', kv_get)
+class TestConsul(unittest.TestCase):
+
+    @patch.object(consul.Consul.Session, 'create', Mock(return_value='fd4f44fe-2cac-bba5-a60b-304b51ff39b7'))
+    @patch.object(consul.Consul.Session, 'renew', Mock(side_effect=NotFound))
+    @patch.object(consul.Consul.KV, 'get', kv_get)
+    @patch.object(consul.Consul.KV, 'delete', Mock())
+    def setUp(self):
+        self.c = Consul('postgresql1', {'ttl': 30, 'scope': 'test', 'host': 'localhost:1'})
+        self.c._base_path = '/service/good'
+        self.c._load_cluster()
+
+    @patch('time.sleep', Mock(side_effect=SleepException))
+    def test_create_or_restore_session(self):
+        self.c._session = None
+        self.assertRaises(SleepException, self.c.create_or_restore_session)
+
+    @patch.object(consul.Consul.Session, 'renew', Mock(side_effect=NotFound))
+    @patch.object(consul.Consul.Session, 'create', Mock(side_effect=ConsulException))
+    def test_referesh_session(self):
+        self.c._session = '1'
+        self.c._name = ''
+        self.assertRaises(ConsulError, self.c.refresh_session)
+
+    @patch.object(consul.Consul.KV, 'delete', Mock())
+    def test_get_cluster(self):
+        self.c._base_path = '/service/test'
+        self.assertIsInstance(self.c.get_cluster(), Cluster)
+        self.assertIsInstance(self.c.get_cluster(), Cluster)
+        self.c._base_path = '/service/fail'
+        self.assertRaises(ConsulError, self.c.get_cluster)
+        self.c._base_path = '/service/good'
+        self.c._session = 'fd4f44fe-2cac-bba5-a60b-304b51ff39b8'
+        self.assertIsInstance(self.c.get_cluster(), Cluster)
+
+    @patch.object(consul.Consul.KV, 'delete', Mock(side_effect=[ConsulException, True, True]))
+    @patch.object(consul.Consul.KV, 'put', Mock(side_effect=[True, ConsulException]))
+    def test_touch_member(self):
+        self.c.refresh_session = Mock(return_value=True)
+        self.c.touch_member('balbla')
+        self.c.touch_member('balbla')
+        self.c.touch_member('balbla')
+        self.c.refresh_session = Mock(return_value=False)
+        self.c.touch_member('balbla')
+
+    @patch.object(consul.Consul.KV, 'put', Mock(return_value=False))
+    def test_take_leader(self):
+        self.c.take_leader()
+
+    @patch.object(consul.Consul.KV, 'put', Mock(return_value=True))
+    def test_set_failover_value(self):
+        self.c.set_failover_value('')
+
+    @patch.object(consul.Consul.KV, 'put', Mock(side_effect=ConsulException))
+    def test_write_leader_optime(self):
+        self.c.write_leader_optime('')
+
+    def test_update_leader(self):
+        self.c.update_leader()
+
+    @patch.object(consul.Consul.KV, 'delete', Mock(return_value=True))
+    def test_delete_leader(self):
+        self.c.delete_leader()
+
+    @patch.object(consul.Consul.KV, 'put', Mock(return_value=True))
+    def test_initialize(self):
+        self.c.initialize()
+
+    @patch.object(consul.Consul.KV, 'delete', Mock(return_value=True))
+    def test_cancel_initialization(self):
+        self.c.cancel_initialization()
+
+    @patch.object(consul.Consul.KV, 'delete', Mock(return_value=True))
+    def test_delete_cluster(self):
+        self.c.delete_cluster()
+
+    @patch.object(AbstractDCS, 'watch', Mock())
+    def test_watch(self):
+        self.c._name = ''
+        self.c.watch(1)
+        with patch.object(consul.Consul.KV, 'get', Mock(side_effect=ConsulException)):
+            self.c.watch(1)

tests/test_ctl.py

@@ -0,0 +1,311 @@
+import etcd
+import os
+import pytest
+import requests.exceptions
+import unittest
+
+from click.testing import CliRunner
+from mock import patch, Mock
+from patroni.ctl import ctl, members, store_config, load_config, output_members, post_patroni, get_dcs, parse_dcs, \
+    wait_for_leader, get_all_members, get_any_member, get_cursor, query_member, configure, PatroniCtlException
+from patroni.etcd import Etcd, Client
+from psycopg2 import OperationalError
+from test_etcd import etcd_read, requests_get, socket_getaddrinfo, MockResponse
+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')
+
+    store_config({'dcs_api': None}, CONFIG_FILE_PATH)
+    load_config(CONFIG_FILE_PATH, None)
+
+
+@patch('patroni.ctl.load_config', Mock(return_value={'etcd': {'host': 'localhost: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'})
+
+    @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_parse_dcs(self):
+        assert parse_dcs(None) is None
+        assert parse_dcs('localhost') == {'etcd': {'host': 'localhost:4001'}}
+        assert parse_dcs('') == {'etcd': {'host': 'localhost:4001'}}
+        assert parse_dcs('localhost:8500') == {'consul': {'host': 'localhost:8500'}}
+        assert parse_dcs('zookeeper://localhost') == {'zookeeper': {'hosts': ['localhost:2181']}}
+        assert parse_dcs('exhibitor://dummy') == {'zookeeper': {'exhibitor': {'hosts': ['dummy'], 'port': 8181}}}
+        assert parse_dcs('consul://localhost') == {'consul': {'host': 'localhost:8500'}}
+        self.assertRaises(PatroniCtlException, parse_dcs, 'invalid://test')
+
+    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.ctl.post_patroni', Mock(return_value=MockResponse()))
+    def test_failover(self):
+        result = self.runner.invoke(ctl, ['failover', 'dummy'], input='''leader\nother\n\ny''')
+        assert 'leader' in result.output
+
+        result = self.runner.invoke(ctl, ['failover', 'dummy'], input='''leader\nother\n2100-01-01T12:23:00\ny''')
+        assert result.exit_code == 0
+
+        result = self.runner.invoke(ctl, ['failover', 'dummy'], input='''leader\nother\n2030-01-01T12:23:00\ny''')
+        assert result.exit_code == 0
+
+        # Aborting failover,as we anser NO to the confirmation
+        result = self.runner.invoke(ctl, ['failover', 'dummy'], input='''leader\nother\n\nN''')
+        assert result.exit_code == 1
+
+        # Target and source are equal
+        result = self.runner.invoke(ctl, ['failover', 'dummy'], input='''leader\nleader\n\ny''')
+        assert result.exit_code == 1
+
+        # Reality is not part of this cluster
+        result = self.runner.invoke(ctl, ['failover', 'dummy'], input='''leader\nReality\n\ny''')
+        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'], 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'], input='''leader\nother\n\ny''')
+            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'], input='''leader\nother\n\ny''')
+            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'], input='''leader\nother\n\ny''')
+            assert 'falling back to DCS' in result.output
+
+        with patch('patroni.ctl.post_patroni') as mocked:
+            mocked.return_value.status_code = 500
+            result = self.runner.invoke(ctl, ['failover', 'dummy'], input='''leader\nother\n\ny''')
+            assert 'Failover failed' in result.output
+
+    def test_get_dcs(self):
+        self.assertRaises(PatroniCtlException, get_dcs, {'dummy': {}}, 'dummy')
+        with patch('patroni.Patroni.get_dcs', Mock(return_value=self.e)):
+            assert get_dcs({'etcd': {'host': 'none'}}, 'dummy').client_path('') == '/service/test/'
+
+    @patch('psycopg2.connect', psycopg2_connect)
+    @patch('patroni.ctl.query_member', Mock(return_value=([['mock column']], None)))
+    @patch.object(etcd.Client, 'read', etcd_read)
+    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'])
+                assert result.exit_code == 0
+
+                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
+
+    def test_query_member(self):
+        with patch('patroni.ctl.get_cursor', Mock(return_value=MockConnect().cursor())):
+            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('test_postgresql.MockCursor.execute', Mock(side_effect=OperationalError('bla'))):
+                rows = query_member(None, None, None, 'replica', 'SELECT pg_is_in_recovery()')
+
+        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()')
+
+    @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'])
+            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.post', requests_get)
+    def test_restart_reinit(self):
+        result = self.runner.invoke(ctl, ['restart', 'alpha'], input='y')
+        assert 'restart failed for' in result.output
+        assert result.exit_code == 0
+
+        result = self.runner.invoke(ctl, ['reinit', 'alpha'], input='y')
+        assert result.exit_code == 1
+
+        # Aborted restart
+        result = self.runner.invoke(ctl, ['restart', 'alpha'], input='N')
+        assert result.exit_code == 1
+
+        # Not a member
+        result = self.runner.invoke(ctl, ['restart', 'alpha', 'dummy', '--any'], input='y')
+        assert result.exit_code == 1
+
+        with patch('requests.post', Mock(return_value=MockResponse())):
+            result = self.runner.invoke(ctl, ['restart', 'alpha'], input='y')
+            assert result.exit_code == 0
+
+    @patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_leader()))
+    @patch.object(etcd.Client, 'delete', Mock(side_effect=etcd.EtcdException))
+    def test_remove(self):
+        with patch('patroni.ctl.get_dcs', Mock(return_value=self.e)):
+            result = self.runner.invoke(ctl, ['remove', 'alpha'], 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'], input='''alpha\nYes I am aware\nslave''')
+            assert result.exit_code == 1
+
+            # cluster specified on cmdline does not match verification prompt
+            result = self.runner.invoke(ctl, ['remove', 'alpha'], input='beta\nleader')
+            assert result.exit_code == 1
+
+            result = self.runner.invoke(ctl, ['remove', 'alpha'], input='''alpha\nYes I am aware\nleader''')
+            assert result.exit_code == 0
+
+    @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):
+        self.assertRaises(PatroniCtlException, wait_for_leader, self.e, 0)
+
+        cluster = wait_for_leader(self.e, timeout=2)
+        assert cluster.leader.member.name == 'leader'
+
+    @patch('requests.post', Mock(side_effect=requests.exceptions.ConnectionError('foo')))
+    def test_post_patroni(self):
+        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))
+    def test_members(self):
+        result = self.runner.invoke(members, ['alpha'])
+        assert '127.0.0.1' in result.output
+        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,63 +1,51 @@
-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 mock import Mock, patch
-from patroni.dcs import Cluster, DCSError, Leader, Member
+from patroni.dcs import Cluster, AbstractDCS
-from patroni.etcd import Client, Etcd
+from patroni.etcd import Client, Etcd, EtcdError
+from patroni.exceptions import DCSError
+from urllib3.exceptions import ReadTimeoutError
 
 
-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):
-        if self.content == 'TimeoutError':
+        return self.content.encode('utf-8')
-            raise urllib3.exceptions.TimeoutError
-        if self.content == 'Exception':
-            raise Exception
-        return self.content
 
     @property
     def status(self):
         return self.status_code
 
+    @staticmethod
     def getheader(*args):
         return ''
 
 
-class MockPostgresql:
-    name = ''
-
-    def last_operation(self):
-        return 0
-
-
 def requests_get(url, **kwargs):
     members = '[{"id":14855829450254237642,"peerURLs":["http://localhost:2380","http://localhost:7001"],' +\
               '"name":"default","clientURLs":["http://localhost:2379","http://localhost:4001"]}]'
     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 = '[{}]' if url.startswith('http://error') else members
-            response.content = '[{}]'
-        else:
-            response.content = members
     elif url.startswith('http://exhibitor'):
         response.content = '{"servers":["127.0.0.1","127.0.0.2","127.0.0.3"],"port":2181}'
     else:
@@ -66,39 +54,33 @@ def requests_get(url, **kwargs):
     return response
 
 
-def etcd_watch(key, index=None, timeout=None, recursive=None):
+def etcd_watch(self, key, index=None, timeout=None, recursive=None):
-    if timeout == 1:
+    if timeout == 2.0:
-        raise urllib3.exceptions.TimeoutError
+        raise etcd.EtcdWatchTimedOut
-    elif timeout == 5:
+    elif timeout == 5.0:
         return etcd.EtcdResult('delete', {})
-    elif timeout == 10:
+    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):
+def etcd_write(self, key, value, **kwargs):
     if key == '/service/exists/leader':
         raise etcd.EtcdAlreadyExist
-    if key == '/service/test/leader':
+    if key in ['/service/test/leader', '/patroni/test/leader'] and \
-        if kwargs.get('prevValue', None) == 'foo' or not kwargs.get('prevExist', True):
+            (kwargs.get('prevValue') == 'foo' or not kwargs.get('prevExist', True)):
-            return True
+        return True
     raise etcd.EtcdException
 
 
-def etcd_delete(key, **kwargs):
+def etcd_read(self, key, **kwargs):
-    raise etcd.EtcdException
+    if key == '/service/noleader/':
-
-
-def etcd_read(key, **kwargs):
-    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",
@@ -110,37 +92,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"
+                     "value": "postgres://replicator:[email protected]:5434/postgres" +
-                        + "?application_name=http://127.0.0.1:8009/patroni",
+                        "?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"
+                     "value": "postgres://replicator:[email protected]:5433/postgres" +
-                        + "?application_name=http://127.0.0.1:8008/patroni",
+                        "?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(_):
-    pass
-
-
 class SleepException(Exception):
     pass
 
 
-def time_sleep_exception(_):
+class MockSRV(object):
-    raise SleepException()
-
-
-class MockSRV:
     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':
@@ -151,66 +125,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 ReadTimeoutError(None, None, None)
     if url == 'http://localhost:2379/':
         return MockResponse()
     raise socket.error
 
 
-class TestMember(unittest.TestCase):
+@patch('dns.resolver.query', dns_query)
-
+@patch('socket.getaddrinfo', socket_getaddrinfo)
-    def __init__(self, method_name='runTest'):
+@patch('requests.get', requests_get)
-        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)
-
-
 class TestClient(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
+    @patch('dns.resolver.query', dns_query)
-        self.setUp = self.set_up
+    @patch('requests.get', requests_get)
-        super(TestClient, self).__init__(method_name)
+    def setUp(self):
-
-    def set_up(self):
-        socket.getaddrinfo = socket_getaddrinfo
-        requests.get = requests_get
-        dns.resolver.query = dns_query
         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__result_from_response(self):
-        response = MockResponse()
-        response.content = 'TimeoutError'
-        self.assertRaises(urllib3.exceptions.TimeoutError, self.client._result_from_response, response)
-        response.content = 'Exception'
-        self.assertRaises(etcd.EtcdException, self.client._result_from_response, response)
-        response.content = b'{}'
-        self.assertRaises(etcd.EtcdException, self.client._result_from_response, response)
-
     def test__get_machines_cache_from_srv(self):
-        self.client.get_srv_record = lambda e: [('localhost', 2380)]
+        self.client.get_srv_record = Mock(return_value=[('localhost', 2380)])
         self.client._get_machines_cache_from_srv('blabla')
 
     def test__get_machines_cache_from_dns(self):
-        self.client._get_machines_cache_from_dns('ok:2379')
+        self.client._get_machines_cache_from_dns('error:2379')
 
     def test__load_machines_cache(self):
         self.client._config = {}
@@ -219,25 +183,26 @@ class TestClient(unittest.TestCase):
         self.assertRaises(etcd.EtcdException, self.client._load_machines_cache)
 
 
+@patch('requests.get', requests_get)
+@patch.object(etcd.Client, 'write', etcd_write)
+@patch.object(etcd.Client, 'read', etcd_read)
+@patch.object(etcd.Client, 'delete', Mock(side_effect=etcd.EtcdException))
 class TestEtcd(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
+    def setUp(self):
-        self.setUp = self.set_up
-        super(TestEtcd, self).__init__(method_name)
-
-    def set_up(self):
-        time.sleep = time_sleep
         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
 
+    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(SleepException, 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)
@@ -245,11 +210,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(), Leader)
         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('', ''))
@@ -257,27 +219,36 @@ class TestEtcd(unittest.TestCase):
     def test_take_leader(self):
         self.assertFalse(self.etcd.take_leader())
 
-    def testattempt_to_acquire_leader(self):
+    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_update_leader(self):
+    def test_write_leader_optime(self):
-        self.assertTrue(self.etcd.update_leader(MockPostgresql()))
+        self.etcd.write_leader_optime('0')
 
-    def test_race(self):
+    def test_update_leader(self):
-        self.assertFalse(self.etcd.race(''))
+        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())
 
+    @patch.object(etcd.Client, 'watch', etcd_watch)
     def test_watch(self):
-        self.etcd.client.watch = etcd_watch
+        self.etcd.watch(0)
-        self.etcd.watch(100)
         self.etcd.get_cluster()
-        self.etcd.watch(1)
+        self.etcd.watch(1.5)
-        self.etcd.watch(5)
+        self.etcd.watch(4.5)
-        self.etcd.watch(10)
+        with patch.object(AbstractDCS, 'watch', Mock()):
-        self.etcd.watch(100)
+            self.etcd.watch(9.5)
+
+    @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 etcd
 import unittest
+import datetime
+import pytz
 
-from mock import Mock, patch
+from mock import Mock, MagicMock, patch
-from patroni.dcs import Cluster, DCSError
+from patroni.dcs import Cluster, Failover, Leader, Member
 from patroni.etcd import Client, Etcd
+from patroni.exceptions import DCSError, PostgresException
 from patroni.ha import Ha
-from test_etcd import etcd_read, etcd_write
+from patroni.postgresql import Postgresql
+from test_etcd import socket_getaddrinfo, etcd_read, etcd_write, requests_get
 
 
 def true(*args, **kwargs):
@@ -15,86 +20,127 @@ def false(*args, **kwargs):
     return False
 
 
-class MockPostgresql:
+def get_cluster(initialize, leader, members, failover):
-
+    return Cluster(initialize, leader, 10, members, failover)
-    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 nop(*args, **kwargs):
+def get_cluster_not_initialized_without_leader():
-    pass
+    return get_cluster(None, None, [], None)
 
 
-def dead_etcd():
+def get_cluster_initialized_without_leader(leader=False, failover=None):
-    raise DCSError('Etcd is not responding properly')
+    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', 'tags': {'clonefrom': True}})
+    return get_cluster(True, l, [m1, m2], failover)
 
 
-def get_unlocked_cluster():
+def get_cluster_initialized_with_leader(failover=None):
-    return Cluster(False, None, 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 = {'foo': 'bar'}
+        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.object(etcd.Client, 'write', etcd_write)
+@patch.object(etcd.Client, 'read', etcd_read)
+@patch.object(etcd.Client, 'delete', Mock(side_effect=etcd.EtcdException))
+@patch('subprocess.call', Mock(return_value=0))
 class TestHa(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
+    @patch('socket.getaddrinfo', socket_getaddrinfo)
-        self.setUp = self.set_up
+    @patch.object(etcd.Client, 'read', etcd_read)
-        super(TestHa, self).__init__(method_name)
+    def setUp(self):
-
-    def set_up(self):
-        self.p = MockPostgresql()
         with patch.object(Client, 'machines') as mock_machines:
             mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
-            self.e = Etcd('foo', {'ttl': 30, 'host': 'remotehost:2379', 'scope': 'test'})
+            self.p = Postgresql({'name': 'postgresql0', 'scope': 'dummy', 'listen': '127.0.0.1:5432',
-            self.e.client.read = etcd_read
+                                 'data_dir': 'data/postgresql0', 'superuser': {}, 'admin': {},
-            self.e.client.write = etcd_write
+                                 'replication': {'username': '', 'password': '', 'network': ''}})
-            self.ha = Ha(self.p, self.e)
+            self.p.set_state('running')
-            self.ha.load_cluster_from_dcs()
+            self.p.set_role('replica')
-            self.ha.cluster = get_unlocked_cluster()
+            self.p.check_replication_lag = true
-            self.ha.load_cluster_from_dcs = nop
+            self.p.can_create_replica_without_replication_connection = MagicMock(return_value=False)
+            self.e = Etcd('foo', {'ttl': 30, 'host': 'ok:2379', 'scope': 'test'})
+            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):
+    def test_update_lock(self):
-        ha = Ha(self.p, self.e)
+        self.p.last_operation = Mock(side_effect=PostgresException(''))
-        ha.load_cluster_from_dcs()
+        self.assertTrue(self.ha.update_lock())
-        self.e.get_cluster = get_unlocked_cluster
-        ha.load_cluster_from_dcs()
 
-    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')
+        self.assertEquals(self.ha.run_cycle(), 'starting 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(), 'starting 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.p.name = 'leader'
+        self.p.set_role('master')
+        self.p.controldata = lambda: {'Database cluster state': 'in production'}
+        self.ha.cluster = get_cluster_initialized_with_leader()
+        self.assertEquals(self.ha.run_cycle(), 'starting as readonly because i had the session lock')
+
+    @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 = false
-        self.p.is_leader = self.p.is_healthy = 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 +148,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 +186,186 @@ 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_another_member(self):
+        self.ha.cluster = get_cluster_initialized_with_leader()
+        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)
+    @patch('time.sleep', Mock())
+    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,7 +1,5 @@
-import datetime
+import etcd
-import patroni.zookeeper
+import os
-import psycopg2
-import subprocess
 import sys
 import time
 import unittest
@@ -9,141 +7,90 @@ import yaml
 
 from mock import Mock, patch
 from patroni.api import RestApiServer
-from patroni.dcs import Cluster, Member
+from patroni.async_executor import AsyncExecutor
+from patroni.consul import Consul
 from patroni.etcd import Etcd
-from patroni import Patroni, main
+from patroni import Patroni, PatroniException, main as _main
 from patroni.zookeeper import ZooKeeper
 from six.moves import BaseHTTPServer
-from test_etcd import Client, etcd_read, etcd_write
+from test_etcd import Client, SleepException, etcd_read, etcd_write
-from test_ha import true, false
+from test_postgresql import Postgresql, psycopg2_connect
-from test_postgresql import Postgresql, subprocess_call, psycopg2_connect
 from test_zookeeper import MockKazooClient
 
 
-def nop(*args, **kwargs):
+@patch('time.sleep', Mock())
-    pass
+@patch('subprocess.call', Mock(return_value=0))
-
+@patch('psycopg2.connect', psycopg2_connect)
-
+@patch.object(Postgresql, 'write_pg_hba', Mock())
-class SleepException(Exception):
+@patch.object(Postgresql, 'write_recovery_conf', Mock())
-    pass
+@patch.object(BaseHTTPServer.HTTPServer, '__init__', Mock())
-
+@patch.object(AsyncExecutor, 'run', Mock())
-
+@patch.object(etcd.Client, 'write', etcd_write)
-def time_sleep(*args):
+@patch.object(etcd.Client, 'read', etcd_read)
-    raise SleepException()
-
-
-class Mock_BaseServer__is_shut_down:
-
-    def set(self):
-        pass
-
-    def clear(self):
-        pass
-
-
 class TestPatroni(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
+    def setUp(self):
-        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'])
-                self.p = Patroni(config)
-
-    def tear_down(self):
-        time.sleep = self.time_sleep
-        Postgresql.write_pg_hba = self.write_pg_hba
-        Postgresql.write_recovery_conf = self.write_recovery_conf
-
-    def test_get_dcs(self):
-        patroni.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()
-        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
+            self.touched = False
-            touch_member = Patroni.touch_member
+            self.init_cancelled = False
-            run = Patroni.run
+            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)
 
-            Patroni.touch_member = self.touch_member
+    @patch('patroni.zookeeper.KazooClient', MockKazooClient())
-            Patroni.run = time_sleep
+    @patch.object(Consul, 'create_or_restore_session', Mock())
+    def test_get_dcs(self):
+        self.assertIsInstance(self.p.get_dcs('', {'zookeeper': {'scope': '', 'hosts': ''}}), ZooKeeper)
+        self.assertIsInstance(self.p.get_dcs('', {'consul': {'scope': '', 'hosts': '127.0.0.1:1'}}), Consul)
+        self.assertRaises(PatroniException, self.p.get_dcs, '', {})
 
-            Etcd.delete_leader = nop
+    @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']
 
-            self.assertRaises(SleepException, main)
+        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()
+        sys.argv = ['patroni.py']
+        # read the content of the yaml configuration file into the environment variable
+        # in order to test how does patroni handle the configuration passed from the environment.
+        with open('postgres0.yml', 'r') as f:
+            os.environ[Patroni.PATRONI_CONFIG_VARIABLE] = f.read()
+        with patch.object(Patroni, 'run', Mock(side_effect=SleepException())):
+            self.assertRaises(SleepException, _main)
+        del os.environ[Patroni.PATRONI_CONFIG_VARIABLE]
 
-            Patroni.run = run
+    def test_run(self):
-            Patroni.touch_member = touch_member
+        self.p.ha.dcs.watch = Mock(side_effect=SleepException)
-
+        self.p.api.start = Mock()
-    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.p.ha.dcs.watch = time_sleep
         self.assertRaises(SleepException, self.p.run)
-        self.p.ha.state_handler.is_leader = lambda: False
-        self.p.api.start = nop
-        self.assertRaises(SleepException, 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()
 
     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_noloadbalance(self):
+        self.p.tags['noloadbalance'] = True
+        self.assertTrue(self.p.noloadbalance)
+
+    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,53 +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 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):
+class MockCursor(object):
-    pass
 
-
+    def __init__(self, connection):
-def subprocess_call(cmd, shell=False, env=None):
+        self.connection = connection
-    return 0
-
-
-def false(*args, **kwargs):
-    return False
-
-
-class MockCursor:
-
-    def __init__(self):
         self.closed = False
         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'):
+        elif sql.startswith('RetryFailedError'):
-            raise psycopg2.InterfaceError()
+            raise RetryFailedError('retry')
         elif sql.startswith('SELECT slot_name'):
             self.results = [('blabla',), ('foobar',)]
-        elif sql.startswith('SELECT pg_current_xlog_location()'):
-            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 pg_xlog_location_diff'):
             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,
@@ -65,155 +51,442 @@ class MockCursor:
     def fetchone(self):
         return self.results[0]
 
-    def close(self):
+    def fetchall(self):
-        pass
+        return self.results
 
     def __iter__(self):
         for i in self.results:
             yield i
 
+    def __enter__(self):
+        return self
 
-class MockConnect:
+    def __exit__(self, *args):
-
-    def __init__(self):
-        self.autocommit = False
-        self.closed = 0
-
-    def cursor(self):
-        return MockCursor()
-
-    def close(self):
         pass
 
 
-def psycopg2_connect(*args, **kwargs):
+class MockConnect(object):
 
+    server_version = '99999'
+    autocommit = False
+    closed = 0
+
+    def cursor(self):
+        return MockCursor(self)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+    @staticmethod
+    def close():
+        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():
+def fake_listdir(path):
-    return False
+    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'):
+    @patch('subprocess.call', Mock(return_value=0))
-        self.setUp = self.set_up
+    @patch('psycopg2.connect', psycopg2_connect)
-        self.tearDown = self.tear_down
+    def setUp(self):
-        super(TestPostgresql, self).__init__(method_name)
-
-    def set_up(self):
-        subprocess.call = subprocess_call
-        shutil.copy = nop
         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'],
+                             'pg_hba': ['host replication replicator 127.0.0.1/32 md5',
-                             'superuser': {'password': ''},
+                                        'hostssl all all 0.0.0.0/0 md5',
+                                        'host all all 0.0.0.0/0 md5'],
+                             'superuser': {'username': 'test', 'password': 'test'},
                              'admin': {'username': 'admin', 'password': 'admin'},
+                             'pg_rewind': {'username': 'admin', 'password': 'admin'},
                              'replication': {'username': 'replicator',
-                                             'password': 'rep-pass',
+                                             'password': 'rep-pass'},
-                                             'network': '127.0.0.1/32'},
                              'parameters': {'foo': 'bar'}, 'recovery_conf': {'foo': 'bar'},
                              'callbacks': {'on_start': 'true', 'on_stop': 'true',
                                            'on_restart': 'true', 'on_role_change': 'true',
                                            'on_reload': 'true'
                                            },
                              'restore': 'true'})
-        psycopg2.connect = psycopg2_connect
         if not os.path.exists(self.p.data_dir):
             os.makedirs(self.p.data_dir)
-        self.leadermem = Member(0, 'leader', 'postgres://replicator:[email protected]:5435/postgres', None, None, 28)
+        self.leadermem = Member(0, 'leader', 28, {'conn_url': 'postgres://replicator:[email protected]:5435/postgres'})
-        self.leader = Leader(-1, None, 28, self.leadermem)
+        self.leader = Leader(-1, 28, self.leadermem)
-        self.other = Member(0, 'test1', 'postgres://replicator:[email protected]:5433/postgres', None, None, 28)
+        self.other = Member(0, 'test1', 28, {'conn_url': 'postgres://replicator:[email protected]:5433/postgres',
-        self.me = Member(0, 'test0', 'postgres://replicator:[email protected]:5434/postgres', None, None, 28)
+                            '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):
+        with open(os.path.join(self.p.data_dir, 'pg_hba.conf')) as f:
-        self.assertFalse(self.p.start())
+            lines = f.readlines()
-        self.p.is_running = is_running
+            assert 'host replication replicator 127.0.0.1/32 md5\n' in lines
-        with open(os.path.join(self.p.data_dir, 'postmaster.pid'), 'w'):
+            assert 'host all all 0.0.0.0/0 md5\n' in lines
-            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):
+    def test_restart(self):
-        self.assertTrue(self.p.sync_from_leader(self.leader))
+        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):
+    @patch.object(builtins, 'open', MagicMock())
-        self.p.demote(self.leader)
+    def test_write_pgpass(self):
-        self.p.follow_the_leader(None)
+        self.p.write_pgpass({'host': 'localhost', 'port': '5432', 'user': 'foo', 'password': 'bar'})
-        self.p.demote(self.leader)
-        self.p.follow_the_leader(self.leader)
-        self.p.follow_the_leader(Leader(-1, None, 28, self.other))
 
-    def test_create_connection_users(self):
+    @patch('patroni.postgresql.Postgresql.write_pgpass', MagicMock(return_value=dict()))
-        cfg = self.p.config
+    def test_sync_replica(self):
-        cfg['superuser']['username'] = 'test'
+        self.assertTrue(self.p.sync_replica(self.leader))
-        p = Postgresql(cfg)
+        self.p.create_replica = Mock(return_value=1)
-        p.create_connection_users()
+        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)
+        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)
+        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.leadermem])
+        cluster = Cluster(True, self.leader, 0, [self.me, self.other, self.leadermem], None)
-        self.p.create_replication_slots(cluster)
+        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.leadermem])
-        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'] = -3
-        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 patroni.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,17 +1,9 @@
-import os
-import time
 import unittest
 
-from patroni.exceptions import DCSError
+from mock import Mock, patch
-from patroni.utils import Retry, RetryFailedError, reap_children, sigchld_handler, sigterm_handler, sleep
+from patroni.exceptions import PatroniException
-
+from patroni.utils import reap_children, Retry, RetryFailedError, set_ignore_sigterm,\
-
+    sigchld_handler, sigterm_handler, sleep
-def nop(*args, **kwargs):
-    pass
-
-
-def os_waitpid(a, b):
-    return (0, 0)
 
 
 def time_sleep(_):
@@ -20,38 +12,27 @@ 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):
+        set_ignore_sigterm(False)
         self.assertRaises(SystemExit, sigterm_handler, None, None)
 
+    @patch('time.sleep', Mock())
     def test_reap_children(self):
-        reap_children()
+        self.assertIsNone(reap_children())
-        os.waitpid = os_waitpid
+        with patch('os.waitpid', Mock(return_value=(0, 0))):
-        sigchld_handler(None, None)
+            sigchld_handler(None, None)
-        reap_children()
+            self.assertIsNone(reap_children())
 
+    @patch('time.sleep', time_sleep)
     def test_sleep(self):
-        time.sleep = time_sleep
+        self.assertIsNone(sleep(0.01))
-        sleep(0.01)
 
 
+@patch('time.sleep', Mock())
 class TestRetrySleeper(unittest.TestCase):
 
-    def _pass(self):
+    @staticmethod
-        pass
+    def _fail(times=1):
-
-    def _fail(self, times=1):
         scope = dict(times=0)
 
         def inner():
@@ -59,43 +40,36 @@ class TestRetrySleeper(unittest.TestCase):
                 pass
             else:
                 scope['times'] += 1
-                raise DCSError('Failed!')
+                raise PatroniException('Failed!')
         return inner
 
-    def _makeOne(self, *args, **kwargs):
-        return Retry(*args, **kwargs)
-
     def test_reset(self):
-        retry = self._makeOne(delay=0, max_tries=2)
+        retry = 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 = self._makeOne(delay=0)
+        retry = Retry(delay=0)
         self.assertRaises(RetryFailedError, retry, self._fail(times=999))
         self.assertEquals(retry._attempts, 1)
 
     def test_maximum_delay(self):
-        def sleep_func(_time):
+        retry = Retry(delay=10, max_tries=100)
-            pass
-
-        retry = self._makeOne(delay=10, max_tries=100, sleep_func=sleep_func)
         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):
-        def sleep_func(_time):
+        retry = Retry(deadline=0.0001)
-            pass
-
-        retry = self._makeOne(deadline=0.0001, sleep_func=sleep_func)
         self.assertRaises(RetryFailedError, retry, self._fail(times=100))
 
     def test_copy(self):
-        _sleep = lambda t: None
+        def _sleep(t):
-        retry = self._makeOne(sleep_func=_sleep)
+            pass
+
+        retry = Retry(sleep_func=_sleep)
         rcopy = retry.copy()
         self.assertTrue(rcopy.sleep_func is _sleep)

tests/test_wale_restore.py

@@ -0,0 +1,80 @@
+import psycopg2
+import subprocess
+import unittest
+
+from mock import MagicMock, patch, PropertyMock
+from patroni.scripts.wale_restore import WALERestore, main as _main
+
+
+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', 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 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,130 +1,109 @@
-import patroni.zookeeper
+import six
-import requests
 import unittest
 
+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):
+    leader = False
-        pass
+    exists = True
-
-    def set(self):
-        pass
-
-    def wait(self, timeout):
-        pass
-
-    def isSet(self):
-        return True
-
-
-class MockEventHandler:
-
-    def event_object(self):
-        return MockEvent()
-
-
-class SleepException(Exception):
-    pass
-
-
-class MockKazooClient:
-
-    def __init__(self, **kwargs):
-        self.handler = MockEventHandler()
-        self.leader = False
-        self.exists = True
-
-    def start(self, timeout):
-        pass
 
     @property
     def client_id(self):
         return (-1, '')
 
-    def add_listener(self, cb):
+    @staticmethod
-        pass
+    def retry(func, *args, **kwargs):
-
-    def retry(self, func, *args, **kwargs):
         func(*args, **kwargs)
 
     def get(self, path, watch=None):
-        if path == '/service/test/no_node':
+        if not isinstance(path, six.string_types):
+            raise TypeError("Invalid type for 'path' (string expected)")
+        if path == '/no_node':
             raise NoNodeError
-        elif path == '/service/test/other_exception':
-            raise Exception()
         elif '/members/' in path:
             return (
-                'postgres://repuser:rep-pass@localhost:5434/postgres?application_name=http://127.0.0.1:8009/patroni',
+                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, 0, 0, 0)
+                ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0 if self.exists else -1, 0, 0, 0)
             )
-        elif path.endswith('/optime/leader'):
-            return '1'
         elif path.endswith('/leader'):
             if self.leader:
-                return ('foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0))
+                return (b'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, 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
-
+        elif path == '/service/test/members/buzz':
-    def set_hosts(self, hosts, randomize_hosts=None):
+            raise Exception
-        pass
+        elif path.endswith('/') or path.endswith('/initialize') or path == '/service/test/members/bar':
-
+            raise NoNodeError
-
-def exhibitor_sleep(_):
-    raise SleepException
 
 
+@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
-        patroni.zookeeper.sleep = exhibitor_sleep
-
     def test_init(self):
         self.assertRaises(SleepException, ExhibitorEnsembleProvider, ['localhost'], 8181)
 
 
 class TestZooKeeper(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
+    @patch('requests.get', requests_get)
-        self.setUp = self.set_up
+    @patch('patroni.zookeeper.KazooClient', MockKazooClient)
-        super(TestZooKeeper, self).__init__(method_name)
+    def setUp(self):
-
-    def set_up(self):
-        requests.get = requests_get
-        patroni.zookeeper.KazooClient = MockKazooClient
         self.zk = ZooKeeper('foo', {'exhibitor': {'hosts': ['localhost', 'exhibitor'], 'port': 8181}, 'scope': 'test'})
 
     def test_session_listener(self):
@@ -132,11 +111,15 @@ 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)
@@ -146,23 +129,50 @@ class TestZooKeeper(unittest.TestCase):
         self.zk.touch_member('foo')
         self.zk.delete_leader()
 
-    def test_race(self):
+    def test_set_failover_value(self):
-        self.assertFalse(self.zk.race('/initialize'))
+        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()
+        with patch.object(MockKazooClient, 'create', Mock(side_effect=Exception)):
+            self.zk.take_leader()
 
     def test_update_leader(self):
-        self.zk.last_leader_operation = -1
+        self.assertTrue(self.zk.update_leader())
-        self.assertTrue(self.zk.update_leader(MockPostgresql()))
+
+    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.last_leader_operation = -1
+        self.zk.write_leader_optime('2')
-        self.assertTrue(self.zk.update_leader(MockPostgresql()))
+
+    def test_delete_cluster(self):
+        self.assertTrue(self.zk.delete_cluster())
 
     def test_watch(self):
         self.zk.watch(0)
+        self.zk.event.isSet = lambda: True
+        self.zk.watch(0)