class AWS::DynamoDB::BatchGet

def convert_unprocessed_keys response
  return nil if response.data['UnprocessedKeys'].empty?
  # convert the json response keys into symbolized keys
  str2sym = lambda do |key_desc|
    type, value = key_desc.to_a.flatten
    case type
    when "S" then { :s => value }
    when "N" then { :n => value }
    else
      raise "unhandled key type: #{type}"
    end
  end
  request_items = {}
  response.data['UnprocessedKeys'].each_pair do |table,keys|
    request_items[table] = {}
    request_items[table][:attributes_to_get] = keys['AttributesToGet'] if
      keys['AttributesToGet']
    request_items[table][:keys] = keys['Keys'].collect do |desc|
      key = {}
      key[:hash_key_element] = str2sym.call(desc['HashKeyElement'])
      key[:range_key_element] = str2sym.call(desc['RangeKeyElement']) if
        desc['RangeKeyElement']
      key
    end
  end
  request_items
end

def each &block
  options = { :request_items => @request_items }
  begin
    response = client.batch_get_item(options)
    response.data['Responses'].each_pair do |table_name,details|
      details['Items'].each do |hash_data|
        attributes = values_from_response_hash(hash_data)
        yield(table_name, attributes)
      end
    end
    options[:request_items] = convert_unprocessed_keys(response)
  end while options[:request_items]
  nil
end

def each_attributes
  each do |table_name, attributes|
    yield(attributes)
  end
end

def initialize options = {}
  super(options)
  @request_items = {}
end

def items attributes, *items
  [items].flatten.each do |item|
    self.table(item.table, attributes, [item])
  end
end

def table table, attributes, items, options = {}
  table = table.is_a?(Table) ? table.name : table.to_s
  attributes = attributes == :all ? nil : [attributes].flatten
  keys = items.collect do |item|
    case item
    when Item then item_key_hash(item)
    when Array
      {
        :hash_key_element => format_attribute_value(item[0]),
        :range_key_element => format_attribute_value(item[1]),
      }
    else
      { :hash_key_element => format_attribute_value(item) }
    end
  end
  # ensure we don't receive 2 different lists of attributes for
  # the same table
  if
    @request_items.has_key?(table) and
    @request_items[table][:attributes_to_get] != attributes
  then
    msg = "When batch getting attributes, you may only provide " +
      "1 list of attributes per table, but the `#{table}` table " +
      "has received reqeusts for 2 different sets of attributes"
    raise ArgumentError, msg
  end
  # merge attributes and items with the request items
  @request_items[table] ||= { :keys => [] }
  @request_items[table][:attributes_to_get] = attributes if attributes
  @request_items[table][:keys] += keys
  @request_items[table].merge!(options)
  nil
end