1#!/usr/bin/env python
2# Copyright (c) 2009 Chris Moyer http://kopertop.blogspot.com/
3#
4# Permission is hereby granted, free of charge, to any person obtaining a
5# copy of this software and associated documentation files (the
6# "Software"), to deal in the Software without restriction, including
7# without limitation the rights to use, copy, modify, merge, publish, dis-
8# tribute, sublicense, and/or sell copies of the Software, and to permit
9# persons to whom the Software is furnished to do so, subject to the fol-
10# lowing conditions:
11#
12# The above copyright notice and this permission notice shall be included
13# in all copies or substantial portions of the Software.
14#
15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
16# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
17# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
18# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 
19# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
21
22#
23# Tools to dump and recover an SDB domain
24#
25VERSION = "%prog version 1.0"
26import boto
27import time
28from boto import sdb
29from boto.compat import json
30
31def choice_input(options, default=None, title=None):
32    """
33    Choice input
34    """
35    if title == None:
36        title = "Please choose"
37    print title
38    objects = []
39    for n, obj in enumerate(options):
40        print "%s: %s" % (n, obj)
41        objects.append(obj)
42    choice = int(raw_input(">>> "))
43    try:
44        choice = objects[choice]
45    except:
46        choice = default
47    return choice
48
49def confirm(message="Are you sure?"):
50    choice = raw_input("%s [yN] " % message)
51    return choice and len(choice) > 0 and choice[0].lower() == "y"
52
53
54def dump_db(domain, file_name, use_json=False, sort_attributes=False):
55    """
56    Dump SDB domain to file
57    """
58    f = open(file_name, "w")
59    if use_json:
60        for item in domain:
61            data = {"name": item.name, "attributes": item}
62            print >> f, json.dumps(data, sort_keys=sort_attributes)
63    else:
64        doc = domain.to_xml(f)
65
66def empty_db(domain):
67    """
68    Remove all entries from domain
69    """
70    for item in domain:
71        item.delete()
72
73def load_db(domain, file, use_json=False):
74    """
75    Load a domain from a file, this doesn't overwrite any existing
76    data in the file so if you want to do a full recovery and restore
77    you need to call empty_db before calling this
78
79    :param domain: The SDB Domain object to load to
80    :param file: The File to load the DB from
81    """
82    if use_json:
83        for line in file.readlines():
84            if line:
85                data = json.loads(line)
86                item = domain.new_item(data['name'])
87                item.update(data['attributes'])
88                item.save()
89                
90    else:
91        domain.from_xml(file)
92
93def check_valid_region(conn, region):
94    if conn is None:
95        print 'Invalid region (%s)' % region
96        sys.exit(1)
97
98def create_db(domain_name, region_name):
99    """Create a new DB
100
101    :param domain: Name of the domain to create
102    :type domain: str
103    """
104    sdb = boto.sdb.connect_to_region(region_name)
105    check_valid_region(sdb, region_name)
106    return sdb.create_domain(domain_name)
107
108if __name__ == "__main__":
109    from optparse import OptionParser
110    parser = OptionParser(version=VERSION, usage="Usage: %prog [--dump|--load|--empty|--list|-l] [options]")
111
112    # Commands
113    parser.add_option("--dump", help="Dump domain to file", dest="dump", default=False, action="store_true")
114    parser.add_option("--load", help="Load domain contents from file", dest="load", default=False, action="store_true")
115    parser.add_option("--empty", help="Empty all contents of domain", dest="empty", default=False, action="store_true")
116    parser.add_option("-l", "--list", help="List All domains", dest="list", default=False, action="store_true")
117    parser.add_option("-c", "--create", help="Create domain", dest="create", default=False, action="store_true")
118
119    parser.add_option("-a", "--all-domains", help="Operate on all domains", action="store_true", default=False, dest="all_domains")
120    if json:
121        parser.add_option("-j", "--use-json", help="Load/Store as JSON instead of XML", action="store_true", default=False, dest="json")
122    parser.add_option("-s", "--sort-attibutes", help="Sort the element attributes", action="store_true", default=False, dest="sort_attributes")
123    parser.add_option("-d", "--domain", help="Do functions on domain (may be more then one)", action="append", dest="domains")
124    parser.add_option("-f", "--file", help="Input/Output file we're operating on", dest="file_name")
125    parser.add_option("-r", "--region", help="Region (e.g. us-east-1[default] or eu-west-1)", default="us-east-1", dest="region_name")
126    (options, args) = parser.parse_args()
127
128    if options.create:
129        for domain_name in options.domains:
130            create_db(domain_name, options.region_name)
131        exit()
132
133    sdb = boto.sdb.connect_to_region(options.region_name)
134    check_valid_region(sdb, options.region_name)
135    if options.list:
136        for db in sdb.get_all_domains():
137            print db
138        exit()
139
140    if not options.dump and not options.load and not options.empty:
141            parser.print_help()
142            exit()
143
144
145
146
147    #
148    # Setup
149    #
150    if options.domains:
151        domains = []
152        for domain_name in options.domains:
153            domains.append(sdb.get_domain(domain_name))
154    elif options.all_domains:
155        domains = sdb.get_all_domains()
156    else:
157        domains = [choice_input(options=sdb.get_all_domains(), title="No domain specified, please choose one")]
158
159
160    #
161    # Execute the commands
162    #
163    stime = time.time()
164    if options.empty:
165        if confirm("WARNING!!! Are you sure you want to empty the following domains?: %s" % domains):
166            stime = time.time()
167            for domain in domains:
168                print "--------> Emptying %s <--------" % domain.name
169                empty_db(domain)
170        else:
171            print "Canceling operations"
172            exit()
173
174    if options.dump:
175        for domain in domains:
176            print "--------> Dumping %s <---------" % domain.name
177            if options.file_name:
178                file_name = options.file_name
179            else:
180                file_name = "%s.db" % domain.name
181            dump_db(domain, file_name, options.json, options.sort_attributes)
182
183    if options.load:
184        for domain in domains:
185            print "---------> Loading %s <----------" % domain.name
186            if options.file_name:
187                file_name = options.file_name
188            else:
189                file_name = "%s.db" % domain.name
190            load_db(domain, open(file_name, "rb"), options.json)
191
192
193    total_time = round(time.time() - stime, 2)
194    print "--------> Finished in %s <--------" % total_time
195