-
Notifications
You must be signed in to change notification settings - Fork 1
/
backup.py
110 lines (109 loc) · 4.06 KB
/
backup.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import os
import sys
import logging
import datetime
import json
from utils import get_consul_session
from utils import get_records_from_consul
from utils import save_records
from utils import get_s3_resource
from utils import s3_upload_file
from utils import s3_download_file
# Configure logging
LOG_FILENAME = '/tmp/consul_backup.log'
logging.basicConfig(filename=LOG_FILENAME,
level=logging.INFO,
)
# Check if the name of the backup folder on S3 was specified,
# otherwise exit
backup_folder = os.environ.get('S3_BACKUP_FOLDER')
# and let's make sure that it has a value
if not backup_folder:
logging.error("No S3 backup folder was specified! Exiting...")
sys.exit(2)
# Check whether necessary environment variables have been set to a non-empty value
# If not, fall back to default values
consul_host = os.environ.get('CONSUL_HOST')
if not consul_host:
consul_host = 'localhost'
consul_port = os.environ.get('CONSUL_PORT')
if not consul_port:
consul_port = '8500'
# Directory to locally store backup file (temporarily)
tmp_dir = os.environ.get('TMP_DIR','/tmp')
# If no action was specified, the default action is backup
action = os.environ.get('ACTION')
if not action:
action = 'backup'
# If we need to restore, we need to be able to identify the backup file to be restored
# and if this has not been set, we need to exit
if action == 'restore':
restore_id = os.environ.get('RESTORE_ID')
if not restore_id:
logging.error('Restore was requested, but no RESTORE_ID was set!')
sys.exit(2)
# See if we are in "overwrite" mode
overwrite = os.environ.get('OVERWRITE',False)
# Register the Consul client with the server
logging.info("Registering consulate with %s on port %s"%(consul_host, consul_port))
try:
session = get_consul_session(consul_host, consul_port)
except Exception,e:
logging.error("Failed to register consulate: %s"%e)
sys.exit(2)
# Request status info to check that we established a connection
try:
leader = session.status.leader()
except Exception, e:
logging.error("Failed to register consulate: %s"%e)
sys.exit(2)
# Now take the appropriate action depending on the value of the ACTION variable
if action == 'backup':
# Get today's date for the time stamp
now = str(datetime.datetime.now().date())
# Construct the name of the backup file
fname = os.environ.get('BACKUP_FILE','adsabs_consul_kv')
version = os.environ.get('VERSION',now)
backup_file = '%s/%s.%s.json' % (tmp_dir,fname,version)
# Get the records from the Consul key/value store
try:
records = get_records_from_consul(session)
except Exception, e:
logging.error('Unable to retrieve records from Consul store: %s'%e)
sys.exit(2)
# Write the records to the backup file
try:
save_records(records, backup_file)
except Exception, e:
logging.error('Unable to write to backup file: %s (%s)'%(backup_file,e))
sys.exit(2)
logging.info('Backup was written to: %s'%backup_file)
# Now copy the backup to S3
s3 = get_s3_resource()
try:
s3_upload_file(s3, backup_file, backup_folder)
except Exception, e:
logging.error('Unable to move backup to S3: %s'%e)
# Finally, remove the local copy
logging.info('Removing local copy of backup file: %s' % backup_file)
os.remove(backup_file)
elif action == 'restore':
# Construct the name of the backup file to retrieve
fname = os.environ.get('BACKUP_FILE','adsabs_consul_kv')
backup_file = '%s/%s.%s.json' % (tmp_dir,fname,restore_id)
# Get the file from S3
s3 = get_s3_resource()
try:
s3_download_file(s3, backup_file, backup_folder)
except Exception, e:
logging.error('Unable to get backup file %s from S3: %s'%(backup_file,e))
sys.exit(2)
# Now do the restore
try:
consul_restore_from_backup(session, backup_file, overwrite)
except Exception, e:
logging.error('Failed restoring Consul key/value store: %s'%e)
sys.exit(2)
else:
logging.error('Unknown action: "%s"'%action)
sys.exit(2)