diff --git a/riak/client.py b/riak/client.py index 91bdb099..756db971 100644 --- a/riak/client.py +++ b/riak/client.py @@ -213,6 +213,14 @@ def set_decoder(self, content_type, decoder): self._decoders[content_type] = decoder return self + def get_buckets(self): + """ + Get the list of buckets. + NOTE: Do not use this in production, as it requires traversing through + all keys stored in a cluster. + """ + return self._transport.get_buckets() + def bucket(self, name): """ Get the bucket by the specified name. Since buckets always exist, diff --git a/riak/mapreduce.py b/riak/mapreduce.py index ae1829c1..702eb1c6 100644 --- a/riak/mapreduce.py +++ b/riak/mapreduce.py @@ -19,6 +19,7 @@ """ import urllib from riak_object import RiakObject +from bucket import RiakBucket class RiakMapReduce(object): """ @@ -34,6 +35,7 @@ def __init__(self, client): self._client = client self._phases = [] self._inputs = [] + self._key_filters = [] self._input_mode = None def add(self, arg1, arg2=None, arg3=None): @@ -72,6 +74,20 @@ def add_bucket(self, bucket) : self._inputs = bucket return self + def add_key_filters(self, key_filters) : + if self._input_mode == 'search': + raise Exception('Key filters are not supported in search query.') + + self._key_filters.extend(key_filters) + return self + + def add_key_filter(self, *args) : + if self._input_mode == 'search': + raise Exception('Key filters are not supported in search query.') + + self._key_filters.append(args) + return self + def search(self, bucket, query): """ Begin a map/reduce operation using a Search. This command will @@ -180,6 +196,17 @@ def run(self, timeout=None): if phase._keep: keep_flag = True query.append(phase.to_array()) + if (len(self._key_filters) > 0): + bucket_name = None + if (type(self._inputs) == str): + bucket_name = self._inputs + elif (type(self._inputs) == RiakBucket): + bucket_name = self._inputs.get_name() + + if (bucket_name is not None): + self._inputs = {'bucket': bucket_name, + 'key_filters': self._key_filters} + t = self._client.get_transport() result = t.mapred(self._inputs, query, timeout) diff --git a/riak/tests/test_all.py b/riak/tests/test_all.py index 06f05272..7e6e16d6 100644 --- a/riak/tests/test_all.py +++ b/riak/tests/test_all.py @@ -310,6 +310,21 @@ def test_javascript_arg_map_reduce(self): .run() self.assertEqual(result, [10]) + def test_key_filters(self): + bucket = self.client.bucket("kftest") + bucket.new("basho-20101215", 1).store() + bucket.new("google-20110103", 2).store() + bucket.new("yahoo-20090613", 3).store() + + result = self.client \ + .add("kftest") \ + .add_key_filters([["tokenize", "-", 2]]) \ + .add_key_filter("ends_with", "0613") \ + .map("function (v, keydata) { return [v.key]; }") \ + .run() + + self.assertEqual(result, ["yahoo-20090613"]) + def test_erlang_map_reduce(self): # Create the object... bucket = self.client.bucket("bucket") diff --git a/riak/transports/http.py b/riak/transports/http.py index d9d0bd3c..0191a849 100644 --- a/riak/transports/http.py +++ b/riak/transports/http.py @@ -138,14 +138,25 @@ def get_keys(self, bucket): host, port, url = self.build_rest_path(bucket, None, None, params) response = self.http_request('GET', host, port, url) - headers = response[0] - encoded_props = response[1] + headers, encoded_props = response[0:2] if (headers['http_code'] == 200): props = json.loads(encoded_props) return props['keys'] else: raise Exception('Error getting bucket properties.') + def get_buckets(self): + params = {'buckets': 'true'} + host, port, url = self.build_rest_path(None, None, None, params) + response = self.http_request('GET', host, port, url) + + headers, encoded_props = response[0:2] + if (headers['http_code'] == 200): + props = json.loads(encoded_props) + return props['buckets'] + else: + raise Exception('Error getting buckets.') + def get_bucket_props(self, bucket, keys=False): # Run the request... params = {'props' : 'True', 'keys' : 'False'} @@ -315,7 +326,10 @@ def build_rest_path(self, bucket, key=None, spec=None, params=None) : # Build 'http://hostname:port/prefix/bucket' path = '' path += '/' + self._prefix - path += '/' + urllib.quote_plus(bucket._name) + + # Add '.../bucket' + if (bucket is not None): + path += '/' + urllib.quote_plus(bucket._name) # Add '.../key' if (key is not None): diff --git a/riak/transports/pbc.py b/riak/transports/pbc.py index 99429a40..74815ce0 100644 --- a/riak/transports/pbc.py +++ b/riak/transports/pbc.py @@ -238,6 +238,19 @@ def get_keys(self, bucket): return keys + def get_buckets(self): + """ + Serialize bucket listing request and deserialize response + """ + req = riakclient_pb2.RpbListBucketsReq() + + self.maybe_connect() + self.send_msg(MSG_CODE_LIST_KEYS_REQ, req) + msg_code, resp = self.recv_msg() + if msg_code != MSG_CODE_LIST_BUCKETS_RESP: + raise RiakError("unexpected protocol buffer message code: ", msg_code) + return resp.buckets + def get_bucket_props(self, bucket): """ Serialize bucket property request and deserialize response diff --git a/riak/transports/transport.py b/riak/transports/transport.py index 889022f3..ebc2f22f 100644 --- a/riak/transports/transport.py +++ b/riak/transports/transport.py @@ -76,6 +76,13 @@ def delete(self, robj, rw = None): """ raise RiakError("not implemented") + def get_buckets(self) : + """ + Serialize get buckets request and deserialize response + @return dict() + """ + raise RiakError("not implemented") + def get_bucket_props(self, bucket) : """ Serialize get bucket property request and deserialize response