Custom bootstrap (#454)

Task of restoring a cluster from backup or cloning existing cluster into a new one was floating around for some time. It was kind of possible to achieve it by doing a lot of manual actions and very error prone. So I come up with the idea of making the way how we bootstrap a new cluster configurable.

In short - we want to run a custom script instead of running initdb.

features/backup_create.sh

@@ -0,0 +1,22 @@
+#!/bin/bash
+
+while getopts ":-:" optchar; do
+    [[ "${optchar}" == "-" ]] || continue
+    case "${OPTARG}" in
+        datadir=* )
+            PGDATA=${OPTARG#*=}
+            ;;
+        dbname=* )
+            DBNAME=${OPTARG#*=}
+            ;;
+        walmethod=* )
+            WALMETHOD=${OPTARG#*=}
+            ;;
+    esac
+done
+
+[[ -z $PGDATA || -z $DBNAME || -z $WALMETHOD ]] && exit 1
+
+[[ $WALMETHOD != "none" ]] && WALMETHOD="-X $WALMETHOD" || WALMETHOD=""
+
+exec pg_basebackup -D $PGDATA $WALMETHOD -c fast -d $DBNAME

features/backup_restore.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+
+set -x
+
+while getopts ":-:" optchar; do
+    [[ "${optchar}" == "-" ]] || continue
+    case "${OPTARG}" in
+        datadir=* )
+            PGDATA=${OPTARG#*=}
+            ;;
+        sourcedir=* )
+            SOURCE=${OPTARG#*=}
+            ;;
+    esac
+done
+
+[[ -z $PGDATA || -z $SOURCE ]] && exit 1
+
+mkdir -p $(dirname $PGDATA)
+
+exec cp -af $SOURCE $PGDATA

features/custom_bootstrap.feature

@@ -0,0 +1,17 @@
+Feature: custom bootstrap
+    We should check that patroni can bootstrap a new cluster from a backup
+
+Scenario: clone existing cluster using pg_basebackup
+    Given I start postgres0
+    Then postgres0 is a leader after 10 seconds
+    When I add the table foo to postgres0
+    And I start postgres1 in a cluster batman1 as a clone of postgres0
+    Then postgres1 is a leader of batman1 after 10 seconds
+    Then table foo is present on postgres1 after 10 seconds
+
+Scenario: make a backup and do a restore into a new cluster
+    Given I add the table bar to postgres1
+    And I do a backup of postgres1
+    When I start postgres2 in a cluster batman2 from backup
+    Then postgres2 is a leader of batman2 after 10 seconds
+    And table bar is present on postgres2 after 10 seconds

features/environment.py

@@ -87,19 +87,18 @@ class PatroniController(AbstractController):
     PATRONI_CONFIG = '{}.yml'
     """ starts and stops individual patronis"""
 
-    def __init__(self, context, name, work_directory, output_dir, tags=None, with_watchdog=False):
+    def __init__(self, context, name, work_directory, output_dir, custom_config=None):
         super(PatroniController, self).__init__(context, 'patroni_' + name, work_directory, output_dir)
         PatroniController.__PORT += 1
         self._data_dir = os.path.join(work_directory, 'data', name)
         self._connstring = None
-        if with_watchdog:
+        if custom_config and 'watchdog' in custom_config:
             self.watchdog = WatchdogMonitor(name, work_directory, output_dir)
-            custom_config = {'watchdog': {'driver': 'testing', 'device': self.watchdog.fifo_path, 'mode': 'required'}}
+            custom_config['watchdog'] = {'driver': 'testing', 'device': self.watchdog.fifo_path, 'mode': 'required'}
         else:
             self.watchdog = None
-            custom_config = None
 
-        self._config = self._make_patroni_test_config(name, tags, custom_config)
+        self._config = self._make_patroni_test_config(name, custom_config)
         self._closables = []
 
         self._conn = None
@@ -142,7 +141,7 @@ class PatroniController(AbstractController):
             cursor.execute("SET synchronous_commit TO 'local'")
             return True
 
-    def _make_patroni_test_config(self, name, tags, custom_config):
+    def _make_patroni_test_config(self, name, custom_config):
         patroni_config_name = self.PATRONI_CONFIG.format(name)
         patroni_config_path = os.path.join(self._output_dir, patroni_config_name)
 
@@ -154,12 +153,9 @@ class PatroniController(AbstractController):
 
         config['postgresql']['listen'] = config['postgresql']['connect_address'] = '{0}:{1}'.format(host, self.__PORT)
 
-        user = config['postgresql'].get('authentication', config['postgresql']).get('superuser', {})
-        self._connkwargs = {k: user[n] for n, k in [('username', 'user'), ('password', 'password')] if n in user}
-        self._connkwargs.update({'host': host, 'port': self.__PORT, 'database': 'postgres'})
-
         config['name'] = name
         config['postgresql']['data_dir'] = self._data_dir
+        config['postgresql']['use_unix_socket'] = True
         config['postgresql']['parameters'].update({
             'logging_collector': 'on', 'log_destination': 'csvlog', 'log_directory': self._output_dir,
             'log_filename': name + '.log', 'log_statement': 'all', 'log_min_messages': 'debug1',
@@ -170,9 +166,6 @@ class PatroniController(AbstractController):
             if 'initdb' in config['bootstrap']:
                 config['bootstrap']['initdb'].extend([{'auth': 'md5'}, {'auth-host': 'md5'}])
 
-        if tags:
-            config['tags'] = tags
-
         if custom_config is not None:
             def recursive_update(dst, src):
                 for k, v in src.items():
@@ -185,6 +178,13 @@ class PatroniController(AbstractController):
         with open(patroni_config_path, 'w') as f:
             yaml.safe_dump(config, f, default_flow_style=False)
 
+        user = config['postgresql'].get('authentication', config['postgresql']).get('superuser', {})
+        self._connkwargs = {k: user[n] for n, k in [('username', 'user'), ('password', 'password')] if n in user}
+        self._connkwargs.update({'host': host, 'port': self.__PORT, 'database': 'postgres'})
+
+        self._replication = config['postgresql'].get('authentication', config['postgresql']).get('replication', {})
+        self._replication.update({'host': host, 'port': self.__PORT, 'database': 'postgres'})
+
         return patroni_config_path
 
     def _connection(self):
@@ -276,6 +276,15 @@ class PatroniController(AbstractController):
             if 'process' not in p.cmdline()[0]:
                 p.terminate()
 
+    @property
+    def backup_source(self):
+        return 'postgres://{username}:{password}@{host}:{port}/{database}'.format(**self._replication)
+
+    def backup(self, dest='basebackup'):
+        subprocess.call([PatroniPoolController.BACKUP_SCRIPT, '--walmethod=none',
+                         '--datadir=' + os.path.join(self._output_dir, dest),
+                         '--dbname=' + self.backup_source])
+
 
 class ProcessHang(object):
 
@@ -304,7 +313,7 @@ class ProcessHang(object):
 
 class AbstractDcsController(AbstractController):
 
-    _CLUSTER_NODE = '/service/batman'
+    _CLUSTER_NODE = '/service/{0}'
 
     def __init__(self, context, mktemp=True):
         work_directory = mktemp and tempfile.mkdtemp() or None
@@ -319,11 +328,11 @@ class AbstractDcsController(AbstractController):
         if self._work_directory:
             shutil.rmtree(self._work_directory)
 
-    def path(self, key=None):
-        return self._CLUSTER_NODE + (key and '/' + key or '')
+    def path(self, key=None, scope='batman'):
+        return self._CLUSTER_NODE.format(scope) + (key and '/' + key or '')
 
     @abc.abstractmethod
-    def query(self, key):
+    def query(self, key, scope='batman'):
         """ query for a value of a given key """
 
     @abc.abstractmethod
@@ -352,18 +361,19 @@ class ConsulController(AbstractDcsController):
         super(ConsulController, self).__init__(context)
         os.environ['PATRONI_CONSUL_HOST'] = 'localhost:8500'
         self._client = consul.Consul()
+        self._config_file = None
 
     def _start(self):
-        config_file = self._work_directory + '.json'
-        with open(config_file, 'wb') as f:
+        self._config_file = self._work_directory + '.json'
+        with open(self._config_file, 'wb') as f:
             f.write(b'{"session_ttl_min":"5s","server":true,"bootstrap":true,"advertise_addr":"127.0.0.1"}')
-        return subprocess.Popen(['consul', 'agent', '-config-file', config_file, '-data-dir', self._work_directory],
-                                stdout=self._log, stderr=subprocess.STDOUT)
+        return subprocess.Popen(['consul', 'agent', '-config-file', self._config_file, '-data-dir',
+                                 self._work_directory], stdout=self._log, stderr=subprocess.STDOUT)
 
     def stop(self, kill=False, timeout=15):
         super(ConsulController, self).stop(kill=kill, timeout=timeout)
-        if self._work_directory:
-            os.unlink(self._work_directory + '.json')
+        if self._config_file:
+            os.unlink(self._config_file)
 
     def _is_running(self):
         try:
@@ -371,18 +381,18 @@ class ConsulController(AbstractDcsController):
         except Exception:
             return False
 
-    def path(self, key=None):
-        return super(ConsulController, self).path(key)[1:]
+    def path(self, key=None, scope='batman'):
+        return super(ConsulController, self).path(key, scope)[1:]
 
-    def query(self, key):
-        _, value = self._client.kv.get(self.path(key))
+    def query(self, key, scope='batman'):
+        _, value = self._client.kv.get(self.path(key, scope))
         return value and value['Value'].decode('utf-8')
 
     def set(self, key, value):
         self._client.kv.put(self.path(key), value)
 
     def cleanup_service_tree(self):
-        self._client.kv.delete(self.path(), recurse=True)
+        self._client.kv.delete(self.path(scope=''), recurse=True)
 
     def start(self, max_wait_limit=15):
         super(ConsulController, self).start(max_wait_limit)
@@ -401,9 +411,9 @@ class EtcdController(AbstractDcsController):
         return subprocess.Popen(["etcd", "--debug", "--data-dir", self._work_directory],
                                 stdout=self._log, stderr=subprocess.STDOUT)
 
-    def query(self, key):
+    def query(self, key, scope='batman'):
         try:
-            return self._client.get(self.path(key)).value
+            return self._client.get(self.path(key, scope)).value
         except etcd.EtcdKeyNotFound:
             return None
 
@@ -412,7 +422,7 @@ class EtcdController(AbstractDcsController):
 
     def cleanup_service_tree(self):
         try:
-            self._client.delete(self.path(), recursive=True)
+            self._client.delete(self.path(scope=''), recursive=True)
         except (etcd.EtcdKeyNotFound, etcd.EtcdConnectionFailed):
             return
         except Exception as e:
@@ -439,9 +449,9 @@ class ZooKeeperController(AbstractDcsController):
     def _start(self):
         pass  # TODO: implement later
 
-    def query(self, key):
+    def query(self, key, scope='batman'):
         try:
-            return self._client.get(self.path(key))[0].decode('utf-8')
+            return self._client.get(self.path(key, scope))[0].decode('utf-8')
         except kazoo.exceptions.NoNodeError:
             return None
 
@@ -450,7 +460,7 @@ class ZooKeeperController(AbstractDcsController):
 
     def cleanup_service_tree(self):
         try:
-            self._client.delete(self.path(), recursive=True)
+            self._client.delete(self.path(scope=''), recursive=True)
         except (kazoo.exceptions.NoNodeError):
             return
         except Exception as e:
@@ -475,6 +485,8 @@ class ExhibitorController(ZooKeeperController):
 
 class PatroniPoolController(object):
 
+    BACKUP_SCRIPT = 'features/backup_create.sh'
+
     def __init__(self, context):
         self._context = context
         self._dcs = None
@@ -499,16 +511,16 @@ class PatroniPoolController(object):
     def output_dir(self):
         return self._output_dir
 
-    def start(self, name, max_wait_limit=20, tags=None, with_watchdog=False):
+    def start(self, name, max_wait_limit=20, custom_config=None):
         if name not in self._processes:
             self._processes[name] = PatroniController(self._context, name, self.patroni_path,
-                                                      self._output_dir, tags, with_watchdog)
+                                                      self._output_dir, custom_config)
         self._processes[name].start(max_wait_limit)
 
     def __getattr__(self, func):
         if func not in ['stop', 'query', 'write_label', 'read_label', 'check_role_has_changed_to', 'add_tag_to_config',
                         'get_watchdog', 'database_is_running', 'checkpoint_hang', 'postmaster_hang',
-                        'terminate_backends']:
+                        'terminate_backends', 'backup']:
             raise AttributeError("PatroniPoolController instance has no attribute '{0}'".format(func))
 
         def wrapper(name, *args, **kwargs):
@@ -528,6 +540,53 @@ class PatroniPoolController(object):
         os.makedirs(feature_dir)
         self._output_dir = feature_dir
 
+    def clone(self, from_name, cluster_name, to_name):
+        f = self._processes[from_name]
+        custom_config = {
+            'scope': cluster_name,
+            'bootstrap': {
+                'method': 'pg_basebackup',
+                'pg_basebackup': {
+                    'command': self.BACKUP_SCRIPT + ' --walmethod=stream --dbname=' + f.backup_source
+                }
+            },
+            'postgresql': {
+                'parameters': {
+                    'archive_mode': 'on',
+                    'archive_command': 'mkdir -p {0} && test ! -f {0}/%f && cp %p {0}/%f'.format(
+                            os.path.join(self._output_dir, 'wal_archive'))
+                },
+                'authentication': {
+                    'superuser': {'password': 'zalando1'},
+                    'replication': {'password': 'rep-pass1'}
+                }
+            }
+        }
+        self.start(to_name, custom_config=custom_config)
+
+    def bootstrap_from_backup(self, name, cluster_name):
+        custom_config = {
+            'scope': cluster_name,
+            'bootstrap': {
+                'method': 'backup_restore',
+                'backup_restore': {
+                    'command': 'features/backup_restore.sh --sourcedir=' + os.path.join(self._output_dir, 'basebackup'),
+                    'recovery_conf': {
+                        'recovery_target_action': 'promote',
+                        'recovery_target_timeline': 'latest',
+                        'restore_command': 'cp {0}/wal_archive/%f %p'.format(self._output_dir)
+                    }
+                }
+            },
+            'postgresql': {
+                'authentication': {
+                    'superuser': {'password': 'zalando2'},
+                    'replication': {'password': 'rep-pass2'}
+                }
+            }
+        }
+        self.start(name, custom_config=custom_config)
+
     @property
     def dcs(self):
         if self._dcs is None:

features/steps/cascading_replication.py

@@ -6,7 +6,7 @@ from behave import step, then
 
 @step('I configure and start {name:w} with a tag {tag_name:w} {tag_value:w}')
 def start_patroni_with_a_name_value_tag(context, name, tag_name, tag_value):
-    return context.pctl.start(name, tags={tag_name: tag_value})
+    return context.pctl.start(name, custom_config={'tags': {tag_name: tag_value}})
 
 
 @then('There is a label with "{content:w}" in {name:w} data directory')

features/steps/custom_bootstrap.py

@@ -0,0 +1,27 @@
+import time
+
+from behave import step, then
+
+
+@step('I start {name:w} in a cluster {cluster_name:w} as a clone of {name2:w}')
+def start_cluster_clone(context, name, cluster_name, name2):
+    context.pctl.clone(name2, cluster_name, name)
+
+
+@step('I start {name:w} in a cluster {cluster_name:w} from backup')
+def start_cluster_from_backup(context, name, cluster_name):
+    context.pctl.bootstrap_from_backup(name, cluster_name)
+
+
+@then('{name:w} is a leader of {cluster_name:w} after {time_limit:d} seconds')
+def is_a_leader(context, name, cluster_name, time_limit):
+    time_limit *= context.timeout_multiplier
+    max_time = time.time() + int(time_limit)
+    while (context.dcs_ctl.query("leader", scope=cluster_name) != name):
+        time.sleep(1)
+        assert time.time() < max_time, "{0} is not a leader in dcs after {1} seconds".format(name, time_limit)
+
+
+@step('I do a backup of {name:w}')
+def do_backup(context, name):
+    context.pctl.backup(name)

features/steps/watchdog.py

@@ -15,7 +15,7 @@ def polling_loop(timeout, interval=1):
 
 @step('I start {name:w} with watchdog')
 def start_patroni_with_watchdog(context, name):
-    return context.pctl.start(name, with_watchdog=True)
+    return context.pctl.start(name, custom_config={'watchdog': True})
 
 
 @step('{name:w} watchdog has been pinged after {timeout:d} seconds')

patroni/__init__.py

@@ -116,7 +116,7 @@ class Patroni(object):
             if cluster and cluster.config and self.config.set_dynamic_configuration(cluster.config):
                 self.reload_config()
 
-            if not self.postgresql.data_directory_empty():
+            if self.postgresql.role != 'uninitialized':
                 self.config.save_cache()
 
             self.schedule_next_run()

patroni/ha.py

@@ -9,8 +9,8 @@ import time
 
 from collections import namedtuple
 from multiprocessing.pool import ThreadPool
-from patroni.async_executor import AsyncExecutor
-from patroni.exceptions import DCSError, PostgresConnectionException
+from patroni.async_executor import AsyncExecutor, CriticalTask
+from patroni.exceptions import DCSError, PostgresConnectionException, PatroniException
 from patroni.postgresql import ACTION_ON_START
 from patroni.utils import polling_loop, null_context, tzutc
 from threading import RLock, Event, Thread
@@ -102,6 +102,8 @@ class Ha(object):
         self.cluster = None
         self.old_cluster = None
         self.recovering = False
+        self._bootstrapping = False
+        self._post_bootstrap_task = None
         self._start_timeout = None
         self._async_executor = AsyncExecutor(self.wakeup)
         self.watchdog = patroni.watchdog
@@ -204,22 +206,11 @@ class Ha(object):
         # no initialize key and node is allowed to be master and has 'bootstrap' section in a configuration file
         elif self.cluster.initialize is None and not self.patroni.nofailover and 'bootstrap' in self.patroni.config:
             if self.dcs.initialize(create_new=True):  # race for initialization
-                with self._background_keepalive_context(wait_for_safepoint=False):
-                    try:
-                        self.state_handler.bootstrap(self.patroni.config['bootstrap'])
-                        self.dcs.initialize(create_new=False, sysid=self.state_handler.sysid)
-                    except:  # initdb or start failed
-                        # remove initialization key and give a chance to other members
-                        logger.info("removing initialize key after failed attempt to initialize the cluster")
-                        self.dcs.cancel_initialization()
-                        self.state_handler.stop('immediate')
-                        self.state_handler.move_data_directory()
-                        raise
-                    self.dcs.set_config_value(json.dumps(self.patroni.config.dynamic_configuration,
-                                                         separators=(',', ':')))
-                    self.dcs.take_leader()
-                    self.load_cluster_from_dcs()
-                    return 'initialized a new cluster'
+                self._bootstrapping = True
+                self._post_bootstrap_task = CriticalTask()
+                self._async_executor.schedule('bootstrap')
+                self._async_executor.run_async(self.state_handler.bootstrap, args=(self.patroni.config['bootstrap'],))
+                return 'trying to bootstrap a new cluster'
             else:
                 return 'failed to acquire initialize lock'
         else:
@@ -899,7 +890,7 @@ class Ha(object):
         try:
             if self.has_lock() and self.update_lock():
                 return 'updated leader lock during ' + self._async_executor.scheduled_action
-            else:
+            elif not self._bootstrapping:
                 # Don't have lock, make sure we are not starting up a master in the background
                 if self.state_handler.role == 'master':
                     logger.info("Demoting master during " + self._async_executor.scheduled_action)
@@ -937,6 +928,35 @@ class Ha(object):
             return 'failed to start postgres'
         return None
 
+    def cancel_initialization(self):
+        logger.info('removing initialize key after failed attempt to bootstrap the cluster')
+        self.dcs.cancel_initialization()
+        self.state_handler.stop('immediate')
+        self.state_handler.move_data_directory()
+        raise PatroniException('Failed to bootstrap cluster')
+
+    def post_bootstrap(self):
+        # bootstrap has failed if postgres is not running
+        if not self.state_handler.is_running() or self._post_bootstrap_task.result is False:
+            self.cancel_initialization()
+
+        self.keepalive()
+
+        if self._post_bootstrap_task.result is None:
+            if not self.state_handler.is_leader():
+                return 'waiting for end of recovery after bootstrap'
+
+            self._async_executor.schedule('post_bootstrap')
+            self._async_executor.run_async(self.state_handler.post_bootstrap,
+                                           args=(self.patroni.config['bootstrap'], self._post_bootstrap_task))
+            return 'running post_bootstrap'
+
+        self._bootstrapping = False
+        self.dcs.set_config_value(json.dumps(self.patroni.config.dynamic_configuration, separators=(',', ':')))
+        self.dcs.take_leader()
+        self.load_cluster_from_dcs()
+        return 'initialized a new cluster'
+
     def handle_starting_instance(self):
         """Starting up PostgreSQL may take a long time. In case we are the leader we may want to
         fail over to."""
@@ -1010,6 +1030,9 @@ class Ha(object):
                 return msg
 
             # we've got here, so any async action has finished.
+            if self._bootstrapping:
+                return self.post_bootstrap()
+
             if self.recovering and not self.state_handler.need_rewind:
                 self.recovering = False
                 # Check if we tried to recover and failed

patroni/postgresql.py

@@ -7,6 +7,7 @@ import re
 import shlex
 import shutil
 import signal
+import socket
 import subprocess
 import tempfile
 import time
@@ -15,7 +16,7 @@ from collections import defaultdict
 from contextlib import contextmanager
 from patroni import call_self
 from patroni.callback_executor import CallbackExecutor
-from patroni.exceptions import PostgresConnectionException, PostgresException
+from patroni.exceptions import PostgresConnectionException
 from patroni.utils import compare_values, parse_bool, parse_int, Retry, RetryFailedError, polling_loop, null_context
 from six import string_types
 from six.moves.urllib.parse import quote_plus
@@ -92,6 +93,8 @@ class Postgresql(object):
         'wal_log_hints': ('on', lambda _: False, 90400)
     }
 
