Bump up minor version.

Call a checkpoint on master before pg_rewind.

.gitignore

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

Dockerfile

@@ -13,8 +13,9 @@ 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 apt-get install python 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-dev postgresql-server-dev-${PGVERSION} -y
+RUN pip install python-etcd psycopg2
 
 ENV PATH /usr/lib/postgresql/${PGVERSION}/bin:$PATH
 

README.rst

@@ -5,10 +5,10 @@ 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
+*There are many ways to run high availability with PostgreSQL. Here, we
-present a template for you to create your own custom fit high
+present a template for you to create your own customized, high-availability 
-availability solution using python and distributed configuration store
+solution using Python and — for maximum accessibility — a distributed 
-(like ZooKeeper or etcd) for maximum accessibility.*
+configuration store like ZooKeeper or etcd.*
 
 Getting Started
 ---------------
@@ -22,8 +22,8 @@ To get started, do the following from different terminals:
     > ./patroni.py postgres1.yml
 
 From there, you will see a high-availability cluster start up. Test
-different settings in the YAML files to see how behavior changes. Kill
+different settings in the YAML files to see how its behavior changes. Kill
-some of the different components to see how the system behaves.
+some of the components to see how the system behaves.
 
 Add more ``postgres*.yml`` files to create an even larger cluster.
 
@@ -39,77 +39,89 @@ run:
 
     > psql --host 127.0.0.1 --port 5000 postgres
 
-How Patroni works
+How Patroni Works
 -----------------
 
-For a diagram of the high availability decision loop, see the included a
+For a diagram of the high availability decision loop, review this PDF:
-PDF:
 `postgres-ha.pdf <https://github.com/zalando/patroni/blob/master/postgres-ha.pdf>`__
 
 YAML Configuration
 ------------------
 
-For an example file, see ``postgres0.yml``. Below is an explanation of
+For an example file, see ``postgres0.yml``. Regarding settings:
-settings:
 
--  *ttl*: the TTL to acquire the leader lock. Think of it as the length of time before automatic failover process is initiated.
+-  *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 provide health-check information for haproxy.
+    -  *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.
+    -  *connect\_address*: IP address + port through which restapi is accessible.
-    -  *auth*: (optional) 'username:password' to protect some dangerous REST API endpoints.
+    -  *auth*: (optional) 'username:password' to protect dangerous REST API endpoints.
-    -  *certfile*: (optional) Specifies a file with the certificate in the PEM format. If certfile is not specified or empty API server will work without SSL.
+    -  *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, thus you can run multiple HA deployments from a single 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 automatic failover process is initiated.
+    -  *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
+    -  *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
+    -  *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 automatic failover process is initiated.
+    -  *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 connection loss. After this timeout we assume that we don't have lock anymore and will restart in read-only mode.
+    -  *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 ZooKeeper cluster under Exhibitor supervisory the following section could be interesting for you
+    -  *exhibitor*: if you are running a ZooKeeper cluster under the Exhibitor supervisory, the following section might interest you:
-        -  *poll\_interval*: how often list of ZooKeeper and Exhibitor nodes should be updated from Exhibitor
+        -  *poll\_interval*: how often the list of ZooKeeper and Exhibitor nodes should be updated from Exhibitor
-        -  *port*: Exhibitor port
+        -  *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.
+        -  *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
+    -  *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.
+    -  *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.
+    -  *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
+    -  *data\_dir*: file path to initialize and store Postgres data files.
-    -  *maximum\_lag\_on\_failover*: the maximum bytes a follower may lag
+    -  *maximum\_lag\_on\_failover*: the maximum bytes a follower may lag.
-    -  *use\_slots*: whether or not to use replication_slots.  Must be False for PostgreSQL 9.3, and you should comment out max_replication_slots. before it is not eligible become leader
+    -  *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.
-    -  *pg\_hba*: list of lines which should be added to pg\_hba.conf
+
-        -  *- host all all 0.0.0.0/0 md5*
+    -  *initdb*:  List options to be passed on to initdb
+        -  *encoding*: default encoding for new databases
+        -  *locale*: default locale for new databases
+        -  *data-checksums*  # When pg_rewind is needed on 9.3, this needs to be enabled
+
+    -  *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
+        -  *username*: replication username; user will be created during initialization.
-        -  *password*: replication password, user will be created during initialization
+        -  *password*: replication password; user will be created during initialization.
-        -  *network*: network setting for replication in pg\_hba.conf
+        -  *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.
+    -  *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\_start*: a script to run when the cluster starts.
-        -  *on\_stop*: a script to run when the cluster stops
+        -  *on\_stop*: a script to run when the cluster stops.
-        -  *on\_restart*: a script to run when the cluster restarts
+        -  *on\_restart*: a script to run when the cluster restarts.
-        -  *on\_reload*: a script to run when configuration reload is triggered
+        -  *on\_reload*: a script to run when configuration reload is triggered.
-        -  *on\_role\_change*: a script to run when the cluster is being promoted or demoted
+        -  *on\_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
+        -  *password*: password for the Postgres user, set during initialization.
 
     -  *admin*:
-        -  *username*: admin username, user will be created during initialization. It would have CREATEDB and CREATEROLE privileges
+        -  *username*: admin username; user is created during initialization. It will have CREATEDB and CREATEROLE privileges.
-        -  *password*: admin password, user will be created during initialization.
+        -  *password*: admin password; user is created during initialization.
 
-    -  *recovery\_conf*: additional configuration settings written to recovery.conf when configuring follower
+    -  *recovery\_conf*: additional configuration settings written to recovery.conf when configuring follower.
         -  *parameters*: list of configuration settings for Postgres.  Many of these are required for replication to work.
 
-Replication choices
+    -  *create\_replica\_methods*: an ordered list of the create methods for turning a patroni node into a new replica.
+       "basebackup" is the default method; other methods are assumed to refer to scripts, each of which is configured
+       as its own config item.
+
+    -  *replica\_method* for each create_replica_method other than basebackup, you would add a configuration section
+       of the same name.  At a minimum, this should include "command" with a full path to the actual script to be
+       executed.  Other configuration parameters will be passed along to the script in the form "parameter=value".
+
+Replication Choices
 -------------------
 
 Patroni uses Postgres' streaming replication. By default, this
@@ -123,43 +135,43 @@ not occur if a follower is more than a certain number of bytes behind
 the follower. This setting should be increased or decreased based on
 business requirements.
 
