gitops/patroni
1
0
Fork 0
mirror of https://github.com/outbackdingo/patroni.git synced 2026-08-25 14:53:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #103 from zalando/pgexperts-restore/movebasebackup

Browse Source 
Replica creation methods are configured via the 'create_replica_method'
parameter in the patroni configuration file. The value should contain a list
of acceptable replica creation method names. Patroni attempts to call them one
by one, until it finds one that returns 0. If there is a configuration section
with the same name as the replica creation method, the values in this section
are converted to the command-line options and appended to the command name.
The command name is constructued either from the 'command' key in the method
name configuration section, or by taking the replica creation method name
verbatim.

The wal_e replica creation method is included with Patroni (wale_restore.py
script), in addition to the 'basebackup', which is a special method: it's taken
by default, if no create_replica_method parameter is specified, or if it is
empty. This merge also reverses the decision to move the basebackup repli
method ouside of the core patroni.

Original pull request by Josh Berkus (@jberkus), with fixes and improvements
from Alexander Kukushkin (CyberDem0n) and Oleksii Kliukin (alexeyklyukin).

Per discussion at https://github.com/zalando/patroni/issues/57
This commit is contained in:
Oleksii Kliukin
2015-11-25 15:04:23 +01:00
committed by Oleksii Kliukin
parent ec1d1c827b d4ab4d1aef
commit 28f1d51cb6
9 changed files with 384 additions and 357 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.rst
+10 -2
View File
@@ -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
+78 -15
View File
@@ -223,24 +223,63 @@ 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)
"""
>>> Postgresql.build_connstring({'host': '127.0.0.1', 'port': '5432'}) == 'host=127.0.0.1 port=5432'
True
"""
return ' '.join('{}={}'.format(param, val) for param, val in sorted(conn.items()))
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.
# If there is no configuration key, or no value is specified, use basebackup
replica_methods = self.config.get('create_replica_method') or ['basebackup']
# go through them in priority order
ret = 1
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:
logger.info("replica has been created using basebackup")
# if basebackup succeeds, exit with success
break
else:
cmd = replica_method
method_config = {}
# user-defined method; check for configuration
# not required, actually
if replica_method in self.config:
method_config = self.config[replica_method].copy()
# look to see if the user has supplied a full command path
# if not, use the method name as the command
cmd = method_config.pop('command', cmd)
# add the default parameters
try:
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()]
# call script with the full set of parameters
ret = subprocess.call(shlex.split(cmd) + params, env=env)
# if we succeeded, stop
if ret == 0:
logger.info("replica has been created using {0}".format(replica_method))
break
except Exception as e:
logger.exception('Error creating replica using method {0}: {1}'.format(replica_method, str(e)))
ret = 1
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
def is_leader(self):
@@ -539,7 +578,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 +702,27 @@ $$""".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
maxfailures = 2
ret = 1
for bbfailures in range(0, 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))
if bbfailures < maxfailures - 1:
logger.error('Trying again in 5 seconds')
time.sleep(5)
return ret
patroni/scripts/restore.py
-216
View File
@@ -1,216 +0,0 @@
#!/usr/bin/env 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]))
patroni/scripts/wale_restore.py
Executable
+164
View File
@@ -0,0 +1,164 @@
#!/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" -p 1
from collections import namedtuple
import logging
import os
import psycopg2
import subprocess
import sys
import argparse
if sys.hexversion >= 0x03000000:
long = int
logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO)
logger = logging.getLogger(__name__)
class WALERestore(object):
def __init__(self, scope, datadir, connstring, env_dir, threshold_mb, threshold_pct, use_iam):
self.scope = scope
self.master_connection = connstring
self.data_dir = datadir
self.wal_e = namedtuple('wale', 'dir,threshold_mb,threshold_pct,iam_string,cmd')
self.wal_e.dir = env_dir
self.wal_e.threshold_mb = threshold_mb
self.wal_e.threshold_pct = threshold_pct
self.wal_e.iam_string = ' --aws-instance-profile ' if use_iam == 1 else ''
self.wal_e.cmd = 'envdir {0} wal-e {1} '.format(self.wal_e.dir, self.wal_e.iam_string)
self.init_error = (not os.path.exists(self.wal_e.dir))
def run(self):
""" creates a new replica using WAL-E """
if not self.init_error and self.should_use_s3_to_create_replica():
return self.create_replica_with_s3()
return 2
def should_use_s3_to_create_replica(self):
""" determine whether it makes sense to use S3 and not pg_basebackup """
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 we're set up, restore the replica using fetch latest
try:
ret = subprocess.call(self.wal_e.cmd.split() + ['backup-fetch', '{}'.format(self.data_dir), 'LATEST'])
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.connstring,
env_dir=args.envdir, threshold_mb=args.threshold_megabytes,
threshold_pct=args.threshold_backup_size_percentage, use_iam=args.use_iam)
ret = restore.run()
if ret == 0:
break
sys.exit(ret)
postgres0.yml
+13 -6
View File
@@ -68,13 +68,20 @@ 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
# - wal_e
# 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" -p 1
parameters:
archive_mode: "on"
wal_level: hot_standby
postgres1.yml
+13 -6
View File
@@ -68,13 +68,20 @@ postgresql:
admin:
username: admin
password: admin
# commented-out example for wal-e provisioning
create_replica_method:
- basebackup
# - wal_e
# commented-out example for wal-e provisioning
#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" -p 1
parameters:
archive_mode: "on"
wal_level: hot_standby
tests/test_postgresql.py
+15 -1
View File
@@ -276,9 +276,23 @@ 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)
self.p.config['create_replica_method'] = ['wale', 'basebackup']
self.p.config['wale'] = {'command': 'foo'}
with patch('subprocess.call', Mock(return_value=0)):
self.assertEquals(self.p.create_replica(self.leader, ''), 0)
del self.p.config['wale']
self.assertEquals(self.p.create_replica(self.leader, ''), 0)
with patch('subprocess.call', Mock(side_effect=Exception("foo"))):
self.assertEquals(self.p.create_replica(self.leader, ''), 1)
def test_create_connection_users(self):
cfg = self.p.config
tests/test_restore.py
-111
View File
@@ -1,111 +0,0 @@
import unittest
from mock import MagicMock, patch
import os
from patroni.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)
tests/test_wale_restore.py
+91
View File
@@ -0,0 +1,91 @@
import unittest
from mock import MagicMock, patch, PropertyMock
import os
import psycopg2
import subprocess
from patroni.scripts.wale_restore import 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
"""
def fake_backup_data_2(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 """
def fake_backup_data_3(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
base_00000001000000000000007F_00000040 2015-05-18T10:13:25.000Z 167772160 00000001000000000000007F 00000040 00000001000000000000007F 00000240
"""
def fake_backup_data_4(self, *args, **kwargs):
""" return the fake result of WAL-E backup-list"""
return """name last_modified expanded_size_foo 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
"""
@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):
self.wale_restore = WALERestore("batman", "/data",
"host=batman port=5432 user=batman", "/etc", 100, 100, 1)
def tearDown(self):
pass
def test_should_use_s3_to_create_replica(self):
with patch('psycopg2.connect', MagicMock(side_effect=psycopg2.Error("foo"))):
self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
with patch('subprocess.check_output', MagicMock(side_effect=subprocess.CalledProcessError(1, "cmd", "foo"))):
self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
with patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data_2)):
self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
with patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data_3)):
self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
with patch('subprocess.check_output', MagicMock(side_effect=fake_backup_data_4)):
self.assertFalse(self.wale_restore.should_use_s3_to_create_replica())
self.wale_restore.should_use_s3_to_create_replica()
def test_create_replica_with_s3(self):
with patch('subprocess.call', MagicMock(return_value=0)):
self.assertEqual(self.wale_restore.create_replica_with_s3(), 0)
with patch('subprocess.call', MagicMock(side_effect=Exception("foo"))):
self.assertEqual(self.wale_restore.create_replica_with_s3(), 1)
def test_run(self):
with patch.object(self.wale_restore, 'init_error', PropertyMock(return_value=True)):
self.assertEqual(self.wale_restore.run(), 2)
with patch.object(self.wale_restore, 'should_use_s3_to_create_replica', MagicMock(return_value=True)):
with patch.object(self.wale_restore, 'create_replica_with_s3', MagicMock(return_value=0)):
self.assertEqual(self.wale_restore.run(), 0)