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