-When asynchronous replication is not best for your use-case, investigate
+When asynchronous replication is not optimal for your use case, investigate
 how Postgres's `synchronous
 replication <http://www.postgresql.org/docs/current/static/warm-standby.html#SYNCHRONOUS-REPLICATION>`__
 works. Synchronous replication ensures consistency across a cluster by
 confirming that writes are written to a secondary before returning to
 the connecting client with a success. The cost of synchronous
-replication will be reduced throughput on writes. This throughput will
+replication: reduced throughput on writes. This throughput will
 be entirely based on network performance. In hosted datacenter
 environments (like AWS, Rackspace, or any network you do not control),
-synchrous replication increases the variability of write performance
+synchrous replication significantly increases the variability of write 
-significantly. If followers become inaccessible from the leader, the
+performance. If followers become inaccessible from the leader, the
-leader will becomes effectively readonly.
+leader effectively becomes readonly.
 
 To enable a simple synchronous replication test, add the follow lines to
-the ``parameters`` section of your YAML configuration files.
+the ``parameters`` section of your YAML configuration files:
 
 .. code:: YAML
 
         synchronous_commit: "on"
         synchronous_standby_names: "*"
 
-When using synchronous replication, use at least a 3-Postgres data nodes
+When using synchronous replication, use at least three Postgres data nodes
 to ensure write availability if one host fails.
 
-Choosing your replication schema is dependent on the many business
+Choosing your replication schema is dependent on your business
-decisions. Investigate both async and sync replication, as well as other
+considerations. Investigate both async and sync replication, as well as other
 HA solutions, to determine which solution is best for you.
 
-Applications should not use superusers
+Applications Should Not Use Superusers
 --------------------------------------
 
 When connecting from an application, always use a non-superuser. Patroni
 requires access to the database to function properly. By using a
-superuser from application, you can potentially use the entire
+superuser from an application, you can potentially use the entire
 connection pool, including the connections reserved for superusers with
 the ``superuser_reserved_connections`` setting. If Patroni cannot access
-the Primary, because the connection pool is full, behavior will be
+the Primary because the connection pool is full, behavior will be
 undesireable.
 
 Requirements on a Mac
@@ -175,14 +187,13 @@ Run the following on a Mac to install requirements:
 Notice
 ------
 
-There are many different ways to do HA with PostgreSQL, see `the
+There are many different ways to do HA with PostgreSQL: See `the
 PostgreSQL
 documentation <https://wiki.postgresql.org/wiki/Replication,_Clustering,_and_Connection_Pooling>`__
 for a complete list.
 
-We call this project a "template" because it is far from a one-size fits
+We call Patroni a "template" because it is far from being a one-size-fits-all
-all, or a plug-and-play replication system. It will have it's own
+or plug-and-play replication system. It will have its own caveats. Use wisely.
-caveats. Use wisely.
 
 .. |Build Status| image:: https://travis-ci.org/zalando/patroni.svg?branch=master
    :target: https://travis-ci.org/zalando/patroni

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

@@ -21,7 +21,7 @@ __EOF__
 }
 
 DOCKER_IP=$(hostname --ip-address)
-PATRONI_SCOPE=batman
+PATRONI_SCOPE=${PATRONI_SCOPE:-batman}
 
 optspec=":vh-:"
 while getopts "$optspec" optchar; do
@@ -83,16 +83,17 @@ 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
+  listen: 0.0.0.0:8008
-  connect_address: 127.0.0.1:8008
+  connect_address: ${DOCKER_IP}:8008
 etcd:
   scope: *scope
   ttl: *ttl
   host: ${ETCD_CLUSTER}
 postgresql:
-  name: postgresql_${DOCKER_IP//./_} ## Replication slots do not allow dots in their name
+  name: ${HOSTNAME}
   scope: *scope
   listen: 0.0.0.0:5432
   connect_address: ${DOCKER_IP}:5432
@@ -115,10 +116,11 @@ 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
+    checkpoint_segments: 64
+    wal_keep_segments: 64
     archive_timeout: 1800s
     max_replication_slots: 20
     hot_standby: "on"

haproxy.cfg

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

patroni/__init__.py

@@ -18,20 +18,24 @@ 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)
-        host, port = config['restapi']['listen'].split(':')
         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 sutable configuration of distributed configuration store')
+        raise Exception('Can not find suitable configuration of distributed configuration store')
 
     def schedule_next_run(self):
         self.next_run += self.nap_time

patroni/api.py

@@ -3,6 +3,8 @@ import fcntl
 import json
 import logging
 import psycopg2
+import socket
+import time
 
 from patroni.exceptions import PostgresConnectionException
 from patroni.utils import Retry, RetryFailedError
@@ -36,16 +38,29 @@ class RestApiHandler(BaseHTTPRequestHandler):
         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_GET(self):
+    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
@@ -69,12 +84,14 @@ class RestApiHandler(BaseHTTPRequestHandler):
             status_code = 503
 
         self.send_response(status_code)
-        self.send_header('Content-Type', 'application/json')
+        if not options:
-        self.end_headers()
+            self.send_header('Content-Type', 'application/json')
-        self.wfile.write(json.dumps(response).encode('utf-8'))
+            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')
@@ -121,6 +138,56 @@ class RestApiHandler(BaseHTTPRequestHandler):
         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
 
@@ -139,6 +206,12 @@ class RestApiHandler(BaseHTTPRequestHandler):
                 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)
@@ -169,11 +242,17 @@ class RestApiHandler(BaseHTTPRequestHandler):
             }
         except (psycopg2.Error, RetryFailedError, PostgresConnectionException):
             state = self.server.patroni.postgresql.state
-            if state in ['stopped', 'starting', 'stopping', 'restarting', 'running']:
+            if state == 'running':
                 logger.exception('get_postgresql_status')
-                state = 'unknown' if state == 'running' else state
+                state = 'unknown'
             return {'state': state}
 
+    def get_tags(self):
+        return {'tags': self.server.patroni.tags}
+
+    def log_message(self, format, *args):
+        logger.debug("API thread: " + format % args)
+
 
 class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
 

patroni/ctl.py

