diff --git a/.travis.yml b/.travis.yml index 3e63c8c1..a47837e0 100644 --- a/.travis.yml +++ b/.travis.yml @@ -2,7 +2,9 @@ language: python python: - "2.6" - "2.7" -install: ./setup.py develop +install: + - ./setup.py develop + - ./setup.py easy_install protobuf script: ./setup.py test before_script: sudo /usr/sbin/search-cmd install searchbucket notifications: diff --git a/riak/__init__.py b/riak/__init__.py index 1e116161..2296f2a1 100644 --- a/riak/__init__.py +++ b/riak/__init__.py @@ -32,19 +32,16 @@ class RiakError(Exception): + """ + Base class for exceptions generated in the Riak API. + """ def __init__(self, value): self.value = value def __str__(self): return repr(self.value) -from riak_object import RiakObject -from bucket import RiakBucket -from client import RiakClient -from mapreduce import RiakMapReduce, RiakMapReducePhase, RiakLinkPhase,\ - RiakKeyFilter -from transports.pbc import RiakPbcTransport -from transports.http import RiakHttpTransport +from mapreduce import RiakKeyFilter ONE = "one" ALL = "all" diff --git a/riak/bucket.py b/riak/bucket.py index 0782a2e1..781aa4f3 100644 --- a/riak/bucket.py +++ b/riak/bucket.py @@ -222,7 +222,11 @@ def _set_r(self, val): def _get_r(self): return self.get_property('r') - r = property(_get_r, _set_r) + r = property(_get_r, _set_r, doc=""" + The default 'read' quorum for this bucket (how many replicas must + reply for a successful read). This should be an integer less than + the 'n_val' property, or a string of 'one', 'quorum', 'all', or + 'default'""") def _set_pr(self, val): return self.set_property('pr', val) @@ -230,7 +234,11 @@ def _set_pr(self, val): def _get_pr(self): return self.get_property('pr') - pr = property(_get_pr, _set_pr) + pr = property(_get_pr, _set_pr, doc=""" + The default 'primary read' quorum for this bucket (how many + primary replicas are required for a successful read). This should + be an integer less than the 'n_val' property, or a string of + 'one', 'quorum', 'all', or 'default'""") def _set_rw(self, val): return self.set_property('rw', val) @@ -238,7 +246,11 @@ def _set_rw(self, val): def _get_rw(self): return self.get_property('rw') - rw = property(_get_rw, _set_rw) + rw = property(_get_rw, _set_rw, doc=""" + The default 'read' and 'write' quorum for this bucket (equivalent + to 'r' and 'w' but for deletes). This should be an integer less + than the 'n_val' property, or a string of 'one', 'quorum', 'all', + or 'default'""") def _set_w(self, val): return self.set_property('w', val) @@ -246,7 +258,11 @@ def _set_w(self, val): def _get_w(self): return self.get_property('w') - w = property(_get_w, _set_w) + w = property(_get_w, _set_w, doc=""" + The default 'write' quorum for this bucket (how many replicas must + acknowledge receipt of a write). This should be an integer less + than the 'n_val' property, or a string of 'one', 'quorum', 'all', + or 'default'""") def _set_dw(self, val): return self.set_property('dw', val) @@ -254,7 +270,11 @@ def _set_dw(self, val): def _get_dw(self): return self.get_property('dw') - dw = property(_get_dw, _set_dw) + dw = property(_get_dw, _set_dw, doc=""" + The default 'durable write' quorum for this bucket (how many + replicas must commit the write). This should be an integer less + than the 'n_val' property, or a string of 'one', 'quorum', 'all', + or 'default'""") def _set_pw(self, val): return self.set_property('pw', val) @@ -262,16 +282,16 @@ def _set_pw(self, val): def _get_pw(self): return self.get_property('pw') - pw = property(_get_pw, _set_pw) + pw = property(_get_pw, _set_pw, doc=""" + The default 'primary write' quorum for this bucket (how many + primary replicas are required for a successful write). This should + be an integer less than the 'n_val' property, or a string of + 'one', 'quorum', 'all', or 'default'""") def set_property(self, key, value): """ Set a bucket property. - .. warning:: - - This should only be used if you know what you are doing. - :param key: Property to set. :type key: string :param value: Property value. @@ -296,24 +316,18 @@ def set_properties(self, props): """ Set multiple bucket properties in one call. - .. warning:: - - This should only be used if you know what you are doing. - - :param props: An associative array of key:value. - :type props: array + :param props: A dictionary of properties + :type props: dict """ - t = self._client.get_transport() - t.set_bucket_props(self, props) + self._client.set_bucket_props(self, props) def get_properties(self): """ - Retrieve an associative array of all bucket properties. + Retrieve a dict of all bucket properties. - :rtype: array + :rtype: dict """ - t = self._client.get_transport() - return t.get_bucket_props(self) + return self._client.get_bucket_props(self) def get_keys(self): """ @@ -323,7 +337,19 @@ def get_keys(self): At current, this is a very expensive operation. Use with caution. """ - return self._client.get_transport().get_keys(self) + return self._client.get_keys(self) + + def stream_keys(self): + """ + Streams all keys within the bucket through an iterator. + + .. warning:: + + At current, this is a very expensive operation. Use with caution. + + :rtype: iterator + """ + return self._client.stream_keys(self) def new_binary_from_file(self, key, filename): """ @@ -332,6 +358,10 @@ def new_binary_from_file(self, key, filename): """ binary_data = open(filename, "rb").read() mimetype, encoding = mimetypes.guess_type(filename) + if encoding: + binary_data = bytearray(binary_data, encoding) + else: + binary_data = bytearray(binary_data) if not mimetype: mimetype = 'application/octet-stream' return self.new_binary(key, binary_data, mimetype) @@ -371,11 +401,13 @@ def search(self, query, **params): """ Queries a search index over objects in this bucket/index. """ - return self._client.solr().search(self.name, query, **params) + return self._client.solr.search(self.name, query, **params) def get_index(self, index, startkey, endkey=None): """ Queries a secondary index over objects in this bucket, returning keys. """ - return self._client._transport.get_index(self.name, index, startkey, - endkey) + return self._client.get_index(self.name, index, startkey, endkey) + + def __str__(self): + return ''.format(self.name) diff --git a/riak/client.py b/riak/client.py deleted file mode 100644 index 483445a0..00000000 --- a/riak/client.py +++ /dev/null @@ -1,256 +0,0 @@ -""" -Copyright 2010 Rusty Klophaus -Copyright 2010 Justin Sheehy -Copyright 2009 Jay Baird - -This file is provided to you under the Apache License, -Version 2.0 (the "License"); you may not use this file -except in compliance with the License. You may obtain -a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, -software distributed under the License is distributed on an -"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -KIND, either express or implied. See the License for the -specific language governing permissions and limitations -under the License. -""" -# Use json as first choice, simplejson as second choice. -try: - import json -except ImportError: - import simplejson as json - -from riak.bucket import RiakBucket -from riak.mapreduce import RiakMapReduce -from riak.search import RiakSearch -from riak.transports import RiakHttpTransport -from riak.util import deprecated -from riak.util import deprecateQuorumAccessors - - -@deprecateQuorumAccessors -class RiakClient(object): - """ - The ``RiakClient`` object holds information necessary to connect to - Riak. The Riak API uses HTTP, so there is no persistent - connection, and the ``RiakClient`` object is extremely lightweight. - """ - def __init__(self, host='127.0.0.1', port=8098, prefix='riak', - mapred_prefix='mapred', transport_class=None, - client_id=None, solr_transport_class=None, - transport_options=None): - """ - Construct a new ``RiakClient`` object. - - :param host: Hostname or IP address - :type host: string - :param port: Port number - :type port: integer - :param prefix: Interface prefix - :type prefix: string - :param mapred_prefix: MapReduce prefix - :type mapred_prefix: string - :param transport_class: transport class to use - :type transport_class: :class:`RiakTransport` - - :param solr_transport_class: HTTP-based transport class for - Solr interface queries - :type solr_transport_class: :class:`RiakHttpTransport` - :param transport_options: Optional key-value args to pass to - the transport constuctor - :type transport_options: dict - """ - if transport_class is None: - transport_class = RiakHttpTransport - - api = getattr(transport_class, 'api', 1) - if api >= 2: - hostports = [(host, port), ] - self._cm = transport_class.default_cm(hostports) - - # If no transport options are provided, then default to the - # empty dict, otherwise just pass through what we are provided. - if transport_options is None: - transport_options = {} - - self._transport = transport_class(self._cm, - prefix=prefix, - mapred_prefix=mapred_prefix, - client_id=client_id, - **transport_options) - else: - deprecated('please upgrade the transport to the new API') - self._cm = None - self._transport = transport_class(host, port, client_id=client_id) - - self._encoders = {'application/json': json.dumps, - 'text/json': json.dumps} - self._decoders = {'application/json': json.loads, - 'text/json': json.loads} - self._solr = None - self._host = host - self._port = port - - def get_transport(self): - """ - Get the transport instance the client is using for it's connection. - """ - return self._transport - - def get_client_id(self): - """ - Get the ``client_id`` for this ``RiakClient`` instance. - - :rtype: string - """ - return self._transport.get_client_id() - - def set_client_id(self, client_id): - """ - Set the client_id for this ``RiakClient`` instance. - - .. warning:: - - Refer to - http://wiki.basho.com/Client-Implementation-Guide.html#Client-IDs - for information on how to set the client_id. - - :param client_id: The new client_id. - :type client_id: string - :rtype: self - """ - self._transport.set_client_id(client_id) - return self - - def get_encoder(self, content_type): - """ - Get the encoding function for the provided content type. - """ - if content_type in self._encoders: - return self._encoders[content_type] - - def set_encoder(self, content_type, encoder): - """ - Set the encoding function for the provided content type. - - :param encoder: - :type encoder: function - """ - self._encoders[content_type] = encoder - return self - - def get_decoder(self, content_type): - """ - Get the decoding function for the provided content type. - """ - if content_type in self._decoders: - return self._decoders[content_type] - - def set_decoder(self, content_type, decoder): - """ - Set the decoding function for the provided content type. - - :param decoder: - :type decoder: function - """ - 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, - this will always return a :class:`RiakBucket `. - - :rtype: :class:`RiakBucket ` - """ - return RiakBucket(self, name) - - def is_alive(self): - """ - Check if the Riak server for this ``RiakClient`` instance is alive. - - :rtype: boolean - """ - return self._transport.ping() - - def add(self, *args): - """ - Start assembling a Map/Reduce operation. A shortcut for - :func:`RiakMapReduce.add`. - - :rtype: :class:`RiakMapReduce` - """ - mr = RiakMapReduce(self) - return apply(mr.add, args) - - def search(self, *args): - """ - Start assembling a Map/Reduce operation based on search - results. This command will return an error unless executed - against a Riak Search cluster. A shortcut for - :func:`RiakMapReduce.search`. - - :rtype: :class:`RiakMapReduce` - """ - mr = RiakMapReduce(self) - return apply(mr.search, args) - - def index(self, *args): - """ - Start assembling a Map/Reduce operation based on secondary - index query results. - - :rtype: :class:`RiakMapReduce` - """ - mr = RiakMapReduce(self) - return apply(mr.index, args) - - def link(self, *args): - """ - Start assembling a Map/Reduce operation. A shortcut for - :func:`RiakMapReduce.link`. - - :rtype: :class:`RiakMapReduce` - """ - mr = RiakMapReduce(self) - return apply(mr.link, args) - - def map(self, *args): - """ - Start assembling a Map/Reduce operation. A shortcut for - :func:`RiakMapReduce.map`. - - :rtype: :class:`RiakMapReduce` - """ - mr = RiakMapReduce(self) - return apply(mr.map, args) - - def reduce(self, *args): - """ - Start assembling a Map/Reduce operation. A shortcut for - :func:`RiakMapReduce.reduce`. - - :rtype: :class:`RiakMapReduce` - """ - mr = RiakMapReduce(self) - return apply(mr.reduce, args) - - def get_index(self, bucket, index, startkey, endkey=None): - return self._transport.get_index(bucket, index, startkey, endkey) - - def solr(self): - if self._solr is None: - self._solr = RiakSearch(self, host=self._host, port=self._port) - - return self._solr diff --git a/riak/client/__init__.py b/riak/client/__init__.py new file mode 100644 index 00000000..4fc285cf --- /dev/null +++ b/riak/client/__init__.py @@ -0,0 +1,229 @@ +""" +Copyright 2011 Basho Technologies, Inc. +Copyright 2010 Rusty Klophaus +Copyright 2010 Justin Sheehy +Copyright 2009 Jay Baird + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" + +import json +import random +from weakref import WeakValueDictionary +from riak.client.operations import RiakClientOperations +from riak.node import RiakNode +from riak.bucket import RiakBucket +from riak.mapreduce import RiakMapReduceChain +from riak.search import RiakSearch +from riak.transports.http import RiakHttpPool +from riak.transports.pbc import RiakPbcPool +from riak.util import deprecated +from riak.util import deprecateQuorumAccessors +from riak.util import lazy_property + + +@deprecateQuorumAccessors +class RiakClient(RiakMapReduceChain, RiakClientOperations): + """ + The ``RiakClient`` object holds information necessary to connect + to Riak. Requests can be made to Riak directly through the client + or by using the methods on related objects. + """ + + PROTOCOLS = ['http', 'https', 'pbc'] + + def __init__(self, protocol='http', transport_options={}, + nodes=None, **unused_args): + """ + Construct a new ``RiakClient`` object. + + :param protocol: the preferred protocol, defaults to 'http' + :type protocol: string + :param nodes: a list of node configurations, + where each configuration is a dict containing the keys + 'host', 'http_port', and 'pb_port' + :type nodes: list + :param transport_options: Optional key-value args to pass to + the transport constuctor + :type transport_options: dict + """ + if 'port' in unused_args: + deprecated("port option is deprecated, use http_port or pb_port," + + " or the nodes option") + + if 'transport_class' in unused_args: + deprecated( + "transport_class is deprecated, use the protocol option") + + if nodes is None: + self.nodes = [self._create_node(unused_args), ] + else: + self.nodes = [self._create_node(n) for n in nodes] + + self.protocol = protocol or 'http' + + self._http_pool = RiakHttpPool(self, **transport_options) + self._pb_pool = RiakPbcPool(self, **transport_options) + + self._encoders = {'application/json': json.dumps, + 'text/json': json.dumps} + self._decoders = {'application/json': json.loads, + 'text/json': json.loads} + self._buckets = WeakValueDictionary() + + def _get_protocol(self): + return self._protocol + + def _set_protocol(self, value): + if value not in self.PROTOCOLS: + raise ValueError("protocol option is invalid, must be one of %s" % + repr(self.PROTOCOLS)) + self._protocol = value + + protocol = property(_get_protocol, _set_protocol, + doc="""Which protocol to prefer, one of PROTOCOLS""") + + def get_transport(self): + """ + Get the transport instance the client is using for it's + connection. DEPRECATED + """ + deprecated("get_transport is deprecated, use client, " + + "bucket, or object methods instead") + return None + + def get_client_id(self): + """ + Get the ``client_id`` for this ``RiakClient`` instance. + DEPRECATED + + :rtype: string + """ + deprecated( + "``get_client_id`` is deprecated, use the ``client_id`` property") + return self.client_id + + def set_client_id(self, client_id): + """ + Set the client_id for this ``RiakClient`` instance. + DEPRECATED + + :param client_id: The new client_id. + :type client_id: string + """ + deprecated( + "``set_client_id`` is deprecated, use the ``client_id`` property") + self.client_id = client_id + return self + + def _get_client_id(self): + with self._transport() as transport: + return transport.client_id + + def _set_client_id(self, client_id): + for http in self._http_pool: + http.client_id = client_id + for pb in self._pb_pool: + pb.client_id = client_id + + client_id = property(_get_client_id, _set_client_id, + doc="""The client ID for this client instance""") + + def get_encoder(self, content_type): + """ + Get the encoding function for the provided content type. + """ + return self._encoders.get(content_type) + + def set_encoder(self, content_type, encoder): + """ + Set the encoding function for the provided content type. + + :param encoder: + :type encoder: function + """ + self._encoders[content_type] = encoder + + def get_decoder(self, content_type): + """ + Get the decoding function for the provided content type. + """ + return self._decoders.get(content_type) + + def set_decoder(self, content_type, decoder): + """ + Set the decoding function for the provided content type. + + :param decoder: + :type decoder: function + """ + self._decoders[content_type] = decoder + + def bucket(self, name): + """ + Get the bucket by the specified name. Since buckets always exist, + this will always return a :class:`RiakBucket `. + + :rtype: :class:`RiakBucket ` + """ + if name in self._buckets: + return self._buckets[name] + else: + bucket = RiakBucket(self, name) + self._buckets[name] = bucket + return bucket + + @lazy_property + def solr(self): + """ + Returns a RiakSearch object which can access search indexes. + """ + return RiakSearch(self) + + def _create_node(self, n): + if isinstance(n, RiakNode): + return n + elif isinstance(n, tuple) and len(n) is 3: + host, http_port, pb_port = n + return RiakNode(host=host, + http_port=http_port, + pb_port=pb_port) + elif isinstance(n, dict): + return RiakNode(**n) + else: + raise TypeError("%s is not a valid node configuration" + % repr(n)) + + def _choose_node(self, nodes=None): + """ + Chooses a random node from the list of nodes in the client, + taking into account each node's recent error rate. + :rtype RiakNode + """ + if not nodes: + nodes = self.nodes + + # Prefer nodes which have gone a reasonable time without + # errors + def _error_rate(node): + return node.error_rate.value() + + good = [n for n in nodes if _error_rate(n) < 0.1] + + if len(good) is 0: + # Fall back to a minimally broken node + return min(nodes, key=_error_rate) + else: + return random.choice(good) diff --git a/riak/client/operations.py b/riak/client/operations.py new file mode 100644 index 00000000..e28c78e3 --- /dev/null +++ b/riak/client/operations.py @@ -0,0 +1,286 @@ +""" +Copyright 2012 Basho Technologies, Inc. + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" + +from transport import RiakClientTransport, retryable, retryableHttpOnly + + +class RiakClientOperations(RiakClientTransport): + """ + Methods for RiakClient that result in requests sent to the Riak + cluster. + + Note that all of these methods have an implicit 'transport' + argument that will be prepended automatically as part of the retry + logic, and does not need to be supplied by the user. + """ + + @retryable + def get_buckets(self, transport): + """ + Get the list of buckets as RiakBucket instances. + NOTE: Do not use this in production, as it requires traversing through + all keys stored in a cluster. + """ + return [self.bucket(name) for name in transport.get_buckets()] + + @retryable + def ping(self, transport): + """ + Check if the Riak server for this ``RiakClient`` instance is alive. + + :rtype: boolean + """ + return transport.ping() + + is_alive = ping + + @retryable + def get_index(self, transport, bucket, index, startkey, endkey=None): + """ + Queries a secondary index, returning matching keys. + + :param bucket: the bucket whose index will be queried + :type bucket: RiakBucket + :param index: the index to query + :type index: string + :param startkey: the sole key to query, or beginning of the query range + :type startkey: string, integer + :param endkey: the end of the query range (optional if equality) + :type endkey: string, integer + :rtype: list + """ + return transport.get_index(bucket, index, startkey, endkey) + + @retryable + def get_bucket_props(self, transport, bucket): + """ + Fetches bucket properties for the given bucket. + + :param bucket: the bucket whose properties will be fetched + :type bucket: RiakBucket + :rtype: dict + """ + return transport.get_bucket_props(bucket) + + @retryable + def set_bucket_props(self, transport, bucket, props): + """ + Sets bucket properties for the given bucket. + + :param bucket: the bucket whose properties will be set + :type bucket: RiakBucket + :param props: the properties to set + :type props: dict + """ + return transport.set_bucket_props(bucket, props) + + @retryable + def get_keys(self, transport, bucket): + """ + Lists all keys in a bucket. + + :param bucket: the bucket whose properties will be set + :type bucket: RiakBucket + :rtype: list + """ + return transport.get_keys(bucket) + + def stream_keys(self, bucket): + """ + Lists all keys in a bucket via a stream. This is a generator + method which should be iterated over. + + + :param bucket: the bucket whose properties will be set + :type bucket: RiakBucket + :rtype: iterator + """ + with self._transport() as transport: + stream = transport.stream_keys(bucket) + try: + for keylist in stream: + if len(keylist) > 0: + yield keylist + finally: + stream.close() + + @retryable + def put(self, transport, robj, w=None, dw=None, pw=None, return_body=None, + if_none_match=None): + """ + Stores an object in the Riak cluster. + + :param robj: the object to store + :type robj: RiakObject + :param w: the write quorum + :type w: integer, string, None + :param dw: the durable write quorum + :type dw: integer, string, None + :param pw: the primary write quorum + :type pw: integer, string, None + :param return_body: whether to return the resulting object + after the write + :type return_body: boolean + :param if_none_match: whether to fail the write if the object + exists + :type if_none_match: boolean + """ + return transport.put(robj, w=w, dw=dw, pw=pw, + return_body=return_body, + if_none_match=if_none_match) + + @retryable + def put_new(self, transport, robj, w=None, dw=None, pw=None, + return_body=None, if_none_match=None): + """ + Stores an object in the Riak cluster with a generated key. + + :param robj: the object to store + :type robj: RiakObject + :param w: the write quorum + :type w: integer, string, None + :param dw: the durable write quorum + :type dw: integer, string, None + :param pw: the primary write quorum + :type pw: integer, string, None + :param return_body: whether to return the resulting object + after the write + :type return_body: boolean + :param if_none_match: whether to fail the write if the object + exists + :type if_none_match: boolean + """ + return transport.put_new(robj, w=w, dw=dw, pw=pw, + return_body=return_body, + if_none_match=if_none_match) + + @retryable + def get(self, transport, robj, r=None, pr=None, vtag=None): + """ + Fetches the contents of a Riak object. + + :param robj: the object to fetch + :type robj: RiakObject + :param r: the read quorum + :type r: integer, string, None + :param pr: the primary read quorum + :type pr: integer, string, None + :param vtag: the specific sibling to fetch + :type vtag: string + """ + return transport.get(robj, r=r, pr=pr, vtag=vtag) + + @retryable + def delete(self, transport, robj, rw=None, r=None, w=None, dw=None, + pr=None, pw=None): + """ + Deletes an object from Riak. + + :param robj: the object to store + :type robj: RiakObject + :param rw: the read/write (delete) quorum + :type rw: integer, string, None + :param r: the read quorum + :type r: integer, string, None + :param pr: the primary read quorum + :type pr: integer, string, None + :param w: the write quorum + :type w: integer, string, None + :param dw: the durable write quorum + :type dw: integer, string, None + :param pw: the primary write quorum + :type pw: integer, string, None + """ + return transport.delete(robj, rw=rw, r=r, w=w, dw=dw, pr=pr, + pw=pw) + + @retryable + def mapred(self, transport, inputs, query, timeout): + """ + Executes a MapReduce query. + + :param inputs: the input list/structure + :type inputs: list, dict + :param query: the list of query phases + :type query: list + :param timeout: the query timeout + :type timeout: integer, None + :rtype: mixed + """ + return transport.mapred(inputs, query, timeout) + + def stream_mapred(self, inputs, query, timeout): + """ + Streams a MapReduce query as (phase, data) pairs. This is a + generator method which should be iterated over. + + :param inputs: the input list/structure + :type inputs: list, dict + :param query: the list of query phases + :type query: list + :param timeout: the query timeout + :type timeout: integer, None + :rtype: iterator + """ + with self._transport() as transport: + stream = transport.stream_mapred(inputs, query, timeout) + try: + for phase, data in stream: + yield phase, data + finally: + stream.close() + + @retryableHttpOnly + def fulltext_search(self, transport, index, query, **params): + """ + Performs a full-text search query. + + :param index: the bucket/index to search over + :type index: string + :param query: the search query + :type query: string + :param params: additional query flags + :type params: dict + """ + return transport.search(index, query, **params) + + @retryableHttpOnly + def fulltext_add(self, transport, index, docs): + """ + Adds documents to the full-text index. + + :param index: the bucket/index in which to index these docs + :type index: string + :param docs: the list of documents + :type docs: list + """ + transport.fulltext_add(index, docs) + + @retryableHttpOnly + def fulltext_delete(self, transport, index, docs=None, queries=None): + """ + Removes documents from the full-text index. + + :param index: the bucket/index from which to delete + :type index: string + :param docs: a list of documents (with ids) + :type docs: list + :param queries: a list of queries to match and delete + :type queries: list + """ + transport.fulltext_delete(index, docs, queries) diff --git a/riak/client/transport.py b/riak/client/transport.py new file mode 100644 index 00000000..5c42c55c --- /dev/null +++ b/riak/client/transport.py @@ -0,0 +1,129 @@ +""" +Copyright 2012 Basho Technologies, Inc. + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" +from contextlib import contextmanager +from riak.transports.pool import BadResource +from riak.transports.pbc import is_retryable as is_pbc_retryable +from riak.transports.http import is_retryable as is_http_retryable +import httplib + + +class RiakClientTransport(object): + """ + Methods for RiakClient related to transport selection and retries. + """ + + RETRY_COUNT = 3 + + # These will be set or redefined by the RiakClient initializer + protocol = 'http' + _http_pool = None + _pb_pool = None + + @contextmanager + def _transport(self): + """ + Yields a single transport to the caller from the default pool, + without retries. + """ + pool = self._choose_pool() + with pool.take() as transport: + yield transport + + def _with_retries(self, pool, fn): + """ + Performs the passed function with retries against the given pool. + + :param pool: the connection pool to use + :type pool: Pool + :param fn: the function to pass a transport + :type fn: function + """ + skip_nodes = [] + + def _skip_bad_nodes(transport): + return transport._node not in skip_nodes + + for retry in range(self.RETRY_COUNT): + try: + with pool.take(_filter=_skip_bad_nodes) as transport: + try: + return fn(transport) + except (IOError, httplib.HTTPException) as e: + if _is_retryable(e): + transport._node.error_rate.incr(1) + skip_nodes.append(transport._node) + raise BadResource(e) + else: + raise e + except BadResource: + continue + + def _choose_pool(self, protocol=None): + """ + Selects a connection pool according to the default protocol + and the passed one. + + :param protocol: the protocol to use + :type protocol: string + :rtype: Pool + """ + if not protocol: + protocol = self.protocol + if protocol in ['http', 'https']: + pool = self._http_pool + elif protocol == 'pbc': + pool = self._pb_pool + else: + raise ValueError("invalid protocol %s" % protocol) + return pool + + +def _is_retryable(error): + """ + Determines whether a given error is retryable according to the + exceptions allowed to be retried by each transport. + + :param error: the error to check + :type error: Exception + :rtype: boolean + """ + return is_pbc_retryable(error) or is_http_retryable(error) + + +def retryable(fn, protocol=None): + """ + Wraps a client operation that can be retried according to the set + RETRY_COUNT. Used internally. + """ + def wrapper(self, *args, **kwargs): + pool = self._choose_pool(protocol) + + def thunk(transport): + return fn(self, transport, *args, **kwargs) + + return self._with_retries(pool, thunk) + + return wrapper + + +def retryableHttpOnly(fn): + """ + Wraps a retryable client operation that is only valid over HTTP. + Used internally. + """ + return retryable(fn, protocol='http') diff --git a/riak/mapreduce.py b/riak/mapreduce.py index 7279b0a0..6031eeb2 100644 --- a/riak/mapreduce.py +++ b/riak/mapreduce.py @@ -18,21 +18,21 @@ under the License. """ import urllib -from riak_object import RiakObject -from bucket import RiakBucket from collections import Iterable class RiakMapReduce(object): """ The RiakMapReduce object allows you to build up and run a - map/reduce operation on Riak. + map/reduce operation on Riak. Most methods return the object on + which it was called, modified with new information, so you can + chain calls together to build the job. """ def __init__(self, client): """ Construct a Map/Reduce object. - @param RiakClient client - A RiakClient object. - @return RiakMapReduce + :param client: A RiakClient object. + :type client: RiakClient """ self._client = client self._phases = [] @@ -46,10 +46,15 @@ def add(self, arg1, arg2=None, arg3=None): different forms, depending on the provided inputs. You can 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 List or blank - @param mixed arg3 - Arg or blank - @return RiakMapReduce + + :param arg1: the object or bucket to add + :type arg1: RiakObject, string + :param arg2: a key or list of keys to add (if a bucket is + given in arg1) + :type arg2: string, list, None + :param arg3: key data for this input (must be convertible to JSON) + :type arg3: string, list, dict, None + :rtype: RiakMapReduce """ if (arg2 is None) and (arg3 is None): if isinstance(arg1, RiakObject): @@ -60,9 +65,27 @@ def add(self, arg1, arg2=None, arg3=None): return self.add_bucket_key_data(arg1, arg2, arg3) def add_object(self, obj): + """ + Adds a RiakObject to the inputs. + + :param obj: the object to add + :type obj: RiakObject + :rtype: RiakMapReduce + """ return self.add_bucket_key_data(obj._bucket._name, obj._key, None) def add_bucket_key_data(self, bucket, key, data): + """ + Adds a bucket/key/keydata triple to the inputs. + + :param bucket: the bucket + :type bucket: string + :param key: the key or list of keys + :type key: string + :param data: the key-specific data + :type data: string, list, dict, None + :rtype: RiakMapReduce + """ if self._input_mode == 'bucket': raise ValueError('Already added a bucket, can\'t add an object.') elif self._input_mode == 'query': @@ -77,11 +100,25 @@ def add_bucket_key_data(self, bucket, key, data): return self def add_bucket(self, bucket): + """ + Adds all keys in a bucket to the inputs. + + :param bucket: the bucket + :type bucket: string + :rtype: RiakMapReduce + """ self._input_mode = 'bucket' self._inputs = bucket return self def add_key_filters(self, key_filters): + """ + Adds key filters to the inputs. + + :param key_filters: a list of filters + :type key_filters: list + :rtype: RiakMapReduce + """ if self._input_mode == 'query': raise ValueError('Key filters are not supported in a query.') @@ -89,6 +126,13 @@ def add_key_filters(self, key_filters): return self def add_key_filter(self, *args): + """ + Add a single key filter to the inputs. + + :param args: a filter + :type args: list + :rtype: RiakMapReduce + """ if self._input_mode == 'query': raise ValueError('Key filters are not supported in a query.') @@ -99,8 +143,12 @@ def search(self, bucket, query): """ Begin a map/reduce operation using a Search. This command will return an error unless executed against a Riak Search cluster. - @param bucket - The bucket over which to perform the search. - @param query - The search query. + + :param bucket: The bucket over which to perform the search + :type bucket: string + :param query: The search query + :type query: string + :rtype: RiakMapReduce """ self._input_mode = 'query' self._inputs = {'module': 'riak_search', @@ -112,10 +160,16 @@ def index(self, bucket, index, startkey, endkey=None): """ Begin a map/reduce operation using a Secondary Index query. - @param bucket - The bucket over which to perform the search. - @param index - The index to use for query - @param startkey - The start key of index range - @param endkey - The end key of index range or blank + + :param bucket: The bucket over which to perform the query + :type bucket: string + :param index: The index to use for query + :type index: string + :param startkey: The start key of index range, or the + value which all entries must equal + :type startkey: string, integer + :param endkey: The end key of index range (if doing a range query) + :type endkey: string, integer, None """ self._input_mode = 'query' @@ -133,13 +187,17 @@ def index(self, bucket, index, startkey, endkey=None): def link(self, bucket='_', tag='_', keep=False): """ Add a link phase to the map/reduce operation. - @param string bucket - Bucket name (default '_', which means all + + :param bucket: Bucket name (default '_', which means all buckets) - @param string tag - Tag (default '_', which means any tag) - @param boolean keep - Flag whether to keep results from this - stage in the map/reduce. (default False, unless this is the last - step in the phase) - @return self + :type bucket: string + :param tag: Tag (default '_', which means any tag) + :type tag: string + :param keep: Flag whether to keep results from this stage in + the map/reduce. (default False, unless this is the last step + in the phase) + :type keep: boolean + :rtype: RiakMapReduce """ self._phases.append(RiakLinkPhase(bucket, tag, keep)) return self @@ -147,13 +205,16 @@ def link(self, bucket='_', tag='_', keep=False): def map(self, function, options=None): """ Add a map phase to the map/reduce operation. - @param mixed function - Either a named Javascript function (ie: - 'Riak.mapValues'), or an anonymous javascript function (ie: - 'function(...) ... ' or an array ['erlang_module', - 'function']. - @param array() options - An optional associative array - containing 'language', 'keep' flag, and/or 'arg'. - @return self + + :param function: Either a named Javascript function (ie: + 'Riak.mapValues'), or an anonymous javascript function (ie: + 'function(...) ... ' or an array ['erlang_module', + 'function']. + :type function: string, list + :param options: phase options, containing 'language', 'keep' + flag, and/or 'arg'. + :type options: dict + :rtype: RiakMapReduce """ if options is None: options = dict() @@ -173,12 +234,15 @@ def map(self, function, options=None): def reduce(self, function, options=None): """ Add a reduce phase to the map/reduce operation. - @param mixed function - Either a named Javascript function (ie. - 'Riak.mapValues'), or an anonymous javascript function(ie: - 'function(...) { ... }' or an array ['erlang_module', 'function']. - @param array() options - An optional associative array - containing 'language', 'keep' flag, and/or 'arg'. - @return self + + :param function: Either a named Javascript function (ie. + 'Riak.reduceSum'), or an anonymous javascript function(ie: + 'function(...) { ... }' or an array ['erlang_module', + 'function']. + :type function: string, list + :param options: phase options, containing 'language', 'keep' + flag, and/or 'arg'. + :rtype: RiakMapReduce """ if options is None: options = dict() @@ -197,17 +261,56 @@ def reduce(self, function, options=None): def run(self, timeout=None): """ - Run the map/reduce operation. Returns an array of results, or an - array of RiakLink objects if the last phase is a link phase. - @param integer timeout - Timeout in milliseconds. - @return array() + Run the map/reduce operation synchronously. Returns a list of + results, or a list of RiakLink objects if the last phase is a + link phase. + + :param timeout: Timeout in milliseconds + :type timeout: integer, None + :rtype: list """ + query, link_results_flag = self._normalize_query() + + result = self._client.mapred(self._inputs, query, timeout) + + # If the last phase is NOT a link phase, then return the result. + if not (link_results_flag + or isinstance(self._phases[-1], RiakLinkPhase)): + return result + + # If there are no results, then return an empty list. + if result == None: + return [] + + # Otherwise, if the last phase IS a link phase, then convert the + # results to RiakLink objects. + a = [] + for r in result: + if (len(r) == 2): + link = RiakLink(r[0], r[1]) + elif (len(r) == 3): + link = RiakLink(r[0], r[1], r[2]) + link._client = self._client + a.append(link) + + return a + + def stream(self, timeout=None): + """ + Streams the MapReduce query (returns an iterator). + + :param timeout: Timeout in milliseconds + :type timeout: integer + :rtype: iterator + """ + query, lrf = self._normalize_query() + return self._client.stream_mapred(self._inputs, query, timeout) + + def _normalize_query(self): num_phases = len(self._phases) - # If there are no phases, then just echo the inputs back to the user. - if (num_phases == 0): - self.reduce(["riak_kv_mapreduce", "reduce_identity"]) - num_phases = 1 + # If there are no phases, return the keys as links + if num_phases is 0: link_results_flag = True else: link_results_flag = False @@ -230,56 +333,84 @@ def run(self, timeout=None): if (type(self._inputs) == str): bucket_name = self._inputs elif (type(self._inputs) == RiakBucket): - bucket_name = self._inputs.get_name() + bucket_name = self._inputs.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) - - # If the last phase is NOT a link phase, then return the result. - if not (link_results_flag - or isinstance(self._phases[-1], RiakLinkPhase)): - return result - - # If there are no results, then return an empty list. - if result == None: - return [] - - # Otherwise, if the last phase IS a link phase, then convert the - # results to RiakLink objects. - a = [] - for r in result: - if (len(r) == 2): - link = RiakLink(r[0], r[1]) - elif (len(r) == 3): - link = RiakLink(r[0], r[1], r[2]) - link._client = self._client - a.append(link) - - return a + return query, link_results_flag ## # Start Shortcuts to built-ins ## def map_values(self, options=None): + """ + Adds the Javascript built-in ``Riak.mapValues`` to the query + as a map phase. + + :param options: phase options, containing 'language', 'keep' + flag, and/or 'arg'. + :type options: dict + """ return self.map("Riak.mapValues", options=options) def map_values_json(self, options=None): + """ + Adds the Javascript built-in ``Riak.mapValuesJson`` to the + query as a map phase. + + :param options: phase options, containing 'language', 'keep' + flag, and/or 'arg'. + :type options: dict + """ return self.map("Riak.mapValuesJson", options=options) def reduce_sum(self, options=None): + """ + Adds the Javascript built-in ``Riak.reduceSum`` to the query + as a reduce phase. + + :param options: phase options, containing 'language', 'keep' + flag, and/or 'arg'. + :type options: dict + """ return self.reduce("Riak.reduceSum", options=options) def reduce_min(self, options=None): + """ + Adds the Javascript built-in ``Riak.reduceMin`` to the query + as a reduce phase. + + :param options: phase options, containing 'language', 'keep' + flag, and/or 'arg'. + :type options: dict + """ return self.reduce("Riak.reduceMin", options=options) def reduce_max(self, options=None): + """ + Adds the Javascript built-in ``Riak.reduceMax`` to the query + as a reduce phase. + + :param options: phase options, containing 'language', 'keep' + flag, and/or 'arg'. + :type options: dict + """ return self.reduce("Riak.reduceMax", options=options) def reduce_sort(self, js_cmp=None, options=None): + """ + Adds the Javascript built-in ``Riak.reduceSort`` to the query + as a reduce phase. + + :param js_cmp: A Javascript comparator function as specified by + Array.sort() + :type js_cmp: string + :param options: phase options, containing 'language', 'keep' + flag, and/or 'arg'. + :type options: dict + """ if options is None: options = dict() @@ -289,9 +420,27 @@ def reduce_sort(self, js_cmp=None, options=None): return self.reduce("Riak.reduceSort", options=options) def reduce_numeric_sort(self, options=None): + """ + Adds the Javascript built-in ``Riak.reduceNumericSort`` to the + query as a reduce phase. + + :param options: phase options, containing 'language', 'keep' + flag, and/or 'arg'. + :type options: dict + """ return self.reduce("Riak.reduceNumericSort", options=options) def reduce_limit(self, limit, options=None): + """ + Adds the Javascript built-in ``Riak.reduceLimit`` to the query + as a reduce phase. + + :param limit: the maximum number of results to return + :type limit: integer + :param options: phase options, containing 'language', 'keep' + flag, and/or 'arg'. + :type options: dict + """ if options is None: options = dict() @@ -303,6 +452,18 @@ def reduce_limit(self, limit, options=None): return self.reduce(code, options=options) def reduce_slice(self, start, end, options=None): + """ + Adds the Javascript built-in ``Riak.reduceSlice`` to the + query as a reduce phase. + + :param start: the beginning of the slice + :type start: integer + :param end: the end of the slice + :type end: integer + :param options: phase options, containing 'language', 'keep' + flag, and/or 'arg'. + :type options: dict + """ if options is None: options = dict() @@ -310,25 +471,42 @@ def reduce_slice(self, start, end, options=None): return self.reduce("Riak.reduceSlice", options=options) def filter_not_found(self, options=None): + """ + Adds the Javascript built-in ``Riak.filterNotFound`` to the query + as a reduce phase. + + :param options: phase options, containing 'language', 'keep' + flag, and/or 'arg'. + :type options: dict + """ return self.reduce("Riak.filterNotFound", options=options) class RiakMapReducePhase(object): """ - The RiakMapReducePhase holds information about a Map phase or - Reduce phase in a RiakMapReduce operation. + The RiakMapReducePhase holds information about a Map or Reduce + phase in a RiakMapReduce operation. + + Normally you won't need to use this object directly, but instead + call methods on RiakMapReduce objects to add instances to the + query. """ def __init__(self, type, function, language, keep, arg): """ Construct a RiakMapReducePhase object. - @param string type - 'map'placeholder149'reduce' - @param mixed function - string or array(): - @param string language - 'javascript'placeholder149'erlang' - @param boolean keep - True to return the output of this phase in - the results. - @param mixed arg - Additional value to pass into the map or - reduce function. + + :param type: the phase type - 'map', 'reduce', 'link' + :type type: string + :param function: the function to execute + :type function: string, list + :param language: 'javascript' or 'erlang' + :type language: string + :param keep: whether to return the output of this phase in the results. + :type keep: boolean + :param arg: Additional static value to pass into the map or + reduce function. + :type arg: string, dict, list """ try: if isinstance(function, basestring): @@ -344,8 +522,10 @@ def __init__(self, type, function, language, keep, arg): def to_array(self): """ - Convert the RiakMapReducePhase to an associative array. Used - internally. + Convert the RiakMapReducePhase to a format that can be output + into JSON. Used internally. + + :rtype: dict """ stepdef = {'keep': self._keep, 'language': self._language, @@ -372,14 +552,21 @@ class RiakLinkPhase(object): """ The RiakLinkPhase object holds information about a Link phase in a map/reduce operation. + + Normally you won't need to use this object directly, but instead + call ``link`` on RiakMapReduce objects to add instances to the + query. """ def __init__(self, bucket, tag, keep): """ Construct a RiakLinkPhase object. - @param string bucket - The bucket name. - @param string tag - The tag. - @param boolean keep - True to return results of this phase. + :param bucket: - The bucket name + :type bucket: string + :param tag: The tag + :type tag: string + :param keep: whether to return results of this phase. + :type keep: boolean """ self._bucket = bucket self._tag = tag @@ -387,8 +574,8 @@ def __init__(self, bucket, tag, keep): def to_array(self): """ - Convert the RiakLinkPhase to an associative array. Used - internally. + Convert the RiakLinkPhase to a format that can be output into + JSON. Used internally. """ stepdef = {'bucket': self._bucket, 'tag': self._tag, @@ -405,9 +592,13 @@ class RiakLink(object): def __init__(self, bucket, key, tag=None): """ Construct a RiakLink object. - @param string bucket - The bucket name. - @param string key - The key. - @param string tag - The tag. + + :param bucket: the bucket name + :type bucket: string + :param key: the key + :type key: string + :param tag: the tag + :type tag: string """ self._bucket = bucket self._key = key @@ -417,31 +608,38 @@ def __init__(self, bucket, key, tag=None): def get(self, r=None): """ Retrieve the RiakObject to which this link points. - @param integer r - The R-value to use. - @return RiakObject + + :param r: the read quorum to use + :type r: string, integer + :rtype: RiakObject """ return self._client.bucket(self._bucket).get(self._key, r) def get_binary(self, r=None): """ Retrieve the RiakObject to which this link points, as a binary. - @param integer r - The R-value to use. - @return RiakObject + + :param r: the read quorum to use + :type r: string, integer + :rtype: RiakObject """ return self._client.bucket(self._bucket).get_binary(self._key, r) def get_bucket(self): """ Get the bucket name of this link. - @return string + + :rtype: string """ return self._bucket def set_bucket(self, name): """ Set the bucket name of this link. - @param string name - The bucket name. - @return self + + :param name: the bucket name + :type name: string + :rtype: RiakLink """ self._bucket = name return self @@ -449,15 +647,18 @@ def set_bucket(self, name): def get_key(self): """ Get the key of this link. - @return string + + :rtype: string """ return self._key def set_key(self, key): """ Set the key of this link. - @param string key - The key. - @return self + + :param key: the key + :type key: string + :rtype: RiakLink """ self._key = key return self @@ -465,7 +666,8 @@ def set_key(self, key): def get_tag(self): """ Get the tag of this link. - @return string + + :rtype: string """ if (self._tag is None): return self._bucket @@ -475,15 +677,20 @@ def get_tag(self): def set_tag(self, tag): """ Set the tag of this link. - @param string tag - The tag. - @return self + + :param tag: the tag + :type tag: string + :rtype: RiakLink """ self._tag = tag return self def to_link_header(self, client): """ - Convert this RiakLink object to a link header string. Used internally. + Convert this RiakLink object to a link header string. Used + internally. + + :rtype: string """ link = '' link += '= 2: - hostports = [(host, port), ] - self._cm = transport_class.default_cm(hostports) - self._transport = transport_class(self._cm, prefix="/solr") - else: - # The old code which attempted to use api==1 would actually - # throw a NameError, so it was obviously never used. We will - # simply raise an error here, intead of a gentle warning. - raise DeprecationWarning('please upgrade the transport to the ' - 'new API') + http://www.apache.org/licenses/LICENSE-2.0 - self._client = client - self._decoders = {"text/xml": ElementTree.fromstring} +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" - def get_decoder(self, content_type): - decoder = (self._client.get_decoder(content_type) - or self._decoders[content_type]) - if not decoder: - decoder = self.decode - return decoder +class RiakSearch(object): + """ + A wrapper around Riak Search-related client operations. See + :func:`RiakClient.solr`. + """ - def decode(self, data): - return data + def __init__(self, client, **unused_args): + self._client = client def add(self, index, *docs): - xml = Document() - root = xml.createElement('add') - for doc in docs: - doc_element = xml.createElement('doc') - for key, value in doc.iteritems(): - field = xml.createElement('field') - field.setAttribute("name", key) - text = xml.createTextNode(value) - field.appendChild(text) - doc_element.appendChild(field) - root.appendChild(doc_element) - xml.appendChild(root) - - url = "/solr/%s/update" % index - self._transport.post_request(uri=url, body=xml.toxml(), - content_type="text/xml") + """ + Adds documents to a fulltext index. Shortcut and backwards + compatibility for :func:`RiakClientOperations.fulltext_add`. + """ + self._client.fulltext_add(index, docs=docs) index = add def delete(self, index, docs=None, queries=None): - xml = Document() - root = xml.createElement('delete') - if docs: - for doc in docs: - doc_element = xml.createElement('id') - text = xml.createTextNode(doc) - doc_element.appendChild(text) - root.appendChild(doc_element) - if queries: - for query in queries: - query_element = xml.createElement('query') - text = xml.createTextNode(query) - query_element.appendChild(text) - root.appendChild(query_element) - - xml.appendChild(root) - - url = "/solr/%s/update" % index - self._transport.post_request(uri=url, body=xml.toxml(), - content_type="text/xml") + """ + Removes documents from a fulltext index. Shortcut and backwards + compatibility for :func:`RiakClientOperations.fulltext_delete`. + """ + self._client.fulltext_delete(index, docs=docs, queries=queries) remove = delete def search(self, index, query, **params): - return self._client._transport.search(index, query, **params) + """ + Searches a fulltext index. Shortcut and backwards + compatibility for :func:`RiakClientOperations.fulltext_search`. + """ + return self._client.fulltext_search(index, query, **params) select = search diff --git a/riak/test_server.py b/riak/test_server.py index 7c6348af..4010fb76 100644 --- a/riak/test_server.py +++ b/riak/test_server.py @@ -92,6 +92,11 @@ class TestServer(object): }, } + _temp_bin = None + _temp_etc = None + _temp_log = None + _temp_pipe = None + def __init__(self, tmp_dir="/tmp/riak/test_server", bin_dir=os.path.expanduser("~/.riak/install/riak-0.14.2/bin"), vm_args=None, **options): diff --git a/riak/tests/suite.py b/riak/tests/suite.py index af74197a..6f96c84d 100644 --- a/riak/tests/suite.py +++ b/riak/tests/suite.py @@ -1,4 +1,3 @@ -import riak.tests.test_server_test import os.path import platform diff --git a/riak/tests/test_all.py b/riak/tests/test_all.py index 5693d797..9433d41f 100644 --- a/riak/tests/test_all.py +++ b/riak/tests/test_all.py @@ -3,34 +3,29 @@ import os import random -import socket import platform if platform.python_version() < '2.7': unittest = __import__('unittest2') else: import unittest -import uuid -import time -from riak import RiakClient -from riak import RiakPbcTransport -from riak import RiakHttpTransport -from riak.mapreduce import RiakLink -from riak import RiakKeyFilter, key_filter +from riak.client import RiakClient +from riak.mapreduce import RiakLink, RiakKeyFilter +from riak import key_filter from riak.test_server import TestServer from riak.tests.test_search import SearchTests, \ EnableSearchTests, SolrSearchTests from riak.tests.test_mapreduce import MapReduceAliasTests, \ - ErlangMapReduceTests, JSMapReduceTests, LinkTests + ErlangMapReduceTests, JSMapReduceTests, LinkTests, MapReduceStreamTests from riak.tests.test_kv import BasicKVTests, KVFileTests, \ HTTPBucketPropsTest, PbcBucketPropsTest from riak.tests.test_2i import TwoITests try: - import riak_pb + __import__('riak_pb') HAVE_PROTO = True except ImportError: HAVE_PROTO = False @@ -56,16 +51,24 @@ class BaseTestCase(object): + host = None + pb_port = None + http_port = None + @staticmethod def randint(): return random.randint(1, 999999) - def create_client(self, host=None, port=None, transport_class=None): - host = host or self.host - port = port or self.port - transport_class = transport_class or self.transport_class - return RiakClient(self.host, self.port, - transport_class=self.transport_class) + def create_client(self, host=None, http_port=None, pb_port=None, + protocol=None, **client_args): + host = host or self.host or HOST + http_port = http_port or self.http_port or HTTP_PORT + pb_port = pb_port or self.pb_port or PB_PORT + protocol = protocol or self.protocol + return RiakClient(protocol=protocol, + host=host, + http_port=http_port, + pb_port=pb_port, **client_args) def setUp(self): self.client = self.create_client() @@ -86,6 +89,7 @@ class RiakPbcTransportTestCase(BasicKVTests, ErlangMapReduceTests, JSMapReduceTests, MapReduceAliasTests, + MapReduceStreamTests, SearchTests, BaseTestCase, unittest.TestCase): @@ -94,72 +98,19 @@ def setUp(self): if not HAVE_PROTO: self.skipTest('protobuf is unavailable') self.host = PB_HOST - self.port = PB_PORT - self.transport_class = RiakPbcTransport + self.pb_port = PB_PORT + self.protocol = 'pbc' super(RiakPbcTransportTestCase, self).setUp() def test_uses_client_id_if_given(self): - self.host = PB_HOST - self.port = PB_PORT zero_client_id = "\0\0\0\0" - c = RiakClient(PB_HOST, PB_PORT, - transport_class=RiakPbcTransport, - client_id=zero_client_id) - self.assertEqual(zero_client_id, c.get_client_id()) - - def test_close_underlying_socket_fails(self): - c = RiakClient(PB_HOST, PB_PORT, transport_class=RiakPbcTransport) - - bucket = c.bucket('bucket_test_close') - rand = self.randint() - obj = bucket.new('foo', rand) - obj.store() - obj = bucket.get('foo') - self.assertTrue(obj.exists) - self.assertEqual(obj.bucket.name, 'bucket_test_close') - self.assertEqual(obj.key, 'foo') - self.assertEqual(obj.data, rand) - - # Close the underlying socket. This gets a bit sketchy, - # since we are reaching into the internals, but there is - # no other way to get at the socket - conns = c._cm.conns - conns[0].sock.close() - - # This shoud fail with a socket error now - self.assertRaises(socket.error, bucket.get, 'foo') - - def test_close_underlying_socket_retry(self): - c = RiakClient(PB_HOST, PB_PORT, transport_class=RiakPbcTransport, - transport_options={"max_attempts": 2}) - - bucket = c.bucket('bucket_test_close') - rand = self.randint() - obj = bucket.new('barbaz', rand) - obj.store() - obj = bucket.get('barbaz') - self.assertTrue(obj.exists) - self.assertEqual(obj.bucket.name, 'bucket_test_close') - self.assertEqual(obj.key, 'barbaz') - self.assertEqual(obj.data, rand) - - # Close the underlying socket. This gets a bit sketchy, - # since we are reaching into the internals, but there is - # no other way to get at the socket - conns = c._cm.conns - conns[0].sock.close() - - # This should work, since we have a retry - obj = bucket.get('barbaz') - self.assertTrue(obj.exists) - self.assertEqual(obj.bucket.name, 'bucket_test_close') - self.assertEqual(obj.key, 'barbaz') - self.assertEqual(obj.data, rand) + c = self.create_client(client_id=zero_client_id) + self.assertEqual(zero_client_id, c.client_id) def test_bucket_search_enabled(self): with self.assertRaises(NotImplementedError): bucket = self.client.bucket("unsearch_bucket") - test = bucket.search_enabled() + bucket.search_enabled() def test_enable_search_commit_hook(self): with self.assertRaises(NotImplementedError): @@ -175,6 +126,7 @@ class RiakHttpTransportTestCase(BasicKVTests, ErlangMapReduceTests, JSMapReduceTests, MapReduceAliasTests, + MapReduceStreamTests, EnableSearchTests, SolrSearchTests, SearchTests, @@ -183,8 +135,8 @@ class RiakHttpTransportTestCase(BasicKVTests, def setUp(self): self.host = HTTP_HOST - self.port = HTTP_PORT - self.transport_class = RiakHttpTransport + self.http_port = HTTP_PORT + self.protocol = 'http' super(RiakHttpTransportTestCase, self).setUp() def test_no_returnbody(self): diff --git a/riak/tests/test_kv.py b/riak/tests/test_kv.py index b477f944..65862410 100644 --- a/riak/tests/test_kv.py +++ b/riak/tests/test_kv.py @@ -2,10 +2,7 @@ import os import cPickle import copy -try: - import json -except ImportError: - import simplejson as json +import json class NotJsonSerializable(object): @@ -75,6 +72,32 @@ def test_generate_key(self): self.assertNotIn(o.key, existing_keys) self.assertEqual(len(bucket.get_keys()), len(existing_keys) + 1) + def test_stream_keys(self): + bucket = self.client.bucket('random_key_bucket') + regular_keys = bucket.get_keys() + self.assertNotEqual(len(regular_keys), 0) + streamed_keys = [] + for keylist in bucket.stream_keys(): + self.assertNotEqual([], keylist) + for key in keylist: + self.assertIsInstance(key, basestring) + streamed_keys += keylist + self.assertEqual(sorted(regular_keys), sorted(streamed_keys)) + + def test_stream_keys_abort(self): + bucket = self.client.bucket('random_key_bucket') + regular_keys = bucket.get_keys() + self.assertNotEqual(len(regular_keys), 0) + try: + for keylist in bucket.stream_keys(): + raise RuntimeError("abort") + except RuntimeError: + pass + + # If the stream was closed correctly, this will not error + robj = bucket.get(regular_keys[0]) + self.assertEqual(True, robj.exists) + def test_binary_store_and_get(self): bucket = self.client.bucket('bucket') # Store as binary, retrieve as binary, then compare... @@ -242,7 +265,7 @@ def test_list_buckets(self): bucket = self.client.bucket("list_bucket") bucket.new("one", {"foo": "one", "bar": "red"}).store() buckets = self.client.get_buckets() - self.assertTrue("list_bucket" in buckets) + self.assertTrue(bucket in buckets) class HTTPBucketPropsTest(object): @@ -288,13 +311,13 @@ class PbcBucketPropsTest(object): def test_rw_settings(self): bucket = self.client.bucket('rwsettings') with self.assertRaises(NotImplementedError): - test = bucket.r + bucket.r with self.assertRaises(NotImplementedError): - test = bucket.w + bucket.w with self.assertRaises(NotImplementedError): - test = bucket.dw + bucket.dw with self.assertRaises(NotImplementedError): - test = bucket.rw + bucket.rw with self.assertRaises(NotImplementedError): bucket.r = 2 @@ -308,9 +331,9 @@ def test_rw_settings(self): def test_primary_quora(self): bucket = self.client.bucket('primary_quora') with self.assertRaises(NotImplementedError): - test = bucket.pr + bucket.pr with self.assertRaises(NotImplementedError): - test = bucket.pw + bucket.pw with self.assertRaises(NotImplementedError): bucket.pr = 2 @@ -321,7 +344,6 @@ def test_primary_quora(self): class KVFileTests(object): def test_store_binary_object_from_file(self): bucket = self.client.bucket('bucket') - rand = str(self.randint()) filepath = os.path.join(os.path.dirname(__file__), 'test_all.py') obj = bucket.new_binary_from_file('foo_from_file', filepath) obj.store() @@ -331,7 +353,6 @@ def test_store_binary_object_from_file(self): def test_store_binary_object_from_file_should_use_default_mimetype(self): bucket = self.client.bucket('bucket') - rand = str(self.randint()) filepath = os.path.join(os.path.dirname(os.path.abspath(__file__)), os.pardir, os.pardir, 'THANKS') obj = bucket.new_binary_from_file('foo_from_file', filepath) @@ -341,7 +362,6 @@ def test_store_binary_object_from_file_should_use_default_mimetype(self): def test_store_binary_object_from_file_should_fail_if_file_not_found(self): bucket = self.client.bucket('bucket') - rand = str(self.randint()) self.assertRaises(IOError, bucket.new_binary_from_file, 'not_found_from_file', 'FILE_NOT_FOUND') obj = bucket.get_binary('not_found_from_file') diff --git a/riak/tests/test_mapreduce.py b/riak/tests/test_mapreduce.py index 23df62e2..b6e2d1b2 100644 --- a/riak/tests/test_mapreduce.py +++ b/riak/tests/test_mapreduce.py @@ -1,7 +1,7 @@ # -*- coding: utf-8 -*- -from riak.mapreduce import RiakLink -from riak import RiakKeyFilter, key_filter +from riak.mapreduce import RiakLink, RiakMapReduce +from riak import key_filter class LinkTests(object): @@ -44,7 +44,7 @@ def test_set_links(self): def test_set_links_all_links(self): bucket = self.client.bucket("bucket") foo1 = bucket.new("foo", 1) - foo2 = bucket.new("foo2", 2).store() + bucket.new("foo2", 2).store() links = [RiakLink("bucket", "foo2")] foo1.set_links(links, True) links = foo1.get_links() @@ -103,7 +103,7 @@ def test_client_exceptional_paths(self): 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): @@ -499,3 +499,37 @@ def test_filter_not_found(self): .run() self.assertEqual(sorted(result), [1, 2]) + + +class MapReduceStreamTests(object): + def test_stream_results(self): + bucket = self.client.bucket('bucket') + bucket.new('one', data=1).store() + bucket.new('two', data=2).store() + + mr = RiakMapReduce(self.client).add('bucket', 'one')\ + .add('bucket', 'two') + mr.map_values_json() + results = [] + for phase, data in mr.stream(): + results.extend(data) + + self.assertEqual(sorted(results), [1, 2]) + + def test_stream_cleanoperationsup(self): + bucket = self.client.bucket('bucket') + bucket.new('one', data=1).store() + bucket.new('two', data=2).store() + + mr = RiakMapReduce(self.client).add('bucket', 'one')\ + .add('bucket', 'two') + mr.map_values_json() + try: + for phase, data in mr.stream(): + raise RuntimeError("woops") + except RuntimeError: + pass + + # This should not raise an exception + obj = bucket.get('one') + self.assertEqual(1, obj.data) diff --git a/riak/tests/test_pool.py b/riak/tests/test_pool.py index f57a2cd3..d552bcaa 100644 --- a/riak/tests/test_pool.py +++ b/riak/tests/test_pool.py @@ -69,7 +69,7 @@ def test_yields_same_object_in_serial_access(self): with pool.take() as element2: self.assertEqual(1, len(pool.elements)) - self.assertEqual([1, 2], element) + self.assertEqual([1, 2], element2) self.assertEqual(1, len(pool.elements)) @@ -93,8 +93,8 @@ def test_unlocks_when_exception_raised(self): """ pool = SimplePool() try: - with pool.take() as x: - with pool.take() as y: + with pool.take(): + with pool.take(): raise RuntimeError except: self.assertEqual(2, len(pool.elements)) @@ -111,7 +111,7 @@ def test_removes_bad_resource(self): self.assertEqual([1], element) element.append(2) try: - with pool.take() as baddie: + with pool.take(): raise BadResource except BadResource: self.assertEqual(0, len(pool.elements)) @@ -127,8 +127,8 @@ def filtereven(numlist): return numlist[0] % 2 == 0 pool = SimplePool() - with pool.take() as x: - with pool.take() as y: + with pool.take(): + with pool.take(): pass with pool.take(_filter=filtereven) as f: @@ -143,7 +143,7 @@ def test_requires_filter_to_be_callable(self): pool = SimplePool() with self.assertRaises(TypeError): - with pool.take(_filter=badfilter) as resource: + with pool.take(_filter=badfilter): pass def test_yields_default_when_empty(self): @@ -245,7 +245,7 @@ def test_clear(self): pool = SimplePool() def worker_run(): - with pool.take() as a: + with pool.take(): startq.put(1) startq.join() sleep(rand.uniform(0, 0.5)) diff --git a/riak/tests/test_search.py b/riak/tests/test_search.py index f3223da0..a56ec63c 100644 --- a/riak/tests/test_search.py +++ b/riak/tests/test_search.py @@ -35,61 +35,55 @@ def test_disable_search_commit_hook(self): class SolrSearchTests(object): @unittest.skipIf(SKIP_SEARCH, 'SKIP_SEARCH is defined') def test_add_document_to_index(self): - self.client.solr().add("searchbucket", + self.client.solr.add("searchbucket", {"id": "doc", "username": "tony"}) - results = self.client.solr().search("searchbucket", "username:tony") + results = self.client.solr.search("searchbucket", "username:tony") self.assertEquals("tony", results['docs'][0]['username']) @unittest.skipIf(SKIP_SEARCH, 'SKIP_SEARCH is defined') def test_add_multiple_documents_to_index(self): - self.client.solr().add("searchbucket", + self.client.solr.add("searchbucket", {"id": "dizzy", "username": "dizzy"}, {"id": "russell", "username": "russell"}) - results = self.client.solr()\ + results = self.client.solr\ .search("searchbucket", "username:russell OR username:dizzy") self.assertEquals(2, len(results['docs'])) @unittest.skipIf(SKIP_SEARCH, 'SKIP_SEARCH is defined') def test_delete_documents_from_search_by_id(self): - self.client.solr().add("searchbucket", + self.client.solr.add("searchbucket", {"id": "dizzy", "username": "dizzy"}, {"id": "russell", "username": "russell"}) - self.client.solr().delete("searchbucket", docs=["dizzy"]) - results = self.client.solr()\ + self.client.solr.delete("searchbucket", docs=["dizzy"]) + results = self.client.solr\ .search("searchbucket", "username:russell OR username:dizzy") self.assertEquals(1, len(results['docs'])) @unittest.skipIf(SKIP_SEARCH, 'SKIP_SEARCH is defined') def test_delete_documents_from_search_by_query(self): - self.client.solr().add("searchbucket", + self.client.solr.add("searchbucket", {"id": "dizzy", "username": "dizzy"}, {"id": "russell", "username": "russell"}) - self.client.solr()\ + self.client.solr\ .delete("searchbucket", queries=["username:dizzy", "username:russell"]) - results = self.client.solr()\ + results = self.client.solr\ .search("searchbucket", "username:russell OR username:dizzy") self.assertEquals(0, len(results['docs'])) @unittest.skipIf(SKIP_SEARCH, 'SKIP_SEARCH is defined') def test_delete_documents_from_search_by_query_and_id(self): - self.client.solr().add("searchbucket", + self.client.solr.add("searchbucket", {"id": "dizzy", "username": "dizzy"}, {"id": "russell", "username": "russell"}) - self.client.solr().delete("searchbucket", + self.client.solr.delete("searchbucket", docs=["dizzy"], queries=["username:russell"]) - results = self.client.solr()\ + results = self.client.solr\ .search("searchbucket", "username:russell OR username:dizzy") self.assertEquals(0, len(results['docs'])) - def test_build_rest_path_excludes_empty_query_params(self): - self.assertEquals( - self.client.get_transport().build_rest_path( - bucket=self.client.bucket("foo"), - key="bar", params={'r': None}), "/riak/foo/bar?") - class SearchTests(object): @unittest.skipIf(SKIP_SEARCH, 'SKIP_SEARCH is defined') @@ -110,16 +104,16 @@ def test_solr_search_with_params_from_bucket(self): def test_solr_search_with_params(self): bucket = self.client.bucket('searchbucket') bucket.new("user", {"username": "roidrage"}).store() - results = self.client.solr().search("searchbucket", - "username:roidrage", wt="xml") + results = self.client.solr.search("searchbucket", + "username:roidrage", wt="xml") self.assertEquals(1, len(results['docs'])) @unittest.skipIf(SKIP_SEARCH, 'SKIP_SEARCH is defined') def test_solr_search(self): bucket = self.client.bucket('searchbucket') bucket.new("user", {"username": "roidrage"}).store() - results = self.client.solr().search("searchbucket", - "username:roidrage") + results = self.client.solr.search("searchbucket", + "username:roidrage") self.assertEquals(1, len(results["docs"])) @unittest.skipIf(SKIP_SEARCH, 'SKIP_SEARCH is defined') @@ -133,7 +127,7 @@ def test_search_integration(self): bucket.new("five", {"foo": "five", "bar": "yellow"}).store() # Run some operations... - results = self.client.solr().search("searchbucket", + results = self.client.solr.search("searchbucket", "foo:one OR foo:two") if (len(results) == 0): print "\n\nNot running test \"testSearchIntegration()\".\n" @@ -144,5 +138,5 @@ def test_search_integration(self): self.assertEqual(len(results['docs']), 2) query = "(foo:one OR foo:two OR foo:three OR foo:four) AND\ (NOT bar:green)" - results = self.client.solr().search("searchbucket", query) + results = self.client.solr.search("searchbucket", query) self.assertEqual(len(results['docs']), 3) diff --git a/riak/transports/__init__.py b/riak/transports/__init__.py index 9c59adb6..e69de29b 100644 --- a/riak/transports/__init__.py +++ b/riak/transports/__init__.py @@ -1,2 +0,0 @@ -from http import RiakHttpTransport -from pbc import RiakPbcTransport diff --git a/riak/transports/connection.py b/riak/transports/connection.py deleted file mode 100644 index d26b85f4..00000000 --- a/riak/transports/connection.py +++ /dev/null @@ -1,192 +0,0 @@ -""" -Copyright 2011 Greg Stein - -This file is provided to you under the Apache License, -Version 2.0 (the "License"); you may not use this file -except in compliance with the License. You may obtain -a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, -software distributed under the License is distributed on an -"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -KIND, either express or implied. See the License for the -specific language governing permissions and limitations -under the License. -""" - -import httplib -import socket -import contextlib -import functools - - -class ConnectionManager(object): - - # Must be constructable with: connection_class(host, port) - # Must have two attribute: host and port - # Must have a close() method - connection_class = None - - def __init__(self, hostports=[]): - # We want a private copy of this list: either to detach the argument - # default, or to detach from the caller's list. - self.hostports = hostports[:] - - # Open a connection to each specified host/port. On single-threaded - # systems, this will create a round-robin across all specified servers. - # When multi-threaded, this will give us an initial set for all the - # threads to work with (and more will be created, according to demand). - self.conns = [self.connection_class(host, port) - for host, port in hostports] - - def add_hostport(self, host, port): - self.hostports.append((host, port)) - - # Open an initial connection. For single-threaded, this adds to the - # round-robin pool. On multi-threaded, it simply gives us an extra - # connectiong for the load-balancing across the servers. - self.conns.append(self.connection_class(host, port)) - - def remove_host(self, host, port=None): - if port is None: - self.hostports = [(h, p) for h, p in self.hostports - if h != host] - else: - self.hostports.remove((host, port)) - - # Now that the host/port pair has been removed from self.hostports, - # no connections on this pair will be added in .giveback(). Thus, the - # existing connections are all that may exist at this time. We'll - # snapshot the list, and look for offending connections, then try and - # remove them, being wary that race conditions may remove them before - # we can remove it. - for conn in self.conns[:]: - if conn.host == host and (port is None or conn.port == port): - try: - self.conns.remove(conn) - except ValueError: - # Another thread removed the connection. It won't - # be coming back, so we have nothing to do here. - pass - else: - # If the connection was still present (no - # ValueError), then we should go ahead and close - # it down. - conn.close() - - # Just in case somebody uses a host/port combo and typos... - remove_hostport = remove_host - - def take(self): - if len(self.conns) == 0: - # RACE: in a multi-threaded environment, a conn might arrive in - # self.conns, but... no biggy. If we're bouncing up against - # needing a new connection, then we'll just create one. - return self._new_connection() - - # RACE: self.conns might empty out right now, so we need to protect - # our access to it. - try: - # round-robin: take from the front, we'll append when it comes back - return self.conns.pop(0) - except IndexError: - return self._new_connection() - - def giveback(self, conn): - # Connections using a host/port pair that is NOT in self.hostports - # should be ignored. Likely, remove_host() was called while this - # connection was borrowed for some work. - if (conn.host, conn.port) in self.hostports: - self.conns.append(conn) - else: - # Proactively close the connection. The caller won't know whether - # we put it into our list, or left the connection for the caller - # to deal with (and close). - conn.close() - - @contextlib.contextmanager - def withconn(self): - conn = self.take() - try: - yield conn - finally: - self.giveback(conn) - - def _new_connection(self): - if len(self.hostports) == 0: - raise NoHostsDefined() - - # Grab the first host/port combo. We'll put this at the end, so that - # we do a round-robin on the host/port pairs. - host, port = self.hostports[0] - conn = self.connection_class(host, port) - - if len(self.hostports) == 1: - # No rotation needed. - return conn - - # Be careful about rotating. We want to append before - # removing, so that we never hit a len==0 race condition - # (which could prevent the creation of needed connections). - self.hostports.append((host, port)) - - # RACE: another thread may have appended the same host/port - # pair. We will add another pair. Each thread will remove - # one (either [0], or one that had been appened), resulting - # in a correct state of a single pair in the list. - # RACE: another thread may get the host/port pair from - # hostports[0] before we have a chance to remove it. We - # don't need precision round-robin behavior; just something - # close. - # RACE: another thread may have removed hostports[0] (which we - # are also trying to remove), but it will have placed - # another copy at the end before doing so. We have also - # added a host/port pair, and will remove one, leaving the - # list in a correct state. - self.hostports.remove((host, port)) - - return conn - - -class Socket(object): - - def __init__(self, host, port): - self.host = host - self.port = port - - self.sock = None - - def maybe_connect(self): - if self.sock is None: - self.sock = s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - - try: - s.connect((self.host, self.port)) - except: - self.close() - raise - - def close(self): - if self.sock is not None: - self.sock.close() - self.sock = None - - -class FactoryConnectionManager(ConnectionManager): - - def __init__(self, connection_class, hostports=[]): - self.connection_class = connection_class - ConnectionManager.__init__(self, hostports) - - -def cm_using(connection_class): - return functools.partial(FactoryConnectionManager, connection_class) - -HTTPConnectionManager = cm_using(httplib.HTTPConnection) -SocketConnectionManager = cm_using(Socket) - - -class NoHostsDefined(Exception): - pass diff --git a/riak/transports/http/__init__.py b/riak/transports/http/__init__.py new file mode 100644 index 00000000..21569113 --- /dev/null +++ b/riak/transports/http/__init__.py @@ -0,0 +1,68 @@ +""" +Copyright 2010 Rusty Klophaus +Copyright 2010 Justin Sheehy +Copyright 2009 Jay Baird + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" + +import httplib +from riak.transports.pool import Pool +from riak.transports.http.transport import RiakHttpTransport + + +class RiakHttpPool(Pool): + """ + A pool of HTTP(S) transport connections. + """ + def __init__(self, client, **options): + self.client = client + self.options = options + if client.protocol == 'https': + self.connection_class = httplib.HTTPSConnection + else: + self.connection_class = httplib.HTTPConnection + super(RiakHttpPool, self).__init__() + + def create_resource(self): + node = self.client._choose_node() + return RiakHttpTransport(node=node, + client=self.client, + connection_class=self.connection_class, + **self.options) + + def destroy_resource(self, transport): + transport.close() + + +CONN_CLOSED_ERRORS = ( + httplib.NotConnected, + httplib.IncompleteRead, + httplib.ImproperConnectionState + ) + + +def is_retryable(err): + """ + Determines if the given exception is something that is + network/socket-related and should thus cause the HTTP connection + to close and the operation retried on another node. + + :rtype: boolean + """ + for errtype in CONN_CLOSED_ERRORS: + if isinstance(err, errtype): + return True + return False diff --git a/riak/transports/http/connection.py b/riak/transports/http/connection.py new file mode 100644 index 00000000..3262b8d3 --- /dev/null +++ b/riak/transports/http/connection.py @@ -0,0 +1,69 @@ +""" +Copyright 2012 Basho Technologies, Inc. + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" + +import httplib + + +class RiakHttpConnection(object): + """ + Connection and low-level request methods for RiakHttpTransport. + """ + + def _request(self, method, uri, headers={}, body='', stream=False): + """ + Given a Method, URL, Headers, and Body, perform and HTTP request, + and return a 2-tuple containing a dictionary of response headers + and the response body. + """ + response = None + try: + self._connection.request(method, uri, body, headers) + response = self._connection.getresponse() + + response_headers = {'http_code': response.status} + for (key, value) in response.getheaders(): + response_headers[key.lower()] = value + + if stream: + # The caller is responsible for fully reading the + # response and closing it when streaming. + response_body = response + else: + response_body = response.read() + finally: + if response and not stream: + response.close() + + return response_headers, response_body + + def _connect(self): + self._connection = self._connection_class(self._node.host, + self._node.http_port) + + def close(self): + """ + Closes the underlying HTTP connection. + """ + try: + self._connection.close() + except httplib.NotConnected: + pass + + # These are set by the RiakHttpTransport initializer + _connection_class = httplib.HTTPConnection + _node = None diff --git a/riak/transports/http/resources.py b/riak/transports/http/resources.py new file mode 100644 index 00000000..b551cbae --- /dev/null +++ b/riak/transports/http/resources.py @@ -0,0 +1,178 @@ +""" +Copyright 2012 Basho Technologies, Inc. + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" + +import re +from urllib import quote_plus, urlencode +from riak import RiakError +from riak.util import lazy_property + + +class RiakHttpResources(object): + """ + Methods for RiakHttpTransport related to URL generation, i.e. + creating the proper paths. + """ + + def ping_path(self): + return mkpath(self.riak_kv_wm_ping) + + def stats_path(self): + return mkpath(self.riak_kv_wm_stats) + + def mapred_path(self, **options): + return mkpath(self.riak_kv_wm_mapred, **options) + + def bucket_list_path(self, **options): + query = options.copy() + query.update(buckets=True) + if self.riak_kv_wm_buckets: + return mkpath(self.riak_kv_wm_buckets, **query) + else: + return mkpath(self.riak_kv_wm_raw, **query) + + def bucket_properties_path(self, bucket, **options): + if self.riak_kv_wm_buckets: + return mkpath(self.riak_kv_wm_buckets, quote_plus(bucket), + "props", **options) + else: + query = options.copy() + query.update(props=True, keys=True) + return mkpath(self.riak_kv_wm_raw, quote_plus(bucket), **query) + + def key_list_path(self, bucket, **options): + query = {'keys': True, 'props': False} + query.update(options) + if self.riak_kv_wm_buckets: + return mkpath(self.riak_kv_wm_buckets, quote_plus(bucket), "keys", + **query) + else: + return mkpath(self.riak_kv_wm_raw, quote_plus(bucket), **query) + + def object_path(self, bucket, key=None, **options): + if key: + key = quote_plus(key) + if self.riak_kv_wm_buckets: + return mkpath(self.riak_kv_wm_buckets, quote_plus(bucket), "keys", + key, **options) + else: + return mkpath(self.riak_kv_wm_raw, quote_plus(bucket), key, + **options) + + # TODO: link_walk_path is undefined here because there is no path + # to it in the client without using MapReduce. + + def index_path(self, bucket, index, start, finish=None, **options): + if not self.riak_kv_wm_buckets: + raise RiakError("Indexes are unsupported by this Riak node") + if finish: + finish = quote_plus(str(finish)) + return mkpath(self.riak_kv_wm_buckets, quote_plus(bucket), + "index", quote_plus(index), quote_plus(str(start)), + finish, **options) + + def solr_select_path(self, index, query, **options): + if not self.riak_solr_searcher_wm: + raise RiakError("Riak Search is unsupported by this Riak node") + qs = {'q': query, 'wt': 'json'} + qs.update(options) + if index: + index = quote_plus(index) + return mkpath(self.riak_solr_searcher_wm, index, "select", **qs) + + def solr_update_path(self, index): + if not self.riak_solr_searcher_wm: + raise RiakError("Riak Search is unsupported by this Riak node") + if index: + index = quote_plus(index) + return mkpath(self.riak_solr_indexer_wm, index, "update") + + def luwak_path(self, key=None): + if not self.luwak_wm_file: + raise RiakError("Luwak is unsupported by this Riak node") + if key: + key = quote_plus(key) + return mkpath(self.luwak_wm_file, key) + + @lazy_property + def riak_kv_wm_buckets(self): + return self.resources.get('riak_kv_wm_index') + + @lazy_property + def riak_kv_wm_raw(self): + return self.resources.get('riak_kv_wm_raw') or "/riak" + + @lazy_property + def riak_kv_wm_link_walker(self): + return self.resources.get('riak_kv_wm_linkwalker') or "/riak" + + @lazy_property + def riak_kv_wm_mapred(self): + return self.resources.get('riak_kv_wm_mapred') or "/mapred" + + @lazy_property + def riak_kv_wm_ping(self): + return self.resources.get('riak_kv_wm_ping') or "/ping" + + @lazy_property + def riak_kv_wm_stats(self): + return self.resources.get('riak_kv_wm_stats') or "/stats" + + @lazy_property + def riak_solr_searcher_wm(self): + return self.resources.get('riak_solr_searcher_wm') + + @lazy_property + def riak_solr_indexer_wm(self): + return self.resources.get('riak_solr_indexer_wm') + + @lazy_property + def luwak_wm_file(self): + return self.resources.get('luwak_wm_file') + + @lazy_property + def resources(self): + return self.get_resources() + + +def mkpath(*segments, **query): + """ + Constructs the path & query portion of a URI from path segments + and a dict. + """ + # Remove empty segments (e.g. no key specified) + segments = [s for s in segments if s is not None] + # Join the segments into a path + pathstring = '/'.join(segments) + # Remove extra slashes + pathstring = re.sub('/+', '/', pathstring) + + # Add the query string if it exists + _query = {} + for key in query: + if query[key] in [False, True]: + _query[key] = str(query[key]).lower() + elif query[key] is not None: + _query[key] = query[key] + + if len(_query) > 0: + pathstring += "?" + urlencode(_query) + + if not pathstring.startswith('/'): + pathstring = '/' + pathstring + + return pathstring diff --git a/riak/transports/http/search.py b/riak/transports/http/search.py new file mode 100644 index 00000000..4e6c69e6 --- /dev/null +++ b/riak/transports/http/search.py @@ -0,0 +1,48 @@ +class XMLSearchResult(object): + # Match tags that are document fields + fieldtags = ['str', 'int', 'date'] + + def __init__(self): + # Results + self.num_found = 0 + self.max_score = 0.0 + self.docs = [] + + # Parser state + self.currdoc = None + self.currfield = None + self.currvalue = None + + def start(self, tag, attrib): + if tag == 'result': + self.num_found = int(attrib['numFound']) + self.max_score = float(attrib['maxScore']) + elif tag == 'doc': + self.currdoc = {} + elif tag in self.fieldtags and self.currdoc is not None: + self.currfield = attrib['name'] + + def end(self, tag): + if tag == 'doc' and self.currdoc is not None: + self.docs.append(self.currdoc) + self.currdoc = None + elif tag in self.fieldtags and self.currdoc is not None: + if tag == 'int': + self.currvalue = int(self.currvalue) + self.currdoc[self.currfield] = self.currvalue + self.currfield = None + self.currvalue = None + + def data(self, data): + if self.currfield: + # riak_solr_output adds NL + 6 spaces + data = data.rstrip() + if self.currvalue: + self.currvalue += data + else: + self.currvalue = data + + def close(self): + return {'num_found': self.num_found, + 'max_score': self.max_score, + 'docs': self.docs} diff --git a/riak/transports/http/stream.py b/riak/transports/http/stream.py new file mode 100644 index 00000000..2f1620f9 --- /dev/null +++ b/riak/transports/http/stream.py @@ -0,0 +1,125 @@ +""" +Copyright 2012 Basho Technologies, Inc. + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" + +import json +import string +import re +from cgi import parse_header +from email import message_from_string + + +class RiakHttpStream(object): + """ + Base class for HTTP streaming iterators. + """ + + BLOCK_SIZE = 2048 + + def __init__(self, response): + self.response = response + self.buffer = '' + self.response_done = False + + def __iter__(self): + return self + + def _read(self): + chunk = self.response.read(self.BLOCK_SIZE) + if chunk == '': + self.response_done = True + self.buffer += chunk + + def next(self): + raise NotImplementedError + + def close(self): + pass + + +class RiakHttpKeyStream(RiakHttpStream): + """ + Streaming iterator for list-keys over HTTP + """ + + def next(self): + while '}' not in self.buffer and not self.response_done: + self._read() + + if '}' in self.buffer: + idx = string.index(self.buffer, '}') + 1 + chunk = self.buffer[:idx] + self.buffer = self.buffer[idx:] + keys = json.loads(chunk)[u'keys'] + return keys + else: + raise StopIteration + + +class RiakHttpMultipartStream(RiakHttpStream): + """ + Streaming iterator for multipart messages over HTTP + """ + def __init__(self, response): + super(RiakHttpMultipartStream, self).__init__(response) + ctypehdr = response.getheader('content-type') + _, params = parse_header(ctypehdr) + self.boundary_re = re.compile('\r?\n--%s(?:--)?\r?\n' % + re.escape(params['boundary'])) + self.next_boundary = None + self.seen_first = False + + def next(self): + # multipart/mixed starts with a boundary, then the first part. + if not self.seen_first: + self.read_until_boundary() + self.advance_buffer() + self.seen_first = True + + self.read_until_boundary() + + if self.next_boundary: + part = self.advance_buffer() + message = message_from_string(part) + return message + else: + raise StopIteration + + def try_match(self): + self.next_boundary = self.boundary_re.search(self.buffer) + return self.next_boundary + + def advance_buffer(self): + part = self.buffer[:self.next_boundary.start()] + self.buffer = self.buffer[self.next_boundary.end():] + self.next_boundary = None + return part + + def read_until_boundary(self): + while not self.try_match() and not self.response_done: + self._read() + + +class RiakHttpMapReduceStream(RiakHttpMultipartStream): + """ + Streaming iterator for MapReduce over HTTP + """ + + def next(self): + message = super(RiakHttpMapReduceStream, self).next() + payload = json.loads(message.get_payload()) + return payload['phase'], payload['data'] diff --git a/riak/transports/http.py b/riak/transports/http/transport.py similarity index 60% rename from riak/transports/http.py rename to riak/transports/http/transport.py index 58ac88b4..9d8269d8 100644 --- a/riak/transports/http.py +++ b/riak/transports/http/transport.py @@ -1,4 +1,5 @@ """ +Copyright 2012 Basho Technologies, Inc. Copyright 2010 Rusty Klophaus Copyright 2010 Justin Sheehy Copyright 2009 Jay Baird @@ -17,96 +18,79 @@ specific language governing permissions and limitations under the License. """ -from __future__ import with_statement +import json import urllib import re import csv -from cStringIO import StringIO import httplib -import socket -import errno -try: - import json -except ImportError: - import simplejson as json - -from transport import RiakTransport -from riak.metadata import * +from riak.transports.transport import RiakTransport +from riak.transports.http.resources import RiakHttpResources +from riak.transports.http.connection import RiakHttpConnection +from riak.transports.http.search import XMLSearchResult +from riak.transports.http.stream import ( + RiakHttpKeyStream, + RiakHttpMapReduceStream + ) +from riak.metadata import ( + MD_CHARSET, + MD_CTYPE, + MD_INDEX, + MD_LASTMOD, + MD_LINKS, + MD_USERMETA, + MD_VTAG, + MD_DELETED + ) from riak.mapreduce import RiakLink from riak import RiakError from riak.riak_index_entry import RiakIndexEntry from riak.multidict import MultiDict -from connection import HTTPConnectionManager -import riak.util from xml.etree import ElementTree +from xml.dom.minidom import Document + # subtract length of "Link: " header string and newline MAX_LINK_HEADER_SIZE = 8192 - 8 -class RiakHttpTransport(RiakTransport): +class RiakHttpTransport(RiakHttpConnection, RiakHttpResources, RiakTransport): """ The RiakHttpTransport object holds information necessary to - connect to Riak. The Riak API uses HTTP, so there is no persistent - connection, and the RiakClient object is extremely lightweight. + connect to Riak via HTTP. """ - # We're using the new RiakTransport API - api = 2 - - # The ConnectionManager class that this transport prefers. - default_cm = HTTPConnectionManager - - # How many times to retry a request - RETRY_COUNT = 3 - - def __init__(self, cm, - prefix='riak', mapred_prefix='mapred', client_id=None, + def __init__(self, node=None, + client=None, + connection_class=httplib.HTTPConnection, + client_id=None, **unused_options): """ - Construct a new RiakClient object. - @param string host - Hostname or IP address (default '127.0.0.1') - @param int port - Port number (default 8098) - @param string prefix - Interface prefix (default 'riak') - @param string mapred_prefix - MapReduce prefix (default 'mapred') - @param string client_id - client id to use for vector clocks + Construct a new HTTP connection to Riak. """ super(RiakHttpTransport, self).__init__() - self._conns = cm - self._prefix = prefix - self._mapred_prefix = mapred_prefix + + self._client = client + self._node = node + self._connection_class = connection_class self._client_id = client_id if not self._client_id: self._client_id = self.make_random_client_id() - - def __copy__(self): - ### not implemented right now - raise Exception('not implemented') - ### we don't have _host and _port. will fix after some refactoring... - return RiakHttpTransport(self._host, self._port, self._prefix, - self._mapred_prefix) - - def set_client_id(self, client_id): - self._client_id = client_id - - def get_client_id(self): - return self._client_id + self._connect() def ping(self): """ Check server is alive over HTTP """ - response = self.http_request('GET', '/ping') + response = self._request('GET', self.ping_path()) return(response is not None) and (response[1] == 'OK') def stats(self): """ Gets performance statistics and server information """ - # TODO: use resource detection - response = self.http_request('GET', '/stats', - {'Accept': 'application/json'}) + response = self._request('GET', self.stats_path(), + {'Accept': 'application/json'}) if response[0]['http_code'] is 200: return json.loads(response[1]) else: @@ -120,7 +104,7 @@ def _server_version(self): # If stats is disabled, we can't assume the Riak version # is >= 1.1. However, we can assume the new URL scheme is # at least version 1.0 - elif 'riak_kv_wm_buckets' in self.get_resources(): + elif self.riak_kv_wm_buckets: return "1.0.0" else: return "0.14.0" @@ -130,8 +114,7 @@ def get_resources(self): Gets a JSON mapping of server-side resource names to paths :rtype dict """ - response = self.http_request('GET', '/', - {'Accept': 'application/json'}) + response = self._request('GET', '/', {'Accept': 'application/json'}) if response[0]['http_code'] is 200: return json.loads(response[1]) else: @@ -143,12 +126,9 @@ def get(self, robj, r=None, pr=None, vtag=None): """ # We could detect quorum_controls here but HTTP ignores # unknown flags/params. - params = {'r': r, 'pr': pr} - if vtag is not None: - params['vtag'] = vtag - url = self.build_rest_path(robj.bucket, robj.key, - params=params) - response = self.http_request('GET', url) + params = {'r': r, 'pr': pr, 'vtag': vtag} + url = self.object_path(robj.bucket.name, robj.key, **params) + response = self._request('GET', url) return self.parse_body(response, [200, 300, 404]) def put(self, robj, w=None, dw=None, pw=None, return_body=True, @@ -158,12 +138,10 @@ def put(self, robj, w=None, dw=None, pw=None, return_body=True, """ # We could detect quorum_controls here but HTTP ignores # unknown flags/params. - params = {'returnbody': str(return_body).lower(), - 'w': w, 'dw': dw, 'pw': pw} - url = self.build_rest_path(bucket=robj.bucket, - key=robj.key, - params=params) + params = {'returnbody': return_body, 'w': w, 'dw': dw, 'pw': pw} + url = self.object_path(robj.bucket.name, robj.key, **params) headers = self.build_put_headers(robj) + # TODO: use a more general 'prevent_stale_writes' semantics, # which is a superset of the if_none_match semantics. if if_none_match: @@ -174,9 +152,9 @@ def put(self, robj, w=None, dw=None, pw=None, return_body=True, def do_put(self, url, headers, content, return_body=False, key=None): if key is None: - response = self.http_request('POST', url, headers, content) + response = self._request('POST', url, headers, content) else: - response = self.http_request('PUT', url, headers, content) + response = self._request('PUT', url, headers, content) if return_body: return self.parse_body(response, [200, 201, 300]) @@ -189,16 +167,15 @@ def put_new(self, robj, w=None, dw=None, pw=None, return_body=True, """Put a new object into the Riak store, returning its (new) key.""" # We could detect quorum_controls here but HTTP ignores # unknown flags/params. - params = {'returnbody': str(return_body).lower(), 'w': w, 'dw': dw, - 'pw': pw} - url = self.build_rest_path(bucket=robj.bucket, params=params) + params = {'returnbody': return_body, 'w': w, 'dw': dw, 'pw': pw} + url = self.object_path(robj.bucket.name, **params) headers = self.build_put_headers(robj) # TODO: use a more general 'prevent_stale_writes' semantics, # which is a superset of the if_none_match semantics. if if_none_match: headers["If-None-Match"] = "*" content = robj.get_encoded_data() - response = self.http_request('POST', url, headers, content) + response = self._request('POST', url, headers, content) location = response[0]['location'] idx = location.rindex('/') key = location[(idx + 1):] @@ -217,11 +194,10 @@ def delete(self, robj, rw=None, r=None, w=None, dw=None, pr=None, pw=None): # unknown flags/params. params = {'rw': rw, 'r': r, 'w': w, 'dw': dw, 'pr': pr, 'pw': pw} headers = {} - url = self.build_rest_path(robj.bucket, robj.key, - params=params) + url = self.object_path(robj.bucket.name, robj.key, **params) if self.tombstone_vclocks() and robj.vclock is not None: headers['X-Riak-Vclock'] = robj.vclock - response = self.http_request('DELETE', url, headers) + response = self._request('DELETE', url, headers) self.check_http_code(response, [204, 404]) return self @@ -229,24 +205,31 @@ def get_keys(self, bucket): """ Fetch a list of keys for the bucket """ - params = {'props': 'True', 'keys': 'true'} - url = self.build_rest_path(bucket, params=params) - response = self.http_request('GET', url) + url = self.key_list_path(bucket.name) + response = self._request('GET', url) 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.') + raise Exception('Error listing keys.') + + def stream_keys(self, bucket): + url = self.key_list_path(bucket.name, keys='stream') + headers, response = self._request('GET', url, stream=True) + + if headers['http_code'] == 200: + return RiakHttpKeyStream(response) + else: + raise Exception('Error listing keys.') def get_buckets(self): """ Fetch a list of all buckets """ - params = {'buckets': 'true'} - url = self.build_rest_path(None, params=params) - response = self.http_request('GET', url) + url = self.bucket_list_path() + response = self._request('GET', url) headers, encoded_props = response[0:2] if headers['http_code'] == 200: @@ -260,9 +243,8 @@ def get_bucket_props(self, bucket): Get properties for a bucket """ # Run the request... - params = {'props': 'true', 'keys': 'false'} - url = self.build_rest_path(bucket, params=params) - response = self.http_request('GET', url) + url = self.bucket_properties_path(bucket.name) + response = self._request('GET', url) headers = response[0] encoded_props = response[1] @@ -276,12 +258,12 @@ def set_bucket_props(self, bucket, props): """ Set the properties on the bucket object given """ - url = self.build_rest_path(bucket) + url = self.bucket_properties_path(bucket.name) headers = {'Content-Type': 'application/json'} content = json.dumps({'props': props}) # Run the request... - response = self.http_request('PUT', url, headers, content) + response = self._request('PUT', url, headers, content) # Handle the response... if response is None: @@ -297,21 +279,13 @@ def mapred(self, inputs, query, timeout=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') - # Construct the job, optionally set the timeout... - job = {'inputs': inputs, 'query': query} - if timeout is not None: - job['timeout'] = timeout - - content = json.dumps(job) + content = self._construct_mapred_json(inputs, query, timeout) # Do the request... - url = "/" + self._mapred_prefix + url = self.mapred_path() headers = {'Content-Type': 'application/json'} - response = self.http_request('POST', url, headers, content) + response = self._request('POST', url, headers, content) # Make sure the expected status code came back... status = response[0]['http_code'] @@ -323,16 +297,28 @@ def mapred(self, inputs, query, timeout=None): result = json.loads(response[1]) return result + def stream_mapred(self, inputs, query, timeout=None): + content = self._construct_mapred_json(inputs, query, timeout) + + url = self.mapred_path(chunked=True) + reqheaders = {'Content-Type': 'application/json'} + headers, response = self._request('POST', url, reqheaders, + content, stream=True) + + if headers['http_code'] is 200: + return RiakHttpMapReduceStream(response) + else: + raise Exception( + 'Error running MapReduce operation. Headers: %s Body: %s' % + (repr(headers), repr(response.read()))) + def get_index(self, bucket, index, startkey, endkey=None): """ Performs a secondary index query. """ - # TODO: use resource detection - segments = ["buckets", bucket, "index", index, str(startkey)] - if endkey: - segments.append(str(endkey)) - uri = '/%s' % ('/'.join(segments)) - headers, data = response = self.get_request(uri) + url = self.index_path(bucket, index, startkey, endkey) + response = self._request('GET', url) + headers, data = response self.check_http_code(response, [200]) jsonData = json.loads(data) return jsonData[u'keys'][:] @@ -344,15 +330,15 @@ def search(self, index, query, **params): if index is None: index = 'search' - options = {'q': query, 'wt': 'json'} + options = {} if 'op' in params: op = params.pop('op') options['q.op'] = op options.update(params) - # TODO: use resource detection - uri = "/solr/%s/select" % index - headers, data = response = self.get_request(uri, options) + url = self.solr_select_path(index, query, **options) + response = self._request('GET', url) + headers, data = response self.check_http_code(response, [200]) if 'json' in headers['content-type']: results = json.loads(data) @@ -362,6 +348,53 @@ def search(self, index, query, **params): else: raise ValueError("Could not decode search response") + def fulltext_add(self, index, docs): + """ + Adds documents to the search index. + """ + xml = Document() + root = xml.createElement('add') + for doc in docs: + doc_element = xml.createElement('doc') + for key in doc: + value = doc[key] + field = xml.createElement('field') + field.setAttribute("name", key) + text = xml.createTextNode(value) + field.appendChild(text) + doc_element.appendChild(field) + root.appendChild(doc_element) + xml.appendChild(root) + + self._request('POST', self.solr_update_path(index), + {'Content-Type': 'text/xml'}, + xml.toxml()) + + def fulltext_delete(self, index, docs=None, queries=None): + """ + Removes documents from the full-text index. + """ + xml = Document() + root = xml.createElement('delete') + if docs: + for doc in docs: + doc_element = xml.createElement('id') + text = xml.createTextNode(doc) + doc_element.appendChild(text) + root.appendChild(doc_element) + if queries: + for query in queries: + query_element = xml.createElement('query') + text = xml.createTextNode(query) + query_element.appendChild(text) + root.appendChild(query_element) + + xml.appendChild(root) + + self._request('POST', self.solr_update_path(index), + {'Content-Type': 'text/xml'}, + xml.toxml()) + def check_http_code(self, response, expected_statuses): status = response[0]['http_code'] if not status in expected_statuses: @@ -370,10 +403,7 @@ def check_http_code(self, response, expected_statuses): def parse_body(self, response, expected_statuses): """ - Given the output of RiakUtils.http_request and a list of - statuses, populate the object. Only for use by the Riak client - library. - @return self + Parse the body of an object response and populate the object. """ # If no response given, then return. if response is None: @@ -444,12 +474,8 @@ def to_link_header(self, link): """ Convert this RiakLink object to a link header string. Used internally. """ - header = '' - header += '; riaktag="' - header += urllib.quote_plus(link.get_tag()) + '"' + url = self.object_path(link.get_bucket(), link.get_key()) + header = '<%s>; riaktag="%s"' % (url, link.get_tag()) return header def parse_links(self, links, linkHeaders): @@ -457,8 +483,8 @@ def parse_links(self, links, linkHeaders): Private. @return self """ - oldform = "; ?riaktag=\"([^\']+)\"" - newform = "; ?riaktag=\"([^\']+)\"" + oldform = "; ?riaktag=\"([^\"]+)\"" + newform = "; ?riaktag=\"([^\"]+)\"" for linkHeader in linkHeaders.strip().split(','): linkHeader = linkHeader.strip() matches = (re.match(oldform, linkHeader) or @@ -488,49 +514,8 @@ def add_links_for_riak_object(self, robject, headers): return headers - def get_request(self, uri=None, params=None): - url = self.build_rest_path(bucket=None, params=params, prefix=uri) - return self.http_request('GET', url) - - def post_request(self, uri=None, body=None, params=None, - content_type="application/json"): - uri = self.build_rest_path(prefix=uri, params=params) - return self.http_request('POST', uri, {'Content-Type': content_type}, - body) - # Utility functions used by Riak library. - def build_rest_path(self, bucket=None, key=None, params=None, prefix=None): - """ - Given a RiakClient, RiakBucket, Key, LinkSpec, and Params, - construct and return a URL. - """ - # Build 'http://hostname:port/prefix/bucket' - path = '' - path += '/' + (prefix or self._prefix) - - # Add '.../bucket' - if bucket is not None: - path += '/' + urllib.quote_plus(bucket.name) - - # Add '.../key' - if key is not None: - path += '/' + urllib.quote_plus(key) - - # Add query parameters. - if params is not None: - s = '' - for key in params.keys(): - if params[key] is not None: - if s != '': - s += '&' - s += (urllib.quote_plus(key) + '=' + - urllib.quote_plus(str(params[key]))) - path += '?' + s - - # Return. - return path - def build_put_headers(self, robj): """Build the headers for a POST/PUT request.""" @@ -544,7 +529,7 @@ def build_put_headers(self, robj): headers['X-Riak-Vclock'] = robj.vclock # Create the header from metadata - links = self.add_links_for_riak_object(robj, headers) + self.add_links_for_riak_object(robj, headers) for key, value in robj.usermeta.iteritems(): headers['X-Riak-Meta-%s' % key] = value @@ -558,50 +543,6 @@ def build_put_headers(self, robj): return headers - def http_request(self, method, uri, headers=None, body=''): - """ - Given a Method, URL, Headers, and Body, perform and HTTP request, - and return a 2-tuple containing a dictionary of response headers - and the response body. - """ - if headers is None: - headers = {} - # Run the request... - for retry in range(self.RETRY_COUNT): - with self._conns.withconn() as conn: - ### should probably build this try/except into a custom - ### contextmanager for the connection. - try: - conn.request(method, uri, body, headers) - response = conn.getresponse() - - try: - # Get the response headers... - response_headers = {'http_code': response.status} - for (key, value) in response.getheaders(): - response_headers[key.lower()] = value - - # Get the body... - response_body = response.read() - finally: - response.close() - - return response_headers, response_body - except socket.error, e: - conn.close() - if e[0] == errno.ECONNRESET: - # Grab another connection and try again. - continue - # Don't know how to handle this. - raise - except httplib.HTTPException: - # Just close the connection and try again. - conn.close() - continue - - # No luck, even with retrying. - raise RiakError("could not get a response") - def _normalize_json_search_response(self, json): """ Normalizes a JSON search response so that PB and HTTP have the @@ -658,53 +599,3 @@ def parse_http_headers(cls, headers): else: retVal[key] = value return retVal - - -class XMLSearchResult(object): - # Match tags that are document fields - fieldtags = ['str', 'int', 'date'] - - def __init__(self): - # Results - self.num_found = 0 - self.max_score = 0.0 - self.docs = [] - - # Parser state - self.currdoc = None - self.currfield = None - self.currvalue = None - - def start(self, tag, attrib): - if tag == 'result': - self.num_found = int(attrib['numFound']) - self.max_score = float(attrib['maxScore']) - elif tag == 'doc': - self.currdoc = {} - elif tag in self.fieldtags and self.currdoc is not None: - self.currfield = attrib['name'] - - def end(self, tag): - if tag == 'doc' and self.currdoc is not None: - self.docs.append(self.currdoc) - self.currdoc = None - elif tag in self.fieldtags and self.currdoc is not None: - if tag == 'int': - self.currvalue = int(self.currvalue) - self.currdoc[self.currfield] = self.currvalue - self.currfield = None - self.currvalue = None - - def data(self, data): - if self.currfield: - # riak_solr_output adds NL + 6 spaces - data = data.rstrip() - if self.currvalue: - self.currvalue += data - else: - self.currvalue = data - - def close(self): - return {'num_found': self.num_found, - 'max_score': self.max_score, - 'docs': self.docs} diff --git a/riak/transports/monitor.py b/riak/transports/monitor.py deleted file mode 100644 index c6223518..00000000 --- a/riak/transports/monitor.py +++ /dev/null @@ -1,47 +0,0 @@ -""" -Copyright 2011 Greg Stein - -This file is provided to you under the Apache License, -Version 2.0 (the "License"); you may not use this file -except in compliance with the License. You may obtain -a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, -software distributed under the License is distributed on an -"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -KIND, either express or implied. See the License for the -specific language governing permissions and limitations -under the License. -""" - -import threading -import time - - -class Monitor(object): - - def __init__(self, cm, transport): - self._cm = cm - self._transport = transport - - self._stop_loop = False - self._thread = threading.Thread(target=self._run) - - self._periodic = 0.050 # 50 msec - - def start(self): - self._thread.start() - - def terminate(self): - self._stop_loop = True - self._thread.join() - - def _run(self): - while not self._stop_loop: - ### look for changes in the ring servers - - ### see if some offline servers came back - - time.sleep(self._periodic) diff --git a/riak/transports/pbc.py b/riak/transports/pbc.py deleted file mode 100644 index cfef37dd..00000000 --- a/riak/transports/pbc.py +++ /dev/null @@ -1,744 +0,0 @@ -""" -Copyright 2010 Rusty Klophaus -Copyright 2010 Justin Sheehy -Copyright 2009 Jay Baird - -This file is provided to you under the Apache License, -Version 2.0 (the "License"); you may not use this file -except in compliance with the License. You may obtain -a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, -software distributed under the License is distributed on an -"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -KIND, either express or implied. See the License for the -specific language governing permissions and limitations -under the License. -""" -from __future__ import with_statement - -import errno -import socket -import struct -try: - import json -except ImportError: - import simplejson as json - -from riak import RiakError -from riak.mapreduce import RiakLink -from riak.metadata import ( - MD_CHARSET, - MD_CTYPE, - MD_ENCODING, - MD_INDEX, - MD_LASTMOD, - MD_LASTMOD_USECS, - MD_LINKS, - MD_USERMETA, - MD_VTAG, - ) -from riak.riak_index_entry import RiakIndexEntry -from riak.transports import connection -from riak.transports.transport import RiakTransport -import riak.util - -try: - import riak_pb -except ImportError: - riak_pb = None - -## Protocol codes -MSG_CODE_ERROR_RESP = 0 -MSG_CODE_PING_REQ = 1 -MSG_CODE_PING_RESP = 2 -MSG_CODE_GET_CLIENT_ID_REQ = 3 -MSG_CODE_GET_CLIENT_ID_RESP = 4 -MSG_CODE_SET_CLIENT_ID_REQ = 5 -MSG_CODE_SET_CLIENT_ID_RESP = 6 -MSG_CODE_GET_SERVER_INFO_REQ = 7 -MSG_CODE_GET_SERVER_INFO_RESP = 8 -MSG_CODE_GET_REQ = 9 -MSG_CODE_GET_RESP = 10 -MSG_CODE_PUT_REQ = 11 -MSG_CODE_PUT_RESP = 12 -MSG_CODE_DEL_REQ = 13 -MSG_CODE_DEL_RESP = 14 -MSG_CODE_LIST_BUCKETS_REQ = 15 -MSG_CODE_LIST_BUCKETS_RESP = 16 -MSG_CODE_LIST_KEYS_REQ = 17 -MSG_CODE_LIST_KEYS_RESP = 18 -MSG_CODE_GET_BUCKET_REQ = 19 -MSG_CODE_GET_BUCKET_RESP = 20 -MSG_CODE_SET_BUCKET_REQ = 21 -MSG_CODE_SET_BUCKET_RESP = 22 -MSG_CODE_MAPRED_REQ = 23 -MSG_CODE_MAPRED_RESP = 24 -MSG_CODE_INDEX_REQ = 25 -MSG_CODE_INDEX_RESP = 26 -MSG_CODE_SEARCH_QUERY_REQ = 27 -MSG_CODE_SEARCH_QUERY_RESP = 28 - -RIAKC_RW_ONE = 4294967294 -RIAKC_RW_QUORUM = 4294967293 -RIAKC_RW_ALL = 4294967292 -RIAKC_RW_DEFAULT = 4294967291 - -# These are a specific set of socket errors -# that could be raised on send/recv that indicate -# that the socket is closed or reset, and is not -# usable. On seeing any of these errors, the socket -# should be closed, and the connection re-established. -CONN_CLOSED_ERRORS = ( - errno.EHOSTUNREACH, - errno.ECONNRESET, - errno.EBADF, - errno.EPIPE - ) - - -class SocketWithId(connection.Socket): - def __init__(self, host, port): - super(SocketWithId, self).__init__(host, port) - self.last_client_id = None - - def maybe_connect(self): - # If we're going to establish a new connection, then reset the last - # client_id used on this connection. - if self.sock is None: - self.last_client_id = None - super(SocketWithId, self).maybe_connect() - - def send(self, pkt): - try: - self.sock.sendall(pkt) - except socket.error, e: - # If the socket is in a bad state, close it and allow it - # to re-connect on the next try - if e[0] in CONN_CLOSED_ERRORS: - self.close() - raise - - def recv(self, want_len): - try: - res = self.sock.recv(want_len) - - # Assume the socket is closed if no data is - # returned on a blocking read. - if len(res) == 0 and want_len > 0: - self.close() - - return res - - except socket.error, e: - # If the socket is in a bad state, close it and allow it - # to re-connect on the next try - if e[0] in CONN_CLOSED_ERRORS: - self.close() - raise - - -class RiakPbcTransport(RiakTransport): - """ - The RiakPbcTransport object holds a connection to the protocol - buffers interface on the riak server. - """ - - # We're using the new RiakTransport API - api = 2 - - rw_names = { - 'default': RIAKC_RW_DEFAULT, - 'all': RIAKC_RW_ALL, - 'quorum': RIAKC_RW_QUORUM, - 'one': RIAKC_RW_ONE - } - - # The ConnectionManager class that this transport prefers. - default_cm = connection.cm_using(SocketWithId) - - def __init__(self, cm, client_id=None, max_attempts=1, **unused_options): - """ - Construct a new RiakPbcTransport object. - """ - if riak_pb is None: - raise RiakError("this transport is not available (no protobuf)") - - super(RiakPbcTransport, self).__init__() - - self._cm = cm - self._client_id = client_id - self._max_attempts = max_attempts - - # FeatureDetection API - def _server_version(self): - return self.get_server_info()['server_version'] - - def translate_rw_val(self, rw): - val = self.rw_names.get(rw) - if val is None: - return rw - return val - - def __copy__(self): - return RiakPbcTransport(self._cm, self._client_id) - - def ping(self): - """ - Ping the remote server - @return boolean - """ - # An expected response code of None implies "any response is valid". - msg_code, msg = self.send_msg_code(MSG_CODE_PING_REQ, None) - if msg_code == MSG_CODE_PING_RESP: - return 1 - else: - return 0 - - def get_server_info(self): - """ - Get information about the server - """ - msg_code, resp = self.send_msg_code(MSG_CODE_GET_SERVER_INFO_REQ, - MSG_CODE_GET_SERVER_INFO_RESP) - return {'node': resp.node, 'server_version': resp.server_version} - - def get_client_id(self): - """ - Get the client id used by this connection - """ - msg_code, resp = self.send_msg_code(MSG_CODE_GET_CLIENT_ID_REQ, - MSG_CODE_GET_CLIENT_ID_RESP) - return resp.client_id - - def set_client_id(self, client_id): - """ - Set the client id used by this connection - """ - req = riak_pb.RpbSetClientIdReq() - req.client_id = client_id - - msg_code, resp = self.send_msg(MSG_CODE_SET_CLIENT_ID_REQ, req, - MSG_CODE_SET_CLIENT_ID_RESP) - - # Using different client_id values across connections is a bad idea - # since you never know which connection you might use for a given - # API call. Setting the client_id manually (rather than as part of - # the transport construction) can be error-prone since the connection - # could drop and be reinstated using self._client_id. - # - # To minimize the potential impact of variant client_id values across - # connections, we'll store this new client_id and use it for all - # future connections. - self._client_id = client_id - - return True - - def get(self, robj, r=None, pr=None, vtag=None): - """ - Serialize get request and deserialize response - """ - if vtag is not None: - raise RiakError("PB transport does not support vtags") - - bucket = robj.bucket - - req = riak_pb.RpbGetReq() - if r: - req.r = self.translate_rw_val(r) - if self.quorum_controls() and pr: - req.pr = self.translate_rw_val(pr) - - if self.tombstone_vclocks(): - req.deletedvclock = 1 - - req.bucket = bucket.name - req.key = robj.key - - # An expected response code of None implies "any response is valid". - msg_code, resp = self.send_msg(MSG_CODE_GET_REQ, req, None) - if msg_code == MSG_CODE_GET_RESP: - contents = [] - for c in resp.content: - contents.append(self.decode_content(c)) - return resp.vclock, contents - else: - return None - - def put(self, robj, w=None, dw=None, pw=None, return_body=True, - if_none_match=False): - """ - Serialize get request and deserialize response - """ - bucket = robj.bucket - - req = riak_pb.RpbPutReq() - if w: - req.w = self.translate_rw_val(w) - if dw: - req.dw = self.translate_rw_val(dw) - if self.quorum_controls() and pw: - req.pw = self.translate_rw_val(pw) - - if return_body: - req.return_body = 1 - if if_none_match: - req.if_none_match = 1 - - req.bucket = bucket.name - req.key = robj.key - vclock = robj.vclock - if vclock: - req.vclock = vclock - - self.pbify_content(robj.metadata, - robj.get_encoded_data(), - req.content) - - msg_code, resp = self.send_msg(MSG_CODE_PUT_REQ, req, - MSG_CODE_PUT_RESP) - if resp is not None: - contents = [] - for c in resp.content: - contents.append(self.decode_content(c)) - return resp.vclock, contents - - def put_new(self, robj, w=None, dw=None, pw=None, return_body=True, - if_none_match=False): - """Put a new object into the Riak store, returning its (new) key. - - If return_meta is False, then the vlock and metadata return values - will be None. - - @return (key, vclock, metadata) - """ - # Note that this won't work on 0.14 nodes. - bucket = robj.bucket - - req = riak_pb.RpbPutReq() - if w: - req.w = self.translate_rw_val(w) - if dw: - req.dw = self.translate_rw_val(dw) - if self.quorum_controls() and pw: - req.pw = self.translate_rw_val(pw) - - if return_body: - req.return_body = 1 - if if_none_match: - req.if_none_match = 1 - - req.bucket = bucket.name - - self.pbify_content(robj.metadata, - robj.get_encoded_data(), - req.content) - - msg_code, resp = self.send_msg(MSG_CODE_PUT_REQ, req, - MSG_CODE_PUT_RESP) - if not resp: - raise RiakError("missing response object") - if len(resp.content) != 1: - raise RiakError("siblings were returned from object creation") - - metadata, content = self.decode_content(resp.content[0]) - return resp.key, resp.vclock, metadata - - def delete(self, robj, rw=None, r=None, w=None, dw=None, pr=None, pw=None): - """ - Serialize get request and deserialize response - """ - bucket = robj.bucket - - req = riak_pb.RpbDelReq() - if rw: - req.rw = self.translate_rw_val(rw) - if r: - req.r = self.translate_rw_val(r) - if w: - req.w = self.translate_rw_val(w) - if dw: - req.dw = self.translate_rw_val(dw) - - if self.quorum_controls(): - if pr: - req.pr = self.translate_rw_val(pr) - if pw: - req.pw = self.translate_rw_val(pw) - - if self.tombstone_vclocks() and robj.vclock: - req.vclock = robj.vclock - - req.bucket = bucket.name - req.key = robj.key - - msg_code, resp = self.send_msg(MSG_CODE_DEL_REQ, req, - MSG_CODE_DEL_RESP) - return self - - def get_keys(self, bucket): - """ - Lists all keys within a bucket. - """ - req = riak_pb.RpbListKeysReq() - req.bucket = bucket.name - - keys = [] - - def _handle_response(resp): - for key in resp.keys: - keys.append(key) - self.send_msg_multi(MSG_CODE_LIST_KEYS_REQ, req, - MSG_CODE_LIST_KEYS_RESP, _handle_response) - - return keys - - def get_buckets(self): - """ - Serialize bucket listing request and deserialize response - """ - msg_code, resp = self.send_msg_code(MSG_CODE_LIST_BUCKETS_REQ, - MSG_CODE_LIST_BUCKETS_RESP) - return resp.buckets - - def get_bucket_props(self, bucket): - """ - Serialize bucket property request and deserialize response - """ - req = riak_pb.RpbGetBucketReq() - req.bucket = bucket.name - - msg_code, resp = self.send_msg(MSG_CODE_GET_BUCKET_REQ, req, - MSG_CODE_GET_BUCKET_RESP) - props = {} - if resp.props.HasField('n_val'): - props['n_val'] = resp.props.n_val - if resp.props.HasField('allow_mult'): - props['allow_mult'] = resp.props.allow_mult - - return props - - def set_bucket_props(self, bucket, props): - """ - Serialize set bucket property request and deserialize response - """ - req = riak_pb.RpbSetBucketReq() - req.bucket = bucket.name - for key in props: - if key not in ['n_val', 'allow_mult']: - raise NotImplementedError - - if 'n_val' in props: - req.props.n_val = props['n_val'] - if 'allow_mult' in props: - req.props.allow_mult = props['allow_mult'] - - msg_code, resp = self.send_msg(MSG_CODE_SET_BUCKET_REQ, req, - MSG_CODE_SET_BUCKET_RESP) - return self - - def mapred(self, inputs, query, timeout=None): - # Construct the job, optionally set the timeout... - job = {'inputs': inputs, 'query': query} - if timeout is not None: - job['timeout'] = timeout - - content = json.dumps(job) - - req = riak_pb.RpbMapRedReq() - req.request = content - req.content_type = "application/json" - - # dictionary of phase results - each content should be an encoded array - # which is appended to the result for that phase. - result = {} - - def _handle_response(resp): - if resp.HasField("phase") and resp.HasField("response"): - content = json.loads(resp.response) - if resp.phase in result: - result[resp.phase] += content - else: - result[resp.phase] = content - self.send_msg_multi(MSG_CODE_MAPRED_REQ, req, MSG_CODE_MAPRED_RESP, - _handle_response) - - # If a single result - return the same as the HTTP interface does - # otherwise return all the phase information - if not len(result): - return None - elif len(result) == 1: - return result[max(result.keys())] - else: - return result - - def get_index(self, bucket, index, startkey, endkey=None): - if not self.pb_indexes(): - return self._get_index_mapred_emu(bucket, index, startkey, endkey) - - req = riak_pb.RpbIndexReq(bucket=bucket, index=index) - if endkey: - req.qtype = riak_pb.RpbIndexReq.range - req.range_min = str(startkey) - req.range_max = str(endkey) - else: - req.qtype = riak_pb.RpbIndexReq.eq - req.key = str(startkey) - - msg_code, resp = self.send_msg(MSG_CODE_INDEX_REQ, req, - MSG_CODE_INDEX_RESP) - return resp.keys - - def search(self, index, query, **params): - if not self.pb_search(): - return self._search_mapred_emu(index, query) - - req = riak_pb.RpbSearchQueryReq(index=index, q=query) - if 'rows' in params: - req.rows = params['rows'] - if 'start' in params: - req.start = params['start'] - if 'sort' in params: - req.sort = params['sort'] - if 'filter' in params: - req.filter = params['filter'] - if 'df' in params: - req.df = params['df'] - if 'op' in params: - req.op = params['op'] - if 'q.op' in params: - req.op = params['q.op'] - if 'fl' in params: - if isinstance(params['fl'], list): - req.fl.extend(params['fl']) - else: - req.fl.append(params['fl']) - if 'presort' in params: - req.presort = params['presort'] - - msg_code, resp = self.send_msg(MSG_CODE_SEARCH_QUERY_REQ, req, - MSG_CODE_SEARCH_QUERY_RESP) - - result = {} - if resp.HasField('max_score'): - result['max_score'] = resp.max_score - if resp.HasField('num_found'): - result['num_found'] = resp.num_found - docs = [] - for doc in resp.docs: - resultdoc = {} - for pair in doc.fields: - resultdoc[pair.key] = pair.value - docs.append(resultdoc) - result['docs'] = docs - return result - - def send_msg_code(self, msg_code, expect): - with self._cm.withconn() as conn: - self.send_pkt(conn, struct.pack("!iB", 1, msg_code)) - return self.recv_msg(conn, expect) - - def encode_msg(self, msg_code, msg): - str = msg.SerializeToString() - slen = len(str) - hdr = struct.pack("!iB", 1 + slen, msg_code) - return hdr + str - - def send_msg(self, msg_code, msg, expect): - with self._cm.withconn() as conn: - self.send_pkt(conn, self.encode_msg(msg_code, msg)) - if msg_code == MSG_CODE_SET_CLIENT_ID_REQ: - conn.last_client_id = self._client_id - return self.recv_msg(conn, expect) - - def send_msg_multi(self, msg_code, msg, expect, handler): - with self._cm.withconn() as conn: - self.send_pkt(conn, self.encode_msg(msg_code, msg)) - while True: - msg_code, resp = self.recv_msg(conn, expect) - handler(resp) - if resp.HasField("done") and resp.done: - break - - def send_pkt(self, conn, pkt): - attempt = 0 - e = None - for attempt in xrange(self._max_attempts): - e = None - try: - conn.maybe_connect() - - # If the last client_id used on this connection is - # different than our client_id, then set a new ID on - # the connection. - if conn.last_client_id != self._client_id: - req = riak_pb.RpbSetClientIdReq() - req.client_id = self._client_id - conn.send(self.encode_msg(MSG_CODE_SET_CLIENT_ID_REQ, req)) - conn.last_client_id = self._client_id - self.recv_msg(conn, MSG_CODE_SET_CLIENT_ID_RESP) - - conn.send(pkt) - break - except socket.error, e: - # If this is some unknown socket error bail out - # instead of retrying - if e[0] not in CONN_CLOSED_ERRORS: - raise - - # Max attempts reached, raise whatever exception we are getting - if attempt + 1 == self._max_attempts and e is not None: - raise e - - def recv_msg(self, conn, expect): - self.recv_pkt(conn) - msg_code, = struct.unpack("B", self._inbuf[:1]) - if msg_code == MSG_CODE_ERROR_RESP: - msg = riak_pb.RpbErrorResp() - msg.ParseFromString(self._inbuf[1:]) - raise Exception(msg.errmsg) - elif msg_code == MSG_CODE_PING_RESP: - msg = None - elif msg_code == MSG_CODE_GET_SERVER_INFO_RESP: - msg = riak_pb.RpbGetServerInfoResp() - msg.ParseFromString(self._inbuf[1:]) - elif msg_code == MSG_CODE_GET_CLIENT_ID_RESP: - msg = riak_pb.RpbGetClientIdResp() - msg.ParseFromString(self._inbuf[1:]) - elif msg_code == MSG_CODE_SET_CLIENT_ID_RESP: - msg = None - elif msg_code == MSG_CODE_GET_RESP: - msg = riak_pb.RpbGetResp() - msg.ParseFromString(self._inbuf[1:]) - elif msg_code == MSG_CODE_PUT_RESP: - msg = riak_pb.RpbPutResp() - msg.ParseFromString(self._inbuf[1:]) - elif msg_code == MSG_CODE_DEL_RESP: - msg = None - elif msg_code == MSG_CODE_LIST_KEYS_RESP: - msg = riak_pb.RpbListKeysResp() - msg.ParseFromString(self._inbuf[1:]) - elif msg_code == MSG_CODE_LIST_BUCKETS_RESP: - msg = riak_pb.RpbListBucketsResp() - msg.ParseFromString(self._inbuf[1:]) - elif msg_code == MSG_CODE_GET_BUCKET_RESP: - msg = riak_pb.RpbGetBucketResp() - msg.ParseFromString(self._inbuf[1:]) - elif msg_code == MSG_CODE_SET_BUCKET_RESP: - msg = None - elif msg_code == MSG_CODE_MAPRED_RESP: - msg = riak_pb.RpbMapRedResp() - msg.ParseFromString(self._inbuf[1:]) - elif msg_code == MSG_CODE_INDEX_RESP: - msg = riak_pb.RpbIndexResp() - msg.ParseFromString(self._inbuf[1:]) - elif msg_code == MSG_CODE_SEARCH_QUERY_RESP: - msg = riak_pb.RpbSearchQueryResp() - msg.ParseFromString(self._inbuf[1:]) - else: - raise Exception("unknown msg code %s" % msg_code) - if expect and msg_code != expect: - raise RiakError("unexpected protocol buffer message code: %d" - % msg_code) - return msg_code, msg - - def recv_pkt(self, conn): - nmsglen = conn.recv(4) - if len(nmsglen) != 4: - raise RiakError( - "Socket returned short packet length %d - expected 4" - % len(nmsglen)) - msglen, = struct.unpack('!i', nmsglen) - self._inbuf_len = msglen - self._inbuf = '' - while len(self._inbuf) < msglen: - want_len = min(8192, msglen - len(self._inbuf)) - recv_buf = conn.recv(want_len) - if not recv_buf: - break - self._inbuf += recv_buf - if len(self._inbuf) != self._inbuf_len: - raise RiakError("Socket returned short packet %d - expected %d" - % (len(self._inbuf), self._inbuf_len)) - - def decode_contents(self, rpb_contents): - contents = [] - for rpb_c in rpb_contents: - contents.append(self.decode_content(rpb_c)) - return contents - - def decode_content(self, rpb_content): - metadata = {} - if rpb_content.HasField("deleted"): - metadata[MD_DELETED] = True - if rpb_content.HasField("content_type"): - metadata[MD_CTYPE] = rpb_content.content_type - if rpb_content.HasField("charset"): - metadata[MD_CHARSET] = rpb_content.charset - if rpb_content.HasField("content_encoding"): - metadata[MD_ENCODING] = rpb_content.content_encoding - if rpb_content.HasField("vtag"): - metadata[MD_VTAG] = rpb_content.vtag - links = [] - for link in rpb_content.links: - if link.HasField("bucket"): - bucket = link.bucket - else: - bucket = None - if link.HasField("key"): - key = link.key - else: - key = None - if link.HasField("tag"): - tag = link.tag - else: - tag = None - links.append(RiakLink(bucket, key, tag)) - if links: - metadata[MD_LINKS] = links - if rpb_content.HasField("last_mod"): - metadata[MD_LASTMOD] = rpb_content.last_mod - if rpb_content.HasField("last_mod_usecs"): - metadata[MD_LASTMOD_USECS] = rpb_content.last_mod_usecs - usermeta = {} - for usermd in rpb_content.usermeta: - usermeta[usermd.key] = usermd.value - if len(usermeta) > 0: - metadata[MD_USERMETA] = usermeta - indexes = [] - for index in rpb_content.indexes: - rie = RiakIndexEntry(index.key, index.value) - indexes.append(rie) - if len(indexes) > 0: - metadata[MD_INDEX] = indexes - return metadata, rpb_content.value - - def pbify_content(self, metadata, data, rpb_content): - # Convert the broken out fields, building up - # pbmetadata for any unknown ones - for k, v in metadata.iteritems(): - if k == MD_CTYPE: - rpb_content.content_type = v - elif k == MD_CHARSET: - rpb_content.charset = v - elif k == MD_ENCODING: - rpb_content.charset = v - elif k == MD_USERMETA: - for uk, uv in v.iteritems(): - pair = rpb_content.usermeta.add() - pair.key = uk - pair.value = uv - elif k == MD_INDEX: - for rie in v: - pair = rpb_content.indexes.add() - pair.key = rie.get_field() - pair.value = rie.get_value() - elif k == MD_LINKS: - for link in v: - pb_link = rpb_content.links.add() - pb_link.bucket = link.get_bucket() - pb_link.key = link.get_key() - pb_link.tag = link.get_tag() - rpb_content.value = data diff --git a/riak/transports/pbc/__init__.py b/riak/transports/pbc/__init__.py new file mode 100644 index 00000000..d49ab6fc --- /dev/null +++ b/riak/transports/pbc/__init__.py @@ -0,0 +1,74 @@ +""" +Copyright 2012 Basho Technologies, Inc. +Copyright 2010 Rusty Klophaus +Copyright 2010 Justin Sheehy +Copyright 2009 Jay Baird + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" + +import errno +import socket +from riak.transports.pool import Pool +from riak.transports.pbc.transport import RiakPbcTransport + + +class RiakPbcPool(Pool): + """ + A resource pool of PBC transports. + """ + def __init__(self, client, **options): + super(RiakPbcPool, self).__init__() + self._client = client + self._options = options + + def create_resource(self): + node = self._client._choose_node() + return RiakPbcTransport(node=node, + client=self._client, + **self._options) + + def destroy_resource(self, pbc): + pbc.close() + +# These are a specific set of socket errors +# that could be raised on send/recv that indicate +# that the socket is closed or reset, and is not +# usable. On seeing any of these errors, the socket +# should be closed, and the connection re-established. +CONN_CLOSED_ERRORS = ( + errno.EHOSTUNREACH, + errno.ECONNRESET, + errno.ECONNREFUSED, + errno.ECONNABORTED, + errno.ETIMEDOUT, + errno.EBADF, + errno.EPIPE + ) + + +def is_retryable(err): + """ + Determines if the given exception is something that is + network/socket-related and should thus cause the PBC connection to + close and the operation retried on another node. + + :rtype: boolean + """ + if isinstance(err, socket.error): + code = err.args[0] + return code in CONN_CLOSED_ERRORS + else: + return False diff --git a/riak/transports/pbc/codec.py b/riak/transports/pbc/codec.py new file mode 100644 index 00000000..2742dc20 --- /dev/null +++ b/riak/transports/pbc/codec.py @@ -0,0 +1,165 @@ +""" +Copyright 2012 Basho Technologies, Inc. + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" +from riak.metadata import ( + MD_CHARSET, + MD_CTYPE, + MD_ENCODING, + MD_INDEX, + MD_LASTMOD, + MD_LASTMOD_USECS, + MD_LINKS, + MD_USERMETA, + MD_VTAG, + MD_DELETED + ) + +import riak_pb +from riak.riak_index_entry import RiakIndexEntry +from riak.mapreduce import RiakLink + +RIAKC_RW_ONE = 4294967294 +RIAKC_RW_QUORUM = 4294967293 +RIAKC_RW_ALL = 4294967292 +RIAKC_RW_DEFAULT = 4294967291 + + +class RiakPbcCodec(object): + """ + Protobuffs Encoding and decoding methods for RiakPbcTransport. + """ + + rw_names = { + 'default': RIAKC_RW_DEFAULT, + 'all': RIAKC_RW_ALL, + 'quorum': RIAKC_RW_QUORUM, + 'one': RIAKC_RW_ONE + } + + def __init__(self, **unused_args): + if riak_pb is None: + raise NotImplementedError("this transport is not available") + super(RiakPbcCodec, self).__init__(**unused_args) + + def translate_rw_val(self, rw): + """ + Converts a symbolic quorum value into its on-the-wire + equivalent. + + :param rw: the quorum + :type rw: string, integer + :rtype: integer + """ + val = self.rw_names.get(rw) + if val is None: + return rw + elif type(rw) is int and rw >= 0: + return val + else: + return None + + def decode_contents(self, rpb_contents): + """ + Decodes the multiple contents (siblings) of a RiakObject from + its protobuf representation. + """ + return [self.decode_content(rpb_c) for rpb_c in rpb_contents] + + def decode_content(self, rpb_content): + """ + Decodes a single sibling from the protobuf representation into + its metadata and value. + + :rtype: (dict, string) + """ + metadata = {} + if rpb_content.HasField("deleted"): + metadata[MD_DELETED] = True + if rpb_content.HasField("content_type"): + metadata[MD_CTYPE] = rpb_content.content_type + if rpb_content.HasField("charset"): + metadata[MD_CHARSET] = rpb_content.charset + if rpb_content.HasField("content_encoding"): + metadata[MD_ENCODING] = rpb_content.content_encoding + if rpb_content.HasField("vtag"): + metadata[MD_VTAG] = rpb_content.vtag + links = [] + for link in rpb_content.links: + if link.HasField("bucket"): + bucket = link.bucket + else: + bucket = None + if link.HasField("key"): + key = link.key + else: + key = None + if link.HasField("tag"): + tag = link.tag + else: + tag = None + links.append(RiakLink(bucket, key, tag)) + if links: + metadata[MD_LINKS] = links + if rpb_content.HasField("last_mod"): + metadata[MD_LASTMOD] = rpb_content.last_mod + if rpb_content.HasField("last_mod_usecs"): + metadata[MD_LASTMOD_USECS] = rpb_content.last_mod_usecs + usermeta = {} + for usermd in rpb_content.usermeta: + usermeta[usermd.key] = usermd.value + if len(usermeta) > 0: + metadata[MD_USERMETA] = usermeta + indexes = [] + for index in rpb_content.indexes: + rie = RiakIndexEntry(index.key, index.value) + indexes.append(rie) + if len(indexes) > 0: + metadata[MD_INDEX] = indexes + return metadata, rpb_content.value + + def encode_content(self, metadata, data, rpb_content): + """ + Fills an RpbContent message with the appropriate data and + metadata from a RiakObject. + """ + # Convert the broken out fields, building up + # pbmetadata for any unknown ones + for k in metadata: + v = metadata[k] + if k == MD_CTYPE: + rpb_content.content_type = v + elif k == MD_CHARSET: + rpb_content.charset = v + elif k == MD_ENCODING: + rpb_content.charset = v + elif k == MD_USERMETA: + for uk in v: + pair = rpb_content.usermeta.add() + pair.key = uk + pair.value = v[uk] + elif k == MD_INDEX: + for rie in v: + pair = rpb_content.indexes.add() + pair.key = rie.get_field() + pair.value = rie.get_value() + elif k == MD_LINKS: + for link in v: + pb_link = rpb_content.links.add() + pb_link.bucket = link.get_bucket() + pb_link.key = link.get_key() + pb_link.tag = link.get_tag() + rpb_content.value = str(data) diff --git a/riak/transports/pbc/connection.py b/riak/transports/pbc/connection.py new file mode 100644 index 00000000..f334fa98 --- /dev/null +++ b/riak/transports/pbc/connection.py @@ -0,0 +1,109 @@ +""" +Copyright 2012 Basho Technologies, Inc. + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" + +import socket +import struct +from riak import RiakError +from messages import ( + MESSAGE_CLASSES, + MSG_CODE_ERROR_RESP + ) + + +class RiakPbcConnection(object): + """ + Connection-related methods for RiakPbcTransport. + """ + + def _encode_msg(self, msg_code, msg=None): + if msg is None: + return struct.pack("!iB", 1, msg_code) + msgstr = msg.SerializeToString() + slen = len(msgstr) + hdr = struct.pack("!iB", 1 + slen, msg_code) + return hdr + msgstr + + def _request(self, msg_code, msg=None, expect=None): + self._send_msg(msg_code, msg) + return self._recv_msg(expect) + + def _send_msg(self, msg_code, msg): + self._socket.send(self._encode_msg(msg_code, msg)) + + def _recv_msg(self, expect=None): + self._recv_pkt() + msg_code, = struct.unpack("B", self._inbuf[:1]) + + if msg_code is MSG_CODE_ERROR_RESP: + err = self._parse_msg(msg_code, self._inbuf[1:]) + raise RiakError(err.errmsg) + elif msg_code in MESSAGE_CLASSES: + msg = self._parse_msg(msg_code, self._inbuf[1:]) + else: + raise Exception("unknown msg code %s" % msg_code) + + if expect and msg_code != expect: + raise RiakError("unexpected protocol buffer message code: %d, %r" + % (msg_code, msg)) + return msg_code, msg + + def _recv_pkt(self): + nmsglen = self._socket.recv(4) + if len(nmsglen) != 4: + raise RiakError( + "Socket returned short packet length %d - expected 4" + % len(nmsglen)) + msglen, = struct.unpack('!i', nmsglen) + self._inbuf_len = msglen + self._inbuf = '' + while len(self._inbuf) < msglen: + want_len = min(8192, msglen - len(self._inbuf)) + recv_buf = self._socket.recv(want_len) + if not recv_buf: + break + self._inbuf += recv_buf + if len(self._inbuf) != self._inbuf_len: + raise RiakError("Socket returned short packet %d - expected %d" + % (len(self._inbuf), self._inbuf_len)) + + def _connect(self): + self._socket = socket.create_connection(self._address, + self._timeouts['connect']) + + def close(self): + """ + Closes the underlying socket of the PB connection. + """ + self._socket.shutdown(socket.SHUT_RDWR) + + def _parse_msg(self, code, packet): + try: + pbclass = MESSAGE_CLASSES[code] + except KeyError: + pbclass = None + + if pbclass is None: + return None + + pbo = pbclass() + pbo.ParseFromString(packet) + return pbo + + # These are set in the RiakPbcTransport initializer + _address = None + _timeouts = {} diff --git a/riak/transports/pbc/messages.py b/riak/transports/pbc/messages.py new file mode 100644 index 00000000..fb9dba5c --- /dev/null +++ b/riak/transports/pbc/messages.py @@ -0,0 +1,92 @@ +""" +Copyright 2012 Basho Technologies, Inc. + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" + +import riak_pb + + +# Protocol codes +MSG_CODE_ERROR_RESP = 0 +MSG_CODE_PING_REQ = 1 +MSG_CODE_PING_RESP = 2 +MSG_CODE_GET_CLIENT_ID_REQ = 3 +MSG_CODE_GET_CLIENT_ID_RESP = 4 +MSG_CODE_SET_CLIENT_ID_REQ = 5 +MSG_CODE_SET_CLIENT_ID_RESP = 6 +MSG_CODE_GET_SERVER_INFO_REQ = 7 +MSG_CODE_GET_SERVER_INFO_RESP = 8 +MSG_CODE_GET_REQ = 9 +MSG_CODE_GET_RESP = 10 +MSG_CODE_PUT_REQ = 11 +MSG_CODE_PUT_RESP = 12 +MSG_CODE_DEL_REQ = 13 +MSG_CODE_DEL_RESP = 14 +MSG_CODE_LIST_BUCKETS_REQ = 15 +MSG_CODE_LIST_BUCKETS_RESP = 16 +MSG_CODE_LIST_KEYS_REQ = 17 +MSG_CODE_LIST_KEYS_RESP = 18 +MSG_CODE_GET_BUCKET_REQ = 19 +MSG_CODE_GET_BUCKET_RESP = 20 +MSG_CODE_SET_BUCKET_REQ = 21 +MSG_CODE_SET_BUCKET_RESP = 22 +MSG_CODE_MAPRED_REQ = 23 +MSG_CODE_MAPRED_RESP = 24 +MSG_CODE_INDEX_REQ = 25 +MSG_CODE_INDEX_RESP = 26 +MSG_CODE_SEARCH_QUERY_REQ = 27 +MSG_CODE_SEARCH_QUERY_RESP = 28 + +# These responses don't include messages +EMPTY_RESPONSES = [ + MSG_CODE_PING_RESP, + MSG_CODE_SET_CLIENT_ID_RESP, + MSG_CODE_DEL_RESP, + MSG_CODE_SET_BUCKET_RESP +] + +# Mapping from code to protobuf class +MESSAGE_CLASSES = { + MSG_CODE_ERROR_RESP: riak_pb.RpbErrorResp, + MSG_CODE_PING_REQ: None, + MSG_CODE_PING_RESP: None, + MSG_CODE_GET_CLIENT_ID_REQ: None, + MSG_CODE_GET_CLIENT_ID_RESP: riak_pb.RpbGetClientIdResp, + MSG_CODE_SET_CLIENT_ID_REQ: riak_pb.RpbSetClientIdReq, + MSG_CODE_SET_CLIENT_ID_RESP: None, + MSG_CODE_GET_SERVER_INFO_REQ: None, + MSG_CODE_GET_SERVER_INFO_RESP: riak_pb.RpbGetServerInfoResp, + MSG_CODE_GET_REQ: riak_pb.RpbGetReq, + MSG_CODE_GET_RESP: riak_pb.RpbGetResp, + MSG_CODE_PUT_REQ: riak_pb.RpbPutReq, + MSG_CODE_PUT_RESP: riak_pb.RpbPutResp, + MSG_CODE_DEL_REQ: riak_pb.RpbDelReq, + MSG_CODE_DEL_RESP: None, + MSG_CODE_LIST_BUCKETS_REQ: None, + MSG_CODE_LIST_BUCKETS_RESP: riak_pb.RpbListBucketsResp, + MSG_CODE_LIST_KEYS_REQ: riak_pb.RpbListKeysReq, + MSG_CODE_LIST_KEYS_RESP: riak_pb.RpbListKeysResp, + MSG_CODE_GET_BUCKET_REQ: riak_pb.RpbGetBucketReq, + MSG_CODE_GET_BUCKET_RESP: riak_pb.RpbGetBucketResp, + MSG_CODE_SET_BUCKET_REQ: riak_pb.RpbSetBucketReq, + MSG_CODE_SET_BUCKET_RESP: None, + MSG_CODE_MAPRED_REQ: riak_pb.RpbMapRedReq, + MSG_CODE_MAPRED_RESP: riak_pb.RpbMapRedResp, + MSG_CODE_INDEX_REQ: riak_pb.RpbIndexReq, + MSG_CODE_INDEX_RESP: riak_pb.RpbIndexResp, + MSG_CODE_SEARCH_QUERY_REQ: riak_pb.RpbSearchQueryReq, + MSG_CODE_SEARCH_QUERY_RESP: riak_pb.RpbSearchQueryResp +} diff --git a/riak/transports/pbc/stream.py b/riak/transports/pbc/stream.py new file mode 100644 index 00000000..ce0bda78 --- /dev/null +++ b/riak/transports/pbc/stream.py @@ -0,0 +1,97 @@ +""" +Copyright 2012 Basho Technologies, Inc. + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" + + +import json +from riak.transports.pbc.messages import ( + MSG_CODE_LIST_KEYS_RESP, MSG_CODE_MAPRED_RESP + ) + + +class RiakPbcStream(object): + """ + Used internally by RiakPbcTransport to implement streaming + operations. Implements the iterator interface. + """ + + _expect = None + + def __init__(self, transport): + self.finished = False + self.transport = transport + + def __iter__(self): + return self + + def next(self): + if self.finished: + raise StopIteration + + msg_code, resp = self.transport._recv_msg(self._expect) + if(self._is_done(resp)): + self.finished = True + + return resp + + def _is_done(self, response): + # This could break if new messages don't name the field the + # same thing. + return response.done + + def close(self): + # We have to drain the socket to make sure that we don't get + # weird responses when some other request comes after a + # failed/prematurely-terminated one. + try: + while self.next(): + pass + except StopIteration: + pass + + +class RiakPbcKeyStream(RiakPbcStream): + """ + Used internally by RiakPbcTransport to implement key-list streams. + """ + + _expect = MSG_CODE_LIST_KEYS_RESP + + def next(self): + response = super(RiakPbcKeyStream, self).next() + + if response.done and len(response.keys) is 0: + raise StopIteration + + return response.keys + + +class RiakPbcMapredStream(RiakPbcStream): + """ + Used internally by RiakPbcTransport to implement MapReduce + streams. + """ + + _expect = MSG_CODE_MAPRED_RESP + + def next(self): + response = super(RiakPbcMapredStream, self).next() + + if response.done and not response.HasField('response'): + raise StopIteration + + return response.phase, json.loads(response.response) diff --git a/riak/transports/pbc/transport.py b/riak/transports/pbc/transport.py new file mode 100644 index 00000000..1ed44a1a --- /dev/null +++ b/riak/transports/pbc/transport.py @@ -0,0 +1,419 @@ +""" +Copyright 2012 Basho Technologies, Inc. +Copyright 2010 Rusty Klophaus +Copyright 2010 Justin Sheehy +Copyright 2009 Jay Baird + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + +http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" + +import riak_pb +from riak import RiakError +from riak.transports.transport import RiakTransport +from connection import RiakPbcConnection +from stream import RiakPbcKeyStream, RiakPbcMapredStream +from codec import RiakPbcCodec +from messages import ( + # MSG_CODE_ERROR_RESP, + MSG_CODE_PING_REQ, + MSG_CODE_PING_RESP, + MSG_CODE_GET_CLIENT_ID_REQ, + MSG_CODE_GET_CLIENT_ID_RESP, + MSG_CODE_SET_CLIENT_ID_REQ, + MSG_CODE_SET_CLIENT_ID_RESP, + MSG_CODE_GET_SERVER_INFO_REQ, + MSG_CODE_GET_SERVER_INFO_RESP, + MSG_CODE_GET_REQ, + MSG_CODE_GET_RESP, + MSG_CODE_PUT_REQ, + MSG_CODE_PUT_RESP, + MSG_CODE_DEL_REQ, + MSG_CODE_DEL_RESP, + MSG_CODE_LIST_BUCKETS_REQ, + MSG_CODE_LIST_BUCKETS_RESP, + MSG_CODE_LIST_KEYS_REQ, + # MSG_CODE_LIST_KEYS_RESP, + MSG_CODE_GET_BUCKET_REQ, + MSG_CODE_GET_BUCKET_RESP, + MSG_CODE_SET_BUCKET_REQ, + MSG_CODE_SET_BUCKET_RESP, + MSG_CODE_MAPRED_REQ, + # MSG_CODE_MAPRED_RESP, + MSG_CODE_INDEX_REQ, + MSG_CODE_INDEX_RESP, + MSG_CODE_SEARCH_QUERY_REQ, + MSG_CODE_SEARCH_QUERY_RESP + ) + + +class RiakPbcTransport(RiakTransport, RiakPbcConnection, RiakPbcCodec): + """ + The RiakPbcTransport object holds a connection to the protocol + buffers interface on the riak server. + """ + + def __init__(self, node=None, client=None, connect_timeout=None, + request_timeout=None, **unused_options): + """ + Construct a new RiakPbcTransport object. + """ + super(RiakPbcTransport, self).__init__() + + self._client = client + self._node = node + self._address = (node.host, node.pb_port) + self._timeouts = {'connect': connect_timeout, + 'request': request_timeout} + self._connect() + + # FeatureDetection API + def _server_version(self): + return self.get_server_info()['server_version'] + + def ping(self): + """ + Ping the remote server + """ + + msg_code, msg = self._request(MSG_CODE_PING_REQ) + if msg_code == MSG_CODE_PING_RESP: + return True + else: + return False + + def get_server_info(self): + """ + Get information about the server + """ + msg_code, resp = self._request(MSG_CODE_GET_SERVER_INFO_REQ, + expect=MSG_CODE_GET_SERVER_INFO_RESP) + return {'node': resp.node, 'server_version': resp.server_version} + + def _get_client_id(self): + msg_code, resp = self._request(MSG_CODE_GET_CLIENT_ID_REQ, + expect=MSG_CODE_GET_CLIENT_ID_RESP) + return resp.client_id + + def _set_client_id(self, client_id): + req = riak_pb.RpbSetClientIdReq() + req.client_id = client_id + + msg_code, resp = self._request(MSG_CODE_SET_CLIENT_ID_REQ, req, + MSG_CODE_SET_CLIENT_ID_RESP) + + self._client_id = client_id + + client_id = property(_get_client_id, _set_client_id, + doc="""the client ID for this connection""") + + def get(self, robj, r=None, pr=None, vtag=None): + """ + Serialize get request and deserialize response + """ + if vtag is not None: + raise RiakError("PB transport does not support vtags") + + bucket = robj.bucket + + req = riak_pb.RpbGetReq() + if r: + req.r = self.translate_rw_val(r) + if self.quorum_controls() and pr: + req.pr = self.translate_rw_val(pr) + + if self.tombstone_vclocks(): + req.deletedvclock = 1 + + req.bucket = bucket.name + req.key = robj.key + + msg_code, resp = self._request(MSG_CODE_GET_REQ, req) + if msg_code == MSG_CODE_GET_RESP: + contents = [] + for c in resp.content: + contents.append(self.decode_content(c)) + return resp.vclock, contents + else: + return None + + def put(self, robj, w=None, dw=None, pw=None, return_body=True, + if_none_match=False): + """ + Serialize get request and deserialize response + """ + bucket = robj.bucket + + req = riak_pb.RpbPutReq() + if w: + req.w = self.translate_rw_val(w) + if dw: + req.dw = self.translate_rw_val(dw) + if self.quorum_controls() and pw: + req.pw = self.translate_rw_val(pw) + + if return_body: + req.return_body = 1 + if if_none_match: + req.if_none_match = 1 + + req.bucket = bucket.name + req.key = robj.key + vclock = robj.vclock + if vclock: + req.vclock = vclock + + self.encode_content(robj.metadata, + robj.get_encoded_data(), + req.content) + + msg_code, resp = self._request(MSG_CODE_PUT_REQ, req, + MSG_CODE_PUT_RESP) + if resp is not None: + contents = [] + for c in resp.content: + contents.append(self.decode_content(c)) + return resp.vclock, contents + + def put_new(self, robj, w=None, dw=None, pw=None, return_body=True, + if_none_match=False): + """Put a new object into the Riak store, returning its (new) key. + + If return_meta is False, then the vlock and metadata return values + will be None. + + @return (key, vclock, metadata) + """ + # Note that this won't work on 0.14 nodes. + bucket = robj.bucket + + req = riak_pb.RpbPutReq() + if w: + req.w = self.translate_rw_val(w) + if dw: + req.dw = self.translate_rw_val(dw) + if self.quorum_controls() and pw: + req.pw = self.translate_rw_val(pw) + + if return_body: + req.return_body = 1 + if if_none_match: + req.if_none_match = 1 + + req.bucket = bucket.name + + self.encode_content(robj.metadata, + robj.get_encoded_data(), + req.content) + + msg_code, resp = self._request(MSG_CODE_PUT_REQ, req, + MSG_CODE_PUT_RESP) + if not resp: + raise RiakError("missing response object") + if len(resp.content) != 1: + raise RiakError("siblings were returned from object creation") + + metadata, content = self.decode_content(resp.content[0]) + return resp.key, resp.vclock, metadata + + def delete(self, robj, rw=None, r=None, w=None, dw=None, pr=None, pw=None): + """ + Serialize get request and deserialize response + """ + bucket = robj.bucket + + req = riak_pb.RpbDelReq() + if rw: + req.rw = self.translate_rw_val(rw) + if r: + req.r = self.translate_rw_val(r) + if w: + req.w = self.translate_rw_val(w) + if dw: + req.dw = self.translate_rw_val(dw) + + if self.quorum_controls(): + if pr: + req.pr = self.translate_rw_val(pr) + if pw: + req.pw = self.translate_rw_val(pw) + + if self.tombstone_vclocks() and robj.vclock: + req.vclock = robj.vclock + + req.bucket = bucket.name + req.key = robj.key + + msg_code, resp = self._request(MSG_CODE_DEL_REQ, req, + MSG_CODE_DEL_RESP) + return self + + def get_keys(self, bucket): + """ + Lists all keys within a bucket. + """ + keys = [] + for keylist in self.stream_keys(bucket): + for key in keylist: + keys.append(key) + + return keys + + def stream_keys(self, bucket): + """ + Streams keys from a bucket, returning an iterator that yields + lists of keys. + """ + req = riak_pb.RpbListKeysReq() + req.bucket = bucket.name + + self._send_msg(MSG_CODE_LIST_KEYS_REQ, req) + + return RiakPbcKeyStream(self) + + def get_buckets(self): + """ + Serialize bucket listing request and deserialize response + """ + msg_code, resp = self._request(MSG_CODE_LIST_BUCKETS_REQ, + expect=MSG_CODE_LIST_BUCKETS_RESP) + return resp.buckets + + def get_bucket_props(self, bucket): + """ + Serialize bucket property request and deserialize response + """ + req = riak_pb.RpbGetBucketReq() + req.bucket = bucket.name + + msg_code, resp = self._request(MSG_CODE_GET_BUCKET_REQ, req, + MSG_CODE_GET_BUCKET_RESP) + props = {} + if resp.props.HasField('n_val'): + props['n_val'] = resp.props.n_val + if resp.props.HasField('allow_mult'): + props['allow_mult'] = resp.props.allow_mult + + return props + + def set_bucket_props(self, bucket, props): + """ + Serialize set bucket property request and deserialize response + """ + req = riak_pb.RpbSetBucketReq() + req.bucket = bucket.name + for key in props: + if key not in ['n_val', 'allow_mult']: + raise NotImplementedError + + if 'n_val' in props: + req.props.n_val = props['n_val'] + if 'allow_mult' in props: + req.props.allow_mult = props['allow_mult'] + + msg_code, resp = self._request(MSG_CODE_SET_BUCKET_REQ, req, + MSG_CODE_SET_BUCKET_RESP) + return self + + def mapred(self, inputs, query, timeout=None): + # dictionary of phase results - each content should be an encoded array + # which is appended to the result for that phase. + result = {} + for phase, content in self.stream_mapred(inputs, query, timeout): + if phase in result: + result[phase] += content + else: + result[phase] = content + + # If a single result - return the same as the HTTP interface does + # otherwise return all the phase information + if not len(result): + return None + elif len(result) == 1: + return result[max(result.keys())] + else: + return result + + def stream_mapred(self, inputs, query, timeout=None): + # Construct the job, optionally set the timeout... + content = self._construct_mapred_json(inputs, query, timeout) + + req = riak_pb.RpbMapRedReq() + req.request = content + req.content_type = "application/json" + + self._send_msg(MSG_CODE_MAPRED_REQ, req) + + return RiakPbcMapredStream(self) + + def get_index(self, bucket, index, startkey, endkey=None): + if not self.pb_indexes(): + return self._get_index_mapred_emu(bucket, index, startkey, endkey) + + req = riak_pb.RpbIndexReq(bucket=bucket, index=index) + if endkey: + req.qtype = riak_pb.RpbIndexReq.range + req.range_min = str(startkey) + req.range_max = str(endkey) + else: + req.qtype = riak_pb.RpbIndexReq.eq + req.key = str(startkey) + + msg_code, resp = self._request(MSG_CODE_INDEX_REQ, req, + MSG_CODE_INDEX_RESP) + return resp.keys + + def search(self, index, query, **params): + if not self.pb_search(): + return self._search_mapred_emu(index, query) + + req = riak_pb.RpbSearchQueryReq(index=index, q=query) + if 'rows' in params: + req.rows = params['rows'] + if 'start' in params: + req.start = params['start'] + if 'sort' in params: + req.sort = params['sort'] + if 'filter' in params: + req.filter = params['filter'] + if 'df' in params: + req.df = params['df'] + if 'op' in params: + req.op = params['op'] + if 'q.op' in params: + req.op = params['q.op'] + if 'fl' in params: + if isinstance(params['fl'], list): + req.fl.extend(params['fl']) + else: + req.fl.append(params['fl']) + if 'presort' in params: + req.presort = params['presort'] + + msg_code, resp = self._request(MSG_CODE_SEARCH_QUERY_REQ, req, + MSG_CODE_SEARCH_QUERY_RESP) + + result = {} + if resp.HasField('max_score'): + result['max_score'] = resp.max_score + if resp.HasField('num_found'): + result['num_found'] = resp.num_found + docs = [] + for doc in resp.docs: + resultdoc = {} + for pair in doc.fields: + resultdoc[pair.key] = pair.value + docs.append(resultdoc) + result['docs'] = docs + return result diff --git a/riak/transports/transport.py b/riak/transports/transport.py index 0b0ff32c..27d1fc52 100644 --- a/riak/transports/transport.py +++ b/riak/transports/transport.py @@ -17,12 +17,12 @@ specific language governing permissions and limitations under the License. """ -from riak import RiakError import base64 import random import threading import platform import os +import json from feature_detect import FeatureDetection @@ -31,11 +31,14 @@ class RiakTransport(FeatureDetection): Class to encapsulate transport details """ - # Subclasses should specify their API level. - # * missing or 1: the API used up and through 1.3.x. - # * 2: the API introduced with 1.4.x - # - # api = 2 + def _get_client_id(self): + return self._client_id + + def _set_client_id(self, value): + self._client_id = value + + client_id = property(_get_client_id, _set_client_id, + doc="""the client ID for this connection""") @classmethod def make_random_client_id(self): @@ -117,9 +120,27 @@ def set_bucket_props(self, bucket, props): """ raise NotImplementedError + def get_keys(self, bucket): + """ + Lists all keys within the given bucket. + """ + raise NotImplementedError + + def stream_keys(self, bucket): + """ + Streams the list of keys for the bucket through an iterator. + """ + raise NotImplementedError + def mapred(self, inputs, query, timeout=None): """ - Serialize map/reduce request + Sends a MapReduce request synchronously. + """ + raise NotImplementedError + + def stream_mapred(self, inputs, query, timeout=None): + """ + Streams the results of a MapReduce request through an iterator. """ raise NotImplementedError @@ -149,6 +170,18 @@ def get_index(self, bucket, index, startkey, endkey=None): """ raise NotImplementedError + def fulltext_add(self, index, *docs): + """ + Adds documents to the full-text index. + """ + raise NotImplementedError + + def fulltext_delete(self, index, docs=None, queries=None): + """ + Removes documents from the full-text index. + """ + raise NotImplementedError + def _search_mapred_emu(self, index, query): """ Emulates a search request via MapReduce. Used in the case @@ -198,3 +231,15 @@ def _get_index_mapred_emu(self, bucket, index, startkey, endkey=None): 'key': startkey}, phases) return [key for bucket, key in result] + + def _construct_mapred_json(self, inputs, query, timeout=None): + 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') + + job = {'inputs': inputs, 'query': query} + if timeout is not None: + job['timeout'] = timeout + + content = json.dumps(job) + return content diff --git a/riak/util.py b/riak/util.py index 6d58dcd8..c590db53 100644 --- a/riak/util.py +++ b/riak/util.py @@ -1,10 +1,23 @@ -import warnings +""" +Copyright 2010 Basho Technologies, Inc. + +This file is provided to you under the Apache License, +Version 2.0 (the "License"); you may not use this file +except in compliance with the License. You may obtain +a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 -try: - from collections import Mapping -except ImportError: - # compatibility with Python 2.5 - Mapping = dict +Unless required by applicable law or agreed to in writing, +software distributed under the License is distributed on an +"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +KIND, either express or implied. See the License for the +specific language governing permissions and limitations +under the License. +""" + +import warnings +from collections import Mapping def quacks_like_dict(object):