Merge pull request #94 from zalando/feature/initdb

Make it possible to specify custom options for initdb

README.rst

@@ -82,6 +82,12 @@ For an example file, see ``postgres0.yml``. Regarding settings:
     -  *data\_dir*: file path to initialize and store Postgres data files.
     -  *maximum\_lag\_on\_failover*: the maximum bytes a follower may lag.
     -  *use\_slots*: whether or not to use replication_slots. Must be False for PostgreSQL 9.3. You should comment out max_replication_slots before it becomes ineligible for leader status.
+
+    -  *initdb*:  List options to be passed on to initdb
+        -  *encoding*: default encoding for new databases
+        -  *locale*: default locale for new databases
+        -  *data-checksums*  # When pg_rewind is needed on 9.3, this needs to be enabled
+
     -  *pg\_hba*: list of lines which should be added to pg\_hba.conf.
         -  *- host all all 0.0.0.0/0 md5*.
 

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
 
@@ -49,6 +51,7 @@ class Postgresql:
         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', {})
@@ -165,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:
@@ -524,21 +557,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.replication['username']), self.replication['password'])
+        self.create_or_update_role(self.replication['username'], self.replication['password'], 'REPLICATION')
 
     def create_connection_users(self):
-        if self.superuser:
-            if 'username' in self.superuser:
-                self.query('CREATE ROLE "{0}" WITH LOGIN SUPERUSER PASSWORD %s'.format(
-                    self.superuser['username']), self.superuser['password'])
-            else:
-                rolsuper = self.query("""SELECT rolname FROM pg_authid WHERE rolsuper = 't'""").fetchone()[0]
-                self.query('ALTER ROLE "{0}" WITH PASSWORD %s'.format(rolsuper), self.superuser['password'])
+        if 'username' in self.superuser:
+            self.create_or_update_role(self.superuser['username'], self.superuser['password'], 'SUPERUSER')
         if self.admin:
-            self.query('CREATE ROLE "{0}" WITH LOGIN CREATEDB CREATEROLE PASSWORD %s'.format(
-                self.admin['username']), self.admin['password'])
+            self.create_or_update_role(self.admin['username'], self.admin['password'], 'CREATEDB CREATEROLE')
 
     def xlog_position(self):
         return self.query("""SELECT pg_xlog_location_diff(CASE WHEN pg_is_in_recovery()

postgres0.yml

@@ -35,6 +35,23 @@ postgresql:
   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

postgres1.yml

@@ -35,6 +35,23 @@ postgresql:
   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

tests/test_postgresql.py

@@ -162,7 +162,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 +187,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')))