diff --git a/riak/bucket.py b/riak/bucket.py index b733e472..f6709fe3 100644 --- a/riak/bucket.py +++ b/riak/bucket.py @@ -340,6 +340,16 @@ def get_keys(self): """ return self._client.get_transport().get_keys(self) + def stream_keys(self, handler): + """ + Stream all keys within the bucket. + This way of key listing is preferable if your bucket contains rather + large amount of keys. + + :param handler: Function of one argument to be called for each key + """ + self._client.get_transport().stream_keys(self, handler) + def new_binary_from_file(self, key, filename): """ Create a new Riak object in the bucket, using the content of diff --git a/riak/mapreduce.py b/riak/mapreduce.py index 7279b0a0..51f8e772 100644 --- a/riak/mapreduce.py +++ b/riak/mapreduce.py @@ -195,11 +195,19 @@ def reduce(self, function, options=None): self._phases.append(mr) return self - def run(self, timeout=None): + def run(self, timeout=None, stream_to=None): """ - Run the map/reduce operation. Returns an array of results, or an + Run the map/reduce operation. + + If `stream_to` is None, returns an array of results, or an array of RiakLink objects if the last phase is a link phase. + + If `stream_to` callable is provided, every result will be passed into + it instead of being accumulated. Thus, when streaming feature is + used, return value will be None. + @param integer timeout - Timeout in milliseconds. + @param callable(phase,data) stream_to - Optional callback @return array() """ num_phases = len(self._phases) @@ -237,7 +245,11 @@ def run(self, timeout=None): 'key_filters': self._key_filters} t = self._client.get_transport() - result = t.mapred(self._inputs, query, timeout) + result = t.mapred(self._inputs, query, timeout, stream_to) + + # If callback is in charge, return value should be None + if stream_to: + return None # If the last phase is NOT a link phase, then return the result. if not (link_results_flag diff --git a/riak/transports/http.py b/riak/transports/http.py index 58ac88b4..8c28b064 100644 --- a/riak/transports/http.py +++ b/riak/transports/http.py @@ -293,13 +293,17 @@ def set_bucket_props(self, bucket, props): raise Exception('Error setting bucket properties.') return True - def mapred(self, inputs, query, timeout=None): + def mapred(self, inputs, query, timeout=None, stream_to=None): """ Run a MapReduce query. """ if not self.phaseless_mapred() and (query is None or len(query) is 0): raise Exception( 'Phase-less MapReduce is not supported by Riak node') + + if stream_to: + raise Exception( + 'Mapred streaming is not supported, use protobuf instead') # Construct the job, optionally set the timeout... job = {'inputs': inputs, 'query': query} diff --git a/riak/transports/pbc.py b/riak/transports/pbc.py index cfef37dd..230cc46b 100644 --- a/riak/transports/pbc.py +++ b/riak/transports/pbc.py @@ -395,6 +395,19 @@ def _handle_response(resp): return keys + def stream_keys(self, bucket, handler): + """ + Stream all keys within a bucket through handler. + """ + req = riak_pb.RpbListKeysReq() + req.bucket = bucket.name + + def _handle_response(resp): + for key in resp.keys: + handler(key) + self.send_msg_multi(MSG_CODE_LIST_KEYS_REQ, req, + MSG_CODE_LIST_KEYS_RESP, _handle_response) + def get_buckets(self): """ Serialize bucket listing request and deserialize response @@ -439,7 +452,7 @@ def set_bucket_props(self, bucket, props): MSG_CODE_SET_BUCKET_RESP) return self - def mapred(self, inputs, query, timeout=None): + def mapred(self, inputs, query, timeout=None, stream_to=None): # Construct the job, optionally set the timeout... job = {'inputs': inputs, 'query': query} if timeout is not None: @@ -458,7 +471,10 @@ def mapred(self, inputs, query, timeout=None): def _handle_response(resp): if resp.HasField("phase") and resp.HasField("response"): content = json.loads(resp.response) - if resp.phase in result: + if stream_to: + for entry in content: + stream_to(resp.phase, entry) + elif resp.phase in result: result[resp.phase] += content else: result[resp.phase] = content