Bumped version to 0.7

Feature/patroni cli or ctl tbd

.gitignore

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

Dockerfile

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

MANIFEST.in

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

README.md

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

README.rst

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

docker/README.md

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

docker/dev_patroni_cluster.sh

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

docker/entrypoint.sh

@@ -3,25 +3,25 @@
 function usage()
 {
     cat <<__EOF__
-Usage: $0 
+Usage: $0
 
 Options:
 
     --etcd      ETCD    Provide an external etcd to connect to
-    --name      NAME    Give the cluster a specific name 
+    --name      NAME    Give the cluster a specific name
     --etcd-only         Do not run Patroni, run a standalone etcd
 
 Examples:
 
     $0 --etcd=127.17.0.84:4001
     $0 --etcd-only
-    $0 
+    $0
     $0 --name=true_scotsman
 __EOF__
 }
 
 DOCKER_IP=$(hostname --ip-address)
-PATRONI_SCOPE=batman
+PATRONI_SCOPE=${PATRONI_SCOPE:-batman}
 
 optspec=":vh-:"
 while getopts "$optspec" optchar; do
@@ -32,7 +32,7 @@ while getopts "$optspec" optchar; do
                     exec etcd --data-dir /tmp/etcd.data \
                         -advertise-client-urls=http://${DOCKER_IP}:4001 \
                         -listen-client-urls=http://0.0.0.0:4001 \
-                        -listen-peer-urls=http://0.0.0.0:2380 
+                        -listen-peer-urls=http://0.0.0.0:2380
                     exit 0
                     ;;
                 cheat)
@@ -83,10 +83,11 @@ cat > /patroni/postgres.yml <<__EOF__
 
 ttl: &ttl 30
 loop_wait: &loop_wait 10
-scope: &scope ${PATRONI_SCOPE}
+scope: &scope '${PATRONI_SCOPE}'
+namespace: 'patroni'
 restapi:
-  listen: 127.0.0.1:8008
-  connect_address: 127.0.0.1:8008
+  listen: 0.0.0.0:8008
+  connect_address: ${DOCKER_IP}:8008
 etcd:
   scope: *scope
   ttl: *ttl
@@ -115,7 +116,7 @@ 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
@@ -133,5 +134,5 @@ then
         sleep 60
     done
 else
-    exec /patroni/patroni.py /patroni/postgres.yml
+    exec python /patroni.py /patroni/postgres.yml
 fi

extras/startup-scripts/README.md

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

extras/startup-scripts/patroni.upstart.conf

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

haproxy.cfg

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

helpers/api.py

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

helpers/dcs.py

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

helpers/ha.py

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

helpers/postgresql.py

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

helpers/utils.py

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

patroni.py

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

patroni/__init__.py

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

patroni/__main__.py

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

patroni/api.py

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

patroni/async_executor.py

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

patroni/ctl.py

