Merge pull request basho#169 from evanmcc/pevm-mr-list-add

add list-consuming convenience function for mr.add

Author: evanmcc

riak/mapreduce.py

@@ -20,6 +20,7 @@
 import urllib
 from riak_object import RiakObject
 from bucket import RiakBucket
+from collections import Iterable
 
 
 class RiakMapReduce(object):
@@ -46,7 +47,7 @@ def add(self, arg1, arg2=None, arg3=None):
         specify either a RiakObject, a string bucket name, or a bucket,
         key, and additional arg.
         @param mixed arg1 - RiakObject or Bucket
-        @param mixed arg2 - Key or blank
+        @param mixed arg2 - Key or List or blank
         @param mixed arg3 - Arg or blank
         @return RiakMapReduce
         """
@@ -63,11 +64,16 @@ def add_object(self, obj):
 
     def add_bucket_key_data(self, bucket, key, data):
         if self._input_mode == 'bucket':
-            raise Exception('Already added a bucket, can\'t add an object.')
+            raise ValueError('Already added a bucket, can\'t add an object.')
         elif self._input_mode == 'query':
-            raise Exception('Already added a query, can\'t add an object.')
+            raise ValueError('Already added a query, can\'t add an object.')
         else:
-            self._inputs.append([bucket, key, data])
+            if isinstance(key, Iterable) and \
+                    not isinstance(key, basestring):
+                for k in key:
+                    self._inputs.append([bucket, k, data])
+            else:
+                self._inputs.append([bucket, key, data])
             return self
 
     def add_bucket(self, bucket):
@@ -77,14 +83,14 @@ def add_bucket(self, bucket):
 
     def add_key_filters(self, key_filters):
         if self._input_mode == 'query':
-            raise Exception('Key filters are not supported in a query.')
+            raise ValueError('Key filters are not supported in a query.')
 
         self._key_filters.extend(key_filters)
         return self
 
     def add_key_filter(self, *args):
         if self._input_mode == 'query':
-            raise Exception('Key filters are not supported in a query.')
+            raise ValueError('Key filters are not supported in a query.')
 
         self._key_filters.append(args)
         return self

riak/tests/test_mapreduce.py

@@ -83,6 +83,27 @@ def test_erlang_map_reduce(self):
             .run()
         self.assertEqual(len(result), 2)
 
+    def test_client_exceptional_paths(self):
+        bucket = self.client.bucket('bucket')
+        bucket.new("foo", 2).store()
+        bucket.new("bar", 2).store()
+        bucket.new("baz", 4).store()
+
+        #adding a b-key pair to a bucket input
+        with self.assertRaises(ValueError):
+            mr = self.client.add('bucket')
+            mr.add('bucket', 'bar')
+
+        #adding a b-key pair to a query input
+        with self.assertRaises(ValueError):
+            mr = self.client.search('bucket', 'fleh')
+            mr.add('bucket', 'bar')
+
+        #adding a key filter to a query input
+        with self.assertRaises(ValueError):
+            mr = self.client.search('bucket', 'fleh')
+            mr.add_key_filter("tokenize", "-", 1)
+ 
 
 class JSMapReduceTests(object):
     def test_javascript_source_map(self):
@@ -229,6 +250,66 @@ def test_map_reduce_from_object(self):
         result = obj.map("Riak.mapValuesJson").run()
         self.assertEqual(result, [2])
 
+    def test_mr_list_add(self):
+        bucket = self.client.bucket("abucket")
+        for x in range(20):
+            bucket.new('baz' + str(x),
+                       'bazval' + str(x)).store()
+        mr = self.client.add('abucket', ['baz' + str(x)
+                                         for x in range(2, 5)])
+        results = mr.map_values().run()
+        results.sort()
+        self.assertEqual(results,
+                         [u'"bazval2"',
+                          u'"bazval3"',
+                          u'"bazval4"'])
+
+    def test_mr_list_add_two_buckets(self):
+        bucket = self.client.bucket("bucket_a")
+        for x in range(10):
+            bucket.new('foo' + str(x),
+                       'fooval' + str(x)).store()
+        bucket = self.client.bucket("bucket_b")
+        for x in range(10):
+            bucket.new('bar' + str(x),
+                       'barval' + str(x)).store()
+
+        mr = self.client.add('bucket_a', ['foo' + str(x)
+                                          for x in range(2, 4)])
+        mr.add('bucket_b', ['bar' + str(x)
+                            for x in range(5, 7)])
+        results = mr.map_values().run()
+        results.sort()
+
+        self.assertEqual(results,
+                         [u'"barval5"',
+                          u'"barval6"',
+                          u'"fooval2"',
+                          u'"fooval3"'])
+
+    def test_mr_list_add_mix(self):
+        bucket = self.client.bucket("bucket_a")
+        for x in range(10):
+            bucket.new('foo' + str(x),
+                       'fooval' + str(x)).store()
+        bucket = self.client.bucket("bucket_b")
+        for x in range(10):
+            bucket.new('bar' + str(x),
+                       'barval' + str(x)).store()
+
+        mr = self.client.add('bucket_a', ['foo' + str(x)
+                                          for x in range(2, 4)])
+        mr.add('bucket_b', 'bar9')
+        mr.add('bucket_b', 'bar2')
+        results = mr.map_values().run()
+        results.sort()
+
+        self.assertEqual(results,
+                         [u'"barval2"',
+                          u'"barval9"',
+                          u'"fooval2"',
+                          u'"fooval3"'])
+
 
 class MapReduceAliasTests(object):
     """This tests the map reduce aliases"""