Add post_init configuration parameter on bootstrap (#296)

* Add bootstrap post_init configuration parameter
* Add documentation

By @zenitraM

docs/SETTINGS.rst

@@ -32,6 +32,7 @@ Bootstrap configuration
         -  **options**: list of options for CREATE USER statement
             -  **- createrole**
             -  **- createdb**
+-  **post_init**: An additional script that will be executed after initializing the cluster. The script receives a connection string URL (with the cluster superuser as a user name). The PGPASSFILE variable is set to the location of pgpass file.
 
 Consul
 ------
@@ -67,7 +68,7 @@ PostgreSQL
 -  **data\_dir**: The location of the Postgres data directory, either existing or to be initialized by Patroni.
 -  **bin\_dir**: Path to PostgreSQL binaries. (pg_ctl, pg_rewind, pg_basebackup, postgres) The  default value is an empty string meaning that PATH environment variable will be used to find the executables.
 -  **listen**: IP address + port that Postgres listens to; must be accessible from other nodes in the cluster, if you're using streaming replication. Multiple comma-separated addresses are permitted, as long as the port component is appended after to the last one with a colon, i.e. ``listen: 127.0.0.1,127.0.0.2:5432``. Patroni will use the first address from this list to establish local connections to the PostgreSQL node.
--  **pgpass**: path to the `.pgpass <https://www.postgresql.org/docs/current/static/libpq-pgpass.html>`__ password file. Patroni creates this file before executing pg\_basebackup and under some other circumstances. The location must be writable by Patroni.
+-  **pgpass**: path to the `.pgpass <https://www.postgresql.org/docs/current/static/libpq-pgpass.html>`__ password file. Patroni creates this file before executing pg\_basebackup, the post_init script and under some other circumstances. The location must be writable by Patroni.
 -  **recovery\_conf**: additional configuration settings written to recovery.conf when configuring follower.
 -  **custom_conf** : path to an optional custom ``postgresql.conf`` file, that will be used in place of ``postgresql.base.conf``. The file must exist on all cluster nodes, be readable by PostgreSQL and will be included from its location on the real ``postgresql.conf``. Note that Patroni will not monitor this file for changes, nor backup it. However, its settings can still be overriden by Patroni's own configuration facilities - see `dynamic configuration <https://github.com/zalando/patroni/blob/master/docs/dynamic_configuration.rst>`__ for details.
 -  **parameters**: list of configuration settings for Postgres. Many of these are required for replication to work.

patroni/postgresql.py

@@ -392,6 +392,30 @@ class Postgresql(object):
             self.set_state('initdb failed')
         return ret
 
+    def run_bootstrap_post_init(self, config):
+        """
+        runs a script after initdb is called and waits until completion.
+        passed: cluster name, parameters
+        """
+        if 'post_init' in config:
+            cmd = config['post_init']
+            r = self._local_connect_kwargs
+            if 'user' in r:
+                connstring = 'postgres://{user}@{host}:{port}/{database}'.format(**r)
+            else:
+                connstring = 'postgres://{host}:{port}/{database}'.format(**r)
+
+            env = self.write_pgpass(r) if 'password' in r else None
+            try:
+                ret = subprocess.call(shlex.split(cmd) + [connstring], env=env)
+            except OSError:
+                logger.error('post_init script %s failed', cmd)
+                return False
+            if ret != 0:
+                logger.error('post_init script %s returned non-zero code %d', cmd, ret)
+                return False
+        return True
+
     def delete_trigger_file(self):
         if os.path.exists(self._trigger_file):
             os.unlink(self._trigger_file)
@@ -981,7 +1005,7 @@ $$""".format(name, ' '.join(options)), name, password, password)
 
     def bootstrap(self, config):
         """ Initialize a new node from scratch and start it. """
-        if self._initialize(config) and self.start():
+        if self._initialize(config) and self.start() and self.run_bootstrap_post_init(config):
             for name, value in config['users'].items():
                 if name not in (self._superuser.get('username'), self._replication['username']):
                     self.create_or_update_role(name, value['password'], value.get('options', []))

postgres0.yml

@@ -48,6 +48,9 @@ bootstrap:
   - host all all 0.0.0.0/0 md5
 #  - hostssl all all 0.0.0.0/0 md5
 
+  # Additional script to be launched after initial cluster creation (will be passed the connection URL as parameter)
+# post_init: /usr/local/bin/setup_cluster.sh
+
   # Some additional users users which needs to be created after initializing new cluster
   users:
     admin:

postgres1.yml

@@ -48,6 +48,9 @@ bootstrap:
   - host all all 0.0.0.0/0 md5
 #  - hostssl all all 0.0.0.0/0 md5
 
+  # Additional script to be launched after initial cluster creation (will be passed the connection URL as parameter)
+# post_init: /usr/local/bin/setup_cluster.sh
+
   # Some additional users users which needs to be created after initializing new cluster
   users:
     admin:

tests/test_postgresql.py

@@ -398,15 +398,34 @@ class TestPostgresql(unittest.TestCase):
         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.p.bootstrap({'users': {'replicator': {'password': 'rep-pass', 'options': ['replication']}},
                           'pg_hba': ['host replication replicator 127.0.0.1/32 md5',
                                      'hostssl all all 0.0.0.0/0 md5',
-                                     'host all all 0.0.0.0/0 md5']})
+                                     'host all all 0.0.0.0/0 md5'],
+                          'post_init': '/bin/false'})
         with open(os.path.join(self.data_dir, 'pg_hba.conf')) as f:
             lines = f.readlines()
             assert 'host replication replicator 127.0.0.1/32 md5\n' in lines
             assert 'host all all 0.0.0.0/0 md5\n' in lines
 
+    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'}))
+
+        with patch('subprocess.call', Mock(side_effect=OSError)):
+            self.assertFalse(self.p.run_bootstrap_post_init({'post_init': '/bin/false'}))
+
+        with patch('subprocess.call', Mock(return_value=0)) as mock_method:
+            self.assertTrue(self.p.run_bootstrap_post_init({'post_init': '/bin/false'}))
+
+        mock_method.assert_called()
+        args, kwargs = mock_method.call_args
+        assert 'PGPASSFILE' in kwargs['env'].keys()
+        self.assertEquals(args[0], ['/bin/false', 'postgres://test@localhost:5432/postgres'])
+
     @patch('patroni.postgresql.Postgresql.create_replica', Mock(return_value=0))
     def test_clone(self):
         self.p.clone(self.leader)