Merge branch 'restore/movebasebackup' of https://github.com/pgexperts/patroni into pgexperts-restore/movebasebackup

README.rst

@@ -110,8 +110,16 @@ For an example file, see ``postgres0.yml``. Regarding settings:
         -  *username*: admin username; user is created during initialization. It will have CREATEDB and CREATEROLE privileges.
         -  *password*: admin password; user is created during initialization.
 
-    -  *recovery\_conf*: additional configuration settings written to recovery.conf when configuring the follower.
-        -  *parameters*: list of configuration settings for Postgres. Many of these are required for replication to work.
+    -  *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.
+
+    -  *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/postgresql.py

@@ -223,25 +223,79 @@ class Postgresql:
         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)
+        mconn = ""
+        for param, val in conn.items():
+            mconn = mconn + "{0}={1} ".format(param, val)
 
-    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
+        return mconn
+
+    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
+        replica_list = self.config.get('create_replica_method', 'basebackup')
+        replica_methods = [rm.strip() for rm in replica_list.split(',')]
+        # go through them in priority order
+        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:
+                    # if basebackup succeeds, exit with success
+                    break
+            else:
+                # user-defined method; check for configuration
+                # not required, actually
+                if replica_method in self.config:
+                    # look to see if the user has supplied a full command path
+                    # if not, use the method name as the command
+                    if "command" in self.config[replica_method]:
+                        cmd = self.config[replica_method]["command"]
+                    else:
+                        cmd = replica_method
+
+                    # get the rest of the replica config
+                    method_config = self.config[replica_method].copy()
+                    # remove the command and turn it into a shlex set
+                    del method_config["command"]
+                    # add the default parameters
+                    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()]
+                else:
+                    cmd = replica_method
+                    method_config = {"scope": self.scope,
+                                     "role": "replica",
+                                     "datadir": self.data_dir,
+                                     "connstring": connstring}
+
+                try:
+                    # call script with the full set of parameters
+                    ret = subprocess.call(shlex.split(cmd) + params, env=env)
+                    # if we succeeded, stop
+                    if ret == 0:
+                        break
+                except Exception as e:
+                    logger.exception('Error creating replica using method {0}: {1}'.format(replica_method, e.str))
+                    ret = 1
+
+        # write the recovery.conf
+        if ret == 0:
+            ret = self.write_recovery_conf(leader)
+            return 0
+
+        # out of methods, return 1
+        return 1
 
     def is_leader(self):
         return not self.query('SELECT pg_is_in_recovery()').fetchone()[0]
@@ -539,7 +593,7 @@ recovery_target_timeline = 'latest'
         """ restore a previously saved postgresql.conf """
         try:
             for f in self.configuration_to_save:
-                not os.path.isfile(f) and os.path.isfile(f+'.backup') and shutil.copy(f + '.backup', f)
+                not os.path.isfile(f) and os.path.isfile(f + '.backup') and shutil.copy(f + '.backup', f)
         except:
             logger.exception('unable to restore configuration files from backup')
 
@@ -663,3 +717,29 @@ $$""".format(name, options), name, password, password)
         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
+        bbfailures = 0
+        maxfailures = 2
+        ret = 1
+        while bbfailures < 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))
+
+            bbfailures += 1
+            if bbfailures < maxfailures:
+                logger.error('Trying again in 5 seconds')
+            time.sleep(5)
+
+        return ret

patroni/scripts/wale_restore.py

@@ -0,0 +1,178 @@
+#!/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"
+
+import logging
+import os
+import psycopg2
+import subprocess
+import sys
+import argparse
+
+
+if sys.hexversion >= 0x03000000:
+    long = int
+
+logger = logging.getLogger(__name__)
+
+
+class WALERestore(object):
+
+    def __init__(self, scope, role, datadir, connstring, env_dir, threshold_mb, threshold_pct, use_iam):
+        self.scope = scope
+        self.role = role
+        self.master_connection = connstring
+        self.data_dir = datadir
+        self.wal_e.dir = env_dir
+        self.wal_e.threshold_mb = threshold_mb
+        self.wal_e.threshold_pct = threshold_pct
+        if use_iam == 1:
+            self.wal_e.iam_string = ' --aws-instance-profile '
+        else:
+            self.wal_e.iam_string = ''
+        if not os.path.exists(self.wal_e.dir):
+            self.init_error = True
+        else:
+            self.init_error = False
+            self.wal_e.cmd = 'envdir {0} wal-e {1} '.\
+                format(self.wal_e.dir, self.wal_e.iam_string)
+
+    def run(self):
+        """ creates a new replica using WAL-E """
+        ret = self.replica_method()
+        return ret
+
+    def replica_method(self):
+        if self.should_use_s3_to_create_replica():
+            return self.create_replica_with_s3
+        return None
+
+    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_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 self.init_error:
+            return 1
+        # if we're set up, restore the replica using fetch latest
+        try:
+            ret = subprocess.call(self.wal_e.cmd + ' backup-fetch {} LATEST'.format(self.data_dir))
+        except Exception as e:
+            logger.error('Error when fetching backup with WAL-E: {0}'.format(e))
+            return 1
+
+        return ret
+
+
+if __name__ == '__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.constring,
+                              env_dir=args.env_dir, threshold_mb=args.threshold_megabytes,
+                              threshold_pct=args.threshold_backup_size_percentage)
+        ret = restore.run()
+        if ret == 0:
+            break
+
+    sys.exit(ret)

postgres0.yml

@@ -68,13 +68,18 @@ postgresql:
   admin:
     username: admin
     password: admin
-  wal_e:
-    env_dir: /home/postgres/etc/wal-e.d/env
-    threshold_megabytes: 10240
-    threshold_backup_size_percentage: 30
-  restore: patroni/scripts/restore.py
+  create_replica_method: basebackup
+# commented-out example for wal-e provisioning
+  #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"
   parameters:
     archive_mode: "on"
     wal_level: hot_standby

postgres1.yml

@@ -68,13 +68,17 @@ postgresql:
   admin:
     username: admin
     password: admin
+# commented-out example for wal-e provisioning
+  #create_replica_method: wal_e, basebackup
+  #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
-  wal_e:
-    env_dir: /home/postgres/etc/wal-e.d/env
-    threshold_megabytes: 10240
-    threshold_backup_size_percentage: 30
-  restore: patroni/scripts/restore.py
+    #restore_command: envdir /etc/wal-e.d/env wal-e wal-fetch "%f" "%p"
   parameters:
     archive_mode: "on"
     wal_level: hot_standby

tests/test_postgresql.py

@@ -273,9 +273,13 @@ 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)
 
     def test_create_connection_users(self):
         cfg = self.p.config