1#!/usr/bin/env python
2# -*-python-*-
3
4import argparse
5import urllib2
6import json
7import subprocess
8import os
9import re
10import base64
11
12"""Written by Daniel Owen owend@couchbase.com on 27 June 2014
13Version 1.4    Last Updated 17 Sept 2014 (Ian McCloy)
14
15This script is to be used to restore backups made using cbbackupwrapper.py script.
16It is a wrapper to the cbrestore command that comes with Couchbase Server 2.5.1
17
18The script uses the same number of processes for each bucket that were used to produce
19original backup.  You can specify a bucket to restore, if no bucket is specified then
20all buckets will be restored.  Note: The destination cluster must have the
21buckets created first.
22An example invocation is as follows:
23
24python cbrestorewrapper.py ../backup/ http://127.0.0.1:8091 -u Administrator \
25-p myPassword --path /opt/couchbbase/bin/ -v
26
27This will restore all the buckets from ../backup onto cluster 127.0.0.1
28Access to the cluster is authenticated using username=Administrator and
29password=myPassword.  Finally, cbrestore will be found in /opt/couchbase/bin
30
31Run python cbrestorewrapper -h for more information."""
32
33bucketList = []
34processes = {}
35bucket_target = ""
36
37
38def argumentParsing():
39    parser = argparse.ArgumentParser(description='cbbackupwrapper')
40    parser.add_argument('backupDir', type=str, help='Root directory of backup')
41    parser.add_argument('node', type=str, help='Server node specified as name or IP address')
42
43    parser.add_argument('-b', '--bucket-source', default='',
44                        help='Specify the bucket to restore.  Defaults to all buckets')
45    parser.add_argument('-B', '--bucket-destination', default='',
46                        help='Target bucket on destination cluster. Defaults to bucket-source name')
47                        #This allows you to transfer to a bucket with a different name
48                        #Only valid if --bucket-source is specified
49    parser.add_argument('-u', '--username', default='Administrator',
50                        help='REST username for source cluster or server node.  Default is Administrator')
51    parser.add_argument('-p', '--password', default='PASSWORD',
52                        help='REST password for source cluster or server node.   Defaults to PASSWORD')
53    parser.add_argument('-v', '--verbose', action='store_true',
54                        default=False, help='Enable verbose messaging')
55    parser.add_argument('--path', default='.',
56                        help='Specify the path to cbrestore.  Defaults to current directory')
57    parser.add_argument('--port', default='11210',
58                        help='Specify the bucket port.  Defaults to 11210')
59    return parser.parse_args()
60
61
62# Get the buckets that exist on the cluster
63def getBuckets(node, rest_port, username, password):
64    request = urllib2.Request(
65        'http://' + node + ':' + rest_port + '/pools/default/buckets')
66    base64string = base64.encodestring(
67        '%s:%s' % (username, password)).replace('\n', '')
68    request.add_header('Authorization', 'Basic %s' % base64string)
69    try:
70        response = urllib2.urlopen(request)
71    except:
72        print('Authorization failed.  Please check username and password.')
73        exit(1)
74    bucketsOnCluster = []
75    data = json.loads(response.read())
76    for item in data:
77        bucket = item['name']
78        bucketsOnCluster.append(bucket)
79    return bucketsOnCluster
80
81
82def getVbucketsToRestore(path, bucket):
83    vBucketList = []
84    # for each file in the directory
85    files = os.listdir(path)
86    regex = re.compile(r'^(\d+)-(\d+)$')
87    cleaned_list = filter(regex.search, files)
88    bucketname = 'bucket-' + bucket
89    okay_items = [
90        x for x in cleaned_list if bucketname in os.listdir(os.path.join(path, x))]
91    return okay_items
92
93if __name__ == '__main__':
94    # Parse the arguments given.
95    args = argumentParsing()
96
97    # Remove any white-spaces from start and end of strings
98    backupDir = args.backupDir.strip()
99    path = args.path.strip()
100
101    # Check to see if root backup directory exists
102    if not os.path.isdir(backupDir):
103        print '\n\nThe directory ' + backupDir + ' does not exist'
104        print 'Please enter a different backup directory\n'
105        exit(1)
106    # Check to see if path is correct
107    if not os.path.isdir(path):
108        print 'The path to cbrestore does not exist'
109        print 'Please run with a different path'
110        exit(1)
111    if not os.path.isfile(os.path.join(path, 'cbrestore')):
112        print 'cbrestore could not be found in ' + path
113        exit(1)
114
115    # Check to see if log directory exists if not create it
116    dir = os.path.join(backupDir, 'logs')
117    try:
118        os.stat(dir)
119    except:
120        try:
121            os.mkdir(dir)
122        except:
123            print('Error trying to create directory ' + dir)
124            exit(1)
125
126    # Separate out node and REST port
127    matchObj = re.match(r'^http://(.*):(\d+)$', args.node, re.I)
128    if matchObj:
129        node = matchObj.group(1)
130        rest = matchObj.group(2)
131    else:
132        print("Please enter the destination as http://hostname:port")
133        print("For example http://localhost:8091 or http://127.0.0.1:8091")
134        exit(1)
135
136    # Check to see if restoring all buckets or just a specified bucket
137    if args.bucket_source == '':
138        if not args.bucket_destination == '':
139            print 'please specify a bucket_source'
140            exit(1)
141        bucketList = getBuckets(
142            node, rest, args.username, args.password)
143    else:
144        # Check that the bucket exists
145        if not args.bucket_destination == '':
146            bucket_target = args.bucket_destination
147        else:
148            bucket_target = args.bucket_source
149        for item in getBuckets(node, rest, args.username, args.password):
150            if item == bucket_target:
151                bucketList.append(bucket_target)
152
153        if len(bucketList) == 0:
154            print 'Bucket ' + bucket_target + ' does not exist'
155            print 'Please enter a different bucket'
156            exit(1)
157
158    # Handle the case when path has spaces
159    if os.name == 'nt':
160        path = re.sub(r' ', '^ ', path)
161    else:
162        path = re.sub(r' ', '\ ', path)
163
164    for bucket in bucketList:
165        if not args.bucket_destination == '':
166            vbucketList = getVbucketsToRestore(backupDir, args.bucket_source.strip())
167            if len(vbucketList) == 0:
168                print 'Error reading source backup vBuckets for bucket', args.bucket_source.strip()
169                exit(1)
170        else:
171            vbucketList = getVbucketsToRestore(backupDir, bucket)
172            if len(vbucketList) == 0:
173                print 'Error reading source backup vBuckets for bucket', bucket
174                exit(1)
175        for vbuckets in vbucketList:
176            # Invoke cbrestore on each of the active vbuckets that reside on
177            # the node
178            if args.verbose:
179                print "vBucket: ", vbuckets
180            if not args.bucket_destination == '':
181                command_line = os.path.join(path, 'cbrestore') + ' -v -t 1 -b ' + args.bucket_source.strip() \
182                    + ' -B ' + args.bucket_destination + ' ' + os.path.join(backupDir, vbuckets) \
183                    + ' http://' + node + ':' + rest \
184                    + ' -u ' + args.username + ' -p ' + args.password + ' 2>' + \
185                    os.path.join(backupDir, 'logs', vbuckets) + \
186                    '-restore-' + bucket + '.err'
187            else:
188                command_line = os.path.join(path, 'cbrestore') + ' -v -t 1 -b ' + bucket \
189                    + ' ' + os.path.join(backupDir, vbuckets) + ' http://' + node + ':' + rest \
190                    + ' -u ' + args.username + ' -p ' + args.password + ' 2>' + \
191                    os.path.join(backupDir, 'logs', vbuckets) + \
192                    '-restore-' + bucket + '.err'
193            if args.verbose:
194                print command_line
195            p = subprocess.Popen(command_line, shell=True)
196            processes[p] = vbuckets + '-restore-' + bucket
197
198    # Did we restore anything?
199    if len(processes) == 0:
200        print 'Did not restore anything'
201        print 'Please check that the backup directory contains data to restore'
202        print 'Also please check that you have the correct buckets created on ' + args.node
203        exit(1)
204    else:
205        print 'Waiting for the restore to complete...'
206        successCount = 0
207        for p in processes:
208            p.wait()
209            if p.returncode == 1:
210                print 'Error with backup - look in ' + os.path.join(backupDir, 'logs', processes[p]) \
211                                                     + '-restore-' + bucket + \
212                    '.err for details'
213            else:
214                successCount += 1
215
216        if successCount == len(processes):
217            print 'SUCCESSFULLY COMPLETED!'
218        else:
219            print 'ERROR!'
220            exit(1)
221