Implemented SimpleDB BatchPutAttributes (by Avdi Grimm)

xtreme · Mar 31, 2010 · d201e2a · d201e2a
1 parent ee66adf
commit d201e2a
Show file tree

Hide file tree

Showing 2 changed files with 107 additions and 6 deletions.
diff --git a/lib/sdb/right_sdb_interface.rb b/lib/sdb/right_sdb_interface.rb
@@ -89,26 +89,42 @@ def request_info(request, parser)  #:nodoc:
 
     # Prepare attributes for putting.
     # (used by put_attributes)
-    def pack_attributes(attributes, replace = false) #:nodoc:
+    def pack_attributes(items_or_attributes, replace = false, batch = false) #:nodoc:
+      if batch
+        index = 0
+        items_or_attributes.inject({}){|result, (item_name, attributes)|
+          item_prefix = "Item.#{index}."
+          result["#{item_prefix}ItemName"] = item_name.to_s
+          result.merge!(
+            pack_single_item_attributes(attributes, replace, item_prefix))
+          index += 1
+          result
+        }
+      else
+        pack_single_item_attributes(items_or_attributes, replace)
+      end
+    end
+
+    def pack_single_item_attributes(attributes, replace, prefix = "")
       result = {}
       if attributes
         idx = 0
         skip_values = attributes.is_a?(Array)
         attributes.each do |attribute, values|
           # set replacement attribute
-          result["Attribute.#{idx}.Replace"] = 'true' if replace
+          result["#{prefix}Attribute.#{idx}.Replace"] = 'true' if replace
           # pack Name/Value
           unless values.nil?
             # Array(values) does not work here:
             #  - Array('') => [] but we wanna get here ['']
             [values].flatten.each do |value|
-              result["Attribute.#{idx}.Name"]  = attribute
-              result["Attribute.#{idx}.Value"] = ruby_to_sdb(value) unless skip_values
+              result["#{prefix}Attribute.#{idx}.Name"]  = attribute
+              result["#{prefix}Attribute.#{idx}.Value"] = ruby_to_sdb(value) unless skip_values
               idx += 1
             end
           else
-            result["Attribute.#{idx}.Name"] = attribute
-            result["Attribute.#{idx}.Value"] = ruby_to_sdb(nil) unless skip_values
+            result["#{prefix}Attribute.#{idx}.Name"] = attribute
+            result["#{prefix}Attribute.#{idx}.Value"] = ruby_to_sdb(nil) unless skip_values
             idx += 1
           end
         end
@@ -309,6 +325,37 @@ def put_attributes(domain_name, item_name, attributes, replace = false)
     rescue Exception
       on_exception
     end
+
+    # Add/Replace attributes for multiple items at a time.
+    #
+    # Params:
+    #   domain_name = DomainName
+    #   items       = {
+    #     'Item1' => {
+    #       'nameA'  => [valueA1, valueA2,..., valueAN],
+    #       ...
+    #       'nameB'  => [valueB1, valueB2,..., valueBN]
+    #     },
+    #     'Item2' => {
+    #       'nameC'  => [valueC1, valueC2,..., valueCN],
+    #       ...
+    #       'nameD'  => [valueD1, valueD2,..., valueDN]
+    #     }
+    #   }
+    #   replace = :replace | any other value to skip replacement
+    #
+    # Usage of batch_put_attributes is similar to put_attributes except that
+    # instead of supplying an item_name and a hash of attributes, you supply a
+    # hash of item names to attributes.
+    #
+    # See: http://docs.amazonwebservices.com/AmazonSimpleDB/latest/DeveloperGuide/index.html?SDB_API_BatchPutAttributes.html
+    def batch_put_attributes(domain_name, items, replace = false)
+      params = { 'DomainName' => domain_name }.merge(pack_attributes(items, replace, true))
+      link = generate_request("BatchPutAttributes", params)
+      request_info( link, QSdbSimpleParser.new)
+    rescue Exception
+      on_exception
+    end
 
     # Retrieve SDB item's attribute(s).
     # 

diff --git a/test/sdb/test_batch_put_attributes.rb b/test/sdb/test_batch_put_attributes.rb
@@ -0,0 +1,54 @@
+# -*- coding: utf-8 -*-
+require File.dirname(__FILE__) + '/test_helper.rb'
+
+class TestSdb < Test::Unit::TestCase
+
+  def setup
+    STDOUT.sync  = true
+    @domain = 'right_sdb_awesome_test_domain'
+    @attributes = {
+      'a' => { 'foo' => '123' },
+      'b' => { 'bar' => '456' }
+    }
+    # Interface instance
+    @sdb = Rightscale::SdbInterface.new
+    @sdb.delete_domain(@domain)
+    wait(SDB_DELAY, "after removing domain")
+    @sdb.create_domain(@domain)
+    wait(SDB_DELAY, "after recreating domain")
+  end
+
+  SDB_DELAY = 2
+
+  def wait(delay, msg='')
+    print "waiting #{delay} seconds #{msg}"
+    while delay>0 do
+      delay -= 1
+      print '.'
+      sleep 1
+    end
+    puts
+  end
+
+  def test_batch_put_attributes
+    @sdb.batch_put_attributes(@domain, @attributes)
+    wait(SDB_DELAY, "after putting attributes")
+    a = @sdb.get_attributes(@domain, 'a')[:attributes]
+    b = @sdb.get_attributes(@domain, 'b')[:attributes]
+    assert_equal( {'foo' => ['123']}, a)
+    assert_equal( {'bar' => ['456']}, b)
+
+    # Replace = false
+    @sdb.batch_put_attributes(@domain, { 'a' => {'foo' => ['789']}})
+    wait(SDB_DELAY, "after putting attributes") 
+    a = @sdb.get_attributes(@domain, 'a')[:attributes]
+    assert_equal ['123', '789'], a['foo'].sort
+
+    # Replace = true
+    @sdb.batch_put_attributes(@domain, {'b' => {'bar' => ['789']}}, true)
+    wait(SDB_DELAY, "after putting attributes") 
+    b = @sdb.get_attributes(@domain, 'b')[:attributes]
+    assert_equal ['789'], b['bar'].sort
+
+  end
+end