+    _CONFIG_WARNING_HEADER = '# Do not edit this file manually!\n# It will be overwritten by Patroni!\n'
+
     def __init__(self, config):
         self.config = config
         self.name = config['name']
@@ -101,12 +104,12 @@ class Postgresql(object):
         self._data_dir = config['data_dir']
         self._config_dir = config.get('config_dir') or self._data_dir
         self._pending_restart = False
+        self._running_custom_bootstrap = False
         self.__thread_ident = current_thread().ident
 
         self._version_file = os.path.join(self._data_dir, 'PG_VERSION')
-        self._major_version = self.get_major_version()
         self._synchronous_standby_names = None
-        self._server_parameters = self.get_server_parameters(config)
+        self._configure_server_parameters()
 
         self._connect_address = config.get('connect_address')
         self._superuser = config['authentication'].get('superuser', {})
@@ -124,6 +127,7 @@ class Postgresql(object):
         self._postgresql_conf = os.path.join(self._config_dir, config_base_name + '.conf')
         self._postgresql_base_conf_name = config_base_name + '.base.conf'
         self._postgresql_base_conf = os.path.join(self._config_dir, self._postgresql_base_conf_name)
+        self._pg_hba_conf = os.path.join(self._config_dir, 'pg_hba.conf')
         self._recovery_conf = os.path.join(self._data_dir, 'recovery.conf')
         self._postmaster_pid = os.path.join(self._data_dir, 'postmaster.pid')
         self._trigger_file = config.get('recovery_conf', {}).get('trigger_file') or 'promote'