@@ -0,0 +1,613 @@
+'''
+Patroni Control
+'''
+
+import click
+import os
+import yaml
+import json
+import time
+import psycopg2
+import random
+import requests
+import datetime
+from prettytable import PrettyTable
+from six.moves.urllib_parse import urlparse
+import logging
+
+from .etcd import Etcd
+from .exceptions import PatroniCtlException
+from .postgresql import parseurl
+
+CONFIG_DIR_PATH = click.get_app_dir('patroni')
+CONFIG_FILE_PATH = os.path.join(CONFIG_DIR_PATH, 'patronictl.yaml')
+LOGLEVEL = 'WARNING'
+
+
+def parse_dcs(dcs):
+    """
+    Break up the provided dcs string
+    >>> parse_dcs('localhost') == {'scheme': 'etcd', 'hostname': 'localhost', 'port': 4001}
+    True
+    >>> parse_dcs('localhost:8500') == {'scheme': 'consul', 'hostname': 'localhost', 'port': 8500}
+    True
+    >>> parse_dcs('zookeeper://localhost') == {'scheme': 'zookeeper', 'hostname': 'localhost', 'port': 2181}
+    True
+    """
+
+    if not dcs:
+        return {}
+
+    parsed = urlparse(dcs)
+    scheme = parsed.scheme
+    if scheme == '' and parsed.netloc == '':
+        parsed = urlparse('//' + dcs)
+
+    if scheme == '':
+        default_schemes = {'2181': 'zookeeper', '8500': 'consul'}
+        scheme = default_schemes.get(str(parsed.port), 'etcd')
+
+    port = parsed.port
+    if port is None:
+        default_ports = {'consul': 8500, 'zookeeper': 2181}
+        port = default_ports.get(str(scheme), 4001)
+
+    return {'scheme': str(scheme), 'hostname': str(parsed.hostname), 'port': int(port)}
+
+
+def load_config(path, dcs):
+    logging.debug('Loading configuration from file {}'.format(path))
+    config = dict()
+    try:
+        with open(path, 'rb') as fd:
+            config = yaml.safe_load(fd)
+    except:
+        logging.exception('Could not load configuration file')
+
+    if dcs:
+        config['dcs'] = parse_dcs(dcs)
+    else:
+        config['dcs'] = parse_dcs(config.get('dcs_api'))
+
+    return config
+
+
+def store_config(config, path):
+    dir_path = os.path.dirname(path)
+    if dir_path:
+        if 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', help='Output format (pretty, json)', default='pretty')
+option_dcs = click.option('--dcs', '-d', help='Use this DCS', envvar='DCS')
+option_watchrefresh = click.option('-w', '--watch', type=float, help='Auto update the screen every X seconds')
+option_watch = click.option('-W', is_flag=True, help='Auto update the screen every 2 seconds')
+option_force = click.option('--force', is_flag=True, help='Do not ask for confirmation at any point')
+
+
+@click.group()
+@click.pass_context
+def ctl(ctx):
+    global LOGLEVEL
+    LOGLEVEL = os.environ.get('LOGLEVEL', LOGLEVEL)
+
+    logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s', level=LOGLEVEL)
+
+
+def get_dcs(config, scope):
+    scheme, hostname, port = map(config.get('dcs', {}).get, ('scheme', 'hostname', 'port'))
+
+    if scheme == 'etcd':
+        return Etcd(name=scope, config={'scope': scope, 'host': '{}:{}'.format(hostname, port)})
+
+    raise PatroniCtlException('Can not find suitable configuration of distributed configuration store')
+
+
+def post_patroni(member, endpoint, content, headers={'Content-Type': 'application/json'}):
+    url = urlparse(member.api_url)
+    logging.debug(url)
+    return requests.post('{}://{}/{}'.format(url.scheme, url.netloc, endpoint), headers=headers,
+                         data=json.dumps(content), timeout=5)
+
+
+def print_output(columns, rows=[], alignment=None, format='pretty', header=True, delimiter='\t'):
+    if format == '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 format == 'json':
+        elements = list()
+        for r in rows:
+            elements.append(dict(zip(columns, r)))
+
+        click.echo(json.dumps(elements))
+
+    if format == 'tsv':
+        if columns is not None and header:
+            click.echo(delimiter.join(columns) + '\n')
+
+        for r in rows or []:
+            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={}):
+    params = connect_parameters.copy()
+    parsed = parseurl(conn_url)
+    params['host'] = parsed['host']
+    params['port'] = parsed['port']
+    params['fallback_application_name'] = 'Patroni ctl'
+    params['connect_timeout'] = '5'
+
+    return params
+
+
+def get_all_members(cluster, role='master'):
+    if role == 'master':
+        if cluster.leader is not None:
+            yield cluster.leader
+        return
+
+    leader_name = (cluster.leader.member.name if cluster.leader else None)
+    for m in cluster.members:
+        if role == 'any' or role == 'replica' and m.name != leader_name:
+            yield m
+
+
+def get_any_member(cluster, role='master', member=None):
+    members = get_all_members(cluster=cluster, role=role)
+    for m in members:
+        if member is None or m.name == member:
+            return m
+
+    return None
+
+
+def get_cursor(cluster, role='master', member=None, connect_parameters={}):
+    member = get_any_member(cluster=cluster, role=role, member=member)
+    if member is None:
+        return None
+
+    params = build_connect_parameters(member.conn_url, connect_parameters=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'
+
+    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+    m = get_any_member(cluster=cluster, role=role, member=member)
+    if m is None:
+        raise PatroniCtlException('Can not find a suitable member')
+
+    params = build_connect_parameters(m.conn_url)
+    click.echo('host={} port={}'.format(params['host'], params['port']))
+
+
+@ctl.command('query', help='Query a Patroni PostgreSQL member')
+@click.argument('cluster_name')
+@option_config_file
+@option_format
+@click.option('--format', help='Output format (pretty, json)', default='tsv')
+@click.option('--file', '-f', help='Execute the SQL commands from this file', type=click.File('rb'))
+@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')
+def query(
+    cluster_name,
+    config_file,
+    dcs,
+    role,
+    member,
+    w,
+    watch,
+    delimiter,
+    command,
+    file,
+    format='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 file is not None and command is not None:
+        raise PatroniCtlException('--file and --command are mutually exclusive options')
+
+    if file is not None:
+        command = file.read()
+
+    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+
+    cursor = None
+    for _ in watching(w, watch, clear=False):
+
+        output, cursor = query_member(cluster=cluster, cursor=cursor, member=member, role=role, command=command)
+        print_output(None, output, format=format, delimiter=delimiter)
+
+        if cursor is None:
+            cluster = dcs.get_cluster()
+
+
+def query_member(cluster, cursor, member, role, command):
+    try:
+        if cursor is None:
+            cursor = get_cursor(cluster, role=role, member=member)
+
+        if cursor is None:
+            if role is None:
+                message = 'No connection to member {} is available'.format(member)
+            else:
+                message = 'No connection to role={} 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: {}'.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, format, dcs):
+    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+
+    if not isinstance(dcs, Etcd):
+        raise PatroniCtlException('We have not implemented this for DCS of type {}'.format(type(dcs)))
+
+    output_members(cluster, format=format)
+
+    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 {}, please type: "{}"'.format(cluster_name,
+                     message), type=str)
+    if message != confirm:
+        raise PatroniCtlException('You did not exactly type "{}"'.format(message))
+
+    if cluster.leader:
+        confirm = click.prompt('This cluster currently is healthy. Please specify the master name to continue')
+        if confirm != cluster.leader.name:
+            raise PatroniCtlException('You did not specify the current master of the cluster')
+
+    dcs.client.delete(dcs._base_path, recursive=True)
+
+
+def wait_for_leader(dcs, timeout=30):
+    t_stop = time.time() + timeout
+    timeout /= 2
+
+    while time.time() < t_stop:
+        dcs.watch(timeout)
+        cluster = dcs.get_cluster()
+
+        if cluster.leader:
+            return cluster
+
+    raise PatroniCtlException('Timeout occured')
+
+
+def empty_post_to_members(cluster, member_names, force, endpoint):
+    candidates = dict()
+    for m in cluster.members:
+        candidates[m.name] = m
+
+    if len(member_names) == 0:
+        member_names = [click.prompt('Which member do you want to {} [{}]?'.format(endpoint,
+                        ', '.join(candidates.keys())), type=str, default='')]
+
+    for mn in member_names:
+        if mn not in candidates.keys():
+            raise PatroniCtlException('{} is not a member of cluster'.format(mn))
+
+    if not force:
+        confirm = click.confirm('Are you sure you want to {} members {}?'.format(endpoint, ', '.join(member_names)))
+        if not confirm:
+            raise PatroniCtlException('Aborted {}'.format(endpoint))
+
+    for mn in member_names:
+        r = post_patroni(candidates[mn], endpoint, '')
+        if r.status_code != 200:
+            click.echo('{} failed for member {}, status code={}, ({})'.format(endpoint, mn, r.status_code, r.text))
+        else:
+            click.echo('Succesful {} on member {}'.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', 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, any):
+    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+
+    role_names = [m.name for m in get_all_members(cluster=cluster, role=role)]
+
+    if len(member_names) > 0:
+        member_names = list(set(member_names) & set(role_names))
+    else:
+        member_names = role_names
+
+    if any:
+        random.shuffle(member_names)
+        member_names = member_names[:1]
+
+    output_members(cluster)
+    empty_post_to_members(cluster, member_names, force, 'restart')
+
+
+@ctl.command('reinit', help='Reinitialize cluster member')
+@click.argument('cluster_name')
+@click.argument('member_names', nargs=-1)
+@option_config_file
+@option_force
+@option_dcs
+def reinit(cluster_name, member_names, config_file, dcs, force):
+    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+    empty_post_to_members(cluster, member_names, force, 'reinitialize')
+
+
+@ctl.command('failover', help='Failover to a replica')
+@click.argument('cluster_name')
+@click.option('--master', help='The name of the current master', default=None)
+@click.option('--candidate', help='The name of the candidate', default=None)
+@click.option('--force', is_flag=True)
+@option_config_file
+@option_dcs
+def failover(config_file, cluster_name, master, candidate, force, dcs):
+    """
+        We want to trigger a failover for the specified cluster name.
+
+        We verify that the cluster name, master name and candidate name are correct.
+        If so, we trigger a failover and keep the client up to date.
+    """
+
+    config, dcs, cluster = ctl_load_config(cluster_name, config_file, dcs)
+
+    if cluster.leader is None:
+        raise PatroniCtlException('This cluster has no master')
+
+    if master is None:
+        if force:
+            master = cluster.leader.member.name
+        else:
+            master = click.prompt('Master', type=str, default=cluster.leader.member.name)
+
+    if cluster.leader.member.name != master:
+        raise PatroniCtlException('Member {} is not the leader of cluster {}'.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 len(candidate_names) == 0:
+        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 {} does not exist in cluster {}'.format(candidate, cluster_name))
+
+    # By now we have established that the leader exists and the candidate exists
+    click.echo('Current cluster topology')
+    output_members(dcs.get_cluster(), name=cluster_name)
+
+    if not force:
+        a = \
+            click.confirm('Are you sure you want to failover cluster {}, demoting current master {}?'.format(
+                cluster_name, master))
+        if not a:
+            raise PatroniCtlException('Aborting failover')
+
+    failover_value = '{}:{}'.format(master, candidate or '')
+
+    t_started = time.time()
+    r = None
+    try:
+        r = post_patroni(cluster.leader.member, 'failover', {'leader': master, 'candidate': candidate or ''})
+        if r.status_code == 200:
+            logging.debug(r)
+            logging.debug(r.text)
+            cluster = dcs.get_cluster()
+            click.echo(timestamp() + ' Failing over to new leader: {}'.format(cluster.leader.member.name))
+        else:
+            click.echo('Failover failed, details: {}, {}'.format(r.status_code, r.text))
+            return
+    except:
+        logging.exception(r)
+        logging.warning('Failing over to DCS')
+        click.echo(timestamp() + ' Could not failover using Patroni api, falling back to DCS')
+        dcs.set_failover_value(failover_value)
+        click.echo(timestamp() + ' Initialized failover from master {}'.format(master))
+        # The failover process should within a minute update the failover key, we will keep watching it until it changes
+        # or we timeout
+        cluster = wait_for_leader(dcs, timeout=60)
+        if cluster.leader.member.name == master:
+            click.echo('Failover failed, master did not change after {:0.1f} seconds'.format(time.time() - t_started))
+            return
+
+    click.echo(timestamp() + ' Failover completed in {:0.1f} seconds, new leader is {}'.format(time.time() - t_started,
+               str(cluster.leader.member.name)))
+    output_members(cluster, name=cluster_name)
+
+
+def output_members(cluster, name=None, format='pretty'):
+    rows = []
+    logging.debug(cluster)
+    leader_name = None
+    if cluster.leader:
+        leader_name = cluster.leader.member.name
+
+    # 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')
+        lag = ''
+        if xlog_location is not None:
+            lag = round(((cluster.last_leader_operation or 0) - m.data.get('xlog_location', 0)) / 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'}
+
+    print_output(columns, rows, alignment, format)
+
+
+@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, format, watch, w, dcs):
+    if len(cluster_names) == 0:
+        logging.warning('Listing members: No cluster names were provided')
+        return
+
+    config = load_config(config_file, dcs)
+    for cn in cluster_names:
+        dcs = get_dcs(config, cn)
+
+        for _ in watching(w, watch):
+            output_members(dcs.get_cluster(), name=cn, format=format)
+
+
+def timestamp(precision=6):
+    return datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:precision - 7]
+
+
+@ctl.command('configure', help='Create configuration file')
+@click.option('--config-file', '-c', help='Configuration file', prompt='Configuration file', default=CONFIG_FILE_PATH)
+@click.option('--dcs', '-d', help='The DCS connect url', prompt='DCS connect url', default='etcd://localhost:4001')
+@click.option('--namespace', '-n', help='The namespace', prompt='Namespace', default='/service/')
+def configure(config_file, dcs, namespace):
+    config = dict()
+    config['dcs_api'] = str(dcs)
+    config['namespace'] = str(namespace)
+    store_config(config, config_file)

patroni/dcs.py

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

patroni/etcd.py

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

patroni/exceptions.py

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

patroni/ha.py

@@ -0,0 +1,439 @@
+import json
+import logging
+import psycopg2
+import requests
+import sys
+
+from patroni.async_executor import AsyncExecutor
+from patroni.exceptions import DCSError, PostgresConnectionException
+from multiprocessing.pool import ThreadPool
+
+logger = logging.getLogger(__name__)
+
+
+class Ha:
+
+    def __init__(self, patroni):
+        self.patroni = patroni
+        self.state_handler = patroni.postgresql
+        self.dcs = patroni.dcs
+        self.cluster = None
+        self.old_cluster = None
+        self._async_executor = AsyncExecutor()
+
+    def load_cluster_from_dcs(self):
+        cluster = self.dcs.get_cluster()
+
+        # We want to keep the state of cluster when it was healhy
+        if not cluster.is_unlocked() or not self.old_cluster:
+            self.old_cluster = cluster
+        self.cluster = cluster
+
+    def acquire_lock(self):
+        return self.dcs.attempt_to_acquire_leader()
+
+    def update_lock(self):
+        ret = self.dcs.update_leader()
+        if ret:
+            try:
+                self.dcs.write_leader_optime(self.state_handler.last_operation())
+            except:
+                pass
+        return ret
+
+    def has_lock(self):
+        lock_owner = self.cluster.leader and self.cluster.leader.name
+        logger.info('Lock owner: %s; I am %s', lock_owner, self.state_handler.name)
+        return lock_owner == self.state_handler.name
+
+    def touch_member(self):
+        data = {
+            'conn_url': self.state_handler.connection_string,
+            'api_url': self.patroni.api.connection_string,
+            'state': self.state_handler.state,
+            'role': self.state_handler.role,
+            'tags': self.patroni.tags
+        }
+        if data['state'] in ['running', 'restarting', 'starting']:
+            try:
+                data['xlog_location'] = self.state_handler.xlog_position()
+            except:
+                pass
+        self.dcs.touch_member(json.dumps(data, separators=(',', ':')))
+
+    def copy_backup_from_leader(self, leader):
+        if self.state_handler.bootstrap(leader):
+            logger.info('bootstrapped from leader')
+        else:
+            self.state_handler.stop('immediate')
+            self.state_handler.remove_data_directory()
+            logger.error('failed to bootstrap from leader')
+
+    def bootstrap(self):
+        if not self.cluster.is_unlocked():  # cluster already has leader
+            self._async_executor.schedule('bootstrap from leader')
+            self._async_executor.run_async(self.copy_backup_from_leader, args=(self.cluster.leader, ))
+            return 'trying to bootstrap from leader'
+        elif not self.cluster.initialize 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()
+                return 'initialized a new cluster'
+            else:
+                return 'failed to acquire initialize lock'
+        else:
+            return 'waiting for leader to bootstrap'
+
+    def recover(self):
+        has_lock = self.has_lock()
+
+        # try to see if we are the former master that crashed. If so - we likely need to run pg_rewind
+        # in order to join the former standby being promoted.
+        pg_controldata = self.state_handler.controldata()
+        if not has_lock and pg_controldata and\
+                pg_controldata.get('Database cluster state', '') == 'in production':  # crashed master
+            self.state_handler.require_rewind()
+
+        # XXX: follow the leader calls stop, which might take quite some time.
+        # perhaps we should run sync asynchronously
+        # (we still need the exit code from follow_the_leader)
+        ret = self.state_handler.follow_the_leader(None if has_lock else self.cluster.leader, recovery=True)
+        if not ret:
+            if not has_lock:
+                return 'failed to start postgres'
+            self.dcs.delete_leader()
+            self.dcs.reset_cluster()
+            return 'removed leader key after trying and failing to start postgres'
+        if not has_lock:
+            return 'started as a secondary'
+        logger.info('started as readonly because i had the session lock')
+        self.load_cluster_from_dcs()
+
+    def follow_the_leader(self, demote_reason, follow_reason, refresh=True):
+        refresh and self.load_cluster_from_dcs()
+        ret = demote_reason if self.state_handler.is_leader() else follow_reason
+        leader = self.cluster.leader
+        leader = None if (leader and leader.name) == self.state_handler.name else leader
+        if not self.state_handler.check_recovery_conf(leader):
+            self._async_executor.schedule('changing primary_conninfo and restarting')
+            self._async_executor.run_async(self.state_handler.follow_the_leader, (leader, ))
+        return ret
+
+    def enforce_master_role(self, message, promote_message):
+        if self.state_handler.is_leader() or self.state_handler.role == 'master':
+            return message
+        else:
+            self.state_handler.promote()
+            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, in_recovery, xlog_location, tags in self.fetch_nodes_statuses(members):
+                if reachable and not tags.get('nofailover', False):
+                    ret = True  # TODO: check xlog_location
+                elif not reachable:
+                    logger.info('Member %s is not reachable', member.name)
+                elif tags.get('nofailover', False):
+                    logger.info('Member %s is not allowed to promote', member.name)
+        else:
+            logger.warning('manual failover: members list is empty')
+        return ret
+
+    def manual_failover_process_no_leader(self):
+        failover = self.cluster.failover
+        if failover.member:  # manual failover to specific member
+            if failover.member == self.state_handler.name:  # manual failover to me
+                return True
+
+            # find specific node and check that it is healthy
+            members = [m for m in self.cluster.members if m.name == failover.member]
+            if members:
+                member, reachable, in_recovery, xlog_location, tags = self.fetch_node_status(members[0])
+                if reachable and not tags.get('nofailover', False):  # node is healthy
+                    logger.info('manual failover: to %s, i am %s', member.name, self.state_handler.name)
+                    return False
+                # we wanted to failover to specific member but it is not healthy
+                if not reachable:
+                    logger.warning('manual failover: member %s is unhealthy', member.name)
+                elif tags.get('nofailover', False):
+                    logger.warning('manual failover: member %s is not allowed to promote', member.name)
+
+            # at this point we should consider all members as a candidates for failover
+            # i.e. we assume that failover.member is None
+
+        # try to pick some other members to failover and check that they are healthy
+        if failover.leader:
+            if self.state_handler.name == failover.leader:  # I was the leader
+                # exclude me and desired member which is unhealthy (failover.member can be None)
+                members = [m for m in self.cluster.members if m.name not in (failover.member, failover.leader)]
+                if self.is_failover_possible(members):  # check that there are healthy members
+                    return False
+                else:  # I was the leader and it looks like currently I am the only healthy member
+                    return True
+
+            # at this point we assume that our node is a candidate for a failover among all nodes except former leader
+
+        # exclude former leader from the list (failover.leader can be None)
+        members = [m for m in self.cluster.members if m.name != failover.leader]
+        return self._is_healthiest_node(members, check_replication_lag=False)
+
+    def is_healthiest_node(self):
+
+        if self.state_handler.is_leader():  # leader is always the healthiest
+            return True
+
+        if self.patroni.nofailover:  # nofailover tag makes node always unhealthy
+            return False
+
+        if self.cluster.failover:
+            return self.manual_failover_process_no_leader()
+
+        # run usual health check
+        members = {m.name: m for m in self.cluster.members + self.old_cluster.members}
+        return self._is_healthiest_node(members.values())
+
+    def demote(self, delete_leader=True):
+        if delete_leader:
+            self.state_handler.stop()
+            self.dcs.delete_leader()
+            self.touch_member()
+            self.dcs.reset_cluster()
+        self.state_handler.follow_the_leader(None)
+
+    def process_manual_failover_from_leader(self):
+        failover = self.cluster.failover
+        if not failover.leader or failover.leader == self.state_handler.name:
+            if not failover.member or failover.member != self.state_handler.name:
+                members = [m for m in self.cluster.members if not failover.member or m.name == failover.member]
+                if self.is_failover_possible(members):  # check that there are healthy members
+                    self._async_executor.schedule('manual failover: demote')
+                    self._async_executor.run_async(self.demote)
+                    return 'manual failover: demoting myself'
+                else:
+                    logger.warning('manual failover: no healthy members found, failover is not possible')
+            else:
+                logger.warning('manual failover: I am already the leader, no need to failover')
+        else:
+            logger.warning('manual failover: leader name does not match: %s != %s',
+                           self.cluster.failover.leader, self.state_handler.name)
+
+        logger.info('Trying to clean up failover key')
+        self.dcs.manual_failover('', '', self.cluster.failover.index)
+
+    def process_unhealthy_cluster(self):
+        if self.is_healthiest_node():
+            if self.acquire_lock():
+                if self.cluster.failover:
+                    logger.info('Cleaning up failover key after acquiring leader lock...')
+                    self.dcs.manual_failover('', '')
+                self.dcs.get_cluster()
+                return self.enforce_master_role('acquired session lock as a leader',
+                                                'promoted self to leader by acquiring session lock')
+            else:
+                return self.follow_the_leader('demoted self due 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_the_leader('demoting self because I am not allowed to become master',
+                                              'following a different leader because I am not allowed to promote')
+            return self.follow_the_leader('demoting self because i am not the healthiest node',
+                                          'following a different leader because i am not the healthiest node')
+
+    def process_healthy_cluster(self):
+        if self.has_lock():
+            if self.cluster.failover:
+                msg = self.process_manual_failover_from_leader()
+                if msg is not None:
+                    return msg
+
+            if self.update_lock():
+                return self.enforce_master_role('no action.  i am the leader with the lock',
+                                                'promoted self to leader because i had the session lock')
+            else:
+                # Either there is no connection to DCS or someone else acquired the lock
+                logger.error('failed to update leader lock')
+                self.load_cluster_from_dcs()
+        else:
+            logger.info('does not have lock')
+        return self.follow_the_leader('demoting self because i do not have the lock and i was a leader',
+                                      'no action.  i am a secondary and i am following a leader', False)
+
+    def schedule(self, action):
+        with self._async_executor:
+            return self._async_executor.schedule(action)
+
+    def restart_scheduled(self):
+        return self._async_executor.scheduled_action == 'restart'
+
+    def schedule_reinitialize(self):
+        return self.schedule('reinitialize')
+
+    def reinitialize_scheduled(self):
+        return self._async_executor.scheduled_action == 'reinitialize'
+
+    def restart(self):
+        with self._async_executor:
+            prev = self._async_executor.schedule('restart', True)
+            if prev is not None:
+                return (False, prev + ' already in progress')
+        if self._async_executor.run(self.state_handler.restart):
+            return (True, 'restarted successfully')
+        else:
+            return (False, 'restart failed')
+
+    def reinitialize(self, cluster):
+        self.state_handler.stop('immediate')
+        self.state_handler.remove_data_directory()
+        self.copy_backup_from_leader(cluster.leader)
+
+    def process_scheduled_action(self):
+        if self.reinitialize_scheduled():
+            if self.cluster.is_unlocked():
+                logger.error('Cluster has no leader, can not reinitialize')
+                self._async_executor.reset_scheduled_action()
+            elif self.has_lock():
+                logger.error('I am the leader, can not reinitialize')
+                self._async_executor.reset_scheduled_action()
+            else:
+                self._async_executor.run_async(self.reinitialize, args=(self.cluster, ))
+                return 'reinitialize started'
+
+    def handle_long_action_in_progress(self):
+        if self.has_lock():
+            if self.update_lock():
+                return 'updated leader lock during ' + self._async_executor.scheduled_action
+            else:
+                return 'failed to update leader lock during ' + self._async_executor.scheduled_action
+        elif self.cluster.is_unlocked():
+            return 'not healthy enough for leader race'
+        else:
+            return self._async_executor.scheduled_action + ' in progress'
+
+    def sysid_valid(self, 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 _run_cycle(self):
+        try:
+            self.load_cluster_from_dcs()
+
+            self.touch_member()
+
+            # cluster has leader key but not initialize key
+            if not self.cluster.is_unlocked() and not self.sysid_valid(self.cluster.initialize) and self.has_lock():
+                self.dcs.initialize(create_new=(self.cluster.initialize is None), sysid=self.state_handler.sysid)
+
+            if self._async_executor.busy:
+                return self.handle_long_action_in_progress()
+
+            # currently it can trigger only reinitialize
+            msg = self.process_scheduled_action()
+            if msg is not None:
+                return msg
+
+            # is data directory empty?
+            if self.state_handler.data_directory_empty():
+                return self.bootstrap()  # new node
+            # "bootstrap", but data directory is not empty
+            elif not self.sysid_valid(self.cluster.initialize) and self.cluster.is_unlocked():
+                self.dcs.initialize(create_new=(self.cluster.initialize is None), sysid=self.state_handler.sysid)
+            else:
+                # check if we are allowed to join
+                if self.sysid_valid(self.cluster.initialize) and self.cluster.initialize != self.state_handler.sysid:
+                    logger.fatal("system ID mismatch, node {0} belongs to a different cluster".
+                                 format(self.state_handler.name))
+                    sys.exit(1)
+
+            # try to start dead postgres
+            if not self.state_handler.is_healthy():
+                msg = self.recover()
+                if msg is not None:
+                    return msg
+
+            try:
+                if self.cluster.is_unlocked():
+                    return self.process_unhealthy_cluster()
+                else:
+                    return self.process_healthy_cluster()
+            finally:
+                self.state_handler.sync_replication_slots(self.cluster)
+        except DCSError:
+            logger.error('Error communicating with DCS')
+            if self.state_handler.is_running() and self.state_handler.is_leader():
+                self.demote(delete_leader=False)
+                return 'demoted self because DCS is not accessible and i was a leader'
+        except (psycopg2.Error, PostgresConnectionException):
+            logger.exception('Error communicating with Postgresql. Will try again later')
+
+    def run_cycle(self):
+        with self._async_executor:
+            return self._run_cycle()

patroni/postgresql.py

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

patroni/scripts/aws.py

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

patroni/scripts/restore.py

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

patroni/utils.py

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

patroni/version.py

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

patroni/zookeeper.py

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

patronictl.py

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

postgres0.yml

@@ -4,6 +4,9 @@ scope: &scope batman
 restapi:
   listen: 127.0.0.1:8008
   connect_address: 127.0.0.1:8008
+  auth: 'username:password'
+#  certfile: /etc/ssl/certs/ssl-cert-snakeoil.pem
+#  keyfile: /etc/ssl/private/ssl-cert-snakeoil.key
 etcd:
   scope: *scope
   ttl: *ttl
@@ -30,6 +33,28 @@ postgresql:
   connect_address: 127.0.0.1:5432
   data_dir: data/postgresql0
   maximum_lag_on_failover: 1048576 # 1 megabyte in bytes
+  use_slots: True
+  pgpass: /tmp/pgpass0
+  initdb:  ## We allow the following options to be passed on to initdb
+  # - auth: authmethod
+  # - auth-host: authmethod
+  # - auth-local: authmethod
+  - encoding: UTF8
+  # - data-checksums  # When pg_rewind is needed on 9.3, this needs to be enabled
+  # - locale: locale
+  # - lc-collate: locale
+  # - lc-ctype: locale
+  # - lc-messages: locale
+  # - lc-monetary: locale
+  # - lc-numeric: locale
+  # - lc-time: locale
+  # - text-search-config: CFG
+  # - xlogdir: directory
+  # - debug
+  # - noclean
+  pg_rewind:
+    username: postgres
+    password: zalando
   pg_hba:
   - host all all 0.0.0.0/0 md5
   - hostssl all all 0.0.0.0/0 md5
@@ -38,6 +63,7 @@ postgresql:
     password: rep-pass
     network:  127.0.0.1/32
   superuser:
+    username: postgres
     password: zalando
   admin:
     username: admin
@@ -46,7 +72,7 @@ postgresql:
     env_dir: /home/postgres/etc/wal-e.d/env
     threshold_megabytes: 10240
     threshold_backup_size_percentage: 30
-  restore: scripts/restore.py
+  restore: patroni/scripts/restore.py
   #recovery_conf:
     #restore_command: cp ../wal_archive/%f %p
   parameters:
@@ -58,3 +84,9 @@ postgresql:
     archive_timeout: 1800s
     max_replication_slots: 5
     hot_standby: "on"
+    wal_log_hints: "on"
+tags:
+    nofailover: False
+    noloadbalance: False
+    clonefrom: False
+    replicatefrom: 127.0.0.1

postgres1.yml

@@ -4,6 +4,9 @@ scope: &scope batman
 restapi:
   listen: 127.0.0.1:8009
   connect_address: 127.0.0.1:8009
+  auth: 'username:password'
+#  certfile: /etc/ssl/certs/ssl-cert-snakeoil.pem
+#  keyfile: /etc/ssl/private/ssl-cert-snakeoil.key
 etcd:
   scope: *scope
   ttl: *ttl
@@ -30,6 +33,28 @@ postgresql:
   connect_address: 127.0.0.1:5433
   data_dir: data/postgresql1
   maximum_lag_on_failover: 1048576 # 1 megabyte in bytes
+  use_slots: True
+  pgpass: /tmp/pgpass1
+  initdb:  ## We allow the following options to be passed on to initdb
+  # - auth: authmethod
+  # - auth-host: authmethod
+  # - auth-local: authmethod
+  - encoding: UTF8
+  # - data-checksums  # When pg_rewind is needed on 9.3, this needs to be enabled
+  # - locale: locale
+  # - lc-collate: locale
+  # - lc-ctype: locale
+  # - lc-messages: locale
+  # - lc-monetary: locale
+  # - lc-numeric: locale
+  # - lc-time: locale
+  # - text-search-config: CFG
+  # - xlogdir: directory
+  # - debug
+  # - noclean
+  pg_rewind:
+    username: postgres
+    password: zalando
   pg_hba:
   - host all all 0.0.0.0/0 md5
   - hostssl all all 0.0.0.0/0 md5
@@ -38,6 +63,7 @@ postgresql:
     password: rep-pass
     network: 127.0.0.1/32
   superuser:
+    user: postgres
     password: zalando
   admin:
     username: admin
@@ -48,7 +74,7 @@ postgresql:
     env_dir: /home/postgres/etc/wal-e.d/env
     threshold_megabytes: 10240
     threshold_backup_size_percentage: 30
-  restore: scripts/restore.py
+  restore: patroni/scripts/restore.py
   parameters:
     archive_mode: "on"
     wal_level: hot_standby
@@ -58,3 +84,9 @@ postgresql:
     archive_timeout: 1800s
     max_replication_slots: 5
     hot_standby: "on"
+    wal_log_hints: "on"
+tags:
+    nofailover: False
+    noloadbalance: False
+    clonefrom: False
+    replicatefrom: 127.0.0.1

release.sh

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

requirements-py2.txt

@@ -1,9 +1,11 @@
 boto
 dnspython
 mock
-psycopg2
+psycopg2>=2.6.1
 PyYAML
 requests
 six >= 1.7
 kazoo>=2.2.1
 python-etcd>=0.4.1
+click>=4.1
+prettytable>=0.7

requirements-py3.txt

@@ -1,9 +1,11 @@
 boto
 mock
 dnspython3
-psycopg2
+psycopg2>=2.6.1
 PyYAML
 requests
 six
 kazoo>=2.2.1
 python-etcd>=0.4.1
+click>=4.1
+prettytable>=0.7

setup.py

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

tests/test_api.py

@@ -1,16 +1,21 @@
 import psycopg2
 import unittest
 
-from helpers.api import RestApiHandler, RestApiServer
+from mock import Mock, patch
+from patroni.api import RestApiHandler, RestApiServer
+from patroni.dcs import Member
 from six import BytesIO as IO
-from test_postgresql import psycopg2_connect
+from six.moves import BaseHTTPServer
+from six.moves.BaseHTTPServer import BaseHTTPRequestHandler
+import socket
+from test_postgresql import psycopg2_connect, MockCursor
 
 
-def throws(*args, **kwargs):
-    raise psycopg2.OperationalError()
+class MockPostgresql(Mock):
 
-
-class MockPostgresql:
+    name = 'test'
+    state = 'running'
+    role = 'master'
 
     def connection(self):
         return psycopg2_connect()
@@ -19,10 +24,33 @@ class MockPostgresql:
         return True
 
 
+class MockHa(Mock):
+
+    dcs = Mock()
+    state_handler = MockPostgresql()
+
+    def schedule_restart(self):
+        return 'restart'
+
+    def schedule_reinitialize(self):
+        return 'reinitialize'
+
+    def restart(self):
+        return (True, '')
+
+    def restart_scheduled(self):
+        return False
+
+    def fetch_nodes_statuses(self, members):
+        return [[None, True, None, None, {}]]
+
+
 class MockPatroni:
 
-    def __init__(self):
-        self.postgresql = MockPostgresql()
+    postgresql = MockPostgresql()
+    ha = MockHa()
+    dcs = Mock()
+    tags = {}
 
 
 class MockRequest:
@@ -36,18 +64,105 @@ class MockRequest:
 
 class MockRestApiServer(RestApiServer):
 
-    def __init__(self, Handler, path, *args):
-        self.patroni = MockPatroni()
-        if len(args) > 0:
-            self.query = args[0]
+    def __init__(self, Handler, path):
+        self.socket = 0
+        BaseHTTPServer.HTTPServer.__init__ = Mock()
+        MockRestApiServer._BaseServer__is_shut_down = Mock()
+        MockRestApiServer._BaseServer__shutdown_request = True
+        config = {'listen': '127.0.0.1:8008', 'auth': 'test:test', 'certfile': 'dumb'}
+        super(MockRestApiServer, self).__init__(MockPatroni(), config)
         Handler(MockRequest(path), ('0.0.0.0', 8080), self)
 
 
+@patch('ssl.wrap_socket', Mock(return_value=0))
 class TestRestApiHandler(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        super(TestRestApiHandler, self).__init__(method_name)
-
     def test_do_GET(self):
-        MockRestApiServer(RestApiHandler, b'GET /')
-        MockRestApiServer(RestApiHandler, b'GET /', throws)
+        MockRestApiServer(RestApiHandler, b'GET /replica')
+        with patch.object(RestApiHandler, 'get_postgresql_status', Mock(return_value={})):
+            MockRestApiServer(RestApiHandler, b'GET /replica')
+        with patch.object(RestApiHandler, 'get_postgresql_status', Mock(return_value={'role': 'master'})):
+            MockRestApiServer(RestApiHandler, b'GET /replica')
+        MockRestApiServer(RestApiHandler, b'GET /master')
+        MockPatroni.dcs.cluster.leader.name = MockPostgresql.name
+        MockRestApiServer(RestApiHandler, b'GET /replica')
+        MockPatroni.dcs.cluster = None
+        with patch.object(RestApiHandler, 'get_postgresql_status', Mock(return_value={'role': 'master'})):
+            MockRestApiServer(RestApiHandler, b'GET /master')
+        with patch.object(MockHa, 'restart_scheduled', Mock(return_value=True)):
+            MockRestApiServer(RestApiHandler, b'GET /master')
+        MockRestApiServer(RestApiHandler, b'GET /master')
+
+    def test_do_OPTIONS(self):
+        MockRestApiServer(RestApiHandler, b'OPTIONS / HTTP/1.0')
+
+        with patch.object(BaseHTTPRequestHandler, 'handle_one_request') as mock_handle_request:
+            mock_handle_request.side_effect = socket.error("foo")
+            MockRestApiServer(RestApiHandler, b'OPTIONS / HTTP/1.0')
+
+        # make sure socket.error gets propagated via wfile object in finalize()
+        with patch.object(MockRequest, 'makefile') as makefile:
+            makefile.return_value.closed = False
+            makefile.return_value.readline.side_effect = lambda x: b"foo"
+            makefile.return_value.flush = Mock(side_effect=socket.error("foo"))
+            MockRestApiServer(RestApiHandler, b'OPTIONS / HTTP/1.0')
+
+    def test_do_GET_patroni(self):
+        MockRestApiServer(RestApiHandler, b'GET /patroni')
+
+    def test_basicauth(self):
+        MockRestApiServer(RestApiHandler, b'POST /restart HTTP/1.0')
+        MockRestApiServer(RestApiHandler, b'POST /restart HTTP/1.0\nAuthorization:')
+
+    def test_do_POST_restart(self):
+        request = b'POST /restart HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0'
+        MockRestApiServer(RestApiHandler, request)
+        with patch.object(MockHa, 'restart', Mock(side_effect=Exception)):
+            MockRestApiServer(RestApiHandler, request)
+
+    @patch.object(MockHa, 'dcs')
+    def test_do_POST_reinitialize(self, dcs):
+        cluster = dcs.get_cluster.return_value
+        request = b'POST /reinitialize HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0'
+        MockRestApiServer(RestApiHandler, request)
+        cluster.is_unlocked.return_value = False
+        MockRestApiServer(RestApiHandler, request)
+        with patch.object(MockHa, 'schedule_reinitialize', Mock(return_value=None)):
+            MockRestApiServer(RestApiHandler, request)
+        cluster.leader.name = 'test'
+        MockRestApiServer(RestApiHandler, request)
+
+    @patch('time.sleep', Mock())
+    def test_RestApiServer_query(self):
+        with patch.object(MockCursor, 'execute', Mock(side_effect=psycopg2.OperationalError)):
+            MockRestApiServer(RestApiHandler, b'GET /patroni')
+        with patch.object(MockPostgresql, 'connection', Mock(side_effect=psycopg2.OperationalError)):
+            MockRestApiServer(RestApiHandler, b'GET /patroni')
+
+    @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: 25\n\n{"leader": "postgresql1"}'
+        MockRestApiServer(RestApiHandler, request)
+        cluster.leader.name = 'postgresql1'
+        MockRestApiServer(RestApiHandler, request)
+        cluster.members = [Member(0, 'postgresql0', 30, {'api_url': 'http'})]
+        MockRestApiServer(RestApiHandler, request)
+        with patch.object(MockPatroni, 'dcs') as d:
+            cluster = d.get_cluster.return_value
+            cluster.leader.name = 'postgresql0'
+            MockRestApiServer(RestApiHandler, request)
+            cluster.leader.name = 'postgresql1'
+            cluster.failover = None
+            MockRestApiServer(RestApiHandler, request)
+            d.get_cluster = Mock(side_effect=Exception())
+            MockRestApiServer(RestApiHandler, request)
+            d.manual_failover.return_value = False
+            MockRestApiServer(RestApiHandler, request)
+        with patch.object(MockHa, 'fetch_nodes_statuses', Mock(return_value=[])):
+            MockRestApiServer(RestApiHandler, request)
+        request = b'POST /failover HTTP/1.0\nAuthorization: Basic dGVzdDp0ZXN0\n' +\
+                  b'Content-Length: 50\n\n{"leader": "postgresql1", "member": "postgresql2"}'
+        MockRestApiServer(RestApiHandler, request)

tests/test_async_executor.py

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

tests/test_aws.py

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

tests/test_ctl.py

@@ -0,0 +1,376 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import os
+import pytest
+import unittest
+import psycopg2
+import requests
+import patroni.exceptions
+import etcd
+from mock import patch, Mock
+
+from click.testing import CliRunner
+from patroni.ctl import ctl, members, store_config, load_config, output_members, post_patroni, get_dcs, \
+    wait_for_leader, get_all_members, get_any_member, get_cursor, query_member, configure
+from patroni.ha import Ha
+from patroni.etcd import Etcd, Client
+from test_ha import get_cluster_initialized_without_leader, get_cluster_initialized_with_leader, \
+    get_cluster_initialized_with_only_leader, MockPostgresql, MockPatroni, run_async, \
+    get_cluster_not_initialized_without_leader
+from test_etcd import etcd_read, etcd_write, requests_get, MockResponse
+from test_postgresql import MockConnect, psycopg2_connect
+
+CONFIG_FILE_PATH = './test-ctl.yaml'
+
+def test_rw_config():
+    runner = CliRunner()
+    config = {'a':'b'}
+    with runner.isolated_filesystem():
+        store_config(config, CONFIG_FILE_PATH + '/dummy')
+        os.remove(CONFIG_FILE_PATH + '/dummy')
+        os.rmdir(CONFIG_FILE_PATH)
+
+        with pytest.raises(Exception):
+            result = load_config(CONFIG_FILE_PATH, None)
+            assert 'Could not load configuration file' in result.output
+
+        os.mkdir(CONFIG_FILE_PATH)
+        with pytest.raises(Exception):
+            store_config(config, CONFIG_FILE_PATH)
+
+        os.rmdir(CONFIG_FILE_PATH)
+
+    store_config(config, CONFIG_FILE_PATH)
+    load_config(CONFIG_FILE_PATH, None)
+    load_config(CONFIG_FILE_PATH, '0.0.0.0')
+
+@patch('patroni.ctl.load_config', Mock(return_value={'dcs': {'scheme': 'etcd', 'hostname': 'localhost', 'port': 4001}}))
+class TestCtl(unittest.TestCase):
+
+    @patch.object(Client, 'machines')
+    def setUp(self, mock_machines):
+        mock_machines.__get__ = Mock(return_value=['http://remotehost:2379'])
+        self.p = MockPostgresql()
+        self.e = Etcd('foo', {'ttl': 30, 'host': 'ok:2379', 'scope': 'test'})
+        self.e.client.read = etcd_read
+        self.e.client.write = etcd_write
+        self.e.client.delete = Mock(side_effect=etcd.EtcdException())
+        self.ha = Ha(MockPatroni(self.p, self.e))
+        self.ha._async_executor.run_async = run_async
+        self.ha.old_cluster = self.e.get_cluster()
+        self.ha.cluster = get_cluster_not_initialized_without_leader()
+        self.ha.load_cluster_from_dcs = Mock()
+
+    @patch('psycopg2.connect', psycopg2_connect)
+    def test_get_cursor(self):
+        c = get_cursor(get_cluster_initialized_without_leader(), role='master')
+        assert c is None
+
+        c = get_cursor(get_cluster_initialized_with_leader(), role='master')
+        assert c is not None
+
+        c = get_cursor(get_cluster_initialized_with_leader(), role='replica')
+        # # MockCursor returns pg_is_in_recovery as false
+        assert c is None
+
+        c = get_cursor(get_cluster_initialized_with_leader(), role='any')
+        assert c is not None
+
+    def test_output_members(self):
+        cluster = get_cluster_initialized_with_leader()
+        output_members(cluster, name='abc', format='pretty')
+        output_members(cluster, name='abc', format='json')
+        output_members(cluster, name='abc', format='tsv')
+
+    @patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_leader()))
+    @patch('patroni.etcd.Etcd.get_etcd_client', Mock(return_value=None))
+    @patch('patroni.etcd.Etcd.set_failover_value', Mock(return_value=None))
+    @patch('patroni.ctl.wait_for_leader', Mock(return_value=get_cluster_initialized_with_leader()))
+    @patch('requests.get', requests_get)
+    @patch('requests.post', requests_get)
+    @patch('patroni.ctl.post_patroni', Mock(return_value=MockResponse()))
+    def test_failover(self):
+        runner = CliRunner()
+
+        with patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_leader())):
+            result = runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+y''')
+            assert 'Failing over to new leader' in result.output
+
+            result = runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+N''')
+            assert 'Aborting failover' in str(result.exception)
+
+            result = runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+leader
+y''')
+            assert 'target and source are the same' in str(result.exception)
+
+            result = runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+Reality
+y''')
+            assert 'Reality does not exist' in str(result.exception)
+
+            result = runner.invoke(ctl, ['failover', 'dummy', '--force'])
+            assert 'Failing over to new leader' in result.output
+
+            result = runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='dummy')
+            assert 'is not the leader of cluster' in str(result.exception)
+
+        with patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_only_leader())):
+            result = runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+y''')
+            assert 'No candidates found to failover to' in str(result.exception)
+
+        with patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_without_leader())):
+            result = runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+y''')
+            assert 'This cluster has no master' in str(result.exception)
+
+        with patch('patroni.ctl.post_patroni', Mock(side_effect=Exception())):
+            result = runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+y''')
+            assert 'falling back to DCS' in result.output
+            assert 'Failover failed' in result.output
+
+        mocked = Mock()
+        mocked.return_value.status_code = 500
+        with patch('patroni.ctl.post_patroni', Mock(return_value=mocked)):
+            result = runner.invoke(ctl, ['failover', 'dummy', '--dcs', '8.8.8.8'], input='''leader
+other
+y''')
+            assert 'Failover failed, details' in result.output
+
+#        with patch('patroni.dcs.AbstractDCS.get_cluster', Mock(return_value=get_cluster_initialized_with_leader())):
+#            result = runner.invoke(ctl, ['failover', 'alpha', '--dcs', '8.8.8.8'], input='nonsense')
+#            assert 'is not the leader of cluster' in str(result.exception)
+
+            # result = runner.invoke(ctl, ['failover', 'alpha', '--dcs', '8.8.8.8', '--master', 'nonsense'])
+            # assert 'is not the leader of cluster' in str(result.exception)
+
+            # result = runner.invoke(ctl, ['failover', 'alpha', '--dcs', '8.8.8.8'], input='leader\nother\nn')
+            # assert 'Aborting failover' in str(result.exception)
+
+            # with patch('patroni.ctl.wait_for_leader', Mock(return_value = get_cluster_initialized_with_leader())):
+                # result = runner.invoke(ctl, ['failover', 'alpha', '--dcs', '8.8.8.8'], input='leader\nother\nY')
+                # assert 'master did not change after' in result.output
+
+                # result = runner.invoke(ctl, ['failover', 'alpha', '--dcs', '8.8.8.8'], input='leader\nother\nY')
+                # assert 'Failover failed' in result.output
+
+    def test_(self):
+        self.assertRaises(patroni.exceptions.PatroniCtlException, get_dcs, {'scheme': 'dummy'}, 'dummy')
+
+    @patch('psycopg2.connect', psycopg2_connect)
+    @patch('patroni.ctl.query_member', Mock(return_value=([['mock column']], None)))
+    def test_query(self):
+        runner = CliRunner()
+
+        with patch('patroni.ctl.get_dcs', Mock(return_value=self.e)):
+            result = runner.invoke(ctl, [
+                'query',
+                'alpha',
+                '--member',
+                'abc',
+                '--role',
+                'master',
+            ])
+            assert 'mutually exclusive' in str(result.exception)
+
+            with runner.isolated_filesystem():
+                dummy_file = open('dummy', 'w')
+                dummy_file.write('SELECT 1')
+                dummy_file.close()
+
+                result = runner.invoke(ctl, [
+                    'query',
+                    'alpha',
+                    '--file',
+                    'dummy',
+                    '--command',
+                    'dummy',
+                ])
+                assert 'mutually exclusive' in str(result.exception)
+
+                result = runner.invoke(ctl, ['query', 'alpha', '--file', 'dummy'])
+
+                os.remove('dummy')
+
+            result = runner.invoke(ctl, ['query', 'alpha', '--command', 'SELECT 1'])
+            assert 'mock column' in result.output
+
+    @patch('patroni.ctl.get_cursor', Mock(return_value=MockConnect().cursor()))
+    def test_query_member(self):
+        rows = query_member(None, None, None, 'master', 'SELECT pg_is_in_recovery()')
+        assert 'False' in str(rows)
+
+        rows = query_member(None, None, None, 'replica', 'SELECT pg_is_in_recovery()')
+        assert rows == (None, None)
+
+        with patch('patroni.ctl.get_cursor', Mock(return_value=None)):
+            rows = query_member(None, None, None, None, 'SELECT pg_is_in_recovery()')
+            assert 'No connection to' in str(rows)
+
+            rows = query_member(None, None, None, 'replica', 'SELECT pg_is_in_recovery()')
+            assert 'No connection to' in str(rows)
+
+        with patch('patroni.ctl.get_cursor', Mock(side_effect=psycopg2.OperationalError('bla'))):
+            rows = query_member(None, None, None, 'replica', 'SELECT pg_is_in_recovery()')
+
+        with patch('test_postgresql.MockCursor.execute', Mock(side_effect=psycopg2.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):
+        runner = CliRunner()
+
+        with patch('patroni.ctl.get_dcs', Mock(return_value=self.e)):
+            result = runner.invoke(ctl, ['dsn', 'alpha', '--dcs', '8.8.8.8'])
+            assert 'host=127.0.0.1 port=5435' in result.output
+
+            result = runner.invoke(ctl, [
+                'dsn',
+                'alpha',
+                '--role',
+                'master',
+                '--member',
+                'dummy',
+            ])
+            assert 'mutually exclusive' in str(result.exception)
+
+            result = runner.invoke(ctl, ['dsn', 'alpha', '--member', 'dummy'])
+            assert 'Can not find' in str(result.exception)
+
+        # result = runner.invoke(ctl, ['dsn', 'alpha', '--dcs', '8.8.8.8', '--role', 'replica'])
+        # assert 'host=127.0.0.1 port=5436' in result.output
+
+    @patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_leader()))
+    @patch('patroni.etcd.Etcd.get_etcd_client', Mock(return_value=None))
+    @patch('requests.get', requests_get)
+    @patch('requests.post', requests_get)
+    def test_restart_reinit(self):
+        runner = CliRunner()
+
+        result = runner.invoke(ctl, ['restart', 'alpha', '--dcs', '8.8.8.8'], input='y')
+        result = runner.invoke(ctl, ['reinit', 'alpha', '--dcs', '8.8.8.8'], input='y')
+
+        result = runner.invoke(ctl, ['restart', 'alpha', '--dcs', '8.8.8.8'], input='N')
+        result = runner.invoke(ctl, [
+            'restart',
+            'alpha',
+            '--dcs',
+            '8.8.8.8',
+            'dummy',
+            '--any',
+        ], input='y')
+        assert 'not a member' in str(result.exception)
+
+        with patch('requests.post', Mock(return_value=MockResponse())):
+            result = runner.invoke(ctl, ['restart', 'alpha', '--dcs', '8.8.8.8'], input='y')
+
+    @patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_leader()))
+    @patch('patroni.etcd.Etcd.get_etcd_client', Mock(return_value=None))
+    def test_remove(self):
+        runner = CliRunner()
+
+        result = runner.invoke(ctl, ['remove', 'alpha', '--dcs', '8.8.8.8'], input='alpha\nslave')
+        assert 'Please confirm' in result.output
+        assert 'You are about to remove all' in result.output
+        assert 'You did not exactly type' in str(result.exception)
+
+        result = runner.invoke(ctl, ['remove', 'alpha', '--dcs', '8.8.8.8'], input='''alpha
+Yes I am aware
+slave''')
+        assert 'You did not specify the current master of the cluster' in str(result.exception)
+
+        result = runner.invoke(ctl, ['remove', 'alpha', '--dcs', '8.8.8.8'], input='beta\nleader')
+        assert 'Cluster names specified do not match' in str(result.exception)
+
+        with patch('patroni.etcd.Etcd.get_cluster', get_cluster_initialized_with_leader):
+            result = runner.invoke(ctl, ['remove', 'alpha', '--dcs', '8.8.8.8'],
+                                   input='''alpha
+Yes I am aware
+leader''')
+            assert 'object has no attribute' in str(result.exception)
+
+        with patch('patroni.ctl.get_dcs', Mock(return_value=Mock())):
+            result = runner.invoke(ctl, ['remove', 'alpha', '--dcs', '8.8.8.8'],
+                                   input='''alpha
+Yes I am aware
+leader''')
+            assert 'We have not implemented this for DCS of type' in str(result.exception)
+
+    @patch('patroni.etcd.Etcd.watch', Mock(return_value=None))
+    @patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_leader()))
+    def test_wait_for_leader(self):
+        dcs = self.e
+        self.assertRaises(patroni.exceptions.PatroniCtlException, wait_for_leader, dcs, 0)
+
+        cluster = wait_for_leader(dcs=dcs, timeout=2)
+        assert cluster.leader.member.name == 'leader'
+
+    def test_post_patroni(self):
+        member = get_cluster_initialized_with_leader().leader.member
+        self.assertRaises(requests.exceptions.ConnectionError, post_patroni, member, 'dummy', {})
+
+    def test_ctl(self):
+        runner = CliRunner()
+
+        runner.invoke(ctl, ['list'])
+
+        result = runner.invoke(ctl, ['--help'])
+        assert 'Usage:' in result.output
+
+    def test_get_any_member(self):
+        m = get_any_member(get_cluster_initialized_without_leader(), role='master')
+        assert m is None
+
+        m = get_any_member(get_cluster_initialized_with_leader(), role='master')
+        assert m.name == 'leader'
+
+    def test_get_all_members(self):
+        r = list(get_all_members(get_cluster_initialized_without_leader(), role='master'))
+        assert len(r) == 0
+
+        r = list(get_all_members(get_cluster_initialized_with_leader(), role='master'))
+        assert len(r) == 1
+        assert r[0].name == 'leader'
+
+        r = list(get_all_members(get_cluster_initialized_with_leader(), role='replica'))
+        assert len(r) == 1
+        assert r[0].name == 'other'
+
+        r = list(get_all_members(get_cluster_initialized_without_leader(), role='replica'))
+        assert len(r) == 2
+
+    @patch('patroni.etcd.Etcd.get_cluster', Mock(return_value=get_cluster_initialized_with_leader()))
+    @patch('patroni.etcd.Etcd.get_etcd_client', Mock(return_value=None))
+    @patch('requests.get', requests_get)
+    @patch('requests.post', requests_get)
+    def test_members(self):
+        runner = CliRunner()
+
+        result = runner.invoke(members, ['alpha'])
+        assert result.exit_code == 0
+
+    def test_configure(self):
+        runner = CliRunner()
+
+        result = runner.invoke(configure, [
+            '--dcs',
+            'abc',
+            '-c',
+            'dummy',
+            '-n',
+            'bla',
+        ])
+
+        assert result.exit_code == 0
+
+

tests/test_etcd.py

@@ -1,16 +1,14 @@
-import datetime
-import dns.resolver
 import etcd
 import json
 import requests
+import urllib3
 import socket
-import time
 import unittest
 
 from dns.exception import DNSException
-from helpers.dcs import Cluster, DCSError, Member
-from helpers.etcd import Client, Etcd
 from mock import Mock, patch
+from patroni.dcs import Cluster, DCSError, Leader
+from patroni.etcd import Client, Etcd, EtcdError
 
 
 class MockResponse:
@@ -19,12 +17,17 @@ class MockResponse:
         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':
+            raise urllib3.exceptions.TimeoutError
+        if self.content == 'Exception':
+            raise Exception
         return self.content
 
     @property
@@ -35,11 +38,10 @@ class MockResponse:
         return ''
 
 
-class MockPostgresql:
-    name = ''
+class MockPostgresql(Mock):
 
     def last_operation(self):
-        return 0
+        return '0'
 
 
 def requests_get(url, **kwargs):
@@ -48,6 +50,8 @@ def requests_get(url, **kwargs):
     response = MockResponse()
     if url.startswith('http://local'):
         raise requests.exceptions.RequestException()
+    elif ':8011/patroni' in url:
+        response.content = '{"role": "replica", "xlog": {"replayed_location": 0}, "tags": {}}'
     elif url.endswith('/members'):
         if url.startswith('http://error'):
             response.content = '[{}]'
@@ -61,24 +65,37 @@ def requests_get(url, **kwargs):
     return response
 
 
+def etcd_watch(key, index=None, timeout=None, recursive=None):
+    if timeout == 2.0:
+        raise urllib3.exceptions.TimeoutError
+    elif timeout == 5.0:
+        return etcd.EtcdResult('delete', {})
+    elif timeout == 10.0:
+        raise etcd.EtcdException
+    elif index == 20729:
+        return etcd.EtcdResult('set', {'value': 'postgresql1', 'modifiedIndex': index + 1})
+    elif index == 20731:
+        return etcd.EtcdResult('set', {'value': 'postgresql2', 'modifiedIndex': index + 1})
+
+
 def etcd_write(key, value, **kwargs):
-    if key == '/service/test/leader':
+    if key == '/service/exists/leader':
+        raise etcd.EtcdAlreadyExist
+    if key == '/service/test/leader' or key == '/patroni/test/leader':
         if kwargs.get('prevValue', None) == 'foo' or not kwargs.get('prevExist', True):
             return True
     raise etcd.EtcdException
 
 
-def etcd_delete(key, **kwargs):
-    raise etcd.EtcdException
-
-
 def etcd_read(key, **kwargs):
-    if key == '/service/noleader':
+    if key == '/service/noleader/':
         raise DCSError('noleader')
-    elif key == '/service/nocluster':
+    elif key == '/service/nocluster/':
         raise etcd.EtcdKeyNotFound
 
     response = {"action": "get", "node": {"key": "/service/batman5", "dir": True, "nodes": [
+                {"key": "/service/batman5/failover", "value": "",
+                 "modifiedIndex": 1582, "createdIndex": 1582},
                 {"key": "/service/batman5/initialize", "value": "postgresql0",
                  "modifiedIndex": 1582, "createdIndex": 1582},
                 {"key": "/service/batman5/leader", "value": "postgresql1",
@@ -90,27 +107,23 @@ def etcd_read(key, **kwargs):
                  "modifiedIndex": 20437, "createdIndex": 20437},
                 {"key": "/service/batman5/members", "dir": True, "nodes": [
                     {"key": "/service/batman5/members/postgresql1",
-                     "value": "postgres://replicator:[email protected]:5434/postgres"
-                        + "?application_name=http://127.0.0.1:8009/patroni",
+                     "value": "postgres://replicator:[email protected]:5434/postgres" +
+                        "?application_name=http://127.0.0.1:8009/patroni",
                      "expiration": "2015-05-15T09:10:59.949384522Z", "ttl": 21,
                      "modifiedIndex": 20727, "createdIndex": 20727},
                     {"key": "/service/batman5/members/postgresql0",
-                     "value": "postgres://replicator:[email protected]:5433/postgres"
-                        + "?application_name=http://127.0.0.1:8008/patroni",
+                     "value": "postgres://replicator:[email protected]:5433/postgres" +
+                        "?application_name=http://127.0.0.1:8008/patroni",
                      "expiration": "2015-05-15T09:11:09.611860899Z", "ttl": 30,
                      "modifiedIndex": 20730, "createdIndex": 20730}],
                  "modifiedIndex": 1581, "createdIndex": 1581}], "modifiedIndex": 1581, "createdIndex": 1581}}
     return etcd.EtcdResult(**response)
 
 
-def time_sleep(_):
+class SleepException(Exception):
     pass
 
 
-def time_sleep_exception(_):
-    raise Exception()
-
-
 class MockSRV:
     port = 2380
     target = '127.0.0.1'
@@ -127,7 +140,7 @@ def dns_query(name, type):
 def socket_getaddrinfo(*args):
     if args[0] == 'ok':
         return [(2, 1, 6, '', ('127.0.0.1', 2379)), (2, 1, 6, '', ('127.0.0.1', 2379))]
-    raise socket.error()
+    raise socket.error
 
 
 def http_request(method, url, **kwargs):
@@ -136,28 +149,14 @@ def http_request(method, url, **kwargs):
     raise socket.error
 
 
-class TestMember(unittest.TestCase):
-
-    def __init__(self, method_name='runTest'):
-        super(TestMember, self).__init__(method_name)
-
-    def test_real_ttl(self):
-        now = datetime.datetime.utcnow()
-        member = Member(0, 'a', 'b', 'c', (now + datetime.timedelta(seconds=2)).strftime('%Y-%m-%dT%H:%M:%S.%fZ'), None)
-        self.assertLess(member.real_ttl(), 2)
-        self.assertEquals(Member(0, 'a', 'b', 'c', '', None).real_ttl(), -1)
-
-
+@patch('dns.resolver.query', dns_query)
+@patch('socket.getaddrinfo', socket_getaddrinfo)
+@patch('requests.get', requests_get)
 class TestClient(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestClient, self).__init__(method_name)
-
-    def set_up(self):
-        socket.getaddrinfo = socket_getaddrinfo
-        requests.get = requests_get
-        dns.resolver.query = dns_query
+    @patch('dns.resolver.query', dns_query)
+    @patch('requests.get', requests_get)
+    def setUp(self):
         with patch.object(etcd.Client, 'machines') as mock_machines:
             mock_machines.__get__ = Mock(return_value=['http://localhost:2379', 'http://localhost:4001'])
             self.client = Client({'discovery_srv': 'test'})
@@ -167,17 +166,31 @@ class TestClient(unittest.TestCase):
         self.client._base_uri = 'http://localhost:4001'
         self.client._machines_cache = ['http://localhost:2379']
         self.client.api_execute('/', 'GET')
+        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)
 
     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 = {}
@@ -186,25 +199,26 @@ class TestClient(unittest.TestCase):
         self.assertRaises(etcd.EtcdException, self.client._load_machines_cache)
 
 
+@patch('requests.get', requests_get)
 class TestEtcd(unittest.TestCase):
 
-    def __init__(self, method_name='runTest'):
-        self.setUp = self.set_up
-        super(TestEtcd, self).__init__(method_name)
-
-    def set_up(self):
-        time.sleep = time_sleep
+    def setUp(self):
         with patch.object(Client, 'machines') as mock_machines:
             mock_machines.__get__ = Mock(return_value=['http://localhost:2379', 'http://localhost:4001'])
-            self.etcd = Etcd('foo', {'ttl': 30, 'host': 'localhost:2379', 'scope': 'test'})
+            self.etcd = Etcd('foo', {'namespace': '/patroni/', 'ttl': 30, 'host': 'localhost:2379', 'scope': 'test'})
             self.etcd.client.write = etcd_write
             self.etcd.client.read = etcd_read
+            self.etcd.client.delete = Mock(side_effect=etcd.EtcdException())
 
+    def test_base_path(self):
+        self.assertEquals(self.etcd._base_path, '/patroni/test')
+
+    @patch('dns.resolver.query', dns_query)
     def test_get_etcd_client(self):
-        time.sleep = time_sleep_exception
         with patch.object(etcd.Client, 'machines') as mock_machines:
             mock_machines.__get__ = Mock(side_effect=etcd.EtcdException)
-            self.assertRaises(Exception, self.etcd.get_etcd_client, {'discovery_srv': 'test'})
+            with patch('time.sleep', Mock(side_effect=SleepException())):
+                self.assertRaises(SleepException, self.etcd.get_etcd_client, {'discovery_srv': 'test'})
 
     def test_get_cluster(self):
         self.assertIsInstance(self.etcd.get_cluster(), Cluster)
@@ -214,7 +228,7 @@ class TestEtcd(unittest.TestCase):
         self.assertIsNone(cluster.leader)
 
     def test_current_leader(self):
-        self.assertIsInstance(self.etcd.current_leader(), Member)
+        self.assertIsInstance(self.etcd.current_leader(), Leader)
         self.etcd._base_path = '/service/noleader'
         self.assertIsNone(self.etcd.current_leader())
 
@@ -224,12 +238,36 @@ class TestEtcd(unittest.TestCase):
     def test_take_leader(self):
         self.assertFalse(self.etcd.take_leader())
 
-    def test_update_leader(self):
-        self.assertTrue(self.etcd.update_leader(MockPostgresql()))
+    def test_attempt_to_acquire_leader(self):
+        self.etcd._base_path = '/service/exists'
+        self.assertFalse(self.etcd.attempt_to_acquire_leader())
+        self.etcd._base_path = '/service/failed'
+        self.assertFalse(self.etcd.attempt_to_acquire_leader())
 
-    def test_race(self):
-        self.assertFalse(self.etcd.race(''))
+    def test_write_leader_optime(self):
+        self.etcd.write_leader_optime('0')
+
+    def test_update_leader(self):
+        self.assertTrue(self.etcd.update_leader())
+
+    def test_initialize(self):
+        self.assertFalse(self.etcd.initialize())
+
+    def test_cancel_initializion(self):
+        self.assertFalse(self.etcd.cancel_initialization())
 
     def test_delete_leader(self):
-        self.etcd.client.delete = etcd_delete
         self.assertFalse(self.etcd.delete_leader())
+
+    def test_watch(self):
+        self.etcd.client.watch = etcd_watch
+        self.etcd.watch(0)
+        self.etcd.get_cluster()
+        self.etcd.watch(1.5)
+        self.etcd.watch(4.5)
+        self.etcd.watch(9.5)
+        self.etcd.watch(100)
+
+    @patch('patroni.etcd.Etcd.retry', Mock(side_effect=AttributeError("foo")))
+    def test_other_exceptions(self):
+        self.assertRaises(EtcdError, self.etcd.cancel_initialization)

tests/test_ha.py

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

tests/test_patroni.py

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

tests/test_postgresql.py

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

tests/test_restore.py

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

tests/test_utils.py

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

tests/test_zookeeper.py

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