1# Copyright (c) 2012 Mitch Garnaat http://garnaat.org/
2# Copyright (c) 2012 Amazon.com, Inc. or its affiliates.  All Rights Reserved
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# IN THE SOFTWARE.
22#
23from boto.compat import six
24
25
26class Batch(object):
27    """
28    Used to construct a BatchGet request.
29
30    :ivar table: The Table object from which the item is retrieved.
31
32    :ivar keys: A list of scalar or tuple values.  Each element in the
33        list represents one Item to retrieve.  If the schema for the
34        table has both a HashKey and a RangeKey, each element in the
35        list should be a tuple consisting of (hash_key, range_key).  If
36        the schema for the table contains only a HashKey, each element
37        in the list should be a scalar value of the appropriate type
38        for the table schema. NOTE: The maximum number of items that
39        can be retrieved for a single operation is 100. Also, the
40        number of items retrieved is constrained by a 1 MB size limit.
41
42    :ivar attributes_to_get: A list of attribute names.
43        If supplied, only the specified attribute names will
44        be returned.  Otherwise, all attributes will be returned.
45
46    :ivar consistent_read: Specify whether or not to use a
47        consistent read. Defaults to False.
48
49    """
50
51    def __init__(self, table, keys, attributes_to_get=None,
52                 consistent_read=False):
53        self.table = table
54        self.keys = keys
55        self.attributes_to_get = attributes_to_get
56        self.consistent_read = consistent_read
57
58    def to_dict(self):
59        """
60        Convert the Batch object into the format required for Layer1.
61        """
62        batch_dict = {}
63        key_list = []
64        for key in self.keys:
65            if isinstance(key, tuple):
66                hash_key, range_key = key
67            else:
68                hash_key = key
69                range_key = None
70            k = self.table.layer2.build_key_from_values(self.table.schema,
71                                                        hash_key, range_key)
72            key_list.append(k)
73        batch_dict['Keys'] = key_list
74        if self.attributes_to_get:
75            batch_dict['AttributesToGet'] = self.attributes_to_get
76        if self.consistent_read:
77            batch_dict['ConsistentRead'] = True
78        else:
79            batch_dict['ConsistentRead'] = False
80        return batch_dict
81
82
83class BatchWrite(object):
84    """
85    Used to construct a BatchWrite request.  Each BatchWrite object
86    represents a collection of PutItem and DeleteItem requests for
87    a single Table.
88
89    :ivar table: The Table object from which the item is retrieved.
90
91    :ivar puts: A list of :class:`boto.dynamodb.item.Item` objects
92        that you want to write to DynamoDB.
93
94    :ivar deletes: A list of scalar or tuple values.  Each element in the
95        list represents one Item to delete.  If the schema for the
96        table has both a HashKey and a RangeKey, each element in the
97        list should be a tuple consisting of (hash_key, range_key).  If
98        the schema for the table contains only a HashKey, each element
99        in the list should be a scalar value of the appropriate type
100        for the table schema.
101    """
102
103    def __init__(self, table, puts=None, deletes=None):
104        self.table = table
105        self.puts = puts or []
106        self.deletes = deletes or []
107
108    def to_dict(self):
109        """
110        Convert the Batch object into the format required for Layer1.
111        """
112        op_list = []
113        for item in self.puts:
114            d = {'Item': self.table.layer2.dynamize_item(item)}
115            d = {'PutRequest': d}
116            op_list.append(d)
117        for key in self.deletes:
118            if isinstance(key, tuple):
119                hash_key, range_key = key
120            else:
121                hash_key = key
122                range_key = None
123            k = self.table.layer2.build_key_from_values(self.table.schema,
124                                                        hash_key, range_key)
125            d = {'Key': k}
126            op_list.append({'DeleteRequest': d})
127        return (self.table.name, op_list)
128
129
130class BatchList(list):
131    """
132    A subclass of a list object that contains a collection of
133    :class:`boto.dynamodb.batch.Batch` objects.
134    """
135
136    def __init__(self, layer2):
137        list.__init__(self)
138        self.unprocessed = None
139        self.layer2 = layer2
140
141    def add_batch(self, table, keys, attributes_to_get=None,
142                  consistent_read=False):
143        """
144        Add a Batch to this BatchList.
145
146        :type table: :class:`boto.dynamodb.table.Table`
147        :param table: The Table object in which the items are contained.
148
149        :type keys: list
150        :param keys: A list of scalar or tuple values.  Each element in the
151            list represents one Item to retrieve.  If the schema for the
152            table has both a HashKey and a RangeKey, each element in the
153            list should be a tuple consisting of (hash_key, range_key).  If
154            the schema for the table contains only a HashKey, each element
155            in the list should be a scalar value of the appropriate type
156            for the table schema. NOTE: The maximum number of items that
157            can be retrieved for a single operation is 100. Also, the
158            number of items retrieved is constrained by a 1 MB size limit.
159
160        :type attributes_to_get: list
161        :param attributes_to_get: A list of attribute names.
162            If supplied, only the specified attribute names will
163            be returned.  Otherwise, all attributes will be returned.
164        """
165        self.append(Batch(table, keys, attributes_to_get, consistent_read))
166
167    def resubmit(self):
168        """
169        Resubmit the batch to get the next result set. The request object is
170        rebuild from scratch meaning that all batch added between ``submit``
171        and ``resubmit`` will be lost.
172
173        Note: This method is experimental and subject to changes in future releases
174        """
175        del self[:]
176
177        if not self.unprocessed:
178            return None
179
180        for table_name, table_req in six.iteritems(self.unprocessed):
181            table_keys = table_req['Keys']
182            table = self.layer2.get_table(table_name)
183
184            keys = []
185            for key in table_keys:
186                h = key['HashKeyElement']
187                r = None
188                if 'RangeKeyElement' in key:
189                    r = key['RangeKeyElement']
190                keys.append((h, r))
191
192            attributes_to_get = None
193            if 'AttributesToGet' in table_req:
194                attributes_to_get = table_req['AttributesToGet']
195
196            self.add_batch(table, keys, attributes_to_get=attributes_to_get)
197
198        return self.submit()
199
200    def submit(self):
201        res = self.layer2.batch_get_item(self)
202        if 'UnprocessedKeys' in res:
203            self.unprocessed = res['UnprocessedKeys']
204        return res
205
206    def to_dict(self):
207        """
208        Convert a BatchList object into format required for Layer1.
209        """
210        d = {}
211        for batch in self:
212            b = batch.to_dict()
213            if b['Keys']:
214                d[batch.table.name] = b
215        return d
216
217
218class BatchWriteList(list):
219    """
220    A subclass of a list object that contains a collection of
221    :class:`boto.dynamodb.batch.BatchWrite` objects.
222    """
223
224    def __init__(self, layer2):
225        list.__init__(self)
226        self.layer2 = layer2
227
228    def add_batch(self, table, puts=None, deletes=None):
229        """
230        Add a BatchWrite to this BatchWriteList.
231
232        :type table: :class:`boto.dynamodb.table.Table`
233        :param table: The Table object in which the items are contained.
234
235        :type puts: list of :class:`boto.dynamodb.item.Item` objects
236        :param puts: A list of items that you want to write to DynamoDB.
237
238        :type deletes: A list
239        :param deletes: A list of scalar or tuple values.  Each element
240            in the list represents one Item to delete.  If the schema
241            for the table has both a HashKey and a RangeKey, each
242            element in the list should be a tuple consisting of
243            (hash_key, range_key).  If the schema for the table
244            contains only a HashKey, each element in the list should
245            be a scalar value of the appropriate type for the table
246            schema.
247        """
248        self.append(BatchWrite(table, puts, deletes))
249
250    def submit(self):
251        return self.layer2.batch_write_item(self)
252
253    def to_dict(self):
254        """
255        Convert a BatchWriteList object into format required for Layer1.
256        """
257        d = {}
258        for batch in self:
259            table_name, batch_dict = batch.to_dict()
260            d[table_name] = batch_dict
261        return d
262