@@ -168,6 +172,8 @@ class Postgresql(object):
             configuration.append(os.path.basename(self._postgresql_base_conf))
         if not self.config['parameters'].get('hba_file'):
             configuration.append('pg_hba.conf')
+        if not self.config['parameters'].get('ident_file'):
+            configuration.append('pg_ident.conf')
         return configuration
 
     @property
@@ -494,7 +500,7 @@ class Postgresql(object):
                 raise Exception('Unknown type of initdb option: {0}'.format(o))
         return options
 
-    def _initialize(self, config):
+    def _initdb(self, config):
         self.set_state('initalizing new cluster')
         options = self.get_initdb_options(config.get('initdb') or [])
         pwfile = None
@@ -521,6 +527,25 @@ class Postgresql(object):
             self.set_state('initdb failed')
         return ret
 
+    def _custom_bootstrap(self, config):
+        params = ['--scope=' + self.scope, '--datadir=' + self._data_dir]
+        try:
+            logger.info('Running custom bootstrap script: %s', config['command'])
+            if subprocess.call(shlex.split(config['command']) + params) != 0:
+                self.set_state('custom bootstrap failed')
+                return False
+        except Exception:
+            logger.exception('Exception during custom bootstrap')
+            return False
+        self._post_restore()
+        self.save_configuration_files()
+
+        if 'recovery_conf' in config:
+            self.write_recovery_conf(config['recovery_conf'])
+        elif os.path.isfile(self._recovery_conf) or os.path.islink(self._recovery_conf):
+            os.unlink(self._recovery_conf)
+        return True
+
     def run_bootstrap_post_init(self, config):
         """
         runs a script after initdb or custom bootstrap script is called and waits until completion.
@@ -802,8 +827,8 @@ class Postgresql(object):
         self._pending_restart = False
 
         self._write_postgresql_conf()
-        self._replace_pg_hba()
         self.resolve_connection_addresses()
+        self._replace_pg_hba()
 
         options = ['--{0}={1}'.format(p, self._server_parameters[p]) for p in self.CMDLINE_OPTIONS
                    if p in self._server_parameters and p != 'wal_keep_segments']
@@ -1001,7 +1026,8 @@ class Postgresql(object):
         elif ready == STATE_NO_RESPONSE:
             self.set_state('start failed')
             self._schedule_load_slots = False  # TODO: can remove this?
-            self.save_configuration_files()  # TODO: maybe remove this?
+            if not self._running_custom_bootstrap:
+                self.save_configuration_files()  # TODO: maybe remove this?
             return True
         else:
             if ready != STATE_RUNNING:
@@ -1011,7 +1037,8 @@ class Postgresql(object):
                                "Unknown" if ready == STATE_UNKNOWN else "Invalid")
             self.set_state('running')
             self._schedule_load_slots = self.use_slots
-            self.save_configuration_files()
+            if not self._running_custom_bootstrap:
+                self.save_configuration_files()
             # TODO: __cb_pending can be None here after PostgreSQL restarts on its own. Do we want to call the callback?
             # Previously we didn't even notice.
             action = self.__cb_pending or ACTION_ON_START
@@ -1056,13 +1083,20 @@ class Postgresql(object):
             os.rename(self._postgresql_conf, self._postgresql_base_conf)
 
         with open(self._postgresql_conf, 'w') as f:
-            f.write('# Do not edit this file manually!\n# It will be overwritten by Patroni!\n')
+            f.write(self._CONFIG_WARNING_HEADER)
             f.write("include '{0}'\n\n".format(self.config.get('custom_conf') or self._postgresql_base_conf_name))
-            f.write("data_directory = '{}'\n".format(self._data_dir))
-            f.write("hba_file = '{}'\n".format(os.path.join(self._config_dir, 'pg_hba.conf')))
-            f.write("ident_file = '{}'\n".format(os.path.join(self._config_dir, 'pg_ident.conf')))
+            f.write("data_directory = '{0}'\n".format(self._data_dir))
             for name, value in sorted(self._server_parameters.items()):
-                f.write("{0} = '{1}'\n".format(name, value))
+                if not self._running_custom_bootstrap or name != 'hba_file':
+                    f.write("{0} = '{1}'\n".format(name, value))
+            # when we are doing custom bootstrap we assume that we don't know superuser password
+            # and in order to be able to change it, we are opening trust access from a certain address
+            # therefore we need to make sure that hba_file is not overriden
+            # after changing superuser password we will "revert" all these "changes"
+            if self._running_custom_bootstrap or 'hba_file' not in self._server_parameters:
+                f.write("hba_file = '{0}'\n".format(self._pg_hba_conf))
+            if 'ident_file' not in self._server_parameters:
+                f.write("ident_file = '{0}'\n".format(os.path.join(self._config_dir, 'pg_ident.conf')))
 
     def is_healthy(self):
         if not self.is_running():
@@ -1071,7 +1105,7 @@ class Postgresql(object):
         return True
 
     def write_pg_hba(self, config):
-        with open(os.path.join(self._config_dir, 'pg_hba.conf'), 'a') as f:
+        with open(self._pg_hba_conf, 'a') as f:
             f.write('\n{}\n'.format('\n'.join(config)))
 
     def _replace_pg_hba(self):
@@ -1081,9 +1115,24 @@ class Postgresql(object):
 
         :returns: True if pg_hba.conf was rewritten.
         """
