Skip to content
This repository
Browse code

Adding support for BatchGetItem to Layer2. Closes #533.

  • Loading branch information...
commit 9d42dea0683117a537046cdae7b0fa5d7145a386 1 parent 4f148ea
Mitch Garnaat authored January 30, 2012
76  boto/dynamodb/batch.py
... ...
@@ -0,0 +1,76 @@
  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
+#
  23
+
  24
+class Batch(object):
  25
+    """
  26
+    :ivar table: The Table object from which the item is retrieved.
  27
+
  28
+    :ivar keys: A list of scalar or tuple values.  Each element in the
  29
+        list represents one Item to retrieve.  If the schema for the
  30
+        table has both a HashKey and a RangeKey, each element in the
  31
+        list should be a tuple consisting of (hash_key, range_key).  If
  32
+        the schema for the table contains only a HashKey, each element
  33
+        in the list should be a scalar value of the appropriate type
  34
+        for the table schema.
  35
+        
  36
+    :ivar attributes_to_get: A list of attribute names.
  37
+        If supplied, only the specified attribute names will
  38
+        be returned.  Otherwise, all attributes will be returned.
  39
+    """
  40
+
  41
+    def __init__(self, table, keys, attributes_to_get=None):
  42
+        self.table = table
  43
+        self.keys = keys
  44
+        self.attributes_to_get = attributes_to_get
  45
+        
  46
+class BatchList(list):
  47
+    """
  48
+    A subclass of a list object that contains a collection of
  49
+    :class:`boto.dynamodb.batch.Batch` objects.
  50
+    """
  51
+
  52
+    def add_batch(self, table, keys, attributes_to_get=None):
  53
+        """
  54
+        Add a Batch to this BatchList.
  55
+        
  56
+        :type table: :class:`boto.dynamodb.table.Table`
  57
+        :param table: The Table object in which the items are contained.
  58
+
  59
+        :type keys: list
  60
+        :param keys: A list of scalar or tuple values.  Each element in the
  61
+            list represents one Item to retrieve.  If the schema for the
  62
+            table has both a HashKey and a RangeKey, each element in the
  63
+            list should be a tuple consisting of (hash_key, range_key).  If
  64
+            the schema for the table contains only a HashKey, each element
  65
+            in the list should be a scalar value of the appropriate type
  66
+            for the table schema.
  67
+        
  68
+        :type attributes_to_get: list
  69
+        :param attributes_to_get: A list of attribute names.
  70
+            If supplied, only the specified attribute names will
  71
+            be returned.  Otherwise, all attributes will be returned.
  72
+        """
  73
+        self.append(Batch(table, keys, attributes_to_get))
  74
+
  75
+        
  76
+
10  boto/dynamodb/layer1.py
@@ -273,7 +273,7 @@ def get_item(self, table_name, key, attributes_to_get=None,
273 273
             )
274 274
         return response
275 275
         
276  
-    def batch_get_item(self, request_items):
  276
+    def batch_get_item(self, request_items, object_hook=None):
277 277
         """
278 278
         Return a set of attributes for a multiple items in
279 279
         multiple tables using their primary keys.
@@ -284,7 +284,8 @@ def batch_get_item(self, request_items):
284 284
         """
285 285
         data = {'RequestItems' : request_items}
286 286
         json_input = json.dumps(data)
287  
-        return self.make_request('BatchGetItem', json_input)
  287
+        return self.make_request('BatchGetItem', json_input,
  288
+                                 object_hook=object_hook)
288 289
 
289 290
     def put_item(self, table_name, item,
290 291
                  expected=None, return_values=None):
@@ -459,7 +460,8 @@ def query(self, table_name, hash_key_value, range_key_conditions=None,
459 460
 
460 461
     def scan(self, table_name, scan_filter=None,
461 462
              attributes_to_get=None, limit=None,
462  
-             count=False, exclusive_start_key=None):
  463
+             count=False, exclusive_start_key=None,
  464
+             object_hook=None):
463 465
         """
464 466
         Perform a scan of DynamoDB.  This version is currently punting
465 467
         and expecting you to provide a full and correct JSON body
@@ -502,6 +504,6 @@ def scan(self, table_name, scan_filter=None,
502 504
         if exclusive_start_key:
503 505
             data['ExclusiveStartKey'] = exclusive_start_key
504 506
         json_input = json.dumps(data)
505  
-        return self.make_request('Scan', json_input)
  507
+        return self.make_request('Scan', json_input, object_hook=object_hook)
506 508
 
507 509
     
51  boto/dynamodb/layer2.py
@@ -25,6 +25,7 @@
25 25
 from boto.dynamodb.table import Table
26 26
 from boto.dynamodb.schema import Schema
27 27
 from boto.dynamodb.item import Item
  28
+from boto.dynamodb.batch import BatchList
28 29
 
29 30
 """
30 31
 Some utility functions to deal with mapping Amazon DynamoDB types to
@@ -120,6 +121,32 @@ def dynamize_expected_value(self, expected_value):
120 121
                 d[attr_name] = attr_value
121 122
         return d
122 123
 
  124
+    def dynamize_request_items(self, batch_list):
  125
+        """
  126
+        Convert a request_items parameter into the data structure
  127
+        required for Layer1.
  128
+        """
  129
+        d = None
  130
+        if batch_list:
  131