@@ -0,0 +1,616 @@
+'''
+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=60)
+
+
+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, 'member': 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
+
+    xlog_location_cluster = cluster.last_leader_operation or 0
+
+    # Mainly for consistent pretty printing and watching we sort the output
+    cluster.members.sort(key=lambda x: x.name)
+    for m in cluster.members:
+        logging.debug(m)
+
+        leader = ''
+        if m.name == leader_name:
+            leader = '*'
+
+        host = build_connect_parameters(m.conn_url)['host']
+
+        xlog_location = m.data.get('xlog_location')
+        if xlog_location is None or (xlog_location_cluster < xlog_location):
+            lag = ''
+        else:
+            lag = round((xlog_location_cluster - xlog_location)/1024/1024)
+
+        rows.append([
+            name,
+            m.name,
+            host,
+            leader,
+            m.data.get('state', ''),
+            lag
+        ])
+
+    columns = [
+        'Cluster',
+        'Member',
+        'Host',
+        'Leader',
+        'State',
+        'Lag in MB',
+    ]
+    alignment = {'Cluster': 'l', 'Member': 'l', 'Host': 'l', 'Lag in MB': 'r'}
+
+    print_output(columns, rows, alignment, 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

@@ -63,6 +63,10 @@ class Member(namedtuple('Member', 'index,name,session,data')):
     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')):
 
@@ -122,8 +126,8 @@ class AbstractDCS:
             i.e.: `zookeeper` for zookeeper, `etcd` for etcd, etc...
         """
         self._name = name
-        self._scope = config['scope']
+        self._namespace = '/{}'.format(config.get('namespace', '/service/').strip('/'))
-        self._base_path = '/service/' + self._scope
+        self._base_path = '/'.join([self._namespace, config['scope']])
 
         self._cluster = None
         self._cluster_thread_lock = Lock()
@@ -240,8 +244,11 @@ class AbstractDCS:
         overwriting the key if necessary."""
 
     @abc.abstractmethod
-    def initialize(self):
+    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`

patroni/etcd.py

@@ -52,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):
@@ -139,6 +143,10 @@ def catch_etcd_errors(func):
             return not func(*args, **kwargs) is None
         except (RetryFailedError, etcd.EtcdException):
             return False
+        except:
+            logger.exception("")
+            raise EtcdError("unexpected error")
+
     return wrapper
 
 
@@ -146,7 +154,7 @@ class Etcd(AbstractDCS):
 
     def __init__(self, name, config):
         super(Etcd, self).__init__(name, config)
-        self.ttl = config['ttl']
+        self.ttl = config.get('ttl', 30)
         self._retry = Retry(deadline=10, max_delay=1, max_tries=-1,
                             retry_exceptions=(etcd.EtcdConnectionFailed,
                                               etcd.EtcdLeaderElectionInProgress,
@@ -177,7 +185,8 @@ class Etcd(AbstractDCS):
             nodes = {os.path.relpath(node.key, result.key): node for node in result.leaves}
 
             # get initialize flag
-            initialize = bool(nodes.get(self._INITIALIZE, False))
+            initialize = nodes.get(self._INITIALIZE, None)
+            initialize = initialize and initialize.value
 
             # get last leader operation
             last_leader_operation = nodes.get(self._LEADER_OPTIME, None)
@@ -235,8 +244,8 @@ class Etcd(AbstractDCS):
         return self.retry(self.client.test_and_set, self.leader_path, self._name, self._name, self.ttl)
 
     @catch_etcd_errors
-    def initialize(self):
+    def initialize(self, create_new=True, sysid=""):
-        return self.retry(self.client.write, self.initialize_path, self._name, prevExist=False)
+        return self.retry(self.client.write, self.initialize_path, sysid, prevExist=(not create_new))
 
     @catch_etcd_errors
     def delete_leader(self):
@@ -244,7 +253,7 @@ class Etcd(AbstractDCS):
 
     @catch_etcd_errors
     def cancel_initialization(self):
-        return self.retry(self.client.delete, self.initialize_path, prevValue=self._name)
+        return self.retry(self.client.delete, self.initialize_path)
 
     def watch(self, timeout):
         cluster = self.cluster

patroni/exceptions.py

@@ -13,6 +13,10 @@ class PatroniException(Exception):
         return repr(self.value)
 
 
+class PatroniCtlException(Exception):
+    pass
+
+
 class PostgresException(PatroniException):
     pass
 

patroni/ha.py

@@ -2,6 +2,7 @@ import json
 import logging
 import psycopg2
 import requests
+import sys
 
 from patroni.async_executor import AsyncExecutor
 from patroni.exceptions import DCSError, PostgresConnectionException
@@ -50,7 +51,8 @@ class Ha:
             'conn_url': self.state_handler.connection_string,
             'api_url': self.patroni.api.connection_string,
             'state': self.state_handler.state,
-            'role': self.state_handler.role
+            'role': self.state_handler.role,
+            'tags': self.patroni.tags
         }
         if data['state'] in ['running', 'restarting', 'starting']:
             try:
@@ -72,10 +74,11 @@ class Ha:
             self._async_executor.schedule('bootstrap from leader')
             self._async_executor.run_async(self.copy_backup_from_leader, args=(self.cluster.leader, ))
             return 'trying to bootstrap from leader'
-        elif not self.cluster.initialize:  # no initialize key
+        elif not self.cluster.initialize and not self.patroni.nofailover:  # no initialize key
-            if self.dcs.initialize():  # race for initialization
+            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")
@@ -130,6 +133,7 @@ class Ha:
             return message
         else:
             self.state_handler.promote()
+            self.touch_member()
             return promote_message
 
     @staticmethod
@@ -139,7 +143,9 @@ class Ha:
 
         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."""
+        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)
@@ -147,10 +153,11 @@ class Ha:
             json = response.json()
             is_master = json['role'] == 'master'
             xlog_location = json['xlog']['location' if is_master else 'replayed_location']
-            return (member, True, not is_master, xlog_location)
+            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)
+        return (member, False, None, 0, {})
 
     def fetch_nodes_statuses(self, members):
         pool = ThreadPool(len(members))
@@ -165,16 +172,19 @@ class Ha:
         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 m.api_url]
+        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 in self.fetch_nodes_statuses(members):
+            for member, reachable, in_recovery, xlog_location, tags in self.fetch_nodes_statuses(members):
-                if reachable:  # If the node is unreachable it's not healhy
+                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
@@ -184,13 +194,15 @@ class Ha:
 
     def is_failover_possible(self, members):
         ret = False
-        members = [m for m in members if m.name != self.state_handler.name and m.api_url]
+        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 in self.fetch_nodes_statuses(members):
+            for member, reachable, in_recovery, xlog_location, tags in self.fetch_nodes_statuses(members):
-                if reachable:
+                if reachable and not tags.get('nofailover', False):
                     ret = True  # TODO: check xlog_location
-                else:
+                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
@@ -204,12 +216,15 @@ class Ha:
             # find specific node and check that it is healthy
             members = [m for m in self.cluster.members if m.name == failover.member]
             if members:
-                member, reachable, in_recovery, xlog_location = self.fetch_node_status(members[0])
+                member, reachable, in_recovery, xlog_location, tags = self.fetch_node_status(members[0])
-                if reachable:  # node is healthy
+                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
-                logger.warning('manual failover: member %s is unhealthy', member.name)
+                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
@@ -218,7 +233,7 @@ class Ha:
         if failover.leader:
             if self.state_handler.name == failover.leader:  # I was the leader
                 # exclude me and desired member which is unhealthy (failover.member can be None)
-                members = [m for m in self.cluster.members if m.name != failover.member]
+                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
@@ -231,6 +246,13 @@ class Ha:
         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()
 
@@ -242,6 +264,7 @@ class Ha:
         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)
 
@@ -269,14 +292,18 @@ class Ha:
         if self.is_healthiest_node():
             if self.acquire_lock():
                 if self.cluster.failover:
-                    logger.info('Cleanning up failover key after acquiring leader lock...')
+                    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')
 
@@ -350,6 +377,11 @@ class Ha:
         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()
@@ -357,8 +389,8 @@ class Ha:
             self.touch_member()
 
             # cluster has leader key but not initialize key
-            if not self.cluster.is_unlocked() and not self.cluster.initialize:
+            if not self.cluster.is_unlocked() and not self.sysid_valid(self.cluster.initialize) and self.has_lock():
-                self.dcs.initialize()  # fix it
+                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()
@@ -372,8 +404,14 @@ class Ha:
             if self.state_handler.data_directory_empty():
                 return self.bootstrap()  # new node
             # "bootstrap", but data directory is not empty
-            elif not self.cluster.initialize and self.cluster.is_unlocked():
+            elif not self.sysid_valid(self.cluster.initialize) and self.cluster.is_unlocked():
-                self.dcs.initialize()
+                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():

patroni/postgresql.py

@@ -4,10 +4,12 @@ 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
 
@@ -42,12 +44,15 @@ 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)
@@ -69,6 +74,7 @@ class Postgresql:
         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)
 
@@ -100,10 +106,14 @@ class Postgresql:
             return False
         # check if the cluster's configuration permits pg_rewind
         data = self.controldata()
-        if data:
+        return data.get('wal_log_hints setting', 'off') == 'on' or data.get('Data page checksum version', '0') != '0'
-            return data.get('wal_log_hints setting', 'off') == 'on' or\
+
-                data.get('Data page checksum version', '0') != '0'
+    @property
-        return False
+    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
@@ -127,9 +137,15 @@ class Postgresql:
 
     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:
@@ -152,9 +168,39 @@ class Postgresql:
     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')
-        ret = subprocess.call(self._pg_ctl + ['initdb', '-o', '--encoding=UTF8']) == 0
+        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:
@@ -165,36 +211,75 @@ class Postgresql:
         os.path.exists(self.trigger_file) and os.unlink(self.trigger_file)
 
     def write_pgpass(self, record):
-        pgpass = 'pgpass'
+        with open(self.pgpass, 'w') as f:
-        with open(pgpass, 'w') as f:
             os.fchmod(f.fileno(), 0o600)
             f.write('{host}:{port}:*:{user}:{password}\n'.format(**record))
+
         env = os.environ.copy()
-        env['PGPASSFILE'] = pgpass
+        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
+        ret = self.create_replica(leader, env) == 0
+        ret and self.delete_trigger_file()
+        return ret
 
     @staticmethod
     def build_connstring(conn):
-        return "host={host} port={port} user={user}".format(**conn)
+        """
+        >>> Postgresql.build_connstring({'host': '127.0.0.1', 'port': '5432'}) == 'host=127.0.0.1 port=5432'
+        True
+        """
+        return ' '.join('{}={}'.format(param, val) for param, val in sorted(conn.items()))
+
+    def create_replica(self, leader, env):
+        # create the replica according to the replica_method
+        # defined by the user.  this is a list, so we need to
+        # loop through all methods the user supplies
+        connstring = leader.conn_url
+        # get list of replica methods from config.
+        # If there is no configuration key, or no value is specified, use basebackup
+        replica_methods = self.config.get('create_replica_method') or ['basebackup']
+        # go through them in priority order
+        ret = 1
+        for replica_method in replica_methods:
+            # if the method is basebackup, then use the built-in
+            if replica_method == "basebackup":
+                ret = self.basebackup(leader, env)
+                if ret == 0:
+                    logger.info("replica has been created using basebackup")
+                    # if basebackup succeeds, exit with success
+                    break
+            else:
+                cmd = replica_method
+                method_config = {}
+                # user-defined method; check for configuration
+                # not required, actually
+                if replica_method in self.config:
+                    method_config = self.config[replica_method].copy()
+                    # look to see if the user has supplied a full command path
+                    # if not, use the method name as the command
+                    cmd = method_config.pop('command', cmd)
+                    # add the default parameters
+                try:
+                    method_config.update({"scope": self.scope,
+                                          "role": "replica",
+                                          "datadir": self.data_dir,
+                                          "connstring": connstring})
+                    params = ["--{0}={1}".format(arg, val) for arg, val in method_config.items()]
+                    # call script with the full set of parameters
+                    ret = subprocess.call(shlex.split(cmd) + params, env=env)
+                    # if we succeeded, stop
+                    if ret == 0:
+                        logger.info("replica has been created using {0}".format(replica_method))
+                        break
+                except Exception as e:
+                    logger.exception('Error creating replica using method {0}: {1}'.format(replica_method, str(e)))
+                    ret = 1
 
-    def create_replica(self, master_connection, env):
-        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):
@@ -257,18 +342,24 @@ class Postgresql:
         ret and not block_callbacks and self.call_nowait(ACTION_ON_START)
         return ret
 
-    def checkpoint(self):
+    def checkpoint(self, connstring=None):
         try:
-            r = parseurl('postgres://{}/postgres'.format(self.local_address))
+            connstring = connstring or 'postgres://{}/postgres'.format(self.local_address)
-            r['options'] = '-c statement_timeout=0'
+            with psycopg2.connect(connstring) as conn:
-            with psycopg2.connect(**r) as conn:
                 conn.autocommit = True
                 with conn.cursor() as cur:
+                    cur.execute("SET statement_timeout = 0")
                     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')
@@ -305,7 +396,7 @@ class Postgresql:
 
     def server_options(self):
         options = "--listen_addresses='{}' --port={}".format(self.listen_addresses, self.port)
-        for setting, value in self.config['parameters'].items():
+        for setting, value in self.server_parameters.items():
             options += " --{}='{}'".format(setting, value)
         return options
 
@@ -316,13 +407,14 @@ class Postgresql:
         return True
 
     def check_replication_lag(self, last_leader_operation):
-        return last_leader_operation - self.xlog_position() <= self.config.get('maximum_lag_on_failover', 0)
+        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.config['parameters'].get('ssl', 'off').lower() != 'on':
+                if line.split()[0].strip() == 'hostssl' and self.server_parameters.get('ssl', 'off').lower() != 'on':
                     continue
                 f.write(line + '\n')
 
@@ -362,6 +454,9 @@ recovery_target_timeline = 'latest'
         r['user'] = r['username']
         env = self.write_pgpass(r)
         pc = "user={user} host={host} port={port} dbname=postgres sslmode=prefer sslcompression=1".format(**r)
+        # first run a checkpoint on a promoted master in order
+        # to make it store the new timeline ([email protected])
+        self.checkpoint(pc)
         logger.info("running pg_rewind from {}".format(pc))
         pg_rewind = ['pg_rewind', '-D', self.data_dir, '--source-server', pc]
         try:
@@ -378,8 +473,8 @@ recovery_target_timeline = 'latest'
         try:
             data = subprocess.check_output(['pg_controldata', self.data_dir])
             if data:
-                data = data.splitlines()
+                data = data.decode().splitlines()
-                result = {l.split(':')[0]: l.split(':')[1].strip() for l in data if l}
+                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:
@@ -472,19 +567,23 @@ recovery_target_timeline = 'latest'
 
     def save_configuration_files(self):
         """
-            copy postgresql.conf to postgresql.conf.backup to preserve it in the WAL-e backup.
+            copy postgresql.conf to postgresql.conf.backup to be able to retrive configuration files
-            see http://comments.gmane.org/gmane.comp.db.postgresql.wal-e/239
+            - originally stored as symlinks, those are normally skipped by pg_basebackup
+            - in case of WAL-E basebackup (see http://comments.gmane.org/gmane.comp.db.postgresql.wal-e/239)
         """
-        for f in self.configuration_to_save:
+        try:
-            shutil.copy(f, f + '.backup')
+            for f in self.configuration_to_save:
+                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:
-                shutil.copy(f + '.backup', f)
+                not os.path.isfile(f) and os.path.isfile(f + '.backup') and shutil.copy(f + '.backup', f)
         except:
-            logger.exception('unable to restore configuration from WAL-E backup')
+            logger.exception('unable to restore configuration files from backup')
 
     def promote(self):
         if self.role == 'master':
@@ -500,21 +599,27 @@ recovery_target_timeline = 'latest'
     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.query('CREATE USER "{}" WITH REPLICATION ENCRYPTED PASSWORD %s'.format(
+        self.create_or_update_role(self.replication['username'], self.replication['password'], 'REPLICATION')
-            self.replication['username']), self.replication['password'])
 
     def create_connection_users(self):
-        if self.superuser:
+        if 'username' in self.superuser:
-            if 'username' in self.superuser:
+            self.create_or_update_role(self.superuser['username'], self.superuser['password'], 'SUPERUSER')
-                self.query('CREATE ROLE "{0}" WITH LOGIN SUPERUSER PASSWORD %s'.format(
-                    self.superuser['username']), self.superuser['password'])
-            else:
-                rolsuper = self.query("""SELECT rolname FROM pg_authid WHERE rolsuper = 't'""").fetchone()[0]
-                self.query('ALTER ROLE "{0}" WITH PASSWORD %s'.format(rolsuper), self.superuser['password'])
         if self.admin:
-            self.query('CREATE ROLE "{0}" WITH LOGIN CREATEDB CREATEROLE PASSWORD %s'.format(
+            self.create_or_update_role(self.admin['username'], self.admin['password'], 'CREATEDB CREATEROLE')
-                self.admin['username']), self.admin['password'])
 
     def xlog_position(self):
         return self.query("""SELECT pg_xlog_location_diff(CASE WHEN pg_is_in_recovery()
@@ -572,6 +677,7 @@ recovery_target_timeline = 'latest'
                 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
@@ -599,3 +705,27 @@ recovery_target_timeline = 'latest'
         except:
             logger.exception('Could not remove data directory %s', self.data_dir)
             self.move_data_directory()
+
+    def basebackup(self, leader, env):
+        # creates a replica data dir using pg_basebackup.
+        # this is the default, built-in create_replica_method
+        # tries twice, then returns failure (as 1)
+        # uses "stream" as the xlog-method to avoid sync issues
+        master_connection = leader.conn_url
+        maxfailures = 2
+        ret = 1
+        for bbfailures in range(0, maxfailures):
+            try:
+                ret = subprocess.call(['pg_basebackup', '--pgdata=' + self.data_dir,
+                                       '--xlog-method=stream', "--dbname=" + master_connection], env=env)
+                if ret == 0:
+                    break
+
+            except Exception as e:
+                logger.error('Error when fetching backup with pg_basebackup: {0}'.format(e))
+
+            if bbfailures < maxfailures - 1:
+                logger.error('Trying again in 5 seconds')
+                time.sleep(5)
+
+        return ret

patroni/scripts/aws.py

@@ -65,8 +65,11 @@ class AWSConnection:
         return self._tag_ebs(new_role) and ret
 
 
-if __name__ == '__main__':
+def main():
     if len(sys.argv) == 4 and sys.argv[1] in ('on_start', 'on_stop', 'on_role_change'):
         AWSConnection(cluster_name=sys.argv[3]).on_role_change(sys.argv[2])
     else:
         sys.exit("Usage: {0} action role name".format(sys.argv[0]))
+
+if __name__ == '__main__':
+    main()

patroni/scripts/restore.py

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

patroni/scripts/wale_restore.py

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

patroni/version.py

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

patroni/zookeeper.py

@@ -139,7 +139,7 @@ class ZooKeeper(AbstractDCS):
             self.fetch_cluster = True
 
         # get initialize flag
-        initialize = self._INITIALIZE in nodes
+        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 []
@@ -198,13 +198,14 @@ class ZooKeeper(AbstractDCS):
             self.client.retry(self.client.set, self.failover_path, value.encode('utf-8'), version=index or -1)
             return True
         except NoNodeError:
-            return value == '' or (not index and self._create(self.failover_path, value.encode('utf-8')))
+            return value == '' or (not index and self._create(self.failover_path, value))
         except:
             logging.exception('set_failover_value')
             return False
 
-    def initialize(self):
+    def initialize(self, create_new=True, sysid=""):
-        return self._create(self.initialize_path, self._name, makepath=True)
+        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
@@ -270,7 +271,7 @@ class ZooKeeper(AbstractDCS):
 
     def _cancel_initialization(self):
         node = self.get_node(self.initialize_path)
-        if node and node[0] == self._name:
+        if node:
             self.client.delete(self.initialize_path, version=node[1].version)
 
     def cancel_initialization(self):

patronictl.py

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

postgres0.yml

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

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

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

@@ -32,7 +32,7 @@ VERSION = read_version(MAIN_PACKAGE)
 DESCRIPTION = 'PostgreSQL High-Available orchestrator and CLI'
 LICENSE = 'The MIT License'
 URL = 'https://github.com/zalando/patroni'
-AUTHOR = 'Alexander Kukushkin, Alexey Klyukin, Feike Steenbergen'
+AUTHOR = 'Alexander Kukushkin, Oleksii Kliukin, Feike Steenbergen'
 AUTHOR_EMAIL = '[email protected], [email protected], [email protected]'
 KEYWORDS = 'etcd governor patroni postgresql postgres ha zookeeper streaming replication'
 
@@ -56,7 +56,10 @@ CLASSIFIERS = [
     'Programming Language :: Python :: Implementation :: CPython',
 ]
 
-CONSOLE_SCRIPTS = ['patroni = patroni:main']
+CONSOLE_SCRIPTS = ['patroni = patroni:main',
+                   'patronictl = patroni.ctl:ctl',
+                   "patroni_wale_restore = patroni.scripts.wale_restore:main",
+                   "patroni_aws = patroni.scripts.aws:main"]
 
 
 class PyTest(TestCommand):

tests/test_api.py

@@ -3,8 +3,11 @@ import unittest
 
 from mock import Mock, patch
 from patroni.api import RestApiHandler, RestApiServer
+from patroni.dcs import Member
 from six import BytesIO as IO
 from six.moves import BaseHTTPServer
+from six.moves.BaseHTTPServer import BaseHTTPRequestHandler
+import socket
 from test_postgresql import psycopg2_connect, MockCursor
 
 
@@ -38,12 +41,16 @@ class MockHa(Mock):
     def restart_scheduled(self):
         return False
 
+    def fetch_nodes_statuses(self, members):
+        return [[None, True, None, None, {}]]
+
 
 class MockPatroni:
 
     postgresql = MockPostgresql()
     ha = MockHa()
     dcs = Mock()
+    tags = {}
 
 
 class MockRequest:
@@ -86,6 +93,20 @@ class TestRestApiHandler(unittest.TestCase):
             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')
 
@@ -117,3 +138,31 @@ class TestRestApiHandler(unittest.TestCase):
             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_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

@@ -8,7 +8,7 @@ import unittest
 from dns.exception import DNSException
 from mock import Mock, patch
 from patroni.dcs import Cluster, DCSError, Leader
-from patroni.etcd import Client, Etcd
+from patroni.etcd import Client, Etcd, EtcdError
 
 
 class MockResponse:
@@ -17,6 +17,7 @@ class MockResponse:
         self.status_code = 200
         self.content = '{}'
         self.ok = True
+        self.text = ''
 
     def json(self):
         return json.loads(self.content)
@@ -50,7 +51,7 @@ def requests_get(url, **kwargs):
     if url.startswith('http://local'):
         raise requests.exceptions.RequestException()
     elif ':8011/patroni' in url:
-        response.content = '{"role": "replica", "xlog": {"replayed_location": 0}}'
+        response.content = '{"role": "replica", "xlog": {"replayed_location": 0}, "tags": {}}'
     elif url.endswith('/members'):
         if url.startswith('http://error'):
             response.content = '[{}]'
@@ -80,7 +81,7 @@ def etcd_watch(key, index=None, timeout=None, recursive=None):
 def etcd_write(key, value, **kwargs):
     if key == '/service/exists/leader':
         raise etcd.EtcdAlreadyExist
-    if key == '/service/test/leader':
+    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
@@ -106,13 +107,13 @@ def etcd_read(key, **kwargs):
                  "modifiedIndex": 20437, "createdIndex": 20437},
                 {"key": "/service/batman5/members", "dir": True, "nodes": [
                     {"key": "/service/batman5/members/postgresql1",
-                     "value": "postgres://replicator:[email protected]:5434/postgres"
+                     "value": "postgres://replicator:[email protected]:5434/postgres" +
-                        + "?application_name=http://127.0.0.1:8009/patroni",
+                        "?application_name=http://127.0.0.1:8009/patroni",
                      "expiration": "2015-05-15T09:10:59.949384522Z", "ttl": 21,
                      "modifiedIndex": 20727, "createdIndex": 20727},
                     {"key": "/service/batman5/members/postgresql0",
-                     "value": "postgres://replicator:[email protected]:5433/postgres"
+                     "value": "postgres://replicator:[email protected]:5433/postgres" +
-                        + "?application_name=http://127.0.0.1:8008/patroni",
+                        "?application_name=http://127.0.0.1:8008/patroni",
                      "expiration": "2015-05-15T09:11:09.611860899Z", "ttl": 30,
                      "modifiedIndex": 20730, "createdIndex": 20730}],
                  "modifiedIndex": 1581, "createdIndex": 1581}], "modifiedIndex": 1581, "createdIndex": 1581}}
@@ -165,6 +166,11 @@ 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'), [])
@@ -199,11 +205,14 @@ class TestEtcd(unittest.TestCase):
     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):
         with patch.object(etcd.Client, 'machines') as mock_machines:
@@ -258,3 +267,7 @@ class TestEtcd(unittest.TestCase):
         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,6 +1,7 @@
+import etcd
 import unittest
 
-from mock import Mock, patch
+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
@@ -26,7 +27,7 @@ def get_cluster_not_initialized_without_leader():
 
 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'})
+                                 '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'})
@@ -36,6 +37,10 @@ def get_cluster_initialized_without_leader(leader=False, failover=None):
 def get_cluster_initialized_with_leader(failover=None):
     return get_cluster_initialized_without_leader(leader=True, failover=failover)
 
+def get_cluster_initialized_with_only_leader(failover=None):
+    l = get_cluster_initialized_without_leader(leader=True, failover=failover).leader
+    return get_cluster(True, l, [l], failover)
+
 
 class MockPostgresql(Mock):
 
@@ -81,6 +86,8 @@ class MockPatroni:
         self.postgresql = p
         self.dcs = d
         self.api = Mock()
+        self.tags = {}
+        self.nofailover = None
         self.api.connection_string = 'http://127.0.0.1:8008'
 
 
@@ -98,6 +105,7 @@ class TestHa(unittest.TestCase):
         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()
@@ -128,6 +136,12 @@ class TestHa(unittest.TestCase):
         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.p.is_leader = self.p.is_healthy = false
@@ -264,6 +278,11 @@ class TestHa(unittest.TestCase):
         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):
@@ -272,24 +291,41 @@ class TestHa(unittest.TestCase):
         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.ha.fetch_node_status = lambda e: (e, True, True, 0, {})  # accessible, in_recovery
         self.assertEquals(self.ha.run_cycle(), 'following a different leader because i am not the healthiest node')
         self.ha.cluster = get_cluster_initialized_without_leader(failover=Failover(0, MockPostgresql.name, ''))
         self.assertEquals(self.ha.run_cycle(), 'following a different leader because i am not the healthiest node')
-        self.ha.fetch_node_status = lambda e: (e, False, True, 0)  # accessible, in_recovery
+        self.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.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.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.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):

tests/test_patroni.py

@@ -74,3 +74,9 @@ class TestPatroni(unittest.TestCase):
         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

@@ -4,12 +4,7 @@ import psycopg2
 import shutil
 import unittest
 
-from sys import version_info
+from six.moves import builtins
-if version_info.major == 2:
-    import __builtin__ as builtins
-else:
-    import builtins
-
 from mock import Mock, MagicMock, PropertyMock, patch, mock_open
 from patroni.dcs import Cluster, Leader, Member
 from patroni.exceptions import PostgresException, PostgresConnectionException
@@ -19,6 +14,11 @@ from test_ha import false
 import subprocess
 
 
+def is_file_raise_on_backup(*args, **kwargs):
+    if args[0].endswith('.backup'):
+        raise Exception("foo")
+
+
 class MockCursor:
 
     def __init__(self, connection):
@@ -56,6 +56,9 @@ class MockCursor:
     def fetchone(self):
         return self.results[0]
 
+    def fetchall(self):
+        return self.results
+
     def close(self):
         pass
 
@@ -86,7 +89,7 @@ class MockConnect(Mock):
 
 
 def pg_controldata_string(*args, **kwargs):
-    return """
+    return b"""
 pg_control version number:            942
 Catalog version number:               201509161
 Database system identifier:           6200971513092291716
@@ -118,12 +121,12 @@ Backup start location:                0/0
 Backup end location:                  0/0
 End-of-backup record required:        no
 wal_level setting:                    hot_standby
-wal_log_hints setting:                on
+Current wal_log_hints setting:                on
-max_connections setting:              100
+Current max_connections setting:              100
-max_worker_processes setting:         8
+Current max_worker_processes setting:         8
-max_prepared_xacts setting:           0
+Current max_prepared_xacts setting:           0
-max_locks_per_xact setting:           64
+Current max_locks_per_xact setting:           64
-track_commit_timestamp setting:       off
+Current track_commit_timestamp setting:       off
 Maximum data alignment:               8
 Database block size:                  8192
 Blocks per segment of large relation: 131072
@@ -141,10 +144,10 @@ 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"'\
+    return '/usr/local/pgsql/bin/postgres "-D" "data/postgresql0" "--listen_addresses=127.0.0.1" \
-        ' "--hot_standby=on" "--wal_keep_segments=8" "--wal_level=hot_standby" "--archive_command=mkdir -p ../wal_archive \n'\
+"--port=5432" "--hot_standby=on" "--wal_keep_segments=8" "--wal_level=hot_standby" \
-        '&& cp %p ../wal_archive/%f" "--wal_log_hints=on" "--max_wal_senders=5" "--archive_timeout=1800s" "--archive_mode=on"'\
+"--archive_command=mkdir -p ../wal_archive && cp %p ../wal_archive/%f" "--wal_log_hints=on" \
-        ' "--max_replication_slots=5"\n'
+"--max_wal_senders=5" "--archive_timeout=1800s" "--archive_mode=on" "--max_replication_slots=5"\n'
 
 
 def psycopg2_connect(*args, **kwargs):
@@ -162,7 +165,7 @@ class TestPostgresql(unittest.TestCase):
         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',
@@ -187,6 +190,16 @@ class TestPostgresql(unittest.TestCase):
     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')))
@@ -210,10 +223,16 @@ class TestPostgresql(unittest.TestCase):
         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))
 
     @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
@@ -222,6 +241,7 @@ class TestPostgresql(unittest.TestCase):
     @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)
@@ -232,8 +252,8 @@ class TestPostgresql(unittest.TestCase):
         self.p.follow_the_leader(self.leader)
         self.p.require_rewind()
         with mock.patch('os.path.islink', MagicMock(return_value=True)):
-            with mock.patch('os.unlink', MagicMock(return_value=True)):
+            with mock.patch('patroni.postgresql.Postgresql.can_rewind', new_callable=PropertyMock(return_value=True)):
-                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)):
@@ -256,9 +276,23 @@ class TestPostgresql(unittest.TestCase):
         self.assertTrue(self.p.can_rewind)
         self.p.controldata = tmp
 
+    @patch('time.sleep', Mock())
     def test_create_replica(self):
         self.p.delete_trigger_file = Mock(side_effect=OSError())
-        self.assertEquals(self.p.create_replica({'host': '', 'port': '', 'user': ''}, ''), 1)
+        with patch('subprocess.call', Mock(side_effect=[1, 0])):
+            self.assertEquals(self.p.create_replica(self.leader, ''), 0)
+        with patch('subprocess.call', Mock(side_effect=[Exception(), 0])):
+            self.assertEquals(self.p.create_replica(self.leader, ''), 0)
+
+        self.p.config['create_replica_method'] = ['wale', 'basebackup']
+        self.p.config['wale'] = {'command': 'foo'}
+        with patch('subprocess.call', Mock(return_value=0)):
+            self.assertEquals(self.p.create_replica(self.leader, ''), 0)
+            del self.p.config['wale']
+            self.assertEquals(self.p.create_replica(self.leader, ''), 0)
+
+        with patch('subprocess.call', Mock(side_effect=Exception("foo"))):
+            self.assertEquals(self.p.create_replica(self.leader, ''), 1)
 
     def test_create_connection_users(self):
         cfg = self.p.config
@@ -327,6 +361,7 @@ class TestPostgresql(unittest.TestCase):
         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)
@@ -429,3 +464,19 @@ class TestPostgresql(unittest.TestCase):
         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,111 +0,0 @@
-import unittest
-from mock import MagicMock, patch
-import os
-from patroni.scripts.restore import Restore, WALERestore
-
-
-def fake_cursor_fetchone(*args, **kwargs):
-    return ('16777216',)
-
-
-def fake_call_fail_for_wal_e(*args, **kwargs):
-    if len(args) > 0 and 'backup-fetch' in args[0]:
-        return 1
-    return 0
-
-
-def fake_call_fail_for_base_backup(*args, **kwargs):
-    if len(args) > 0 and 'backup-fetch' in args[0]:
-        return 0
-    return 1
-
-
-def fake_backup_data(self, *args, **kwargs):
-    """ return the fake result of WAL-E backup-list"""
-    return """name    last_modified   expanded_size_bytes wal_segment_backup_start    wal_segment_offset_backup_start wal_segment_backup_stop wal_segment_offset_backup_stop
-base_00000001000000000000007F_00000040  2015-05-18T10:13:25.000Z 167772160   00000001000000000000007F    00000040 00000001000000000000007F    00000240
-"""
-
-
-class TestRestore(unittest.TestCase):
-
-    def setUp(self):
-        self.restore = Restore("batman", "master", "/data", "host=batman port=5432 user=batman")
-        pass
-
-    def tearDown(self):
-        pass
-
-    def test_parse_connstring(self):
-        self.assertDictEqual(self.restore.master_connection, {'host': 'batman', 'port': '5432', 'user': 'batman'})
-
-    @patch('subprocess.call', MagicMock(return_value=0))
-    def test_run(self):
-        ret = self.restore.run()
-        self.assertEqual(ret, 0)
-
-    @patch('subprocess.call', MagicMock(return_value=1))
-    def test_run_fail(self):
-        ret = self.restore.run()
-        self.assertEqual(ret, 1)
-
-
-@patch('os.access', MagicMock(return_value=True))
-@patch('os.makedirs', MagicMock(return_value=True))
-@patch('os.path.exists', MagicMock(return_value=True))
-@patch('os.path.isdir', MagicMock(return_value=True))
-@patch('psycopg2.extensions.cursor.fetchone', MagicMock(side_effect=fake_cursor_fetchone))
-@patch('psycopg2.extensions.cursor', MagicMock(autospec=True))
-@patch('psycopg2.extensions.connection', MagicMock(autospec=True))
-@patch('psycopg2.connect', MagicMock(autospec=True))
-@patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data))
-class TestWALERestore(unittest.TestCase):
-
-    def setUp(self):
-        env = {}
-        env['WAL_S3_BUCKET'] = 'batman'
-        env['WALE_BACKUP_THRESHOLD_PERCENTAGE'] = 100
-        env['WALE_BACKUP_THRESHOLD_MEGABYTES'] = 100
-        self.wale_restore = WALERestore("batman", "master", "/data", "host=batman port=5432 user=batman", env=env)
-
-    def tearDown(self):
-        pass
-
-    def test_setup(self):
-        self.wale_restore.setup()
-        self.assertFalse(self.wale_restore.init_error)
-
-    # have to redefine the class-level os.access mock inside the function
-    # since the class-level mock will be applied after the function level one.
-    @patch('os.access', return_value=False)
-    def test_setup_fail(self, mock_no_access):
-        os.access = mock_no_access
-        self.wale_restore.setup()
-        self.assertTrue(self.wale_restore.init_error)
-
-    # The 3 tests above only differ with the mock function instead of a subprocess call
-    # in the first one, subprocess call should return success only for wal-e command,
-    # checking the primary use-case of restoring from WAL-E backup.
-    # In the second one, we test fallbacks by failing at WAL-E, but succeeding at
-    # pg_basebackup.
-    # Finally, the last use case is when all subprocess.call fails. resulting in a
-    # failure to restore from replica
-    @patch('subprocess.call',
-           MagicMock(side_effect=lambda *args, **kwargs: 0 if 'wal-e' in args[0] else 1))
-    def test_run(self):
-        self.wale_restore.setup()
-        ret = self.wale_restore.run()
-        self.assertEqual(ret, 0)
-
-    @patch('subprocess.call',
-           MagicMock(side_effect=lambda *args, **kwargs: 0 if 'pg_basebackup' in args[0] else 1))
-    def test_run_fallback(self):
-        self.wale_restore.setup()
-        ret = self.wale_restore.run()
-        self.assertEqual(ret, 0)
-
-    @patch('subprocess.call', MagicMock(return_value=1))
-    def test_run_all_fail(self):
-        self.wale_restore.setup()
-        ret = self.wale_restore.run()
-        self.assertEqual(ret, 1)

tests/test_utils.py

@@ -67,7 +67,9 @@ class TestRetrySleeper(unittest.TestCase):
         self.assertRaises(RetryFailedError, retry, self._fail(times=100))
 
     def test_copy(self):
-        _sleep = lambda t: None
+        def _sleep(t):
+            None
+
         retry = self._makeOne(sleep_func=_sleep)
         rcopy = retry.copy()
         self.assertTrue(rcopy.sleep_func is _sleep)

tests/test_wale_restore.py

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