-        if not self.config['parameters'].get('hba_file') and self.config.get('pg_hba'):
-            with open(os.path.join(self._config_dir, 'pg_hba.conf'), 'w') as f:
-                f.write('# Do not edit this file manually!\n# It will be overwritten by Patroni!\n')
+
+        # when we are doing custom bootstrap we assume that we don't know superuser password
+        # and in order to be able to change it, we are opening trust access from a certain address
+        if self._running_custom_bootstrap:
+            addresses = {'': 'local'}
+            if 'host' in self._local_address and not self._local_address['host'].startswith('/'):
+                for _, _, _, _, sa in socket.getaddrinfo(self._local_address['host'], self._local_address['port'],
+                                                         0, socket.SOCK_STREAM, socket.IPPROTO_TCP):
+                    addresses[sa[0] + '/32'] = 'host'
+
+            with open(self._pg_hba_conf, 'w') as f:
+                f.write(self._CONFIG_WARNING_HEADER)
+                for address, t in addresses.items():
+                    f.write('{0}\t{1}\t{2}\t{3}\ttrust\n'.format(t, self._database,
+                                                                 self._superuser.get('username') or 'all', address))
+        elif not self.config['parameters'].get('hba_file') and self.config.get('pg_hba'):
+            with open(self._pg_hba_conf, 'w') as f:
+                f.write(self._CONFIG_WARNING_HEADER)
                 for line in self.config['pg_hba']:
                     f.write('{0}\n'.format(line))
             return True