+            d = {}
  132
+            for batch in batch_list:
  133
+                batch_dict = {}
  134
+                key_list = []
  135
+                for key in batch.keys:
  136
+                    if isinstance(key, tuple):
  137
+                        hash_key, range_key = key
  138
+                    else:
  139
+                        hash_key = key
  140
+                        range_key = None
  141
+                    k = self.build_key_from_values(batch.table.schema,
  142
+                                                   hash_key, range_key)
  143
+                    key_list.append(k)
  144
+                batch_dict['Keys'] = key_list
  145
+                if batch.attributes_to_get:
  146
+                    batch_dict['AttributesToGet'] = batch.attributes_to_get
  147
+            d[batch.table.name] = batch_dict
  148
+        return d
  149
+
123 150
     def get_dynamodb_type(self, val):
124 151
         """
125 152
         Take a scalar Python value and return a string representing
@@ -201,6 +228,9 @@ def build_key_from_values(self, schema, hash_key, range_key=None):
201 228
             dynamodb_key['RangeKeyElement'] = dynamodb_value
202 229
         return dynamodb_key
203 230
 
  231
+    def new_batch_list(self):
  232
+        return BatchList()
  233
+
204 234
     def list_tables(self, limit=None, start_table=None):
205 235
         """
206 236
         Return a list of the names of all Tables associated with the
@@ -376,6 +406,22 @@ def get_item(self, table, hash_key, range_key=None,
376 406
             item.consumed_units = response['ConsumedCapacityUnits']
377 407
         return item
378 408
 
  409
+    def batch_get_item(self, batch_list):
  410
+        """
  411
+        Return a set of attributes for a multiple items in
  412
+        multiple tables using their primary keys.
  413
+
  414
+        :type batch_list: :class:`boto.dynamodb.batch.BatchList`
  415
+        :param batch_list: A BatchList object which consists of a
  416
+            list of :class:`boto.dynamoddb.batch.Batch` objects.
  417
+            Each Batch object contains the information about one
  418
+            batch of objects that you wish to retrieve in this
  419
+            request.
  420
+        """
  421
+        request_items = self.dynamize_request_items(batch_list)
  422
+        return self.layer1.batch_get_item(request_items,
  423
+                                          object_hook=item_object_hook)
  424
+
379 425
     def put_item(self, item, expected_value=None, return_values=None):
380 426
         """
381 427
         Store a new item or completely replace an existing item
@@ -551,8 +597,9 @@ def scan(self, table, scan_filter=None,
551 597
                 break
552 598
 
553 599
             response = self.layer1.scan(table.name, scan_filter,
554  
-                attributes_to_get,limit,
555  
-                count, exclusive_start_key)
  600
+                                        attributes_to_get,limit,
  601
+                                        count, exclusive_start_key,
  602
+                                        object_hook=item_object_hook)
556 603
             if response:
557 604
                 for item in response['Items']:
558 605
                     yield item_class(table, attrs=item)
10  boto/dynamodb/table.py
@@ -179,9 +179,10 @@ def get_item(self, hash_key, range_key=None,
179 179
             :class:`boto.dynamodb.item.Item`
180 180
         """
181 181
         return self.layer2.get_item(self, hash_key, range_key,
182  
-                                    attributes_to_get, consistent_read, item_class)
183  
-
  182
+                                    attributes_to_get, consistent_read,
  183
+                                    item_class)
184 184
     lookup = get_item
  185
+    
185 186
     def has_item(self, hash_key, range_key=None, consistent_read=False):
186 187
         """
187 188
         Checks the table to see if the Item with the specified ``hash_key``
@@ -288,9 +289,10 @@ def scan(self, scan_filter=None,
288 289
              attributes_to_get=None, limit=None,
289 290
              count=False, exclusive_start_key=None,
290 291
              item_class=Item):
291  
-        """Scan through this table, this is a very long
  292
+        """
  293
+        Scan through this table, this is a very long
292 294
         and expensive operation, and should be avoided if
293  
-        at all possible
  295
+        at all possible.
294 296
 
295 297
         :type scan_filter: dict
296 298
         :param scan_filter: A Python version of the
7  docs/source/ref/dynamodb.rst
Source Rendered
@@ -46,4 +46,11 @@ boto.dynamodb.item
46 46
    :members:
47 47
    :undoc-members:
48 48
 
  49
+boto.dynamodb.batch
  50
+-------------------
  51
+
  52
+.. automodule:: boto.dynamodb.batch
  53
+   :members:
  54
+   :undoc-members:
  55
+
49 56
 
7  tests/dynamodb/test_layer2.py
@@ -248,6 +248,13 @@ def test_layer2_basic(self):
248 248
             assert i in mixed_set
249 249
         for i in item4['StrSetAttr']:
250 250
             assert i in str_set
  251
+
  252
+        # Try a batch get
  253
+        batch_list = c.new_batch_list()
  254
+        batch_list.add_batch(table, [(item2_key, item2_range),
  255
+                                     (item3_key, item3_range)])
  256
+        response = c.batch_get_item(batch_list)
  257
+        assert len(response['Responses'][table.name]['Items']) == 2
251 258
         
252 259
         # Try to delete the item with the right Expected value
253 260
         expected = {'Views': 0}

0 notes on commit 9d42dea

Please sign in to comment.
Something went wrong with that request. Please try again.