Merge pull request #8 from zalando/feature/external_backup

Move the functionality of restoring a replica from backup into an external script

helpers/postgresql.py

@@ -59,10 +59,6 @@ class Postgresql:
         self.is_promoted = False
 
         self._pg_ctl = ['pg_ctl', '-w', '-D', self.data_dir]
-        self.wal_e = config.get('wal_e', None)
-        if self.wal_e:
-            self.wal_e_path = 'envdir {} wal-e --aws-instance-profile '.\
-                format(self.wal_e.get('env_dir', '/home/postgres/etc/wal-e.d/env'))
 
         self.local_address = self.get_local_address()
         connect_address = config.get('connect_address', None) or self.local_address
@@ -143,99 +139,21 @@ class Postgresql:
         env['PGPASSFILE'] = pgpass
         return self.create_replica(r, env) == 0
 
+    @staticmethod
+    def build_connstring(conn):
+        return "host={host} port={port} user={user}".format(**conn)
+
     def create_replica(self, master_connection, env):
-        """ creates a new replica using either pg_basebackup or WAL-E """
-        if self.should_use_s3_to_create_replica(master_connection):
-            result = self.create_replica_with_s3()
-            # if restore from the backup on S3 failed - try with the pg_basebackup
-            if result == 0:
-                return result
-        return self.create_replica_with_pg_basebackup(master_connection, env)
-
-    def create_replica_with_pg_basebackup(self, master_connection, env):
-        ret = subprocess.call(['pg_basebackup', '-R', '-D', self.data_dir, '--host=' + master_connection['host'],
-                               '--port=' + str(master_connection['port']), '-U', master_connection['user']], env=env)
-        self.delete_trigger_file()
-        return ret
-
-    def create_replica_with_s3(self):
-        if not self.wal_e or not self.wal_e_path:
+        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')
             return 1
-
-        ret = subprocess.call(self.wal_e_path + ' backup-fetch {} LATEST'.format(self.data_dir), shell=True)
-        self.restore_configuration_files()
         return ret
 
-    def should_use_s3_to_create_replica(self, master_connection):
-        """ determine whether it makes sense to use S3 and not pg_basebackup """
-        if not self.wal_e or not self.wal_e_path:
-            return False
-
-        threshold_megabytes = self.wal_e.get('threshold_megabytes', 10240)
-        threshold_backup_size_percentage = self.wal_e.get('threshold_backup_size_percentage', 30)
-
-        try:
-            latest_backup = subprocess.check_output(self.wal_e_path.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 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(**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 is_leader(self, check_only=False):
         ret = not self.query('SELECT pg_is_in_recovery()').fetchone()[0]
         if ret and self.is_promoted and not check_only:
@@ -257,12 +175,11 @@ class Postgresql:
             except psycopg2.OperationalError as e:
                 logger.warning("unable to perform {0} action, cannot obtain the cluster role: {1}".format(cb_name, e))
                 return False
-        scope = self.scope
         try:
             role = "master" if is_leader else "replica"
-            subprocess.Popen(shlex.split(os.path.abspath(cmd))+[cb_name, role, scope])
-        except Exception as e:
-            logger.warning("callback {0} {1} {2} {3} failed: {4}".format(os.path.abspath(cmd), cb_name, role, scope, e))
+            subprocess.Popen(shlex.split(cmd) + [cb_name, role, self.scope])
+        except:
+            logger.exception('callback %s %s %s %s failed', cmd, cb_name, role, self.scope)
             return False
         return True
 
@@ -415,8 +332,8 @@ primary_conninfo = '{}'
         try:
             for f in self.configuration_to_save:
                 shutil.copy(f + '.backup', f)
-        except Exception as e:
-            logger.error("unable to restore configuration from WAL-E backup: {}".format(e))
+        except:
+            logger.exception('unable to restore configuration from WAL-E backup')
 
     def promote(self):
         self.is_promoted = subprocess.call(self._pg_ctl + ['promote']) == 0

postgres0.yml

@@ -46,6 +46,7 @@ postgresql:
     env_dir: /home/postgres/etc/wal-e.d/env
     threshold_megabytes: 10240
     threshold_backup_size_percentage: 30
+  restore: "true"
   #recovery_conf:
     #restore_command: cp ../wal_archive/%f %p
   parameters:

requirements-py2.txt

@@ -1,7 +1,8 @@
 boto
 dnspython
+mock
 psycopg2
 PyYAML
 requests
-six
+six >= 1.7
 kazoo>=2.2.1

requirements-py3.txt

@@ -1,4 +1,5 @@
 boto
+mock
 dnspython3
 psycopg2
 PyYAML

scripts/restore.py

@@ -0,0 +1,216 @@
+#!/usr/bin/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]))

tests/test_postgresql.py

@@ -118,10 +118,11 @@ class TestPostgresql(unittest.TestCase):
                                              'password': 'rep-pass',
                                              'network': '127.0.0.1/32'},
                              'parameters': {'foo': 'bar'}, 'recovery_conf': {'foo': 'bar'},
-                             'callbacks': {'on_start': '/usr/bin/true', 'on_stop': '/usr/bin/true',
-                                           'on_restart': '/usr/bin/true', 'on_role_change': '/bin/true',
-                                           'on_reload': '/usr/bin/true'
-                                           }})
+                             'callbacks': {'on_start': 'true', 'on_stop': 'true',
+                                           'on_restart': 'true', 'on_role_change': 'true',
+                                           'on_reload': 'true'
+                                           },
+                             'restore': '/usr/bin/true'})
         psycopg2.connect = psycopg2_connect
         if not os.path.exists(self.p.data_dir):
             os.makedirs(self.p.data_dir)

tests/test_restore.py

@@ -0,0 +1,111 @@
+import unittest
+from mock import MagicMock, patch
+import os
+from 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)