@@ -1109,16 +1158,10 @@ class Postgresql(object):
                     return primary_conninfo and (primary_conninfo in line)
         return not primary_conninfo
 
-    def write_recovery_conf(self, primary_conninfo):
+    def write_recovery_conf(self, recovery_params):
         with open(self._recovery_conf, 'w') as f:
-            f.write("standby_mode = 'on'\nrecovery_target_timeline = 'latest'\n")
-            if primary_conninfo:
-                f.write("primary_conninfo = '{0}'\n".format(primary_conninfo))
-                if self.use_slots:
-                    f.write("primary_slot_name = '{0}'\n".format(slot_name_from_member_name(self.name)))
-            for name, value in self.config.get('recovery_conf', {}).items():
-                if name not in ('standby_mode', 'recovery_target_timeline', 'primary_conninfo', 'primary_slot_name'):
-                    f.write("{0} = '{1}'\n".format(name, value))
+            for name, value in recovery_params.items():
+                f.write("{0} = '{1}'\n".format(name, value))
 
     def pg_rewind(self, r):
         # prepare pg_rewind connection
@@ -1303,7 +1346,15 @@ class Postgresql(object):
         primary_conninfo = self.primary_conninfo(member)
         change_role = self.role in ('master', 'demoted')
 
-        self.write_recovery_conf(primary_conninfo)
+        recovery_params = self.config.get('recovery_conf', {}).copy()
+        recovery_params.update({'standby_mode': 'on', 'recovery_target_timeline': 'latest'})
+        if primary_conninfo:
+            recovery_params['primary_conninfo'] = primary_conninfo
+        if self.use_slots:
+            recovery_params['primary_slot_name'] = slot_name_from_member_name(self.name)
+
+        self.write_recovery_conf(recovery_params)
+
         if self.is_running():
             self.restart()
         else:
@@ -1449,6 +1500,15 @@ $$""".format(name, ' '.join(options)), name, password, password)
     def last_operation(self):
         return str(self.wal_position())
 
+    def _post_restore(self):
+        self.delete_trigger_file()
+        self.restore_configuration_files()
+
+    def _configure_server_parameters(self):
+        self._major_version = self.get_major_version()
+        self._server_parameters = self.get_server_parameters(self.config)
+        return True
+
     def clone(self, clone_member):
         """
              - initialize the replica from an existing member (master or replica)
@@ -1459,21 +1519,50 @@ $$""".format(name, ' '.join(options)), name, password, password)
 
         ret = self.create_replica(clone_member) == 0
         if ret:
-            self._major_version = self.get_major_version()
-            self._server_parameters = self.get_server_parameters(self.config)
-            self.delete_trigger_file()
-            self.restore_configuration_files()
+            self._post_restore()
+            self._configure_server_parameters()
         return ret
 
     def bootstrap(self, config):
         """ Initialize a new node from scratch and start it. """
-        if self._initialize(config) and self.start() and self.run_bootstrap_post_init(config):
-            for name, value in (config.get('users') or {}).items():
-                if name not in (self._superuser.get('username'), self._replication['username']):
-                    self.create_or_update_role(name, value['password'], value.get('options', []))
-            self.create_or_update_role(self._replication['username'], self._replication['password'], ['REPLICATION'])
+        method = config.get('method') or 'initdb'
+        self._running_custom_bootstrap = method != 'initdb' and method in config and 'command' in config[method]
+        if self._running_custom_bootstrap:
+            do_initialize = self._custom_bootstrap
+            config = config[method]
         else:
-            raise PostgresException("Could not bootstrap master PostgreSQL")
+            do_initialize = self._initdb
+        return do_initialize(config) and self._configure_server_parameters() and self.start()
+
+    def post_bootstrap(self, config, task):
+        try:
+            self.create_or_update_role(self._superuser['username'], self._superuser['password'], ['SUPERUSER'])
+
+            # We were doing a custom bootstrap instead of running initdb, therefore we opened trust
+            # access from certain addresses to be able to reach cluster and change password
+            if self._running_custom_bootstrap:
+                self._running_custom_bootstrap = False
+                # If we don't have custom configuration for pg_hba.conf we need to restore original file
+                if not self.config.get('pg_hba'):
+                    os.unlink(self._pg_hba_conf)
+                    self.restore_configuration_files()
+                self._write_postgresql_conf()
+                self._replace_pg_hba()
+                self.reload()
+                time.sleep(1)  # give a time to postgres to "reload" configuration files
+                self.close_connection()  # close connection to reconnect with a new password
+
+            task.complete(self.run_bootstrap_post_init(config))
+            if task.result:
+                self.create_or_update_role(self._replication['username'],
+                                           self._replication['password'], ['REPLICATION'])
+                for name, value in (config.get('users') or {}).items():
+                    if name not in (self._superuser.get('username'), self._replication['username']):
+                        self.create_or_update_role(name, value['password'], value.get('options', []))
+        except Exception:
+            logger.exception('post_bootstrap')
+            task.complete(False)
+        return task.result
 
     def move_data_directory(self):
         if os.path.isdir(self._data_dir) and not self.is_running():
@@ -1517,19 +1606,13 @@ $$""".format(name, ' '.join(options)), name, password, password)
                 self.remove_data_directory()
 
             try:
-                version = 0
-                with psycopg2.connect(conn_url + '?replication=1') as c:
-                    version = c.server_version
-
                 ret = subprocess.call([self._pgcommand('pg_basebackup'), '--pgdata=' + self._data_dir,
-                                       '--{0}-method=stream'.format(self._wal_name(version)), '--dbname=' + conn_url])
+                                       '-X', 'stream', '--dbname=' + conn_url])
                 if ret == 0:
                     break
                 else:
                     logger.error('Error when fetching backup: pg_basebackup exited with code=%s', ret)
 
-            except psycopg2.Error:
-                logger.error('Can not connect to %s', conn_url)
             except Exception as e:
                 logger.error('Error when fetching backup with pg_basebackup: %s', e)
 

tests/test_ha.py

@@ -8,7 +8,7 @@ from mock import Mock, MagicMock, PropertyMock, patch
 from patroni.config import Config
 from patroni.dcs import Cluster, ClusterConfig, Failover, Leader, Member, get_dcs, SyncState
 from patroni.dcs.etcd import Client
-from patroni.exceptions import DCSError, PostgresException
+from patroni.exceptions import DCSError, PostgresConnectionException, PatroniException
 from patroni.ha import Ha, _MemberStatus, BackgroundKeepaliveSender
 from patroni.postgresql import Postgresql
 from patroni.watchdog import Watchdog
@@ -162,7 +162,7 @@ class TestHa(unittest.TestCase):
             self.ha.is_synchronous_mode = false
 
     def test_update_lock(self):
-        self.p.last_operation = Mock(side_effect=PostgresException(''))
+        self.p.last_operation = Mock(side_effect=PostgresConnectionException(''))
         self.assertTrue(self.ha.update_lock(True))
 
     def test_touch_member(self):
@@ -305,13 +305,19 @@ class TestHa(unittest.TestCase):
     def test_bootstrap_initialized_new_cluster(self):
         self.ha.cluster = get_cluster_not_initialized_without_leader()
         self.e.initialize = true
-        self.assertEquals(self.ha.bootstrap(), 'initialized a new cluster')
+        self.assertEquals(self.ha.bootstrap(), 'trying to bootstrap a new cluster')
+        self.p.is_leader = false
+        self.assertEquals(self.ha.run_cycle(), 'waiting for end of recovery after bootstrap')
+        self.p.is_leader = true
+        self.assertEquals(self.ha.run_cycle(), 'running post_bootstrap')
+        self.assertEquals(self.ha.run_cycle(), 'initialized a new cluster')
 
     def test_bootstrap_release_initialize_key_on_failure(self):
         self.ha.cluster = get_cluster_not_initialized_without_leader()
         self.e.initialize = true
-        self.p.bootstrap = Mock(side_effect=PostgresException("Could not bootstrap master PostgreSQL"))
-        self.assertRaises(PostgresException, self.ha.bootstrap)
+        self.ha.bootstrap()
+        self.p.is_running = false
+        self.assertRaises(PatroniException, self.ha.post_bootstrap)
 
     @patch('psycopg2.connect', psycopg2_connect)
     def test_reinitialize(self):

tests/test_patroni.py

@@ -105,6 +105,7 @@ class TestPatroni(unittest.TestCase):
     @patch('patroni.config.Config.reload_local_configuration', Mock(return_value=True))
     @patch.object(Postgresql, 'state', PropertyMock(return_value='running'))
     def test_run(self):
+        self.p.postgresql.set_role('replica')
         self.p.sighup_handler()
         self.p.ha.dcs.watch = Mock(side_effect=SleepException)
         self.p.api.start = Mock()

tests/test_postgresql.py

@@ -1,6 +1,7 @@
 import errno
 import mock  # for the mock.call method, importing it without a namespace breaks python3
 import os
+import psutil
 import psycopg2
 import shutil
 import subprocess
@@ -9,7 +10,7 @@ import unittest
 from mock import Mock, MagicMock, PropertyMock, patch, mock_open
 from patroni.async_executor import CriticalTask
 from patroni.dcs import Cluster, Leader, Member, SyncState
-from patroni.exceptions import PostgresException, PostgresConnectionException
+from patroni.exceptions import PostgresConnectionException
 from patroni.postgresql import Postgresql, STATE_REJECT, STATE_NO_RESPONSE
 from patroni.utils import RetryFailedError
 from six.moves import builtins
@@ -47,7 +48,7 @@ class MockCursor(object):
                             ('port', '5433', None, 'integer', 'postmaster'),
                             ('listen_addresses', '*', None, 'string', 'postmaster'),
                             ('autovacuum', 'on', None, 'bool', 'sighup'),
-                            ('unix_socket_directories', '.', None, 'string', 'postmaster')]
+                            ('unix_socket_directories', '/tmp', None, 'string', 'postmaster')]
         elif sql.startswith('IDENTIFY_SYSTEM'):
             self.results = [('1', 2, '0/402EEC0', '')]
         elif sql.startswith('TIMELINE_HISTORY '):
@@ -183,11 +184,8 @@ class TestPostgresql(unittest.TestCase):
                              'parameters': self._PARAMETERS,
                              'recovery_conf': {'foo': 'bar'},
                              'pg_hba': ['host all all 0.0.0.0/0 md5'],
-                             'callbacks': {'on_start': 'true', 'on_stop': 'true',
-                                           'on_restart': 'true', 'on_role_change': 'true',
-                                           'on_reload': 'true'
-                                           },
-                             'use_unix_socket': True})
+                             'callbacks': {'on_start': 'true', 'on_stop': 'true', 'on_reload': 'true',
+                                           'on_restart': 'true', 'on_role_change': 'true'}})
         self.p._callback_executor = Mock()
         self.leadermem = Member(0, 'leader', 28, {'conn_url': 'postgres://replicator:[email protected]:5435/postgres'})
         self.leader = Leader(-1, 28, self.leadermem)
@@ -281,7 +279,8 @@ class TestPostgresql(unittest.TestCase):
         mock_get_pid.return_value = -1
         self.assertFalse(self.p.stop())
         mock_get_pid.return_value = 123
-        with patch('os.kill', Mock(side_effect=[OSError(errno.ESRCH, ''), OSError, None])):
+        with patch('os.kill', Mock(side_effect=[OSError(errno.ESRCH, ''), OSError, None])),\
+             patch('psutil.Process', Mock(side_effect=psutil.NoSuchProcess(123))):
             self.assertTrue(self.p.stop())
             self.assertFalse(self.p.stop())
             self.p.stop_safepoint_reached.clear()
@@ -316,9 +315,9 @@ class TestPostgresql(unittest.TestCase):
         self.assertFalse(self.p.pg_rewind(r))
 
     def test_check_recovery_conf(self):
-        self.p.write_recovery_conf('foo')
+        self.p.write_recovery_conf({'primary_conninfo': 'foo'})
         self.assertFalse(self.p.check_recovery_conf(None))
-        self.p.write_recovery_conf(None)
+        self.p.write_recovery_conf({})
         self.assertTrue(self.p.check_recovery_conf(None))
 
     @patch.object(Postgresql, 'start', Mock())
@@ -512,10 +511,7 @@ class TestPostgresql(unittest.TestCase):
     @patch.object(Postgresql, 'is_running', Mock(return_value=True))
     def test_bootstrap(self):
         with patch('subprocess.call', Mock(return_value=1)):
-            self.assertRaises(PostgresException, self.p.bootstrap, {})
-
-        with patch.object(Postgresql, 'run_bootstrap_post_init', Mock(return_value=False)):
-            self.assertRaises(PostgresException, self.p.bootstrap, {})
+            self.assertFalse(self.p.bootstrap({}))
 
         config = {'users': {'replicator': {'password': 'rep-pass', 'options': ['replication']}}}
 
@@ -534,6 +530,49 @@ class TestPostgresql(unittest.TestCase):
             lines = f.readlines()
             self.assertTrue('host replication replicator 127.0.0.1/32 md5\n' in lines)
 
+    def test_custom_bootstrap(self):
+        config = {'method': 'foo', 'foo': {'command': 'bar'}}
+        with patch('subprocess.call', Mock(return_value=1)):
+            self.assertFalse(self.p.bootstrap(config))
+        with patch('subprocess.call', Mock(side_effect=Exception)):
+            self.assertFalse(self.p.bootstrap(config))
+        with patch('subprocess.call', Mock(return_value=0)),\
+                patch('subprocess.Popen', Mock(side_effect=Exception("42"))),\
+                patch('os.path.isfile', Mock(return_value=True)),\
+                patch('os.unlink', Mock()),\
+                patch.object(Postgresql, 'save_configuration_files', Mock()),\
+                patch.object(Postgresql, 'restore_configuration_files', Mock()),\
+             patch.object(Postgresql, 'write_recovery_conf', Mock()):
+            with self.assertRaises(Exception) as e:
+                self.p.bootstrap(config)
+            self.assertEqual(str(e.exception), '42')
+
+            config['foo']['recovery_conf'] = {'foo': 'bar'}
+
+            with self.assertRaises(Exception) as e:
+                self.p.bootstrap(config)
+            self.assertEqual(str(e.exception), '42')
+
+    @patch('time.sleep', Mock())
+    @patch.object(Postgresql, 'run_bootstrap_post_init', Mock(side_effect=Exception))
+    def test_post_bootstrap(self):
+        config = {'method': 'foo', 'foo': {'command': 'bar'}}
+        with patch('subprocess.call', Mock(return_value=0)), \
+                patch('subprocess.Popen', Mock(side_effect=Exception("42"))), \
+                patch('os.path.isfile', Mock(return_value=True)),\
+                patch('os.unlink', Mock()), \
+                patch.object(Postgresql, 'save_configuration_files', Mock()), \
+                patch.object(Postgresql, 'restore_configuration_files', Mock()), \
+             patch.object(Postgresql, 'write_recovery_conf', Mock()):
+            with self.assertRaises(Exception) as e:
+                self.p.bootstrap(config)
+            self.assertEqual(str(e.exception), '42')
+
+        self.p.config.pop('pg_hba')
+        task = CriticalTask()
+        self.p.post_bootstrap({}, task)
+        self.assertFalse(task.result)
+
     def test_run_bootstrap_post_init(self):
         with patch('subprocess.call', Mock(return_value=1)):
             self.assertFalse(self.p.run_bootstrap_post_init({'post_init': '/bin/false'}))
@@ -547,7 +586,7 @@ class TestPostgresql(unittest.TestCase):
             mock_method.assert_called()
             args, kwargs = mock_method.call_args
             self.assertTrue('PGPASSFILE' in kwargs['env'])
-            self.assertEquals(args[0], ['/bin/false', 'postgres://%2Ftmp:5432/postgres'])
+            self.assertEquals(args[0], ['/bin/false', 'postgres://127.0.0.2:5432/postgres'])
 
             mock_method.reset_mock()
             self.p._local_address.pop('host')
@@ -629,8 +668,10 @@ class TestPostgresql(unittest.TestCase):
         parameters['autovacuum'] = 'off'
         parameters.pop('search_path')
         config['listen'] = '*:5433'
+        self.p.reload_config(config)
         parameters['unix_socket_directories'] = '.'
         self.p.reload_config(config)
+        self.p.resolve_connection_addresses()
 
     @patch.object(Postgresql, '_version_file_exists', Mock(return_value=True))
     def test_get_major_version(self):