diff --git a/.gitignore b/.gitignore index 81fd6c28..e20b3037 100644 --- a/.gitignore +++ b/.gitignore @@ -1,22 +1,17 @@ +README.rst *.pyc .python-version __pycache__/ - .tox/ - -.tox/ - docs/_build - .*.swp .coverage - +riak-*/ py-build/ dist/ - riak.egg-info/ *.egg .eggs/ - #*# *~ +*.ps1 diff --git a/MANIFEST.in b/MANIFEST.in index d9f9a3fd..ddf59c00 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,6 +1,7 @@ include docs/* include riak/erl_src/* include README.md +include README.rst include LICENSE include RELNOTES.md include version.py diff --git a/buildbot/Makefile b/buildbot/Makefile index 8272fdf1..ecc933b8 100644 --- a/buildbot/Makefile +++ b/buildbot/Makefile @@ -26,7 +26,6 @@ compile: lint: @pip install --upgrade pep8 flake8 - @cd ..; pep8 --exclude=riak/pb riak *.py @cd ..; flake8 --exclude=riak/pb riak *.py test: setup test_normal test_security @@ -46,7 +45,7 @@ test_security: test_timeseries: @echo "Testing Riak Python Client (timeseries)" @$(RIAK_ADMIN) security disable - @RIAK_TEST_PROTOCOL='pbc' RUN_YZ=0 RUN_DATATYPES=0 RUN_INDEXES=1 RUN_TIMESERIES=1 ./tox_runner.sh .. + @RIAK_TEST_PROTOCOL='pbc' RIAK_TEST_PB_PORT=8087 RUN_YZ=0 RUN_DATATYPES=0 RUN_INDEXES=1 RUN_TIMESERIES=1 ./tox_runner.sh .. setup: ./tox_setup.sh diff --git a/docs/advanced.rst b/docs/advanced.rst index 523b465d..45475dd8 100644 --- a/docs/advanced.rst +++ b/docs/advanced.rst @@ -93,7 +93,7 @@ Transports .. currentmodule:: riak.transports.transport -.. autoclass:: RiakTransport +.. autoclass:: Transport :members: :private-members: @@ -124,20 +124,24 @@ HTTP Transport .. currentmodule:: riak.transports.http -.. autoclass:: RiakHttpPool +.. autoclass:: HttpPool .. autofunction:: is_retryable -.. autoclass:: RiakHttpTransport +.. autoclass:: HttpTransport :members: -^^^^^^^^^^^^^^^^^^^^^^^^^^ -Protocol Buffers Transport -^^^^^^^^^^^^^^^^^^^^^^^^^^ +^^^^^^^^^^^^^ +TCP Transport +^^^^^^^^^^^^^ + +.. currentmodule:: riak.transports.tcp -.. currentmodule:: riak.transports.pbc +.. autoclass:: TcpPool + +.. autofunction:: is_retryable -.. autoclass:: RiakPbcTransport +.. autoclass:: TcpTransport :members: --------- diff --git a/riak/benchmark.py b/riak/benchmark.py index 13286100..cfb220c1 100644 --- a/riak/benchmark.py +++ b/riak/benchmark.py @@ -1,24 +1,9 @@ -""" -Copyright 2013 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 __future__ import print_function + import os import gc +import sys +import traceback __all__ = ['measure', 'measure_with_rehearsal'] @@ -172,5 +157,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): elif exc_type is KeyboardInterrupt: return False else: - print("EXCEPTION! %r" % ((exc_type, exc_val, exc_tb),)) - return True + msg = "EXCEPTION! type: %r val: %r" % (exc_type, exc_val) + print(msg, file=sys.stderr) + traceback.print_tb(exc_tb) + return True if exc_type is None else False diff --git a/riak/benchmarks/timeseries.py b/riak/benchmarks/timeseries.py new file mode 100644 index 00000000..84960962 --- /dev/null +++ b/riak/benchmarks/timeseries.py @@ -0,0 +1,74 @@ +import datetime +import random +import sys + +import riak.benchmark as benchmark + +from multiprocessing import cpu_count +from riak import RiakClient + +# logger = logging.getLogger() +# logger.level = logging.DEBUG +# logger.addHandler(logging.StreamHandler(sys.stdout)) + +# batch sizes 8, 16, 32, 64, 128, 256 +if len(sys.argv) != 3: + raise AssertionError( + 'first arg is batch size, second arg is true / false' + 'for use_ttb') + +rowcount = 32768 +batchsz = int(sys.argv[1]) +if rowcount % batchsz != 0: + raise AssertionError('rowcount must be divisible by batchsz') +use_ttb = sys.argv[2].lower() == 'true' + +epoch = datetime.datetime.utcfromtimestamp(0) +onesec = datetime.timedelta(0, 1) + +weather = ['typhoon', 'hurricane', 'rain', 'wind', 'snow'] +rows = [] +for i in range(rowcount): + ts = datetime.datetime(2016, 1, 1, 12, 0, 0) + \ + datetime.timedelta(seconds=i) + family_idx = i % batchsz + series_idx = i % batchsz + family = 'hash{:d}'.format(family_idx) + series = 'user{:d}'.format(series_idx) + w = weather[i % len(weather)] + temp = (i % 100) + random.random() + row = [family, series, ts, w, temp] + key = [family, series, ts] + rows.append(row) + +print("Benchmarking timeseries:") +print(" Use TTB: {}".format(use_ttb)) +print("Batch Size: {}".format(batchsz)) +print(" CPUs: {}".format(cpu_count())) +print(" Rows: {}".format(len(rows))) +print() + +tbl = 'GeoCheckin' +h = 'riak-test' +n = [ + {'host': h, 'pb_port': 10017}, + {'host': h, 'pb_port': 10027}, + {'host': h, 'pb_port': 10037}, + {'host': h, 'pb_port': 10047}, + {'host': h, 'pb_port': 10057} +] +client = RiakClient(nodes=n, protocol='pbc', + transport_options={'use_ttb': use_ttb}) +table = client.table(tbl) + +with benchmark.measure() as b: + for i in (1, 2, 3): + with b.report('populate-%d' % i): + for i in range(0, rowcount, batchsz): + x = i + y = i + batchsz + r = rows[x:y] + ts_obj = table.new(r) + result = ts_obj.store() + if result is not True: + raise AssertionError("expected success") diff --git a/riak/client/__init__.py b/riak/client/__init__.py index d623f970..ea9abaca 100644 --- a/riak/client/__init__.py +++ b/riak/client/__init__.py @@ -11,8 +11,8 @@ from riak.mapreduce import RiakMapReduceChain from riak.resolver import default_resolver from riak.table import Table -from riak.transports.http import RiakHttpPool -from riak.transports.pbc import RiakPbcPool +from riak.transports.http import HttpPool +from riak.transports.tcp import TcpPool from riak.security import SecurityCreds from riak.util import lazy_property, bytes_to_str, str_to_bytes from six import string_types, PY2 @@ -68,7 +68,7 @@ class RiakClient(RiakMapReduceChain, RiakClientOperations): PROTOCOLS = ['http', 'pbc'] def __init__(self, protocol='pbc', transport_options={}, nodes=None, - credentials=None, multiget_pool_size=None, **unused_args): + credentials=None, multiget_pool_size=None, **kwargs): """ Construct a new ``RiakClient`` object. @@ -88,10 +88,10 @@ def __init__(self, protocol='pbc', transport_options={}, nodes=None, CPUs in the system :type multiget_pool_size: int """ - unused_args = unused_args.copy() + kwargs = kwargs.copy() if nodes is None: - self.nodes = [self._create_node(unused_args), ] + self.nodes = [self._create_node(kwargs), ] else: self.nodes = [self._create_node(n) for n in nodes] @@ -99,8 +99,8 @@ def __init__(self, protocol='pbc', transport_options={}, nodes=None, self.protocol = protocol or 'pbc' self._resolver = None self._credentials = self._create_credentials(credentials) - self._http_pool = RiakHttpPool(self, **transport_options) - self._pb_pool = RiakPbcPool(self, **transport_options) + self._http_pool = HttpPool(self, **transport_options) + self._tcp_pool = TcpPool(self, **transport_options) if PY2: self._encoders = {'application/json': default_encoder, @@ -167,7 +167,7 @@ def _get_client_id(self): def _set_client_id(self, client_id): for http in self._http_pool: http.client_id = client_id - for pb in self._pb_pool: + for pb in self._tcp_pool: pb.client_id = client_id client_id = property(_get_client_id, _set_client_id, @@ -298,8 +298,8 @@ def close(self): """ if self._http_pool is not None: self._http_pool.clear() - if self._pb_pool is not None: - self._pb_pool.clear() + if self._tcp_pool is not None: + self._tcp_pool.clear() def _create_node(self, n): if isinstance(n, RiakNode): diff --git a/riak/client/multiget.py b/riak/client/multiget.py index 64a0e3b4..9b5d7522 100644 --- a/riak/client/multiget.py +++ b/riak/client/multiget.py @@ -1,26 +1,9 @@ -""" -Copyright 2013 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 __future__ import print_function from collections import namedtuple from threading import Thread, Lock, Event from multiprocessing import cpu_count from six import PY2 + if PY2: from Queue import Queue else: @@ -177,8 +160,8 @@ def multiget(client, keys, **options): :meth:`RiakBucket.get ` :type options: dict :rtype: list - """ + outq = Queue() if 'pool' in options: diff --git a/riak/client/transport.py b/riak/client/transport.py index 027951d6..6aca7f24 100644 --- a/riak/client/transport.py +++ b/riak/client/transport.py @@ -1,23 +1,6 @@ -""" -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.tcp import is_retryable as is_pbc_retryable from riak.transports.http import is_retryable as is_http_retryable import threading from six import PY2 @@ -49,7 +32,7 @@ class RiakClientTransport(object): # These will be set or redefined by the RiakClient initializer protocol = 'pbc' _http_pool = None - _pb_pool = None + _tcp_pool = None _locals = _client_locals() def _get_retry_count(self): @@ -163,8 +146,8 @@ def _choose_pool(self, protocol=None): protocol = self.protocol if protocol == 'http': pool = self._http_pool - elif protocol == 'pbc': - pool = self._pb_pool + elif protocol == 'tcp' or protocol == 'pbc': + pool = self._tcp_pool else: raise ValueError("invalid protocol %s" % protocol) return pool diff --git a/riak/codecs/__init__.py b/riak/codecs/__init__.py new file mode 100644 index 00000000..e356b5f9 --- /dev/null +++ b/riak/codecs/__init__.py @@ -0,0 +1,25 @@ +import collections + +from riak import RiakError + +Msg = collections.namedtuple('Msg', + ['msg_code', 'data', 'resp_code'], + verbose=False) + + +class Codec(object): + def parse_msg(self): + raise NotImplementedError('parse_msg not implemented') + + def maybe_incorrect_code(self, resp_code, expect=None): + if expect and resp_code != expect: + raise RiakError("unexpected message code: %d, expected %d" + % (resp_code, expect)) + + def maybe_riak_error(self, err_code, msg_code, data=None): + if msg_code == err_code: + if data is None: + raise RiakError('no error provided!') + return data + else: + return None diff --git a/riak/transports/http/codec.py b/riak/codecs/http.py similarity index 93% rename from riak/transports/http/codec.py rename to riak/codecs/http.py index 9f040220..1078dd43 100644 --- a/riak/transports/http/codec.py +++ b/riak/codecs/http.py @@ -1,26 +1,6 @@ -""" -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 re import csv + from six import PY2, PY3 from cgi import parse_header from email import message_from_string @@ -32,6 +12,7 @@ from riak.multidict import MultiDict from riak.transports.http.search import XMLSearchResult from riak.util import decode_index_value, bytes_to_str + if PY2: from urllib import unquote_plus else: @@ -42,7 +23,7 @@ MAX_LINK_HEADER_SIZE = 8192 - 8 -class RiakHttpCodec(object): +class HttpCodec(object): """ Methods for HTTP transport that marshals and unmarshals HTTP messages. diff --git a/riak/codecs/pbuf.py b/riak/codecs/pbuf.py new file mode 100644 index 00000000..a976d1ff --- /dev/null +++ b/riak/codecs/pbuf.py @@ -0,0 +1,1242 @@ +import datetime +import six + +import riak.pb.messages +import riak.pb.riak_pb2 +import riak.pb.riak_dt_pb2 +import riak.pb.riak_kv_pb2 +import riak.pb.riak_ts_pb2 + +from riak import RiakError +from riak.codecs import Codec, Msg +from riak.content import RiakContent +from riak.pb.riak_ts_pb2 import TsColumnType +from riak.riak_object import VClock +from riak.ts_object import TsColumns +from riak.util import decode_index_value, str_to_bytes, bytes_to_str, \ + unix_time_millis, datetime_from_unix_time_millis +from riak.multidict import MultiDict + + +def _invert(d): + out = {} + for key in d: + value = d[key] + out[value] = key + return out + +REPL_TO_PY = { + riak.pb.riak_pb2.RpbBucketProps.FALSE: False, + riak.pb.riak_pb2.RpbBucketProps.TRUE: True, + riak.pb.riak_pb2.RpbBucketProps.REALTIME: 'realtime', + riak.pb.riak_pb2.RpbBucketProps.FULLSYNC: 'fullsync' +} + +REPL_TO_PB = _invert(REPL_TO_PY) + +RIAKC_RW_ONE = 4294967294 +RIAKC_RW_QUORUM = 4294967293 +RIAKC_RW_ALL = 4294967292 +RIAKC_RW_DEFAULT = 4294967291 + +QUORUM_TO_PB = {'default': RIAKC_RW_DEFAULT, + 'all': RIAKC_RW_ALL, + 'quorum': RIAKC_RW_QUORUM, + 'one': RIAKC_RW_ONE} + +QUORUM_TO_PY = _invert(QUORUM_TO_PB) + +NORMAL_PROPS = ['n_val', 'allow_mult', 'last_write_wins', 'old_vclock', + 'young_vclock', 'big_vclock', 'small_vclock', 'basic_quorum', + 'notfound_ok', 'search', 'backend', 'search_index', 'datatype', + 'write_once'] +COMMIT_HOOK_PROPS = ['precommit', 'postcommit'] +MODFUN_PROPS = ['chash_keyfun', 'linkfun'] +QUORUM_PROPS = ['r', 'pr', 'w', 'pw', 'dw', 'rw'] + +MAP_FIELD_TYPES = { + riak.pb.riak_dt_pb2.MapField.COUNTER: 'counter', + riak.pb.riak_dt_pb2.MapField.SET: 'set', + riak.pb.riak_dt_pb2.MapField.REGISTER: 'register', + riak.pb.riak_dt_pb2.MapField.FLAG: 'flag', + riak.pb.riak_dt_pb2.MapField.MAP: 'map', + 'counter': riak.pb.riak_dt_pb2.MapField.COUNTER, + 'set': riak.pb.riak_dt_pb2.MapField.SET, + 'register': riak.pb.riak_dt_pb2.MapField.REGISTER, + 'flag': riak.pb.riak_dt_pb2.MapField.FLAG, + 'map': riak.pb.riak_dt_pb2.MapField.MAP +} + +DT_FETCH_TYPES = { + riak.pb.riak_dt_pb2.DtFetchResp.COUNTER: 'counter', + riak.pb.riak_dt_pb2.DtFetchResp.SET: 'set', + riak.pb.riak_dt_pb2.DtFetchResp.MAP: 'map' +} + + +class PbufCodec(Codec): + ''' + Protobuffs Encoding and decoding methods for TcpTransport. + ''' + + def __init__(self, + client_timeouts=False, quorum_controls=False, + tombstone_vclocks=False, bucket_types=False): + if riak.pb is None: + raise NotImplementedError("this codec is not available") + self._client_timeouts = client_timeouts + self._quorum_controls = quorum_controls + self._tombstone_vclocks = tombstone_vclocks + self._bucket_types = bucket_types + + def parse_msg(self, msg_code, data): + pbclass = riak.pb.messages.MESSAGE_CLASSES.get(msg_code, None) + if pbclass is None: + return None + pbo = pbclass() + pbo.ParseFromString(data) + return pbo + + def maybe_riak_error(self, msg_code, data=None): + err_code = riak.pb.messages.MSG_CODE_ERROR_RESP + err_data = super(PbufCodec, self).maybe_riak_error( + err_code, msg_code, data) + if err_data: + err = self.parse_msg(msg_code, err_data) + raise RiakError(bytes_to_str(err.errmsg)) + + def encode_auth(self, username, password): + req = riak.pb.riak_pb2.RpbAuthReq() + req.user = str_to_bytes(username) + req.password = str_to_bytes(password) + mc = riak.pb.messages.MSG_CODE_AUTH_REQ + rc = riak.pb.messages.MSG_CODE_AUTH_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_ping(self): + return Msg(riak.pb.messages.MSG_CODE_PING_REQ, None, + riak.pb.messages.MSG_CODE_PING_RESP) + + def encode_quorum(self, rw): + """ + Converts a symbolic quorum value into its on-the-wire + equivalent. + + :param rw: the quorum + :type rw: string, integer + :rtype: integer + """ + if rw in QUORUM_TO_PB: + return QUORUM_TO_PB[rw] + elif type(rw) is int and rw >= 0: + return rw + else: + return None + + def decode_quorum(self, rw): + """ + Converts a protobuf quorum value to a symbolic value if + necessary. + + :param rw: the quorum + :type rw: int + :rtype int or string + """ + if rw in QUORUM_TO_PY: + return QUORUM_TO_PY[rw] + else: + return rw + + def decode_contents(self, contents, obj): + """ + Decodes the list of siblings from the protobuf representation + into the object. + + :param contents: a list of RpbContent messages + :type contents: list + :param obj: a RiakObject + :type obj: RiakObject + :rtype RiakObject + """ + obj.siblings = [self.decode_content(c, RiakContent(obj)) + for c in contents] + # Invoke sibling-resolution logic + if len(obj.siblings) > 1 and obj.resolver is not None: + obj.resolver(obj) + return obj + + def decode_content(self, rpb_content, sibling): + """ + Decodes a single sibling from the protobuf representation into + a RiakObject. + + :param rpb_content: a single RpbContent message + :type rpb_content: riak.pb.riak_pb2.RpbContent + :param sibling: a RiakContent sibling container + :type sibling: RiakContent + :rtype: RiakContent + """ + + if rpb_content.HasField("deleted") and rpb_content.deleted: + sibling.exists = False + else: + sibling.exists = True + if rpb_content.HasField("content_type"): + sibling.content_type = bytes_to_str(rpb_content.content_type) + if rpb_content.HasField("charset"): + sibling.charset = bytes_to_str(rpb_content.charset) + if rpb_content.HasField("content_encoding"): + sibling.content_encoding = \ + bytes_to_str(rpb_content.content_encoding) + if rpb_content.HasField("vtag"): + sibling.etag = bytes_to_str(rpb_content.vtag) + + sibling.links = [self.decode_link(link) + for link in rpb_content.links] + if rpb_content.HasField("last_mod"): + sibling.last_modified = float(rpb_content.last_mod) + if rpb_content.HasField("last_mod_usecs"): + sibling.last_modified += rpb_content.last_mod_usecs / 1000000.0 + + sibling.usermeta = dict([(bytes_to_str(usermd.key), + bytes_to_str(usermd.value)) + for usermd in rpb_content.usermeta]) + sibling.indexes = set([(bytes_to_str(index.key), + decode_index_value(index.key, index.value)) + for index in rpb_content.indexes]) + sibling.encoded_data = rpb_content.value + + return sibling + + def encode_content(self, robj, rpb_content): + """ + Fills an RpbContent message with the appropriate data and + metadata from a RiakObject. + + :param robj: a RiakObject + :type robj: RiakObject + :param rpb_content: the protobuf message to fill + :type rpb_content: riak.pb.riak_pb2.RpbContent + """ + if robj.content_type: + rpb_content.content_type = str_to_bytes(robj.content_type) + if robj.charset: + rpb_content.charset = str_to_bytes(robj.charset) + if robj.content_encoding: + rpb_content.content_encoding = str_to_bytes(robj.content_encoding) + for uk in robj.usermeta: + pair = rpb_content.usermeta.add() + pair.key = str_to_bytes(uk) + pair.value = str_to_bytes(robj.usermeta[uk]) + for link in robj.links: + pb_link = rpb_content.links.add() + try: + bucket, key, tag = link + except ValueError: + raise RiakError("Invalid link tuple %s" % link) + + pb_link.bucket = str_to_bytes(bucket) + pb_link.key = str_to_bytes(key) + if tag: + pb_link.tag = str_to_bytes(tag) + else: + pb_link.tag = str_to_bytes('') + + for field, value in robj.indexes: + pair = rpb_content.indexes.add() + pair.key = str_to_bytes(field) + pair.value = str_to_bytes(str(value)) + + # Python 2.x data is stored in a string + if six.PY2: + rpb_content.value = str(robj.encoded_data) + else: + rpb_content.value = robj.encoded_data + + def decode_link(self, link): + """ + Decodes an RpbLink message into a tuple + + :param link: an RpbLink message + :type link: riak.pb.riak_pb2.RpbLink + :rtype tuple + """ + + if link.HasField("bucket"): + bucket = bytes_to_str(link.bucket) + else: + bucket = None + if link.HasField("key"): + key = bytes_to_str(link.key) + else: + key = None + if link.HasField("tag"): + tag = bytes_to_str(link.tag) + else: + tag = None + + return (bucket, key, tag) + + def decode_index_value(self, index, value): + """ + Decodes a secondary index value into the correct Python type. + :param index: the name of the index + :type index: str + :param value: the value of the index entry + :type value: str + :rtype str or int + """ + if index.endswith("_int"): + return int(value) + else: + return bytes_to_str(value) + + def encode_bucket_props(self, props, msg): + """ + Encodes a dict of bucket properties into the protobuf message. + + :param props: bucket properties + :type props: dict + :param msg: the protobuf message to fill + :type msg: riak.pb.riak_pb2.RpbSetBucketReq + """ + for prop in NORMAL_PROPS: + if prop in props and props[prop] is not None: + if isinstance(props[prop], six.string_types): + setattr(msg.props, prop, str_to_bytes(props[prop])) + else: + setattr(msg.props, prop, props[prop]) + for prop in COMMIT_HOOK_PROPS: + if prop in props: + setattr(msg.props, 'has_' + prop, True) + self.encode_hooklist(props[prop], getattr(msg.props, prop)) + for prop in MODFUN_PROPS: + if prop in props and props[prop] is not None: + self.encode_modfun(props[prop], getattr(msg.props, prop)) + for prop in QUORUM_PROPS: + if prop in props and props[prop] not in (None, 'default'): + value = self.encode_quorum(props[prop]) + if value is not None: + if isinstance(value, six.string_types): + setattr(msg.props, prop, str_to_bytes(value)) + else: + setattr(msg.props, prop, value) + if 'repl' in props: + msg.props.repl = REPL_TO_PY[props['repl']] + + return msg + + def decode_bucket_props(self, msg): + """ + Decodes the protobuf bucket properties message into a dict. + + :param msg: the protobuf message to decode + :type msg: riak.pb.riak_pb2.RpbBucketProps + :rtype dict + """ + props = {} + for prop in NORMAL_PROPS: + if msg.HasField(prop): + props[prop] = getattr(msg, prop) + if isinstance(props[prop], bytes): + props[prop] = bytes_to_str(props[prop]) + for prop in COMMIT_HOOK_PROPS: + if getattr(msg, 'has_' + prop): + props[prop] = self.decode_hooklist(getattr(msg, prop)) + for prop in MODFUN_PROPS: + if msg.HasField(prop): + props[prop] = self.decode_modfun(getattr(msg, prop)) + for prop in QUORUM_PROPS: + if msg.HasField(prop): + props[prop] = self.decode_quorum(getattr(msg, prop)) + if msg.HasField('repl'): + props['repl'] = REPL_TO_PY[msg.repl] + return props + + def decode_modfun(self, modfun): + """ + Decodes a protobuf modfun pair into a dict with 'mod' and + 'fun' keys. Used in bucket properties. + + :param modfun: the protobuf message to decode + :type modfun: riak.pb.riak_pb2.RpbModFun + :rtype dict + """ + return {'mod': bytes_to_str(modfun.module), + 'fun': bytes_to_str(modfun.function)} + + def encode_modfun(self, props, msg=None): + """ + Encodes a dict with 'mod' and 'fun' keys into a protobuf + modfun pair. Used in bucket properties. + + :param props: the module/function pair + :type props: dict + :param msg: the protobuf message to fill + :type msg: riak.pb.riak_pb2.RpbModFun + :rtype riak.pb.riak_pb2.RpbModFun + """ + if msg is None: + msg = riak.pb.riak_pb2.RpbModFun() + msg.module = str_to_bytes(props['mod']) + msg.function = str_to_bytes(props['fun']) + return msg + + def decode_hooklist(self, hooklist): + """ + Decodes a list of protobuf commit hooks into their python + equivalents. Used in bucket properties. + + :param hooklist: a list of protobuf commit hooks + :type hooklist: list + :rtype list + """ + return [self.decode_hook(hook) for hook in hooklist] + + def encode_hooklist(self, hooklist, msg): + """ + Encodes a list of commit hooks into their protobuf equivalent. + Used in bucket properties. + + :param hooklist: a list of commit hooks + :type hooklist: list + :param msg: a protobuf field that is a list of commit hooks + """ + for hook in hooklist: + pbhook = msg.add() + self.encode_hook(hook, pbhook) + + def decode_hook(self, hook): + """ + Decodes a protobuf commit hook message into a dict. Used in + bucket properties. + + :param hook: the hook to decode + :type hook: riak.pb.riak_pb2.RpbCommitHook + :rtype dict + """ + if hook.HasField('modfun'): + return self.decode_modfun(hook.modfun) + else: + return {'name': bytes_to_str(hook.name)} + + def encode_hook(self, hook, msg): + """ + Encodes a commit hook dict into the protobuf message. Used in + bucket properties. + + :param hook: the hook to encode + :type hook: dict + :param msg: the protobuf message to fill + :type msg: riak.pb.riak_pb2.RpbCommitHook + :rtype riak.pb.riak_pb2.RpbCommitHook + """ + if 'name' in hook: + msg.name = str_to_bytes(hook['name']) + else: + self.encode_modfun(hook, msg.modfun) + return msg + + def encode_index_req(self, bucket, index, startkey, endkey=None, + return_terms=None, max_results=None, + continuation=None, timeout=None, term_regex=None, + streaming=False): + """ + Encodes a secondary index request into the protobuf message. + + :param bucket: the bucket whose index to query + :type bucket: string + :param index: the index to query + :type index: string + :param startkey: the value or beginning of the range + :type startkey: integer, string + :param endkey: the end of the range + :type endkey: integer, string + :param return_terms: whether to return the index term with the key + :type return_terms: bool + :param max_results: the maximum number of results to return (page size) + :type max_results: integer + :param continuation: the opaque continuation returned from a + previous paginated request + :type continuation: string + :param timeout: a timeout value in milliseconds, or 'infinity' + :type timeout: int + :param term_regex: a regular expression used to filter index terms + :type term_regex: string + :param streaming: encode as streaming request + :type streaming: bool + :rtype riak.pb.riak_kv_pb2.RpbIndexReq + """ + req = riak.pb.riak_kv_pb2.RpbIndexReq( + bucket=str_to_bytes(bucket.name), + index=str_to_bytes(index)) + self._add_bucket_type(req, bucket.bucket_type) + if endkey is not None: + req.qtype = riak.pb.riak_kv_pb2.RpbIndexReq.range + req.range_min = str_to_bytes(str(startkey)) + req.range_max = str_to_bytes(str(endkey)) + else: + req.qtype = riak.pb.riak_kv_pb2.RpbIndexReq.eq + req.key = str_to_bytes(str(startkey)) + if return_terms is not None: + req.return_terms = return_terms + if max_results: + req.max_results = max_results + if continuation: + req.continuation = str_to_bytes(continuation) + if timeout: + if timeout == 'infinity': + req.timeout = 0 + else: + req.timeout = timeout + if term_regex: + req.term_regex = str_to_bytes(term_regex) + req.stream = streaming + mc = riak.pb.messages.MSG_CODE_INDEX_REQ + rc = riak.pb.messages.MSG_CODE_INDEX_RESP + return Msg(mc, req.SerializeToString(), rc) + + def decode_index_req(self, resp, index, + return_terms=None, max_results=None): + if return_terms and resp.results: + results = [(decode_index_value(index, pair.key), + bytes_to_str(pair.value)) + for pair in resp.results] + else: + results = resp.keys[:] + if six.PY3: + results = [bytes_to_str(key) for key in resp.keys] + + if max_results is not None and resp.HasField('continuation'): + return (results, bytes_to_str(resp.continuation)) + else: + return (results, None) + + def decode_search_index(self, index): + """ + Fills an RpbYokozunaIndex message with the appropriate data. + + :param index: a yz index message + :type index: riak.pb.riak_yokozuna_pb2.RpbYokozunaIndex + :rtype dict + """ + result = {} + result['name'] = bytes_to_str(index.name) + if index.HasField('schema'): + result['schema'] = bytes_to_str(index.schema) + if index.HasField('n_val'): + result['n_val'] = index.n_val + return result + + def _add_bucket_type(self, req, bucket_type): + if bucket_type and not bucket_type.is_default(): + if not self._bucket_types: + raise NotImplementedError( + 'Server does not support bucket-types') + req.type = str_to_bytes(bucket_type.name) + + def encode_search_query(self, req, **kwargs): + if 'rows' in kwargs: + req.rows = kwargs['rows'] + if 'start' in kwargs: + req.start = kwargs['start'] + if 'sort' in kwargs: + req.sort = str_to_bytes(kwargs['sort']) + if 'filter' in kwargs: + req.filter = str_to_bytes(kwargs['filter']) + if 'df' in kwargs: + req.df = str_to_bytes(kwargs['df']) + if 'op' in kwargs: + req.op = str_to_bytes(kwargs['op']) + if 'q.op' in kwargs: + req.op = kwargs['q.op'] + if 'fl' in kwargs: + if isinstance(kwargs['fl'], list): + req.fl.extend(kwargs['fl']) + else: + req.fl.append(kwargs['fl']) + if 'presort' in kwargs: + req.presort = kwargs['presort'] + + def decode_search_doc(self, doc): + resultdoc = MultiDict() + for pair in doc.fields: + if six.PY2: + ukey = unicode(pair.key, 'utf-8') # noqa + uval = unicode(pair.value, 'utf-8') # noqa + else: + ukey = bytes_to_str(pair.key) + uval = bytes_to_str(pair.value) + resultdoc.add(ukey, uval) + return resultdoc.mixed() + + def decode_dt_fetch(self, resp): + dtype = DT_FETCH_TYPES.get(resp.type) + if dtype is None: + raise ValueError("Unknown datatype on wire: {}".format(resp.type)) + + value = self.decode_dt_value(dtype, resp.value) + + if resp.HasField('context'): + context = resp.context[:] + else: + context = None + + return dtype, value, context + + def decode_dt_value(self, dtype, msg): + if dtype == 'counter': + return msg.counter_value + elif dtype == 'set': + return self.decode_set_value(msg.set_value) + elif dtype == 'map': + return self.decode_map_value(msg.map_value) + + def encode_dt_options(self, req, **kwargs): + for q in ['r', 'pr', 'w', 'dw', 'pw']: + if q in kwargs and kwargs[q] is not None: + setattr(req, q, self.encode_quorum(kwargs[q])) + + for o in ['basic_quorum', 'notfound_ok', 'timeout', 'return_body', + 'include_context']: + if o in kwargs and kwargs[o] is not None: + setattr(req, o, kwargs[o]) + + def decode_map_value(self, entries): + out = {} + for entry in entries: + name = bytes_to_str(entry.field.name[:]) + dtype = MAP_FIELD_TYPES[entry.field.type] + if dtype == 'counter': + value = entry.counter_value + elif dtype == 'set': + value = self.decode_set_value(entry.set_value) + elif dtype == 'register': + value = bytes_to_str(entry.register_value[:]) + elif dtype == 'flag': + value = entry.flag_value + elif dtype == 'map': + value = self.decode_map_value(entry.map_value) + out[(name, dtype)] = value + return out + + def decode_set_value(self, set_value): + return [bytes_to_str(string[:]) for string in set_value] + + def encode_dt_op(self, dtype, req, op): + if dtype == 'counter': + req.op.counter_op.increment = op[1] + elif dtype == 'set': + self.encode_set_op(req.op, op) + elif dtype == 'map': + self.encode_map_op(req.op.map_op, op) + else: + raise TypeError("Cannot send operation on datatype {!r}". + format(dtype)) + + def encode_set_op(self, msg, op): + if 'adds' in op: + msg.set_op.adds.extend(str_to_bytes(op['adds'])) + if 'removes' in op: + msg.set_op.removes.extend(str_to_bytes(op['removes'])) + + def encode_map_op(self, msg, ops): + for op in ops: + name, dtype = op[1] + ftype = MAP_FIELD_TYPES[dtype] + if op[0] == 'add': + add = msg.adds.add() + add.name = str_to_bytes(name) + add.type = ftype + elif op[0] == 'remove': + remove = msg.removes.add() + remove.name = str_to_bytes(name) + remove.type = ftype + elif op[0] == 'update': + update = msg.updates.add() + update.field.name = str_to_bytes(name) + update.field.type = ftype + self.encode_map_update(dtype, update, op[2]) + + def encode_map_update(self, dtype, msg, op): + if dtype == 'counter': + # ('increment', some_int) + msg.counter_op.increment = op[1] + elif dtype == 'set': + self.encode_set_op(msg, op) + elif dtype == 'map': + self.encode_map_op(msg.map_op, op) + elif dtype == 'register': + # ('assign', some_str) + msg.register_op = str_to_bytes(op[1]) + elif dtype == 'flag': + if op == 'enable': + msg.flag_op = riak.pb.riak_dt_pb2.MapUpdate.ENABLE + else: + msg.flag_op = riak.pb.riak_dt_pb2.MapUpdate.DISABLE + + def encode_to_ts_cell(self, cell, ts_cell): + if cell is not None: + if isinstance(cell, datetime.datetime): + ts_cell.timestamp_value = unix_time_millis(cell) + elif isinstance(cell, bool): + ts_cell.boolean_value = cell + elif isinstance(cell, six.binary_type): + ts_cell.varchar_value = cell + elif isinstance(cell, six.text_type): + ts_cell.varchar_value = str_to_bytes(cell) + elif isinstance(cell, six.string_types): + ts_cell.varchar_value = str_to_bytes(cell) + elif (isinstance(cell, six.integer_types)): + ts_cell.sint64_value = cell + elif isinstance(cell, float): + ts_cell.double_value = cell + else: + t = type(cell) + raise RiakError("can't serialize type '{}', value '{}'" + .format(t, cell)) + + def encode_timeseries_keyreq(self, table, key, is_delete=False): + key_vals = None + if isinstance(key, list): + key_vals = key + else: + raise ValueError("key must be a list") + + req = riak.pb.riak_ts_pb2.TsGetReq() + mc = riak.pb.messages.MSG_CODE_TS_GET_REQ + rc = riak.pb.messages.MSG_CODE_TS_GET_RESP + if is_delete: + req = riak.pb.riak_ts_pb2.TsDelReq() + mc = riak.pb.messages.MSG_CODE_TS_DEL_REQ + rc = riak.pb.messages.MSG_CODE_TS_DEL_RESP + + req.table = str_to_bytes(table.name) + for cell in key_vals: + ts_cell = req.key.add() + self.encode_to_ts_cell(cell, ts_cell) + return Msg(mc, req.SerializeToString(), rc) + + def encode_timeseries_listkeysreq(self, table, timeout=None): + req = riak.pb.riak_ts_pb2.TsListKeysReq() + req.table = str_to_bytes(table.name) + if self._client_timeouts and timeout: + req.timeout = timeout + mc = riak.pb.messages.MSG_CODE_TS_LIST_KEYS_REQ + rc = riak.pb.messages.MSG_CODE_TS_LIST_KEYS_RESP + return Msg(mc, req.SerializeToString(), rc) + + def validate_timeseries_put_resp(self, resp_code, resp): + if resp is not None: + return True + else: + raise RiakError("missing response object") + + def encode_timeseries_put(self, tsobj): + """ + Fills an TsPutReq message with the appropriate data and + metadata from a TsObject. + + :param tsobj: a TsObject + :type tsobj: TsObject + :param req: the protobuf message to fill + :type req: riak.pb.riak_ts_pb2.TsPutReq + """ + req = riak.pb.riak_ts_pb2.TsPutReq() + req.table = str_to_bytes(tsobj.table.name) + + if tsobj.columns: + raise NotImplementedError("columns are not implemented yet") + + if tsobj.rows and isinstance(tsobj.rows, list): + for row in tsobj.rows: + tsr = req.rows.add() # NB: type TsRow + if not isinstance(row, list): + raise ValueError("TsObject row must be a list of values") + for cell in row: + tsc = tsr.cells.add() # NB: type TsCell + self.encode_to_ts_cell(cell, tsc) + else: + raise RiakError("TsObject requires a list of rows") + + mc = riak.pb.messages.MSG_CODE_TS_PUT_REQ + rc = riak.pb.messages.MSG_CODE_TS_PUT_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_timeseries_query(self, table, query, interpolations=None): + req = riak.pb.riak_ts_pb2.TsQueryReq() + q = query + if '{table}' in q: + q = q.format(table=table.name) + req.query.base = str_to_bytes(q) + mc = riak.pb.messages.MSG_CODE_TS_QUERY_REQ + rc = riak.pb.messages.MSG_CODE_TS_QUERY_RESP + return Msg(mc, req.SerializeToString(), rc) + + def decode_timeseries(self, resp, tsobj): + """ + Fills an TsObject with the appropriate data and + metadata from a TsGetResp / TsQueryResp. + + :param resp: the protobuf message from which to process data + :type resp: riak.pb.riak_ts_pb2.TsQueryRsp or + riak.pb.riak_ts_pb2.TsGetResp + :param tsobj: a TsObject + :type tsobj: TsObject + """ + if resp.columns is not None: + col_names = [] + col_types = [] + for col in resp.columns: + col_names.append(bytes_to_str(col.name)) + col_type = self.decode_timeseries_col_type(col.type) + col_types.append(col_type) + tsobj.columns = TsColumns(col_names, col_types) + + tsobj.rows = [] + if resp.rows is not None: + for row in resp.rows: + tsobj.rows.append( + self.decode_timeseries_row( + row, resp.columns)) + + def decode_timeseries_col_type(self, col_type): + # NB: these match the atom names for column types + if col_type == TsColumnType.Value('VARCHAR'): + return 'varchar' + elif col_type == TsColumnType.Value('SINT64'): + return 'sint64' + elif col_type == TsColumnType.Value('DOUBLE'): + return 'double' + elif col_type == TsColumnType.Value('TIMESTAMP'): + return 'timestamp' + elif col_type == TsColumnType.Value('BOOLEAN'): + return 'boolean' + else: + msg = 'could not decode column type: {}'.format(col_type) + raise RiakError(msg) + + def decode_timeseries_row(self, tsrow, tscols=None): + """ + Decodes a TsRow into a list + + :param tsrow: the protobuf TsRow to decode. + :type tsrow: riak.pb.riak_ts_pb2.TsRow + :param tscols: the protobuf TsColumn data to help decode. + :type tscols: list + :rtype list + """ + row = [] + for i, cell in enumerate(tsrow.cells): + col = None + if tscols is not None: + col = tscols[i] + if cell.HasField('varchar_value'): + if col and col.type != TsColumnType.Value('VARCHAR'): + raise TypeError('expected VARCHAR column') + else: + row.append(cell.varchar_value) + elif cell.HasField('sint64_value'): + if col and col.type != TsColumnType.Value('SINT64'): + raise TypeError('expected SINT64 column') + else: + row.append(cell.sint64_value) + elif cell.HasField('double_value'): + if col and col.type != TsColumnType.Value('DOUBLE'): + raise TypeError('expected DOUBLE column') + else: + row.append(cell.double_value) + elif cell.HasField('timestamp_value'): + if col and col.type != TsColumnType.Value('TIMESTAMP'): + raise TypeError('expected TIMESTAMP column') + else: + dt = datetime_from_unix_time_millis( + cell.timestamp_value) + row.append(dt) + elif cell.HasField('boolean_value'): + if col and col.type != TsColumnType.Value('BOOLEAN'): + raise TypeError('expected BOOLEAN column') + else: + row.append(cell.boolean_value) + else: + row.append(None) + return row + + def decode_preflist(self, item): + """ + Decodes a preflist response + + :param preflist: a bucket/key preflist + :type preflist: list of + riak.pb.riak_kv_pb2.RpbBucketKeyPreflistItem + :rtype dict + """ + result = {'partition': item.partition, + 'node': bytes_to_str(item.node), + 'primary': item. primary} + return result + + def encode_get(self, robj, r=None, pr=None, timeout=None, + basic_quorum=None, notfound_ok=None): + bucket = robj.bucket + req = riak.pb.riak_kv_pb2.RpbGetReq() + if r: + req.r = self.encode_quorum(r) + if self._quorum_controls: + if pr: + req.pr = self.encode_quorum(pr) + if basic_quorum is not None: + req.basic_quorum = basic_quorum + if notfound_ok is not None: + req.notfound_ok = notfound_ok + if self._client_timeouts and timeout: + req.timeout = timeout + if self._tombstone_vclocks: + req.deletedvclock = True + req.bucket = str_to_bytes(bucket.name) + self._add_bucket_type(req, bucket.bucket_type) + req.key = str_to_bytes(robj.key) + mc = riak.pb.messages.MSG_CODE_GET_REQ + rc = riak.pb.messages.MSG_CODE_GET_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_put(self, robj, w=None, dw=None, pw=None, + return_body=True, if_none_match=False, + timeout=None): + bucket = robj.bucket + req = riak.pb.riak_kv_pb2.RpbPutReq() + if w: + req.w = self.encode_quorum(w) + if dw: + req.dw = self.encode_quorum(dw) + if self._quorum_controls and pw: + req.pw = self.encode_quorum(pw) + if return_body: + req.return_body = 1 + if if_none_match: + req.if_none_match = 1 + if self._client_timeouts and timeout: + req.timeout = timeout + req.bucket = str_to_bytes(bucket.name) + self._add_bucket_type(req, bucket.bucket_type) + if robj.key: + req.key = str_to_bytes(robj.key) + if robj.vclock: + req.vclock = robj.vclock.encode('binary') + self.encode_content(robj, req.content) + mc = riak.pb.messages.MSG_CODE_PUT_REQ + rc = riak.pb.messages.MSG_CODE_PUT_RESP + return Msg(mc, req.SerializeToString(), rc) + + def decode_get(self, robj, resp): + if resp is not None: + if resp.HasField('vclock'): + robj.vclock = VClock(resp.vclock, 'binary') + # We should do this even if there are no contents, i.e. + # the object is tombstoned + self.decode_contents(resp.content, robj) + else: + # "not found" returns an empty message, + # so let's make sure to clear the siblings + robj.siblings = [] + return robj + + def decode_put(self, robj, resp): + if resp is not None: + if resp.HasField('key'): + robj.key = bytes_to_str(resp.key) + if resp.HasField("vclock"): + robj.vclock = VClock(resp.vclock, 'binary') + if resp.content: + self.decode_contents(resp.content, robj) + elif not robj.key: + raise RiakError("missing response object") + return robj + + def encode_delete(self, robj, rw=None, r=None, + w=None, dw=None, pr=None, pw=None, + timeout=None): + req = riak.pb.riak_kv_pb2.RpbDelReq() + if rw: + req.rw = self.encode_quorum(rw) + if r: + req.r = self.encode_quorum(r) + if w: + req.w = self.encode_quorum(w) + if dw: + req.dw = self.encode_quorum(dw) + + if self._quorum_controls: + if pr: + req.pr = self.encode_quorum(pr) + if pw: + req.pw = self.encode_quorum(pw) + + if self._client_timeouts and timeout: + req.timeout = timeout + + use_vclocks = (self._tombstone_vclocks and + hasattr(robj, 'vclock') and robj.vclock) + if use_vclocks: + req.vclock = robj.vclock.encode('binary') + + bucket = robj.bucket + req.bucket = str_to_bytes(bucket.name) + self._add_bucket_type(req, bucket.bucket_type) + req.key = str_to_bytes(robj.key) + mc = riak.pb.messages.MSG_CODE_DEL_REQ + rc = riak.pb.messages.MSG_CODE_DEL_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_stream_keys(self, bucket, timeout=None): + req = riak.pb.riak_kv_pb2.RpbListKeysReq() + req.bucket = str_to_bytes(bucket.name) + if self._client_timeouts and timeout: + req.timeout = timeout + self._add_bucket_type(req, bucket.bucket_type) + mc = riak.pb.messages.MSG_CODE_LIST_KEYS_REQ + rc = riak.pb.messages.MSG_CODE_LIST_KEYS_RESP + return Msg(mc, req.SerializeToString(), rc) + + def decode_get_keys(self, stream): + keys = [] + for keylist in stream: + for key in keylist: + keys.append(bytes_to_str(key)) + return keys + + def decode_get_server_info(self, resp): + return {'node': bytes_to_str(resp.node), + 'server_version': bytes_to_str(resp.server_version)} + + def encode_get_client_id(self): + mc = riak.pb.messages.MSG_CODE_GET_CLIENT_ID_REQ + rc = riak.pb.messages.MSG_CODE_GET_CLIENT_ID_RESP + return Msg(mc, None, rc) + + def decode_get_client_id(self, resp): + return bytes_to_str(resp.client_id) + + def encode_set_client_id(self, client_id): + req = riak.pb.riak_kv_pb2.RpbSetClientIdReq() + req.client_id = str_to_bytes(client_id) + mc = riak.pb.messages.MSG_CODE_SET_CLIENT_ID_REQ + rc = riak.pb.messages.MSG_CODE_SET_CLIENT_ID_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_get_buckets(self, bucket_type, + timeout=None, streaming=False): + # Bucket streaming landed in the same release as timeouts, so + # we don't need to check the capability. + req = riak.pb.riak_kv_pb2.RpbListBucketsReq() + req.stream = streaming + self._add_bucket_type(req, bucket_type) + if self._client_timeouts and timeout: + req.timeout = timeout + mc = riak.pb.messages.MSG_CODE_LIST_BUCKETS_REQ + rc = riak.pb.messages.MSG_CODE_LIST_BUCKETS_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_get_bucket_props(self, bucket): + req = riak.pb.riak_pb2.RpbGetBucketReq() + req.bucket = str_to_bytes(bucket.name) + self._add_bucket_type(req, bucket.bucket_type) + mc = riak.pb.messages.MSG_CODE_GET_BUCKET_REQ + rc = riak.pb.messages.MSG_CODE_GET_BUCKET_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_set_bucket_props(self, bucket, props): + req = riak.pb.riak_pb2.RpbSetBucketReq() + req.bucket = str_to_bytes(bucket.name) + self._add_bucket_type(req, bucket.bucket_type) + self.encode_bucket_props(props, req) + mc = riak.pb.messages.MSG_CODE_SET_BUCKET_REQ + rc = riak.pb.messages.MSG_CODE_SET_BUCKET_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_clear_bucket_props(self, bucket): + req = riak.pb.riak_pb2.RpbResetBucketReq() + req.bucket = str_to_bytes(bucket.name) + self._add_bucket_type(req, bucket.bucket_type) + mc = riak.pb.messages.MSG_CODE_RESET_BUCKET_REQ + rc = riak.pb.messages.MSG_CODE_RESET_BUCKET_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_get_bucket_type_props(self, bucket_type): + req = riak.pb.riak_pb2.RpbGetBucketTypeReq() + req.type = str_to_bytes(bucket_type.name) + mc = riak.pb.messages.MSG_CODE_GET_BUCKET_TYPE_REQ + rc = riak.pb.messages.MSG_CODE_GET_BUCKET_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_set_bucket_type_props(self, bucket_type, props): + req = riak.pb.riak_pb2.RpbSetBucketTypeReq() + req.type = str_to_bytes(bucket_type.name) + self.encode_bucket_props(props, req) + mc = riak.pb.messages.MSG_CODE_SET_BUCKET_TYPE_REQ + rc = riak.pb.messages.MSG_CODE_SET_BUCKET_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_stream_mapred(self, content): + req = riak.pb.riak_kv_pb2.RpbMapRedReq() + req.request = str_to_bytes(content) + req.content_type = str_to_bytes("application/json") + mc = riak.pb.messages.MSG_CODE_MAP_RED_REQ + rc = riak.pb.messages.MSG_CODE_MAP_RED_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_create_search_index(self, index, schema=None, + n_val=None, timeout=None): + index = str_to_bytes(index) + idx = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndex(name=index) + if schema: + idx.schema = str_to_bytes(schema) + if n_val: + idx.n_val = n_val + req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexPutReq(index=idx) + if timeout is not None: + req.timeout = timeout + mc = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_PUT_REQ + rc = riak.pb.messages.MSG_CODE_PUT_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_get_search_index(self, index): + req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexGetReq( + name=str_to_bytes(index)) + mc = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_REQ + rc = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_list_search_indexes(self): + req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexGetReq() + mc = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_REQ + rc = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_delete_search_index(self, index): + req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexDeleteReq( + name=str_to_bytes(index)) + mc = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_DELETE_REQ + rc = riak.pb.messages.MSG_CODE_DEL_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_create_search_schema(self, schema, content): + scma = riak.pb.riak_yokozuna_pb2.RpbYokozunaSchema( + name=str_to_bytes(schema), + content=str_to_bytes(content)) + req = riak.pb.riak_yokozuna_pb2.RpbYokozunaSchemaPutReq( + schema=scma) + mc = riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_PUT_REQ + rc = riak.pb.messages.MSG_CODE_PUT_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_get_search_schema(self, schema): + req = riak.pb.riak_yokozuna_pb2.RpbYokozunaSchemaGetReq( + name=str_to_bytes(schema)) + mc = riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_GET_REQ + rc = riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_GET_RESP + return Msg(mc, req.SerializeToString(), rc) + + def decode_get_search_schema(self, resp): + result = {} + result['name'] = bytes_to_str(resp.schema.name) + result['content'] = bytes_to_str(resp.schema.content) + return result + + def encode_search(self, index, query, **kwargs): + req = riak.pb.riak_search_pb2.RpbSearchQueryReq( + index=str_to_bytes(index), + q=str_to_bytes(query)) + self.encode_search_query(req, **kwargs) + mc = riak.pb.messages.MSG_CODE_SEARCH_QUERY_REQ + rc = riak.pb.messages.MSG_CODE_SEARCH_QUERY_RESP + return Msg(mc, req.SerializeToString(), rc) + + def decode_search(self, resp): + result = {} + if resp.HasField('max_score'): + result['max_score'] = resp.max_score + if resp.HasField('num_found'): + result['num_found'] = resp.num_found + result['docs'] = [self.decode_search_doc(doc) for doc in resp.docs] + return result + + def encode_get_counter(self, bucket, key, **kwargs): + req = riak.pb.riak_kv_pb2.RpbCounterGetReq() + req.bucket = str_to_bytes(bucket.name) + req.key = str_to_bytes(key) + if kwargs.get('r') is not None: + req.r = self.encode_quorum(kwargs['r']) + if kwargs.get('pr') is not None: + req.pr = self.encode_quorum(kwargs['pr']) + if kwargs.get('basic_quorum') is not None: + req.basic_quorum = kwargs['basic_quorum'] + if kwargs.get('notfound_ok') is not None: + req.notfound_ok = kwargs['notfound_ok'] + mc = riak.pb.messages.MSG_CODE_COUNTER_GET_REQ + rc = riak.pb.messages.MSG_CODE_COUNTER_GET_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_update_counter(self, bucket, key, value, **kwargs): + req = riak.pb.riak_kv_pb2.RpbCounterUpdateReq() + req.bucket = str_to_bytes(bucket.name) + req.key = str_to_bytes(key) + req.amount = value + if kwargs.get('w') is not None: + req.w = self.encode_quorum(kwargs['w']) + if kwargs.get('dw') is not None: + req.dw = self.encode_quorum(kwargs['dw']) + if kwargs.get('pw') is not None: + req.pw = self.encode_quorum(kwargs['pw']) + if kwargs.get('returnvalue') is not None: + req.returnvalue = kwargs['returnvalue'] + mc = riak.pb.messages.MSG_CODE_COUNTER_UPDATE_REQ + rc = riak.pb.messages.MSG_CODE_COUNTER_UPDATE_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_fetch_datatype(self, bucket, key, **kwargs): + req = riak.pb.riak_dt_pb2.DtFetchReq() + req.type = str_to_bytes(bucket.bucket_type.name) + req.bucket = str_to_bytes(bucket.name) + req.key = str_to_bytes(key) + self.encode_dt_options(req, **kwargs) + mc = riak.pb.messages.MSG_CODE_DT_FETCH_REQ + rc = riak.pb.messages.MSG_CODE_DT_FETCH_RESP + return Msg(mc, req.SerializeToString(), rc) + + def encode_update_datatype(self, datatype, **kwargs): + op = datatype.to_op() + type_name = datatype.type_name + if not op: + raise ValueError("No operation to send on datatype {!r}". + format(datatype)) + req = riak.pb.riak_dt_pb2.DtUpdateReq() + req.bucket = str_to_bytes(datatype.bucket.name) + req.type = str_to_bytes(datatype.bucket.bucket_type.name) + if datatype.key: + req.key = str_to_bytes(datatype.key) + if datatype._context: + req.context = datatype._context + self.encode_dt_options(req, **kwargs) + self.encode_dt_op(type_name, req, op) + mc = riak.pb.messages.MSG_CODE_DT_UPDATE_REQ + rc = riak.pb.messages.MSG_CODE_DT_UPDATE_RESP + return Msg(mc, req.SerializeToString(), rc) + + def decode_update_datatype(self, datatype, resp, **kwargs): + type_name = datatype.type_name + if resp.HasField('key'): + datatype.key = resp.key[:] + if resp.HasField('context'): + datatype._context = resp.context[:] + if kwargs.get('return_body'): + datatype._set_value(self.decode_dt_value(type_name, resp)) + + def encode_get_preflist(self, bucket, key): + req = riak.pb.riak_kv_pb2.RpbGetBucketKeyPreflistReq() + req.bucket = str_to_bytes(bucket.name) + req.key = str_to_bytes(key) + req.type = str_to_bytes(bucket.bucket_type.name) + mc = riak.pb.messages.MSG_CODE_GET_BUCKET_KEY_PREFLIST_REQ + rc = riak.pb.messages.MSG_CODE_GET_BUCKET_KEY_PREFLIST_RESP + return Msg(mc, req.SerializeToString(), rc) diff --git a/riak/codecs/ttb.py b/riak/codecs/ttb.py new file mode 100644 index 00000000..6bb81f0d --- /dev/null +++ b/riak/codecs/ttb.py @@ -0,0 +1,206 @@ +import datetime +import six + +from erlastic import encode, decode +from erlastic.types import Atom + +from riak import RiakError +from riak.codecs import Codec, Msg +from riak.ts_object import TsColumns +from riak.util import bytes_to_str, unix_time_millis, \ + datetime_from_unix_time_millis + +udef_a = Atom('undefined') + +rpberrorresp_a = Atom('rpberrorresp') +tsgetreq_a = Atom('tsgetreq') +tsgetresp_a = Atom('tsgetresp') +tsqueryreq_a = Atom('tsqueryreq') +tsqueryresp_a = Atom('tsqueryresp') +tsinterpolation_a = Atom('tsinterpolation') +tsputreq_a = Atom('tsputreq') +tsputresp_a = Atom('tsputresp') +tsdelreq_a = Atom('tsdelreq') +timestamp_a = Atom('timestamp') + +# TODO RTS-842 +MSG_CODE_TS_TTB = 104 + + +class TtbCodec(Codec): + ''' + Erlang term-to-binary Encoding and decoding methods for TcpTransport + ''' + + def __init__(self, **unused_args): + super(TtbCodec, self).__init__(**unused_args) + + def parse_msg(self, msg_code, data): + if msg_code != MSG_CODE_TS_TTB: + raise RiakError("TTB can't parse code: {}".format(msg_code)) + if len(data) > 0: + decoded = decode(data) + self.maybe_err_ttb(decoded) + return decoded + else: + return None + + def maybe_err_ttb(self, err_ttb): + resp_a = err_ttb[0] + if resp_a == rpberrorresp_a: + errmsg = err_ttb[1] + # errcode = err_ttb[2] + raise RiakError(bytes_to_str(errmsg)) + + def maybe_riak_error(self, msg_code, data=None): + pass + + def encode_to_ts_cell(self, cell): + if cell is None: + return [] + else: + if isinstance(cell, datetime.datetime): + ts = unix_time_millis(cell) + return ts + elif isinstance(cell, bool): + return cell + elif isinstance(cell, six.text_type) or \ + isinstance(cell, six.binary_type) or \ + isinstance(cell, six.string_types): + return cell + elif (isinstance(cell, six.integer_types)): + return cell + elif isinstance(cell, float): + return cell + else: + t = type(cell) + raise RiakError("can't serialize type '{}', value '{}'" + .format(t, cell)) + + def encode_timeseries_keyreq(self, table, key, is_delete=False): + key_vals = None + if isinstance(key, list): + key_vals = key + else: + raise ValueError("key must be a list") + + mc = MSG_CODE_TS_TTB + rc = MSG_CODE_TS_TTB + req_atom = tsgetreq_a + if is_delete: + req_atom = tsdelreq_a + + # TODO RTS-842 timeout is last + req = req_atom, table.name, \ + [self.encode_to_ts_cell(k) for k in key_vals], udef_a + return Msg(mc, encode(req), rc) + + def validate_timeseries_put_resp(self, resp_code, resp): + if resp is None and resp_code == MSG_CODE_TS_TTB: + return True + if resp is not None: + return True + else: + raise RiakError("missing response object") + + def encode_timeseries_put(self, tsobj): + ''' + Returns an Erlang-TTB encoded tuple with the appropriate data and + metadata from a TsObject. + + :param tsobj: a TsObject + :type tsobj: TsObject + :rtype: term-to-binary encoded object + ''' + if tsobj.columns: + raise NotImplementedError('columns are not used') + + if tsobj.rows and isinstance(tsobj.rows, list): + req_rows = [] + for row in tsobj.rows: + req_r = [] + for cell in row: + req_r.append(self.encode_to_ts_cell(cell)) + req_rows.append(tuple(req_r)) + req = tsputreq_a, tsobj.table.name, [], req_rows + mc = MSG_CODE_TS_TTB + rc = MSG_CODE_TS_TTB + return Msg(mc, encode(req), rc) + else: + raise RiakError("TsObject requires a list of rows") + + def encode_timeseries_query(self, table, query, interpolations=None): + q = query + if '{table}' in q: + q = q.format(table=table.name) + tsi = tsinterpolation_a, q, [] + req = tsqueryreq_a, tsi, False, [] + mc = MSG_CODE_TS_TTB + rc = MSG_CODE_TS_TTB + return Msg(mc, encode(req), rc) + + def decode_timeseries(self, resp_ttb, tsobj): + """ + Fills an TsObject with the appropriate data and + metadata from a TTB-encoded TsGetResp / TsQueryResp. + + :param resp_ttb: the decoded TTB data + :type resp_ttb: TTB-encoded tsqueryrsp or tsgetresp + :param tsobj: a TsObject + :type tsobj: TsObject + """ + if resp_ttb is None: + return tsobj + + self.maybe_err_ttb(resp_ttb) + + resp_a = resp_ttb[0] + if resp_a == tsputresp_a: + return + elif resp_a == tsgetresp_a or resp_a == tsqueryresp_a: + resp_data = resp_ttb[1] + if len(resp_data) == 0: + return + elif len(resp_data) == 3: + resp_colnames = resp_data[0] + resp_coltypes = resp_data[1] + tsobj.columns = self.decode_timeseries_cols( + resp_colnames, resp_coltypes) + resp_rows = resp_data[2] + tsobj.rows = [] + for resp_row in resp_rows: + tsobj.rows.append( + self.decode_timeseries_row(resp_row, resp_coltypes)) + else: + raise RiakError( + "Expected 3-tuple in response, got: {}".format(resp_data)) + else: + raise RiakError("Unknown TTB response type: {}".format(resp_a)) + + def decode_timeseries_cols(self, cnames, ctypes): + cnames = [bytes_to_str(cname) for cname in cnames] + ctypes = [str(ctype) for ctype in ctypes] + return TsColumns(cnames, ctypes) + + def decode_timeseries_row(self, tsrow, tsct): + """ + Decodes a TTB-encoded TsRow into a list + + :param tsrow: the TTB decoded TsRow to decode. + :type tsrow: TTB dncoded row + :param tsct: the TTB decoded column types (atoms). + :type tsct: list + :rtype list + """ + row = [] + for i, cell in enumerate(tsrow): + if cell is None: + row.append(None) + elif isinstance(cell, list) and len(cell) == 0: + row.append(None) + else: + if tsct[i] == timestamp_a: + row.append(datetime_from_unix_time_millis(cell)) + else: + row.append(cell) + return row diff --git a/riak/node.py b/riak/node.py index 332dc654..9a999ece 100644 --- a/riak/node.py +++ b/riak/node.py @@ -1,22 +1,6 @@ -""" -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 math import time + from threading import RLock diff --git a/riak/pb/messages.py b/riak/pb/messages.py index 0fea1f49..9bbf284a 100644 --- a/riak/pb/messages.py +++ b/riak/pb/messages.py @@ -77,6 +77,11 @@ MSG_CODE_TS_GET_RESP = 97 MSG_CODE_TS_LIST_KEYS_REQ = 98 MSG_CODE_TS_LIST_KEYS_RESP = 99 +MSG_CODE_TS_COVERAGE_REQ = 100 +MSG_CODE_TS_COVERAGE_RESP = 101 +MSG_CODE_TS_COVERAGE_ENTRY = 102 +MSG_CODE_TS_RANGE = 103 +MSG_CODE_TS_TTB_PUT_REQ = 104 MSG_CODE_TOGGLE_ENCODING_REQ = 110 MSG_CODE_TOGGLE_ENCODING_RESP = 111 MSG_CODE_AUTH_REQ = 253 @@ -159,6 +164,11 @@ MSG_CODE_TS_GET_RESP: riak.pb.riak_ts_pb2.TsGetResp, MSG_CODE_TS_LIST_KEYS_REQ: riak.pb.riak_ts_pb2.TsListKeysReq, MSG_CODE_TS_LIST_KEYS_RESP: riak.pb.riak_ts_pb2.TsListKeysResp, + MSG_CODE_TS_COVERAGE_REQ: riak.pb.riak_ts_pb2.TsCoverageReq, + MSG_CODE_TS_COVERAGE_RESP: riak.pb.riak_ts_pb2.TsCoverageResp, + MSG_CODE_TS_COVERAGE_ENTRY: riak.pb.riak_ts_pb2.TsCoverageEntry, + MSG_CODE_TS_RANGE: riak.pb.riak_ts_pb2.TsRange, + MSG_CODE_TS_TTB_PUT_REQ: riak.pb.riak_ts_pb2.TsTtbPutReq, MSG_CODE_TOGGLE_ENCODING_REQ: riak.pb.riak_pb2.RpbToggleEncodingReq, MSG_CODE_TOGGLE_ENCODING_RESP: riak.pb.riak_pb2.RpbToggleEncodingResp, MSG_CODE_AUTH_REQ: riak.pb.riak_pb2.RpbAuthReq, diff --git a/riak/pb/riak_ts_pb2.py b/riak/pb/riak_ts_pb2.py index b371fdea..ce9b250f 100644 --- a/riak/pb/riak_ts_pb2.py +++ b/riak/pb/riak_ts_pb2.py @@ -16,7 +16,7 @@ DESCRIPTOR = _descriptor.FileDescriptor( name='riak_ts.proto', package='', - serialized_pb='\n\rriak_ts.proto\x1a\nriak.proto\"D\n\nTsQueryReq\x12\x1f\n\x05query\x18\x01 \x01(\x0b\x32\x10.TsInterpolation\x12\x15\n\x06stream\x18\x02 \x01(\x08:\x05\x66\x61lse\"^\n\x0bTsQueryResp\x12%\n\x07\x63olumns\x18\x01 \x03(\x0b\x32\x14.TsColumnDescription\x12\x14\n\x04rows\x18\x02 \x03(\x0b\x32\x06.TsRow\x12\x12\n\x04\x64one\x18\x03 \x01(\x08:\x04true\"@\n\x08TsGetReq\x12\r\n\x05table\x18\x01 \x02(\x0c\x12\x14\n\x03key\x18\x02 \x03(\x0b\x32\x07.TsCell\x12\x0f\n\x07timeout\x18\x03 \x01(\r\"H\n\tTsGetResp\x12%\n\x07\x63olumns\x18\x01 \x03(\x0b\x32\x14.TsColumnDescription\x12\x14\n\x04rows\x18\x02 \x03(\x0b\x32\x06.TsRow\"V\n\x08TsPutReq\x12\r\n\x05table\x18\x01 \x02(\x0c\x12%\n\x07\x63olumns\x18\x02 \x03(\x0b\x32\x14.TsColumnDescription\x12\x14\n\x04rows\x18\x03 \x03(\x0b\x32\x06.TsRow\"\x0b\n\tTsPutResp\"P\n\x08TsDelReq\x12\r\n\x05table\x18\x01 \x02(\x0c\x12\x14\n\x03key\x18\x02 \x03(\x0b\x32\x07.TsCell\x12\x0e\n\x06vclock\x18\x03 \x01(\x0c\x12\x0f\n\x07timeout\x18\x04 \x01(\r\"\x0b\n\tTsDelResp\"A\n\x0fTsInterpolation\x12\x0c\n\x04\x62\x61se\x18\x01 \x02(\x0c\x12 \n\x0einterpolations\x18\x02 \x03(\x0b\x32\x08.RpbPair\"@\n\x13TsColumnDescription\x12\x0c\n\x04name\x18\x01 \x02(\x0c\x12\x1b\n\x04type\x18\x02 \x02(\x0e\x32\r.TsColumnType\"\x1f\n\x05TsRow\x12\x16\n\x05\x63\x65lls\x18\x01 \x03(\x0b\x32\x07.TsCell\"{\n\x06TsCell\x12\x15\n\rvarchar_value\x18\x01 \x01(\x0c\x12\x14\n\x0csint64_value\x18\x02 \x01(\x12\x12\x17\n\x0ftimestamp_value\x18\x03 \x01(\x12\x12\x15\n\rboolean_value\x18\x04 \x01(\x08\x12\x14\n\x0c\x64ouble_value\x18\x05 \x01(\x01\"/\n\rTsListKeysReq\x12\r\n\x05table\x18\x01 \x02(\x0c\x12\x0f\n\x07timeout\x18\x02 \x01(\r\"4\n\x0eTsListKeysResp\x12\x14\n\x04keys\x18\x01 \x03(\x0b\x32\x06.TsRow\x12\x0c\n\x04\x64one\x18\x02 \x01(\x08*O\n\x0cTsColumnType\x12\x0b\n\x07VARCHAR\x10\x00\x12\n\n\x06SINT64\x10\x01\x12\n\n\x06\x44OUBLE\x10\x02\x12\r\n\tTIMESTAMP\x10\x03\x12\x0b\n\x07\x42OOLEAN\x10\x04\x42#\n\x17\x63om.basho.riak.protobufB\x08RiakTsPB') + serialized_pb='\n\rriak_ts.proto\x1a\nriak.proto\"[\n\nTsQueryReq\x12\x1f\n\x05query\x18\x01 \x01(\x0b\x32\x10.TsInterpolation\x12\x15\n\x06stream\x18\x02 \x01(\x08:\x05\x66\x61lse\x12\x15\n\rcover_context\x18\x03 \x01(\x0c\"^\n\x0bTsQueryResp\x12%\n\x07\x63olumns\x18\x01 \x03(\x0b\x32\x14.TsColumnDescription\x12\x14\n\x04rows\x18\x02 \x03(\x0b\x32\x06.TsRow\x12\x12\n\x04\x64one\x18\x03 \x01(\x08:\x04true\"@\n\x08TsGetReq\x12\r\n\x05table\x18\x01 \x02(\x0c\x12\x14\n\x03key\x18\x02 \x03(\x0b\x32\x07.TsCell\x12\x0f\n\x07timeout\x18\x03 \x01(\r\"H\n\tTsGetResp\x12%\n\x07\x63olumns\x18\x01 \x03(\x0b\x32\x14.TsColumnDescription\x12\x14\n\x04rows\x18\x02 \x03(\x0b\x32\x06.TsRow\"V\n\x08TsPutReq\x12\r\n\x05table\x18\x01 \x02(\x0c\x12%\n\x07\x63olumns\x18\x02 \x03(\x0b\x32\x14.TsColumnDescription\x12\x14\n\x04rows\x18\x03 \x03(\x0b\x32\x06.TsRow\"Y\n\x0bTsTtbPutReq\x12\r\n\x05table\x18\x01 \x02(\x0c\x12%\n\x07\x63olumns\x18\x02 \x03(\x0b\x32\x14.TsColumnDescription\x12\x14\n\x04rows\x18\x03 \x03(\x0b\x32\x06.TsRow\"\x0b\n\tTsPutResp\"P\n\x08TsDelReq\x12\r\n\x05table\x18\x01 \x02(\x0c\x12\x14\n\x03key\x18\x02 \x03(\x0b\x32\x07.TsCell\x12\x0e\n\x06vclock\x18\x03 \x01(\x0c\x12\x0f\n\x07timeout\x18\x04 \x01(\r\"\x0b\n\tTsDelResp\"A\n\x0fTsInterpolation\x12\x0c\n\x04\x62\x61se\x18\x01 \x02(\x0c\x12 \n\x0einterpolations\x18\x02 \x03(\x0b\x32\x08.RpbPair\"@\n\x13TsColumnDescription\x12\x0c\n\x04name\x18\x01 \x02(\x0c\x12\x1b\n\x04type\x18\x02 \x02(\x0e\x32\r.TsColumnType\"\x1f\n\x05TsRow\x12\x16\n\x05\x63\x65lls\x18\x01 \x03(\x0b\x32\x07.TsCell\"{\n\x06TsCell\x12\x15\n\rvarchar_value\x18\x01 \x01(\x0c\x12\x14\n\x0csint64_value\x18\x02 \x01(\x12\x12\x17\n\x0ftimestamp_value\x18\x03 \x01(\x12\x12\x15\n\rboolean_value\x18\x04 \x01(\x08\x12\x14\n\x0c\x64ouble_value\x18\x05 \x01(\x01\"/\n\rTsListKeysReq\x12\r\n\x05table\x18\x01 \x02(\x0c\x12\x0f\n\x07timeout\x18\x02 \x01(\r\"4\n\x0eTsListKeysResp\x12\x14\n\x04keys\x18\x01 \x03(\x0b\x32\x06.TsRow\x12\x0c\n\x04\x64one\x18\x02 \x01(\x08\"q\n\rTsCoverageReq\x12\x1f\n\x05query\x18\x01 \x01(\x0b\x32\x10.TsInterpolation\x12\r\n\x05table\x18\x02 \x02(\x0c\x12\x15\n\rreplace_cover\x18\x03 \x01(\x0c\x12\x19\n\x11unavailable_cover\x18\x04 \x03(\x0c\"3\n\x0eTsCoverageResp\x12!\n\x07\x65ntries\x18\x01 \x03(\x0b\x32\x10.TsCoverageEntry\"[\n\x0fTsCoverageEntry\x12\n\n\x02ip\x18\x01 \x02(\x0c\x12\x0c\n\x04port\x18\x02 \x02(\r\x12\x15\n\rcover_context\x18\x03 \x02(\x0c\x12\x17\n\x05range\x18\x04 \x01(\x0b\x32\x08.TsRange\"\x93\x01\n\x07TsRange\x12\x12\n\nfield_name\x18\x01 \x02(\x0c\x12\x13\n\x0blower_bound\x18\x02 \x02(\x12\x12\x1d\n\x15lower_bound_inclusive\x18\x03 \x02(\x08\x12\x13\n\x0bupper_bound\x18\x04 \x02(\x12\x12\x1d\n\x15upper_bound_inclusive\x18\x05 \x02(\x08\x12\x0c\n\x04\x64\x65sc\x18\x06 \x02(\x0c*O\n\x0cTsColumnType\x12\x0b\n\x07VARCHAR\x10\x00\x12\n\n\x06SINT64\x10\x01\x12\n\n\x06\x44OUBLE\x10\x02\x12\r\n\tTIMESTAMP\x10\x03\x12\x0b\n\x07\x42OOLEAN\x10\x04\x42#\n\x17\x63om.basho.riak.protobufB\x08RiakTsPB') _TSCOLUMNTYPE = _descriptor.EnumDescriptor( name='TsColumnType', @@ -47,8 +47,8 @@ ], containing_type=None, options=None, - serialized_start=925, - serialized_end=1004, + serialized_start=1450, + serialized_end=1529, ) TsColumnType = enum_type_wrapper.EnumTypeWrapper(_TSCOLUMNTYPE) @@ -81,6 +81,13 @@ message_type=None, enum_type=None, containing_type=None, is_extension=False, extension_scope=None, options=None), + _descriptor.FieldDescriptor( + name='cover_context', full_name='TsQueryReq.cover_context', index=2, + number=3, type=12, cpp_type=9, label=1, + has_default_value=False, default_value="", + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), ], extensions=[ ], @@ -91,7 +98,7 @@ is_extendable=False, extension_ranges=[], serialized_start=29, - serialized_end=97, + serialized_end=120, ) @@ -132,8 +139,8 @@ options=None, is_extendable=False, extension_ranges=[], - serialized_start=99, - serialized_end=193, + serialized_start=122, + serialized_end=216, ) @@ -174,8 +181,8 @@ options=None, is_extendable=False, extension_ranges=[], - serialized_start=195, - serialized_end=259, + serialized_start=218, + serialized_end=282, ) @@ -209,8 +216,8 @@ options=None, is_extendable=False, extension_ranges=[], - serialized_start=261, - serialized_end=333, + serialized_start=284, + serialized_end=356, ) @@ -251,8 +258,50 @@ options=None, is_extendable=False, extension_ranges=[], - serialized_start=335, - serialized_end=421, + serialized_start=358, + serialized_end=444, +) + + +_TSTTBPUTREQ = _descriptor.Descriptor( + name='TsTtbPutReq', + full_name='TsTtbPutReq', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='table', full_name='TsTtbPutReq.table', index=0, + number=1, type=12, cpp_type=9, label=2, + has_default_value=False, default_value="", + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='columns', full_name='TsTtbPutReq.columns', index=1, + number=2, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='rows', full_name='TsTtbPutReq.rows', index=2, + number=3, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + extension_ranges=[], + serialized_start=446, + serialized_end=535, ) @@ -272,8 +321,8 @@ options=None, is_extendable=False, extension_ranges=[], - serialized_start=423, - serialized_end=434, + serialized_start=537, + serialized_end=548, ) @@ -321,8 +370,8 @@ options=None, is_extendable=False, extension_ranges=[], - serialized_start=436, - serialized_end=516, + serialized_start=550, + serialized_end=630, ) @@ -342,8 +391,8 @@ options=None, is_extendable=False, extension_ranges=[], - serialized_start=518, - serialized_end=529, + serialized_start=632, + serialized_end=643, ) @@ -377,8 +426,8 @@ options=None, is_extendable=False, extension_ranges=[], - serialized_start=531, - serialized_end=596, + serialized_start=645, + serialized_end=710, ) @@ -412,8 +461,8 @@ options=None, is_extendable=False, extension_ranges=[], - serialized_start=598, - serialized_end=662, + serialized_start=712, + serialized_end=776, ) @@ -440,8 +489,8 @@ options=None, is_extendable=False, extension_ranges=[], - serialized_start=664, - serialized_end=695, + serialized_start=778, + serialized_end=809, ) @@ -496,8 +545,8 @@ options=None, is_extendable=False, extension_ranges=[], - serialized_start=697, - serialized_end=820, + serialized_start=811, + serialized_end=934, ) @@ -531,8 +580,8 @@ options=None, is_extendable=False, extension_ranges=[], - serialized_start=822, - serialized_end=869, + serialized_start=936, + serialized_end=983, ) @@ -566,8 +615,197 @@ options=None, is_extendable=False, extension_ranges=[], - serialized_start=871, - serialized_end=923, + serialized_start=985, + serialized_end=1037, +) + + +_TSCOVERAGEREQ = _descriptor.Descriptor( + name='TsCoverageReq', + full_name='TsCoverageReq', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='query', full_name='TsCoverageReq.query', index=0, + number=1, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='table', full_name='TsCoverageReq.table', index=1, + number=2, type=12, cpp_type=9, label=2, + has_default_value=False, default_value="", + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='replace_cover', full_name='TsCoverageReq.replace_cover', index=2, + number=3, type=12, cpp_type=9, label=1, + has_default_value=False, default_value="", + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='unavailable_cover', full_name='TsCoverageReq.unavailable_cover', index=3, + number=4, type=12, cpp_type=9, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + extension_ranges=[], + serialized_start=1039, + serialized_end=1152, +) + + +_TSCOVERAGERESP = _descriptor.Descriptor( + name='TsCoverageResp', + full_name='TsCoverageResp', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='entries', full_name='TsCoverageResp.entries', index=0, + number=1, type=11, cpp_type=10, label=3, + has_default_value=False, default_value=[], + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + extension_ranges=[], + serialized_start=1154, + serialized_end=1205, +) + + +_TSCOVERAGEENTRY = _descriptor.Descriptor( + name='TsCoverageEntry', + full_name='TsCoverageEntry', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='ip', full_name='TsCoverageEntry.ip', index=0, + number=1, type=12, cpp_type=9, label=2, + has_default_value=False, default_value="", + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='port', full_name='TsCoverageEntry.port', index=1, + number=2, type=13, cpp_type=3, label=2, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='cover_context', full_name='TsCoverageEntry.cover_context', index=2, + number=3, type=12, cpp_type=9, label=2, + has_default_value=False, default_value="", + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='range', full_name='TsCoverageEntry.range', index=3, + number=4, type=11, cpp_type=10, label=1, + has_default_value=False, default_value=None, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + extension_ranges=[], + serialized_start=1207, + serialized_end=1298, +) + + +_TSRANGE = _descriptor.Descriptor( + name='TsRange', + full_name='TsRange', + filename=None, + file=DESCRIPTOR, + containing_type=None, + fields=[ + _descriptor.FieldDescriptor( + name='field_name', full_name='TsRange.field_name', index=0, + number=1, type=12, cpp_type=9, label=2, + has_default_value=False, default_value="", + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='lower_bound', full_name='TsRange.lower_bound', index=1, + number=2, type=18, cpp_type=2, label=2, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='lower_bound_inclusive', full_name='TsRange.lower_bound_inclusive', index=2, + number=3, type=8, cpp_type=7, label=2, + has_default_value=False, default_value=False, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='upper_bound', full_name='TsRange.upper_bound', index=3, + number=4, type=18, cpp_type=2, label=2, + has_default_value=False, default_value=0, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='upper_bound_inclusive', full_name='TsRange.upper_bound_inclusive', index=4, + number=5, type=8, cpp_type=7, label=2, + has_default_value=False, default_value=False, + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + _descriptor.FieldDescriptor( + name='desc', full_name='TsRange.desc', index=5, + number=6, type=12, cpp_type=9, label=2, + has_default_value=False, default_value="", + message_type=None, enum_type=None, containing_type=None, + is_extension=False, extension_scope=None, + options=None), + ], + extensions=[ + ], + nested_types=[], + enum_types=[ + ], + options=None, + is_extendable=False, + extension_ranges=[], + serialized_start=1301, + serialized_end=1448, ) _TSQUERYREQ.fields_by_name['query'].message_type = _TSINTERPOLATION @@ -578,16 +816,22 @@ _TSGETRESP.fields_by_name['rows'].message_type = _TSROW _TSPUTREQ.fields_by_name['columns'].message_type = _TSCOLUMNDESCRIPTION _TSPUTREQ.fields_by_name['rows'].message_type = _TSROW +_TSTTBPUTREQ.fields_by_name['columns'].message_type = _TSCOLUMNDESCRIPTION +_TSTTBPUTREQ.fields_by_name['rows'].message_type = _TSROW _TSDELREQ.fields_by_name['key'].message_type = _TSCELL _TSINTERPOLATION.fields_by_name['interpolations'].message_type = riak.pb.riak_pb2._RPBPAIR _TSCOLUMNDESCRIPTION.fields_by_name['type'].enum_type = _TSCOLUMNTYPE _TSROW.fields_by_name['cells'].message_type = _TSCELL _TSLISTKEYSRESP.fields_by_name['keys'].message_type = _TSROW +_TSCOVERAGEREQ.fields_by_name['query'].message_type = _TSINTERPOLATION +_TSCOVERAGERESP.fields_by_name['entries'].message_type = _TSCOVERAGEENTRY +_TSCOVERAGEENTRY.fields_by_name['range'].message_type = _TSRANGE DESCRIPTOR.message_types_by_name['TsQueryReq'] = _TSQUERYREQ DESCRIPTOR.message_types_by_name['TsQueryResp'] = _TSQUERYRESP DESCRIPTOR.message_types_by_name['TsGetReq'] = _TSGETREQ DESCRIPTOR.message_types_by_name['TsGetResp'] = _TSGETRESP DESCRIPTOR.message_types_by_name['TsPutReq'] = _TSPUTREQ +DESCRIPTOR.message_types_by_name['TsTtbPutReq'] = _TSTTBPUTREQ DESCRIPTOR.message_types_by_name['TsPutResp'] = _TSPUTRESP DESCRIPTOR.message_types_by_name['TsDelReq'] = _TSDELREQ DESCRIPTOR.message_types_by_name['TsDelResp'] = _TSDELRESP @@ -597,6 +841,10 @@ DESCRIPTOR.message_types_by_name['TsCell'] = _TSCELL DESCRIPTOR.message_types_by_name['TsListKeysReq'] = _TSLISTKEYSREQ DESCRIPTOR.message_types_by_name['TsListKeysResp'] = _TSLISTKEYSRESP +DESCRIPTOR.message_types_by_name['TsCoverageReq'] = _TSCOVERAGEREQ +DESCRIPTOR.message_types_by_name['TsCoverageResp'] = _TSCOVERAGERESP +DESCRIPTOR.message_types_by_name['TsCoverageEntry'] = _TSCOVERAGEENTRY +DESCRIPTOR.message_types_by_name['TsRange'] = _TSRANGE @add_metaclass(_reflection.GeneratedProtocolMessageType) class TsQueryReq(_message.Message): @@ -628,6 +876,12 @@ class TsPutReq(_message.Message): # @@protoc_insertion_point(class_scope:TsPutReq) +@add_metaclass(_reflection.GeneratedProtocolMessageType) +class TsTtbPutReq(_message.Message): + DESCRIPTOR = _TSTTBPUTREQ + + # @@protoc_insertion_point(class_scope:TsTtbPutReq) + @add_metaclass(_reflection.GeneratedProtocolMessageType) class TsPutResp(_message.Message): DESCRIPTOR = _TSPUTRESP @@ -682,6 +936,30 @@ class TsListKeysResp(_message.Message): # @@protoc_insertion_point(class_scope:TsListKeysResp) +@add_metaclass(_reflection.GeneratedProtocolMessageType) +class TsCoverageReq(_message.Message): + DESCRIPTOR = _TSCOVERAGEREQ + + # @@protoc_insertion_point(class_scope:TsCoverageReq) + +@add_metaclass(_reflection.GeneratedProtocolMessageType) +class TsCoverageResp(_message.Message): + DESCRIPTOR = _TSCOVERAGERESP + + # @@protoc_insertion_point(class_scope:TsCoverageResp) + +@add_metaclass(_reflection.GeneratedProtocolMessageType) +class TsCoverageEntry(_message.Message): + DESCRIPTOR = _TSCOVERAGEENTRY + + # @@protoc_insertion_point(class_scope:TsCoverageEntry) + +@add_metaclass(_reflection.GeneratedProtocolMessageType) +class TsRange(_message.Message): + DESCRIPTOR = _TSRANGE + + # @@protoc_insertion_point(class_scope:TsRange) + DESCRIPTOR.has_options = True DESCRIPTOR._options = _descriptor._ParseOptions(descriptor_pb2.FileOptions(), '\n\027com.basho.riak.protobufB\010RiakTsPB') diff --git a/riak/tests/base.py b/riak/tests/base.py index ec0f397c..97b9607c 100644 --- a/riak/tests/base.py +++ b/riak/tests/base.py @@ -9,7 +9,6 @@ class IntegrationTestBase(object): - host = None pb_port = None http_port = None @@ -28,7 +27,7 @@ def randname(length=12): @classmethod def create_client(cls, host=None, http_port=None, pb_port=None, - protocol=None, credentials=None, **client_args): + protocol=None, credentials=None, **kwargs): host = host or HOST http_port = http_port or HTTP_PORT pb_port = pb_port or PB_PORT @@ -43,22 +42,26 @@ def create_client(cls, host=None, http_port=None, pb_port=None, credentials = credentials or SECURITY_CREDS + if hasattr(cls, 'client_options'): + kwargs.update(cls.client_options) + if hasattr(cls, 'logging_enabled') and cls.logging_enabled: cls.logger.debug("RiakClient(protocol='%s', host='%s', " + "pb_port='%d', http_port='%d', " + - "credentials='%s', client_args='%s')", + "credentials='%s', kwargs='%s')", protocol, host, pb_port, http_port, credentials, - client_args) + kwargs) return RiakClient(protocol=protocol, host=host, http_port=http_port, credentials=credentials, - pb_port=pb_port, **client_args) + pb_port=pb_port, + **kwargs) @classmethod def setUpClass(cls): diff --git a/riak/tests/suite.py b/riak/tests/suite.py index 97f3532c..e317213a 100644 --- a/riak/tests/suite.py +++ b/riak/tests/suite.py @@ -1,10 +1,5 @@ import os.path -import platform - -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest +import unittest def additional_tests(): diff --git a/riak/tests/test_2i.py b/riak/tests/test_2i.py index d7b254e3..6db10602 100644 --- a/riak/tests/test_2i.py +++ b/riak/tests/test_2i.py @@ -1,14 +1,10 @@ # -*- coding: utf-8 -*- -import platform +import unittest + from riak import RiakError from riak.tests import RUN_INDEXES from riak.tests.base import IntegrationTestBase -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - class TwoITests(IntegrationTestBase, unittest.TestCase): def is_2i_supported(self): diff --git a/riak/tests/test_btypes.py b/riak/tests/test_btypes.py index 30c9d6ac..d0fe728b 100644 --- a/riak/tests/test_btypes.py +++ b/riak/tests/test_btypes.py @@ -1,15 +1,11 @@ -import platform +import unittest + from riak import RiakError, RiakObject from riak.bucket import RiakBucket, BucketType from riak.tests import RUN_BTYPES from riak.tests.base import IntegrationTestBase from riak.tests.comparison import Comparison -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - @unittest.skipUnless(RUN_BTYPES, "RUN_BTYPES is 0") class BucketTypeTests(IntegrationTestBase, unittest.TestCase, Comparison): diff --git a/riak/tests/test_client.py b/riak/tests/test_client.py index 46700a61..19379d06 100644 --- a/riak/tests/test_client.py +++ b/riak/tests/test_client.py @@ -1,4 +1,5 @@ -import platform +import unittest + from six import PY2 from threading import Thread from riak.riak_object import RiakObject @@ -10,11 +11,6 @@ else: from queue import Queue -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - class ClientTests(IntegrationTestBase, unittest.TestCase): def test_uses_client_id_if_given(self): @@ -201,10 +197,10 @@ def test_pool_close(self): # Do something to add to the connection pool self.test_multiget_bucket() if self.client.protocol == 'pbc': - self.assertGreater(len(self.client._pb_pool.resources), 1) + self.assertGreater(len(self.client._tcp_pool.resources), 1) else: self.assertGreater(len(self.client._http_pool.resources), 1) # Now close them all up self.client.close() self.assertEqual(len(self.client._http_pool.resources), 0) - self.assertEqual(len(self.client._pb_pool.resources), 0) + self.assertEqual(len(self.client._tcp_pool.resources), 0) diff --git a/riak/tests/test_comparison.py b/riak/tests/test_comparison.py index 446bc031..86fb9f8b 100644 --- a/riak/tests/test_comparison.py +++ b/riak/tests/test_comparison.py @@ -1,14 +1,10 @@ # -*- coding: utf-8 -*- -import platform +import unittest + from riak.riak_object import RiakObject from riak.bucket import RiakBucket, BucketType from riak.tests.base import IntegrationTestBase -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - class BucketTypeRichComparisonTest(unittest.TestCase): def test_btype_eq(self): diff --git a/riak/tests/test_datatypes.py b/riak/tests/test_datatypes.py index 39166069..3e945920 100644 --- a/riak/tests/test_datatypes.py +++ b/riak/tests/test_datatypes.py @@ -1,16 +1,12 @@ # -*- coding: utf-8 -*- -import platform -from riak import RiakBucket, BucketType, RiakObject +import unittest import riak.datatypes as datatypes + +from riak import RiakBucket, BucketType, RiakObject from riak.tests import RUN_DATATYPES from riak.tests.base import IntegrationTestBase from riak.tests.comparison import Comparison -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - class DatatypeUnitTestBase(object): dtype = None diff --git a/riak/tests/test_feature_detection.py b/riak/tests/test_feature_detection.py index d88334aa..894c8b17 100644 --- a/riak/tests/test_feature_detection.py +++ b/riak/tests/test_feature_detection.py @@ -1,11 +1,7 @@ # -*- coding: utf-8 -*- -import platform -from riak.transports.feature_detect import FeatureDetection +import unittest -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest +from riak.transports.feature_detect import FeatureDetection class IncompleteTransport(FeatureDetection): diff --git a/riak/tests/test_filters.py b/riak/tests/test_filters.py index c821ce95..e41eea6c 100644 --- a/riak/tests/test_filters.py +++ b/riak/tests/test_filters.py @@ -1,13 +1,9 @@ # -*- coding: utf-8 -*- -import platform +import unittest + from riak.mapreduce import RiakKeyFilter from riak import key_filter -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - class FilterTests(unittest.TestCase): def test_simple(self): diff --git a/riak/tests/test_kv.py b/riak/tests/test_kv.py index bfa2b888..5513c603 100644 --- a/riak/tests/test_kv.py +++ b/riak/tests/test_kv.py @@ -1,9 +1,10 @@ # -*- coding: utf-8 -*- +import copy import os -import platform -from six import string_types, PY2, PY3 +import sys +import unittest -import copy +from six import string_types, PY2, PY3 from time import sleep from riak import ConflictError, RiakBucket, RiakError from riak.resolver import default_resolver, last_written_resolver @@ -16,11 +17,6 @@ except ImportError: import json -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - if PY2: import cPickle test_pickle_dumps = cPickle.dumps @@ -695,7 +691,9 @@ def test_store_binary_object_from_file(self): obj.store() obj = bucket.get(self.key_name) self.assertNotEqual(obj.encoded_data, None) + is_win32 = sys.platform == 'win32' self.assertTrue(obj.content_type == 'text/x-python' or + (is_win32 and obj.content_type == 'text/plain') or obj.content_type == 'application/x-python-code') def test_store_binary_object_from_file_should_use_default_mimetype(self): diff --git a/riak/tests/test_mapreduce.py b/riak/tests/test_mapreduce.py index a1827398..b22a70ba 100644 --- a/riak/tests/test_mapreduce.py +++ b/riak/tests/test_mapreduce.py @@ -1,7 +1,7 @@ # -*- coding: utf-8 -*- from __future__ import print_function -import platform +import unittest from six import PY2 from riak.mapreduce import RiakMapReduce @@ -12,11 +12,6 @@ from riak.tests import RUN_SECURITY from riak.tests.yz_setup import yzSetUp, yzTearDown -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - testrun_yz_mr = {'btype': 'mr', 'bucket': 'mrbucket', diff --git a/riak/tests/test_pool.py b/riak/tests/test_pool.py index f1088244..a5f8ffd5 100644 --- a/riak/tests/test_pool.py +++ b/riak/tests/test_pool.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- +import unittest + from six import PY2 -import platform from threading import Thread, currentThread from riak.transports.pool import Pool, BadResource from random import SystemRandom @@ -8,11 +9,6 @@ from riak.tests import RUN_POOL from riak.tests.comparison import Comparison -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - if PY2: from Queue import Queue else: diff --git a/riak/tests/test_search.py b/riak/tests/test_search.py index 7cc369b6..17e2ea6a 100644 --- a/riak/tests/test_search.py +++ b/riak/tests/test_search.py @@ -1,14 +1,11 @@ # -*- coding: utf-8 -*- from __future__ import print_function -import platform + +import unittest + from riak.tests import RUN_SEARCH, RUN_YZ from riak.tests.base import IntegrationTestBase -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - testrun_search_bucket = 'searchbucket' diff --git a/riak/tests/test_security.py b/riak/tests/test_security.py index 056c1b48..8a3db8f7 100644 --- a/riak/tests/test_security.py +++ b/riak/tests/test_security.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -import platform import sys +import unittest from riak.tests import RUN_SECURITY, SECURITY_USER, SECURITY_PASSWD, \ SECURITY_CACERT, SECURITY_KEY, SECURITY_CERT, SECURITY_REVOKED, \ @@ -8,11 +8,6 @@ from riak.security import SecurityCreds from riak.tests.base import IntegrationTestBase -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - class SecurityTests(IntegrationTestBase, unittest.TestCase): @unittest.skipIf(RUN_SECURITY, 'RUN_SECURITY is 1') diff --git a/riak/tests/test_server_test.py b/riak/tests/test_server_test.py index d02debe6..45dcbd55 100644 --- a/riak/tests/test_server_test.py +++ b/riak/tests/test_server_test.py @@ -1,7 +1,10 @@ -from riak.test_server import TestServer +import sys import unittest +from riak.test_server import TestServer + +@unittest.skipIf(sys.platform == 'win32', 'Windows is not supported') class TestServerTestCase(unittest.TestCase): def setUp(self): self.test_server = TestServer() diff --git a/riak/tests/test_timeseries.py b/riak/tests/test_timeseries_pbuf.py similarity index 67% rename from riak/tests/test_timeseries.py rename to riak/tests/test_timeseries_pbuf.py index 600c67db..50c0818c 100644 --- a/riak/tests/test_timeseries.py +++ b/riak/tests/test_timeseries_pbuf.py @@ -1,25 +1,21 @@ # -*- coding: utf-8 -*- import datetime -import platform import random import string +import unittest import riak.pb.riak_ts_pb2 +from riak.pb.riak_ts_pb2 import TsColumnType from riak import RiakError +from riak.codecs.pbuf import PbufCodec from riak.table import Table +from riak.tests import RUN_TIMESERIES +from riak.tests.base import IntegrationTestBase from riak.ts_object import TsObject -from riak.transports.pbc.codec import RiakPbcCodec from riak.util import str_to_bytes, bytes_to_str, \ + unix_time_millis, datetime_from_unix_time_millis, \ is_timeseries_supported -from riak.tests import RUN_TIMESERIES -from riak.tests.base import IntegrationTestBase -from riak.pb.riak_ts_pb2 import TsColumnType - -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest table_name = 'GeoCheckin' @@ -36,17 +32,15 @@ @unittest.skipUnless(is_timeseries_supported(), "Timeseries not supported") -class TimeseriesUnitTests(unittest.TestCase): +class TimeseriesPbufUnitTests(unittest.TestCase): @classmethod def setUpClass(cls): - cls.c = RiakPbcCodec() - - cls.ts0ms = cls.c._unix_time_millis(ts0) + cls.ts0ms = unix_time_millis(ts0) if cls.ts0ms != ex0ms: raise AssertionError( 'expected {:d} to equal {:d}'.format(cls.ts0ms, ex0ms)) - cls.ts1ms = cls.c._unix_time_millis(ts1) + cls.ts1ms = unix_time_millis(ts1) if cls.ts1ms != ex1ms: raise AssertionError( 'expected {:d} to equal {:d}'.format(cls.ts1ms, ex1ms)) @@ -56,7 +50,7 @@ def setUpClass(cls): [bd1, 3, 4.5, ts1, False] ] cls.test_key = ['hash1', 'user2', ts0] - cls.table = Table(None, 'test-table') + cls.table = Table(None, table_name) def validate_keyreq(self, req): self.assertEqual(self.table.name, bytes_to_str(req.table)) @@ -66,31 +60,39 @@ def validate_keyreq(self, req): self.assertEqual(self.ts0ms, req.key[2].timestamp_value) def test_encode_decode_timestamp(self): - ts0ms = self.c._unix_time_millis(ts0) + ts0ms = unix_time_millis(ts0) self.assertEqual(ts0ms, ex0ms) - ts0_d = self.c._datetime_from_unix_time_millis(ts0ms) + ts0_d = datetime_from_unix_time_millis(ts0ms) self.assertEqual(ts0, ts0_d) def test_encode_data_for_get(self): + c = PbufCodec() + msg = c.encode_timeseries_keyreq( + self.table, self.test_key, is_delete=False) req = riak.pb.riak_ts_pb2.TsGetReq() - self.c._encode_timeseries_keyreq(self.table, self.test_key, req) + req.ParseFromString(msg.data) self.validate_keyreq(req) def test_encode_data_for_delete(self): + c = PbufCodec() + msg = c.encode_timeseries_keyreq( + self.table, self.test_key, is_delete=True) req = riak.pb.riak_ts_pb2.TsDelReq() - self.c._encode_timeseries_keyreq(self.table, self.test_key, req) + req.ParseFromString(msg.data) self.validate_keyreq(req) def test_encode_data_for_put(self): + c = PbufCodec() tsobj = TsObject(None, self.table, self.rows, None) - ts_put_req = riak.pb.riak_ts_pb2.TsPutReq() - self.c._encode_timeseries_put(tsobj, ts_put_req) + msg = c.encode_timeseries_put(tsobj) + req = riak.pb.riak_ts_pb2.TsPutReq() + req.ParseFromString(msg.data) # NB: expected, actual - self.assertEqual(self.table.name, bytes_to_str(ts_put_req.table)) - self.assertEqual(len(self.rows), len(ts_put_req.rows)) + self.assertEqual(self.table.name, bytes_to_str(req.table)) + self.assertEqual(len(self.rows), len(req.rows)) - r0 = ts_put_req.rows[0] + r0 = req.rows[0] self.assertEqual(bytes_to_str(r0.cells[0].varchar_value), self.rows[0][0]) self.assertEqual(r0.cells[1].sint64_value, self.rows[0][1]) @@ -98,7 +100,7 @@ def test_encode_data_for_put(self): self.assertEqual(r0.cells[3].timestamp_value, self.ts0ms) self.assertEqual(r0.cells[4].boolean_value, self.rows[0][4]) - r1 = ts_put_req.rows[1] + r1 = req.rows[1] self.assertEqual(bytes_to_str(r1.cells[0].varchar_value), self.rows[1][0]) self.assertEqual(r1.cells[1].sint64_value, self.rows[1][1]) @@ -107,8 +109,10 @@ def test_encode_data_for_put(self): self.assertEqual(r1.cells[4].boolean_value, self.rows[1][4]) def test_encode_data_for_listkeys(self): + c = PbufCodec(client_timeouts=True) + msg = c.encode_timeseries_listkeysreq(self.table, 1234) req = riak.pb.riak_ts_pb2.TsListKeysReq() - self.c._encode_timeseries_listkeysreq(self.table, req, 1234) + req.ParseFromString(msg.data) self.assertEqual(self.table.name, bytes_to_str(req.table)) self.assertEqual(1234, req.timeout) @@ -155,34 +159,35 @@ def test_decode_data_from_query(self): r1c4 = r1.cells.add() r1c4.boolean_value = self.rows[1][4] - tsobj = TsObject(None, self.table, [], []) - c = RiakPbcCodec() - c._decode_timeseries(tqr, tsobj) - - self.assertEqual(len(self.rows), len(tsobj.rows)) - self.assertEqual(len(tqr.columns), len(tsobj.columns)) - - c = tsobj.columns - self.assertEqual(c[0][0], 'col_varchar') - self.assertEqual(c[0][1], TsColumnType.Value('VARCHAR')) - self.assertEqual(c[1][0], 'col_integer') - self.assertEqual(c[1][1], TsColumnType.Value('SINT64')) - self.assertEqual(c[2][0], 'col_double') - self.assertEqual(c[2][1], TsColumnType.Value('DOUBLE')) - self.assertEqual(c[3][0], 'col_timestamp') - self.assertEqual(c[3][1], TsColumnType.Value('TIMESTAMP')) - self.assertEqual(c[4][0], 'col_boolean') - self.assertEqual(c[4][1], TsColumnType.Value('BOOLEAN')) + tsobj = TsObject(None, self.table) + c = PbufCodec() + c.decode_timeseries(tqr, tsobj) + + self.assertEqual(len(tsobj.rows), len(self.rows)) + self.assertEqual(len(tsobj.columns.names), len(tqr.columns)) + self.assertEqual(len(tsobj.columns.types), len(tqr.columns)) + + cn, ct = tsobj.columns + self.assertEqual(cn[0], 'col_varchar') + self.assertEqual(ct[0], 'varchar') + self.assertEqual(cn[1], 'col_integer') + self.assertEqual(ct[1], 'sint64') + self.assertEqual(cn[2], 'col_double') + self.assertEqual(ct[2], 'double') + self.assertEqual(cn[3], 'col_timestamp') + self.assertEqual(ct[3], 'timestamp') + self.assertEqual(cn[4], 'col_boolean') + self.assertEqual(ct[4], 'boolean') r0 = tsobj.rows[0] - self.assertEqual(r0[0], self.rows[0][0]) + self.assertEqual(bytes_to_str(r0[0]), self.rows[0][0]) self.assertEqual(r0[1], self.rows[0][1]) self.assertEqual(r0[2], self.rows[0][2]) self.assertEqual(r0[3], ts0) self.assertEqual(r0[4], self.rows[0][4]) r1 = tsobj.rows[1] - self.assertEqual(r1[0], self.rows[1][0]) + self.assertEqual(bytes_to_str(r1[0]), self.rows[1][0]) self.assertEqual(r1[1], self.rows[1][1]) self.assertEqual(r1[2], self.rows[1][2]) self.assertEqual(r1[3], ts1) @@ -191,10 +196,12 @@ def test_decode_data_from_query(self): @unittest.skipUnless(is_timeseries_supported() and RUN_TIMESERIES, 'Timeseries not supported or RUN_TIMESERIES is 0') -class TimeseriesTests(IntegrationTestBase, unittest.TestCase): +class TimeseriesPbufTests(IntegrationTestBase, unittest.TestCase): + client_options = {'transport_options': {'use_ttb': False}} + @classmethod def setUpClass(cls): - super(TimeseriesTests, cls).setUpClass() + super(TimeseriesPbufTests, cls).setUpClass() cls.now = datetime.datetime.utcfromtimestamp(144379690.987000) fiveMinsAgo = cls.now - fiveMins tenMinsAgo = fiveMinsAgo - fiveMins @@ -217,26 +224,44 @@ def setUpClass(cls): raise AssertionError("expected success") client.close() - codec = RiakPbcCodec() - cls.nowMsec = codec._unix_time_millis(cls.now) + cls.nowMsec = unix_time_millis(cls.now) cls.fiveMinsAgo = fiveMinsAgo cls.twentyMinsAgo = twentyMinsAgo cls.twentyFiveMinsAgo = twentyFiveMinsAgo - cls.tenMinsAgoMsec = codec._unix_time_millis(tenMinsAgo) - cls.twentyMinsAgoMsec = codec._unix_time_millis(twentyMinsAgo) + cls.tenMinsAgoMsec = unix_time_millis(tenMinsAgo) + cls.twentyMinsAgoMsec = unix_time_millis(twentyMinsAgo) cls.numCols = len(rows[0]) cls.rows = rows + encoded_rows = [ + [str_to_bytes('hash1'), str_to_bytes('user2'), + twentyFiveMinsAgo, str_to_bytes('typhoon'), 90.3], + [str_to_bytes('hash1'), str_to_bytes('user2'), + twentyMinsAgo, str_to_bytes('hurricane'), 82.3], + [str_to_bytes('hash1'), str_to_bytes('user2'), + fifteenMinsAgo, str_to_bytes('rain'), 79.0], + [str_to_bytes('hash1'), str_to_bytes('user2'), + fiveMinsAgo, str_to_bytes('wind'), None], + [str_to_bytes('hash1'), str_to_bytes('user2'), + cls.now, str_to_bytes('snow'), 20.1] + ] + cls.encoded_rows = encoded_rows + + def validate_len(self, ts_obj, expected_len): + self.assertEqual(len(ts_obj.columns.names), expected_len) + self.assertEqual(len(ts_obj.columns.types), expected_len) + self.assertEqual(len(ts_obj.rows), expected_len) def validate_data(self, ts_obj): if ts_obj.columns is not None: - self.assertEqual(len(ts_obj.columns), self.numCols) + self.assertEqual(len(ts_obj.columns.names), self.numCols) + self.assertEqual(len(ts_obj.columns.types), self.numCols) self.assertEqual(len(ts_obj.rows), 1) row = ts_obj.rows[0] - self.assertEqual(row[0], 'hash1') - self.assertEqual(row[1], 'user2') + self.assertEqual(bytes_to_str(row[0]), 'hash1') + self.assertEqual(bytes_to_str(row[1]), 'user2') self.assertEqual(row[2], self.fiveMinsAgo) self.assertEqual(row[2].microsecond, 987000) - self.assertEqual(row[3], 'wind') + self.assertEqual(bytes_to_str(row[3]), 'wind') self.assertIsNone(row[4]) def test_query_that_creates_table_using_interpolation(self): @@ -254,44 +279,38 @@ def test_query_that_creates_table_using_interpolation(self): """ ts_obj = self.client.ts_query(table, query) self.assertIsNotNone(ts_obj) - self.assertEqual(len(ts_obj.columns), 0) - self.assertEqual(len(ts_obj.rows), 0) + self.validate_len(ts_obj, 0) def test_query_that_returns_table_description(self): fmt = 'DESCRIBE {table}' query = fmt.format(table=table_name) - ts_obj = self.client.ts_query('GeoCheckin', query) + ts_obj = self.client.ts_query(table_name, query) self.assertIsNotNone(ts_obj) - self.assertEqual(len(ts_obj.columns), 5) - self.assertEqual(len(ts_obj.rows), 5) + self.validate_len(ts_obj, 5) def test_query_that_returns_table_description_using_interpolation(self): query = 'Describe {table}' - ts_obj = self.client.ts_query('GeoCheckin', query) + ts_obj = self.client.ts_query(table_name, query) self.assertIsNotNone(ts_obj) - self.assertEqual(len(ts_obj.columns), 5) - self.assertEqual(len(ts_obj.rows), 5) + self.validate_len(ts_obj, 5) def test_query_description_via_table(self): query = 'describe {table}' - table = Table(self.client, 'GeoCheckin') + table = Table(self.client, table_name) ts_obj = table.query(query) self.assertIsNotNone(ts_obj) - self.assertEqual(len(ts_obj.columns), 5) - self.assertEqual(len(ts_obj.rows), 5) + self.validate_len(ts_obj, 5) def test_get_description(self): - ts_obj = self.client.ts_describe('GeoCheckin') + ts_obj = self.client.ts_describe(table_name) self.assertIsNotNone(ts_obj) - self.assertEqual(len(ts_obj.columns), 5) - self.assertEqual(len(ts_obj.rows), 5) + self.validate_len(ts_obj, 5) def test_get_description_via_table(self): - table = Table(self.client, 'GeoCheckin') + table = Table(self.client, table_name) ts_obj = table.describe() self.assertIsNotNone(ts_obj) - self.assertEqual(len(ts_obj.columns), 5) - self.assertEqual(len(ts_obj.rows), 5) + self.validate_len(ts_obj, 5) def test_query_that_returns_no_data(self): fmt = """ @@ -301,9 +320,8 @@ def test_query_that_returns_no_data(self): user = 'user1' """ query = fmt.format(table=table_name) - ts_obj = self.client.ts_query('GeoCheckin', query) - self.assertEqual(len(ts_obj.columns), 0) - self.assertEqual(len(ts_obj.rows), 0) + ts_obj = self.client.ts_query(table_name, query) + self.validate_len(ts_obj, 0) def test_query_that_returns_no_data_using_interpolation(self): query = """ @@ -312,9 +330,8 @@ def test_query_that_returns_no_data_using_interpolation(self): geohash = 'hash1' and user = 'user1' """ - ts_obj = self.client.ts_query('GeoCheckin', query) - self.assertEqual(len(ts_obj.columns), 0) - self.assertEqual(len(ts_obj.rows), 0) + ts_obj = self.client.ts_query(table_name, query) + self.validate_len(ts_obj, 0) def test_query_that_matches_some_data(self): fmt = """ @@ -327,7 +344,7 @@ def test_query_that_matches_some_data(self): table=table_name, t1=self.tenMinsAgoMsec, t2=self.nowMsec) - ts_obj = self.client.ts_query('GeoCheckin', query) + ts_obj = self.client.ts_query(table_name, query) self.validate_data(ts_obj) def test_query_that_matches_some_data_using_interpolation(self): @@ -340,7 +357,7 @@ def test_query_that_matches_some_data_using_interpolation(self): query = fmt.format( t1=self.tenMinsAgoMsec, t2=self.nowMsec) - ts_obj = self.client.ts_query('GeoCheckin', query) + ts_obj = self.client.ts_query(table_name, query) self.validate_data(ts_obj) def test_query_that_matches_more_data(self): @@ -354,9 +371,9 @@ def test_query_that_matches_more_data(self): table=table_name, t1=self.twentyMinsAgoMsec, t2=self.nowMsec) - ts_obj = self.client.ts_query('GeoCheckin', query) + ts_obj = self.client.ts_query(table_name, query) j = 0 - for i, want in enumerate(self.rows): + for i, want in enumerate(self.encoded_rows): if want[2] == self.twentyFiveMinsAgo: continue got = ts_obj.rows[j] @@ -366,23 +383,23 @@ def test_query_that_matches_more_data(self): def test_get_with_invalid_key(self): key = ['hash1', 'user2'] with self.assertRaises(RiakError): - self.client.ts_get('GeoCheckin', key) + self.client.ts_get(table_name, key) def test_get_single_value(self): key = ['hash1', 'user2', self.fiveMinsAgo] - ts_obj = self.client.ts_get('GeoCheckin', key) + ts_obj = self.client.ts_get(table_name, key) self.assertIsNotNone(ts_obj) self.validate_data(ts_obj) def test_get_single_value_via_table(self): key = ['hash1', 'user2', self.fiveMinsAgo] - table = Table(self.client, 'GeoCheckin') + table = Table(self.client, table_name) ts_obj = table.get(key) self.assertIsNotNone(ts_obj) self.validate_data(ts_obj) def test_stream_keys(self): - table = Table(self.client, 'GeoCheckin') + table = Table(self.client, table_name) streamed_keys = [] for keylist in table.stream_keys(): self.assertNotEqual([], keylist) @@ -390,14 +407,23 @@ def test_stream_keys(self): for key in keylist: self.assertIsInstance(key, list) self.assertEqual(len(key), 3) - self.assertEqual('hash1', key[0]) - self.assertEqual('user2', key[1]) + self.assertEqual(bytes_to_str(key[0]), 'hash1') + self.assertEqual(bytes_to_str(key[1]), 'user2') self.assertIsInstance(key[2], datetime.datetime) self.assertGreater(len(streamed_keys), 0) def test_delete_single_value(self): key = ['hash1', 'user2', self.twentyFiveMinsAgo] - rslt = self.client.ts_delete('GeoCheckin', key) + rslt = self.client.ts_delete(table_name, key) self.assertTrue(rslt) - ts_obj = self.client.ts_get('GeoCheckin', key) + ts_obj = self.client.ts_get(table_name, key) + self.assertIsNotNone(ts_obj) self.assertEqual(len(ts_obj.rows), 0) + self.assertEqual(len(ts_obj.columns.names), 0) + self.assertEqual(len(ts_obj.columns.types), 0) + + def test_create_error_via_put(self): + table = Table(self.client, table_name) + ts_obj = table.new([]) + with self.assertRaises(RiakError): + ts_obj.store() diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py new file mode 100644 index 00000000..339dd040 --- /dev/null +++ b/riak/tests/test_timeseries_ttb.py @@ -0,0 +1,196 @@ +# -*- coding: utf-8 -*- +import datetime +import logging +import six +import unittest + +from erlastic import decode, encode +from erlastic.types import Atom + +from riak import RiakError +from riak.table import Table +from riak.ts_object import TsObject +from riak.codecs.ttb import TtbCodec +from riak.util import str_to_bytes, \ + unix_time_millis, datetime_from_unix_time_millis, \ + is_timeseries_supported +from riak.tests import RUN_TIMESERIES +from riak.tests.base import IntegrationTestBase + +rpberrorresp_a = Atom('rpberrorresp') +tsgetreq_a = Atom('tsgetreq') +tsgetresp_a = Atom('tsgetresp') +tsputreq_a = Atom('tsputreq') + +udef_a = Atom('undefined') +varchar_a = Atom('varchar') +sint64_a = Atom('sint64') +double_a = Atom('double') +timestamp_a = Atom('timestamp') +boolean_a = Atom('boolean') + +table_name = 'GeoCheckin' + +str0 = 'ascii-0' +str1 = 'ascii-1' + +bd0 = six.u('时间序列') +bd1 = six.u('временные ряды') + +fiveMins = datetime.timedelta(0, 300) +ts0 = datetime.datetime(2015, 1, 1, 12, 0, 0) +ts1 = ts0 + fiveMins + + +@unittest.skipUnless(is_timeseries_supported(), "Timeseries not supported") +class TimeseriesTtbUnitTests(unittest.TestCase): + def setUp(self): + self.table = Table(None, table_name) + + def test_encode_data_for_get(self): + keylist = [ + str_to_bytes('hash1'), str_to_bytes('user2'), unix_time_millis(ts0) + ] + req = tsgetreq_a, str_to_bytes(table_name), keylist, udef_a + req_test = encode(req) + + test_key = ['hash1', 'user2', ts0] + c = TtbCodec() + msg = c.encode_timeseries_keyreq(self.table, test_key) + self.assertEqual(req_test, msg.data) + + # {tsgetresp, + # { + # [<<"geohash">>, <<"user">>, <<"time">>, + # <<"weather">>, <<"temperature">>], + # [varchar, varchar, timestamp, varchar, double], + # [(<<"hash1">>, <<"user2">>, 144378190987, <<"typhoon">>, 90.3)] + # } + # } + def test_decode_data_from_get(self): + colnames = ["varchar", "sint64", "double", "timestamp", + "boolean", "varchar", "varchar"] + coltypes = [varchar_a, sint64_a, double_a, timestamp_a, + boolean_a, varchar_a, varchar_a] + r0 = (bd0, 0, 1.2, unix_time_millis(ts0), True, + [], str1, None) + r1 = (bd1, 3, 4.5, unix_time_millis(ts1), False, + [], str1, None) + rows = [r0, r1] + # { tsgetresp, { [colnames], [coltypes], [rows] } } + data_t = colnames, coltypes, rows + rsp_data = tsgetresp_a, data_t + rsp_ttb = encode(rsp_data) + + tsobj = TsObject(None, self.table) + c = TtbCodec() + c.decode_timeseries(decode(rsp_ttb), tsobj) + + for i in range(0, 1): + dr = rows[i] + r = tsobj.rows[i] # encoded + self.assertEqual(r[0], dr[0].encode('utf-8')) + self.assertEqual(r[1], dr[1]) + self.assertEqual(r[2], dr[2]) + dt = datetime_from_unix_time_millis(dr[3]) + self.assertEqual(r[3], dt) + if i == 0: + self.assertEqual(r[4], True) + else: + self.assertEqual(r[4], False) + self.assertEqual(r[5], None) + self.assertEqual(r[6], dr[6].encode('ascii')) + self.assertEqual(r[7], None) + + def test_encode_data_for_put(self): + r0 = (bd0, 0, 1.2, unix_time_millis(ts0), True, []) + r1 = (bd1, 3, 4.5, unix_time_millis(ts1), False, []) + rows = [r0, r1] + req = tsputreq_a, str_to_bytes(table_name), [], rows + req_test = encode(req) + + rows_to_encode = [ + [bd0, 0, 1.2, ts0, True, None], + [bd1, 3, 4.5, ts1, False, None] + ] + + tsobj = TsObject(None, self.table, rows_to_encode, None) + c = TtbCodec() + msg = c.encode_timeseries_put(tsobj) + self.assertEqual(req_test, msg.data) + + +@unittest.skipUnless(is_timeseries_supported() and RUN_TIMESERIES, + 'Timeseries not supported or RUN_TIMESERIES is 0') +class TimeseriesTtbTests(IntegrationTestBase, unittest.TestCase): + client_options = {'transport_options': {'use_ttb': True}} + + @classmethod + def setUpClass(cls): + super(TimeseriesTtbTests, cls).setUpClass() + + def test_query_that_returns_table_description(self): + fmt = 'DESCRIBE {table}' + query = fmt.format(table=table_name) + ts_obj = self.client.ts_query(table_name, query) + self.assertIsNotNone(ts_obj) + ts_cols = ts_obj.columns + self.assertEqual(len(ts_cols.names), 5) + self.assertEqual(len(ts_cols.types), 5) + row = ts_obj.rows[0] + self.assertEqual(len(row), 5) + + def test_store_and_fetch(self): + now = datetime.datetime.utcfromtimestamp(144379690.987000) + fiveMinsAgo = now - fiveMins + tenMinsAgo = fiveMinsAgo - fiveMins + fifteenMinsAgo = tenMinsAgo - fiveMins + twentyMinsAgo = fifteenMinsAgo - fiveMins + twentyFiveMinsAgo = twentyMinsAgo - fiveMins + + table = self.client.table(table_name) + rows = [ + ['hash1', 'user2', twentyFiveMinsAgo, 'typhoon', 90.3], + ['hash1', 'user2', twentyMinsAgo, 'hurricane', 82.3], + ['hash1', 'user2', fifteenMinsAgo, 'rain', 79.0], + ['hash1', 'user2', fiveMinsAgo, 'wind', None], + ['hash1', 'user2', now, 'snow', 20.1] + ] + # NB: response data is binary + exp_rows = [ + [six.b('hash1'), six.b('user2'), twentyFiveMinsAgo, + six.b('typhoon'), 90.3], + [six.b('hash1'), six.b('user2'), twentyMinsAgo, + six.b('hurricane'), 82.3], + [six.b('hash1'), six.b('user2'), fifteenMinsAgo, + six.b('rain'), 79.0], + [six.b('hash1'), six.b('user2'), fiveMinsAgo, + six.b('wind'), None], + [six.b('hash1'), six.b('user2'), now, + six.b('snow'), 20.1] + ] + ts_obj = table.new(rows) + result = ts_obj.store() + self.assertTrue(result) + + for i, r in enumerate(rows): + k = r[0:3] + ts_obj = self.client.ts_get(table_name, k) + self.assertIsNotNone(ts_obj) + ts_cols = ts_obj.columns + self.assertEqual(len(ts_cols.names), 5) + self.assertEqual(len(ts_cols.types), 5) + self.assertEqual(len(ts_obj.rows), 1) + row = ts_obj.rows[0] + exp = exp_rows[i] + self.assertEqual(len(row), 5) + self.assertEqual(row, exp) + + def test_create_error_via_put(self): + table = Table(self.client, table_name) + ts_obj = table.new([]) + with self.assertRaises(RiakError) as cm: + ts_obj.store() + logging.debug( + "[test_timeseries_ttb] saw exception: {}" + .format(cm.exception)) diff --git a/riak/tests/test_util.py b/riak/tests/test_util.py index 3cc69e95..af704516 100644 --- a/riak/tests/test_util.py +++ b/riak/tests/test_util.py @@ -1,12 +1,7 @@ -import platform +import unittest from riak.util import is_timeseries_supported -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - class UtilUnitTests(unittest.TestCase): def test_is_timeseries_supported(self): diff --git a/riak/tests/test_yokozuna.py b/riak/tests/test_yokozuna.py index 52f9af88..a4f325f1 100644 --- a/riak/tests/test_yokozuna.py +++ b/riak/tests/test_yokozuna.py @@ -1,15 +1,11 @@ # -*- coding: utf-8 -*- -import platform +import unittest + from riak.tests import RUN_YZ from riak.tests.base import IntegrationTestBase from riak.tests.comparison import Comparison from riak.tests.yz_setup import yzSetUp, yzTearDown -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - def wait_for_yz_index(bucket, key, index=None): """ diff --git a/riak/transports/feature_detect.py b/riak/transports/feature_detect.py index c73ba37d..8f5808ac 100644 --- a/riak/transports/feature_detect.py +++ b/riak/transports/feature_detect.py @@ -40,7 +40,7 @@ class FeatureDetection(object): should return the server's version as a string. :class:`FeatureDetection` is a parent class of - :class:`RiakTransport `. + :class:`Transport `. """ def _server_version(self): diff --git a/riak/transports/http/__init__.py b/riak/transports/http/__init__.py index d7e69c3d..69c7de8c 100644 --- a/riak/transports/http/__init__.py +++ b/riak/transports/http/__init__.py @@ -1,27 +1,11 @@ -""" -Copyright 2015 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 select + from six import PY2 from riak.security import SecurityError, USE_STDLIB_SSL from riak.transports.pool import Pool -from riak.transports.http.transport import RiakHttpTransport +from riak.transports.http.transport import HttpTransport + if USE_STDLIB_SSL: import ssl from riak.transports.security import configure_ssl_context @@ -29,6 +13,7 @@ import OpenSSL.SSL from riak.transports.security import RiakWrappedSocket,\ configure_pyopenssl_context + if PY2: from httplib import HTTPConnection, \ NotConnected, \ @@ -149,7 +134,7 @@ def connect(self): self.sock.context = ssl_ctx -class RiakHttpPool(Pool): +class HttpPool(Pool): """ A pool of HTTP(S) transport connections. """ @@ -160,14 +145,14 @@ def __init__(self, client, **options): if self.client._credentials: self.connection_class = RiakHTTPSConnection - super(RiakHttpPool, self).__init__() + super(HttpPool, 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) + return HttpTransport(node=node, + client=self.client, + connection_class=self.connection_class, + **self.options) def destroy_resource(self, transport): transport.close() diff --git a/riak/transports/http/connection.py b/riak/transports/http/connection.py index e1f570e0..87a5716d 100644 --- a/riak/transports/http/connection.py +++ b/riak/transports/http/connection.py @@ -1,33 +1,17 @@ -""" -Copyright 2015 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 base64 from six import PY2 -import base64 from riak.util import str_to_bytes + if PY2: from httplib import NotConnected, HTTPConnection else: from http.client import NotConnected, HTTPConnection -class RiakHttpConnection(object): +class HttpConnection(object): """ - Connection and low-level request methods for RiakHttpTransport. + Connection and low-level request methods for HttpTransport. """ def _request(self, method, uri, headers={}, body='', stream=False): @@ -93,7 +77,7 @@ def close(self): except NotConnected: pass - # These are set by the RiakHttpTransport initializer + # These are set by the HttpTransport initializer _connection_class = HTTPConnection _node = None diff --git a/riak/transports/http/resources.py b/riak/transports/http/resources.py index c13925bc..2017f420 100644 --- a/riak/transports/http/resources.py +++ b/riak/transports/http/resources.py @@ -1,34 +1,18 @@ -""" -Copyright 2015 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 six import PY2 from riak import RiakError from riak.util import lazy_property, bytes_to_str + if PY2: from urllib import quote_plus, urlencode else: from urllib.parse import quote_plus, urlencode -class RiakHttpResources(object): +class HttpResources(object): """ - Methods for RiakHttpTransport related to URL generation, i.e. + Methods for HttpTransport related to URL generation, i.e. creating the proper paths. """ @@ -204,7 +188,7 @@ def index_term_regex(self): if self.riak_kv_wm_bucket_type is not None: return True else: - return super(RiakHttpResources, self).index_term_regex() + return super(HttpResources, self).index_term_regex() # Resource root paths @lazy_property diff --git a/riak/transports/http/stream.py b/riak/transports/http/stream.py index edb1c818..b5ec00d7 100644 --- a/riak/transports/http/stream.py +++ b/riak/transports/http/stream.py @@ -1,23 +1,6 @@ -""" -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 re + from cgi import parse_header from email import message_from_string from riak.util import decode_index_value @@ -26,7 +9,7 @@ from six import PY2 -class RiakHttpStream(object): +class HttpStream(object): """ Base class for HTTP streaming iterators. """ @@ -66,7 +49,7 @@ def close(self): self.resource.release() -class RiakHttpJsonStream(RiakHttpStream): +class HttpJsonStream(HttpStream): _json_field = None def next(self): @@ -92,26 +75,26 @@ def __next__(self): return self.next() -class RiakHttpKeyStream(RiakHttpJsonStream): +class HttpKeyStream(HttpJsonStream): """ Streaming iterator for list-keys over HTTP """ _json_field = u'keys' -class RiakHttpBucketStream(RiakHttpJsonStream): +class HttpBucketStream(HttpJsonStream): """ Streaming iterator for list-buckets over HTTP """ _json_field = u'buckets' -class RiakHttpMultipartStream(RiakHttpStream): +class HttpMultipartStream(HttpStream): """ Streaming iterator for multipart messages over HTTP """ def __init__(self, response): - super(RiakHttpMultipartStream, self).__init__(response) + super(HttpMultipartStream, self).__init__(response) ctypehdr = response.getheader('content-type') _, params = parse_header(ctypehdr) self.boundary_re = re.compile('\r?\n--%s(?:--)?\r?\n' % @@ -154,13 +137,13 @@ def read_until_boundary(self): self._read() -class RiakHttpMapReduceStream(RiakHttpMultipartStream): +class HttpMapReduceStream(HttpMultipartStream): """ Streaming iterator for MapReduce over HTTP """ def next(self): - message = super(RiakHttpMapReduceStream, self).next() + message = super(HttpMapReduceStream, self).next() payload = json.loads(message.get_payload()) return payload['phase'], payload['data'] @@ -169,18 +152,18 @@ def __next__(self): return self.next() -class RiakHttpIndexStream(RiakHttpMultipartStream): +class HttpIndexStream(HttpMultipartStream): """ Streaming iterator for secondary indexes over HTTP """ def __init__(self, response, index, return_terms): - super(RiakHttpIndexStream, self).__init__(response) + super(HttpIndexStream, self).__init__(response) self.index = index self.return_terms = return_terms def next(self): - message = super(RiakHttpIndexStream, self).next() + message = super(HttpIndexStream, self).next() payload = json.loads(message.get_payload()) if u'error' in payload: raise RiakError(payload[u'error']) diff --git a/riak/transports/http/transport.py b/riak/transports/http/transport.py index c139b3ea..10238d32 100644 --- a/riak/transports/http/transport.py +++ b/riak/transports/http/transport.py @@ -1,24 +1,3 @@ -""" -Copyright 2015 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. -""" - try: import simplejson as json except ImportError: @@ -26,28 +5,30 @@ from six import PY2 from xml.dom.minidom import Document -from riak.transports.transport import RiakTransport -from riak.transports.http.resources import RiakHttpResources -from riak.transports.http.connection import RiakHttpConnection -from riak.transports.http.codec import RiakHttpCodec -from riak.transports.http.stream import ( - RiakHttpKeyStream, - RiakHttpMapReduceStream, - RiakHttpBucketStream, - RiakHttpIndexStream) + from riak import RiakError +from riak.codecs.http import HttpCodec +from riak.transports.transport import Transport +from riak.transports.http.resources import HttpResources +from riak.transports.http.connection import HttpConnection +from riak.transports.http.stream import ( + HttpKeyStream, + HttpMapReduceStream, + HttpBucketStream, + HttpIndexStream) from riak.security import SecurityError from riak.util import decode_index_value, bytes_to_str, str_to_long + if PY2: from httplib import HTTPConnection else: from http.client import HTTPConnection -class RiakHttpTransport(RiakHttpConnection, RiakHttpResources, RiakHttpCodec, - RiakTransport): +class HttpTransport(Transport, + HttpConnection, HttpResources, HttpCodec): """ - The RiakHttpTransport object holds information necessary to + The HttpTransport object holds information necessary to connect to Riak via HTTP. """ @@ -59,7 +40,7 @@ def __init__(self, node=None, """ Construct a new HTTP connection to Riak. """ - super(RiakHttpTransport, self).__init__() + super(HttpTransport, self).__init__() self._client = client self._node = node @@ -219,7 +200,7 @@ def stream_keys(self, bucket, timeout=None): status, headers, response = self._request('GET', url, stream=True) if status == 200: - return RiakHttpKeyStream(response) + return HttpKeyStream(response) else: raise RiakError('Error listing keys.') @@ -252,7 +233,7 @@ def stream_buckets(self, bucket_type=None, timeout=None): status, headers, response = self._request('GET', url, stream=True) if status == 200: - return RiakHttpBucketStream(response) + return HttpBucketStream(response) else: raise RiakError('Error listing buckets.') @@ -371,7 +352,7 @@ def stream_mapred(self, inputs, query, timeout=None): content, stream=True) if status == 200: - return RiakHttpMapReduceStream(response) + return HttpMapReduceStream(response) else: raise RiakError( 'Error running MapReduce operation. Headers: %s Body: %s' % @@ -441,7 +422,7 @@ def stream_index(self, bucket, index, startkey, endkey=None, status, headers, response = self._request('GET', url, stream=True) if status == 200: - return RiakHttpIndexStream(response, index, return_terms) + return HttpIndexStream(response, index, return_terms) else: raise RiakError('Error streaming secondary index.') diff --git a/riak/transports/pbc/__init__.py b/riak/transports/pbc/__init__.py deleted file mode 100644 index fc8914b6..00000000 --- a/riak/transports/pbc/__init__.py +++ /dev/null @@ -1,74 +0,0 @@ -""" -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 deleted file mode 100644 index 97629752..00000000 --- a/riak/transports/pbc/codec.py +++ /dev/null @@ -1,773 +0,0 @@ -import datetime -import logging -import riak.pb -import riak.pb.riak_pb2 -import riak.pb.riak_dt_pb2 -import riak.pb.riak_kv_pb2 -import riak.pb.riak_ts_pb2 - -from riak import RiakError -from riak.content import RiakContent -from riak.util import decode_index_value, str_to_bytes, bytes_to_str -from riak.multidict import MultiDict -from riak.pb.riak_ts_pb2 import TsColumnType - -from six import string_types, PY2 - -epoch = datetime.datetime.utcfromtimestamp(0) - - -def _invert(d): - out = {} - for key in d: - value = d[key] - out[value] = key - return out - -REPL_TO_PY = { - riak.pb.riak_pb2.RpbBucketProps.FALSE: False, - riak.pb.riak_pb2.RpbBucketProps.TRUE: True, - riak.pb.riak_pb2.RpbBucketProps.REALTIME: 'realtime', - riak.pb.riak_pb2.RpbBucketProps.FULLSYNC: 'fullsync' -} - -REPL_TO_PB = _invert(REPL_TO_PY) - -RIAKC_RW_ONE = 4294967294 -RIAKC_RW_QUORUM = 4294967293 -RIAKC_RW_ALL = 4294967292 -RIAKC_RW_DEFAULT = 4294967291 - -QUORUM_TO_PB = {'default': RIAKC_RW_DEFAULT, - 'all': RIAKC_RW_ALL, - 'quorum': RIAKC_RW_QUORUM, - 'one': RIAKC_RW_ONE} - -QUORUM_TO_PY = _invert(QUORUM_TO_PB) - -NORMAL_PROPS = ['n_val', 'allow_mult', 'last_write_wins', 'old_vclock', - 'young_vclock', 'big_vclock', 'small_vclock', 'basic_quorum', - 'notfound_ok', 'search', 'backend', 'search_index', 'datatype', - 'write_once'] -COMMIT_HOOK_PROPS = ['precommit', 'postcommit'] -MODFUN_PROPS = ['chash_keyfun', 'linkfun'] -QUORUM_PROPS = ['r', 'pr', 'w', 'pw', 'dw', 'rw'] - -MAP_FIELD_TYPES = { - riak.pb.riak_dt_pb2.MapField.COUNTER: 'counter', - riak.pb.riak_dt_pb2.MapField.SET: 'set', - riak.pb.riak_dt_pb2.MapField.REGISTER: 'register', - riak.pb.riak_dt_pb2.MapField.FLAG: 'flag', - riak.pb.riak_dt_pb2.MapField.MAP: 'map', - 'counter': riak.pb.riak_dt_pb2.MapField.COUNTER, - 'set': riak.pb.riak_dt_pb2.MapField.SET, - 'register': riak.pb.riak_dt_pb2.MapField.REGISTER, - 'flag': riak.pb.riak_dt_pb2.MapField.FLAG, - 'map': riak.pb.riak_dt_pb2.MapField.MAP -} - -DT_FETCH_TYPES = { - riak.pb.riak_dt_pb2.DtFetchResp.COUNTER: 'counter', - riak.pb.riak_dt_pb2.DtFetchResp.SET: 'set', - riak.pb.riak_dt_pb2.DtFetchResp.MAP: 'map' -} - - -class RiakPbcCodec(object): - """ - Protobuffs Encoding and decoding methods for RiakPbcTransport. - """ - - def __init__(self, **unused_args): - if riak.pb is None: - raise NotImplementedError("this transport is not available") - super(RiakPbcCodec, self).__init__(**unused_args) - - def _unix_time_millis(self, dt): - td = dt - epoch - return int(td.total_seconds() * 1000.0) - - def _datetime_from_unix_time_millis(self, ut): - return datetime.datetime.utcfromtimestamp(ut / 1000.0) - - def _encode_quorum(self, rw): - """ - Converts a symbolic quorum value into its on-the-wire - equivalent. - - :param rw: the quorum - :type rw: string, integer - :rtype: integer - """ - if rw in QUORUM_TO_PB: - return QUORUM_TO_PB[rw] - elif type(rw) is int and rw >= 0: - return rw - else: - return None - - def _decode_quorum(self, rw): - """ - Converts a protobuf quorum value to a symbolic value if - necessary. - - :param rw: the quorum - :type rw: int - :rtype int or string - """ - if rw in QUORUM_TO_PY: - return QUORUM_TO_PY[rw] - else: - return rw - - def _decode_contents(self, contents, obj): - """ - Decodes the list of siblings from the protobuf representation - into the object. - - :param contents: a list of RpbContent messages - :type contents: list - :param obj: a RiakObject - :type obj: RiakObject - :rtype RiakObject - """ - obj.siblings = [self._decode_content(c, RiakContent(obj)) - for c in contents] - # Invoke sibling-resolution logic - if len(obj.siblings) > 1 and obj.resolver is not None: - obj.resolver(obj) - return obj - - def _decode_content(self, rpb_content, sibling): - """ - Decodes a single sibling from the protobuf representation into - a RiakObject. - - :param rpb_content: a single RpbContent message - :type rpb_content: riak.pb.riak_pb2.RpbContent - :param sibling: a RiakContent sibling container - :type sibling: RiakContent - :rtype: RiakContent - """ - - if rpb_content.HasField("deleted") and rpb_content.deleted: - sibling.exists = False - else: - sibling.exists = True - if rpb_content.HasField("content_type"): - sibling.content_type = bytes_to_str(rpb_content.content_type) - if rpb_content.HasField("charset"): - sibling.charset = bytes_to_str(rpb_content.charset) - if rpb_content.HasField("content_encoding"): - sibling.content_encoding = \ - bytes_to_str(rpb_content.content_encoding) - if rpb_content.HasField("vtag"): - sibling.etag = bytes_to_str(rpb_content.vtag) - - sibling.links = [self._decode_link(link) - for link in rpb_content.links] - if rpb_content.HasField("last_mod"): - sibling.last_modified = float(rpb_content.last_mod) - if rpb_content.HasField("last_mod_usecs"): - sibling.last_modified += rpb_content.last_mod_usecs / 1000000.0 - - sibling.usermeta = dict([(bytes_to_str(usermd.key), - bytes_to_str(usermd.value)) - for usermd in rpb_content.usermeta]) - sibling.indexes = set([(bytes_to_str(index.key), - decode_index_value(index.key, index.value)) - for index in rpb_content.indexes]) - sibling.encoded_data = rpb_content.value - - return sibling - - def _encode_content(self, robj, rpb_content): - """ - Fills an RpbContent message with the appropriate data and - metadata from a RiakObject. - - :param robj: a RiakObject - :type robj: RiakObject - :param rpb_content: the protobuf message to fill - :type rpb_content: riak.pb.riak_pb2.RpbContent - """ - if robj.content_type: - rpb_content.content_type = str_to_bytes(robj.content_type) - if robj.charset: - rpb_content.charset = str_to_bytes(robj.charset) - if robj.content_encoding: - rpb_content.content_encoding = str_to_bytes(robj.content_encoding) - for uk in robj.usermeta: - pair = rpb_content.usermeta.add() - pair.key = str_to_bytes(uk) - pair.value = str_to_bytes(robj.usermeta[uk]) - for link in robj.links: - pb_link = rpb_content.links.add() - try: - bucket, key, tag = link - except ValueError: - raise RiakError("Invalid link tuple %s" % link) - - pb_link.bucket = str_to_bytes(bucket) - pb_link.key = str_to_bytes(key) - if tag: - pb_link.tag = str_to_bytes(tag) - else: - pb_link.tag = str_to_bytes('') - - for field, value in robj.indexes: - pair = rpb_content.indexes.add() - pair.key = str_to_bytes(field) - pair.value = str_to_bytes(str(value)) - - # Python 2.x data is stored in a string - if PY2: - rpb_content.value = str(robj.encoded_data) - else: - rpb_content.value = robj.encoded_data - - def _decode_link(self, link): - """ - Decodes an RpbLink message into a tuple - - :param link: an RpbLink message - :type link: riak.pb.riak_pb2.RpbLink - :rtype tuple - """ - - if link.HasField("bucket"): - bucket = bytes_to_str(link.bucket) - else: - bucket = None - if link.HasField("key"): - key = bytes_to_str(link.key) - else: - key = None - if link.HasField("tag"): - tag = bytes_to_str(link.tag) - else: - tag = None - - return (bucket, key, tag) - - def _decode_index_value(self, index, value): - """ - Decodes a secondary index value into the correct Python type. - :param index: the name of the index - :type index: str - :param value: the value of the index entry - :type value: str - :rtype str or int - """ - if index.endswith("_int"): - return int(value) - else: - return bytes_to_str(value) - - def _encode_bucket_props(self, props, msg): - """ - Encodes a dict of bucket properties into the protobuf message. - - :param props: bucket properties - :type props: dict - :param msg: the protobuf message to fill - :type msg: riak.pb.riak_pb2.RpbSetBucketReq - """ - for prop in NORMAL_PROPS: - if prop in props and props[prop] is not None: - if isinstance(props[prop], string_types): - setattr(msg.props, prop, str_to_bytes(props[prop])) - else: - setattr(msg.props, prop, props[prop]) - for prop in COMMIT_HOOK_PROPS: - if prop in props: - setattr(msg.props, 'has_' + prop, True) - self._encode_hooklist(props[prop], getattr(msg.props, prop)) - for prop in MODFUN_PROPS: - if prop in props and props[prop] is not None: - self._encode_modfun(props[prop], getattr(msg.props, prop)) - for prop in QUORUM_PROPS: - if prop in props and props[prop] not in (None, 'default'): - value = self._encode_quorum(props[prop]) - if value is not None: - if isinstance(value, string_types): - setattr(msg.props, prop, str_to_bytes(value)) - else: - setattr(msg.props, prop, value) - if 'repl' in props: - msg.props.repl = REPL_TO_PY[props['repl']] - - return msg - - def _decode_bucket_props(self, msg): - """ - Decodes the protobuf bucket properties message into a dict. - - :param msg: the protobuf message to decode - :type msg: riak.pb.riak_pb2.RpbBucketProps - :rtype dict - """ - props = {} - - for prop in NORMAL_PROPS: - if msg.HasField(prop): - props[prop] = getattr(msg, prop) - if isinstance(props[prop], bytes): - props[prop] = bytes_to_str(props[prop]) - for prop in COMMIT_HOOK_PROPS: - if getattr(msg, 'has_' + prop): - props[prop] = self._decode_hooklist(getattr(msg, prop)) - for prop in MODFUN_PROPS: - if msg.HasField(prop): - props[prop] = self._decode_modfun(getattr(msg, prop)) - for prop in QUORUM_PROPS: - if msg.HasField(prop): - props[prop] = self._decode_quorum(getattr(msg, prop)) - if msg.HasField('repl'): - props['repl'] = REPL_TO_PY[msg.repl] - - return props - - def _decode_modfun(self, modfun): - """ - Decodes a protobuf modfun pair into a dict with 'mod' and - 'fun' keys. Used in bucket properties. - - :param modfun: the protobuf message to decode - :type modfun: riak.pb.riak_pb2.RpbModFun - :rtype dict - """ - return {'mod': bytes_to_str(modfun.module), - 'fun': bytes_to_str(modfun.function)} - - def _encode_modfun(self, props, msg=None): - """ - Encodes a dict with 'mod' and 'fun' keys into a protobuf - modfun pair. Used in bucket properties. - - :param props: the module/function pair - :type props: dict - :param msg: the protobuf message to fill - :type msg: riak.pb.riak_pb2.RpbModFun - :rtype riak.pb.riak_pb2.RpbModFun - """ - if msg is None: - msg = riak.pb.riak_pb2.RpbModFun() - msg.module = str_to_bytes(props['mod']) - msg.function = str_to_bytes(props['fun']) - return msg - - def _decode_hooklist(self, hooklist): - """ - Decodes a list of protobuf commit hooks into their python - equivalents. Used in bucket properties. - - :param hooklist: a list of protobuf commit hooks - :type hooklist: list - :rtype list - """ - return [self._decode_hook(hook) for hook in hooklist] - - def _encode_hooklist(self, hooklist, msg): - """ - Encodes a list of commit hooks into their protobuf equivalent. - Used in bucket properties. - - :param hooklist: a list of commit hooks - :type hooklist: list - :param msg: a protobuf field that is a list of commit hooks - """ - for hook in hooklist: - pbhook = msg.add() - self._encode_hook(hook, pbhook) - - def _decode_hook(self, hook): - """ - Decodes a protobuf commit hook message into a dict. Used in - bucket properties. - - :param hook: the hook to decode - :type hook: riak.pb.riak_pb2.RpbCommitHook - :rtype dict - """ - if hook.HasField('modfun'): - return self._decode_modfun(hook.modfun) - else: - return {'name': bytes_to_str(hook.name)} - - def _encode_hook(self, hook, msg): - """ - Encodes a commit hook dict into the protobuf message. Used in - bucket properties. - - :param hook: the hook to encode - :type hook: dict - :param msg: the protobuf message to fill - :type msg: riak.pb.riak_pb2.RpbCommitHook - :rtype riak.pb.riak_pb2.RpbCommitHook - """ - if 'name' in hook: - msg.name = str_to_bytes(hook['name']) - else: - self._encode_modfun(hook, msg.modfun) - return msg - - def _encode_index_req(self, bucket, index, startkey, endkey=None, - return_terms=None, max_results=None, - continuation=None, timeout=None, term_regex=None): - """ - Encodes a secondary index request into the protobuf message. - - :param bucket: the bucket whose index to query - :type bucket: string - :param index: the index to query - :type index: string - :param startkey: the value or beginning of the range - :type startkey: integer, string - :param endkey: the end of the range - :type endkey: integer, string - :param return_terms: whether to return the index term with the key - :type return_terms: bool - :param max_results: the maximum number of results to return (page size) - :type max_results: integer - :param continuation: the opaque continuation returned from a - previous paginated request - :type continuation: string - :param timeout: a timeout value in milliseconds, or 'infinity' - :type timeout: int - :param term_regex: a regular expression used to filter index terms - :type term_regex: string - :rtype riak.pb.riak_kv_pb2.RpbIndexReq - """ - req = riak.pb.riak_kv_pb2.RpbIndexReq( - bucket=str_to_bytes(bucket.name), - index=str_to_bytes(index)) - self._add_bucket_type(req, bucket.bucket_type) - if endkey is not None: - req.qtype = riak.pb.riak_kv_pb2.RpbIndexReq.range - req.range_min = str_to_bytes(str(startkey)) - req.range_max = str_to_bytes(str(endkey)) - else: - req.qtype = riak.pb.riak_kv_pb2.RpbIndexReq.eq - req.key = str_to_bytes(str(startkey)) - if return_terms is not None: - req.return_terms = return_terms - if max_results: - req.max_results = max_results - if continuation: - req.continuation = str_to_bytes(continuation) - if timeout: - if timeout == 'infinity': - req.timeout = 0 - else: - req.timeout = timeout - if term_regex: - req.term_regex = str_to_bytes(term_regex) - return req - - def _decode_search_index(self, index): - """ - Fills an RpbYokozunaIndex message with the appropriate data. - - :param index: a yz index message - :type index: riak.pb.riak_yokozuna_pb2.RpbYokozunaIndex - :rtype dict - """ - result = {} - result['name'] = bytes_to_str(index.name) - if index.HasField('schema'): - result['schema'] = bytes_to_str(index.schema) - if index.HasField('n_val'): - result['n_val'] = index.n_val - return result - - def _add_bucket_type(self, req, bucket_type): - if bucket_type and not bucket_type.is_default(): - if not self.bucket_types(): - raise NotImplementedError( - 'Server does not support bucket-types') - req.type = str_to_bytes(bucket_type.name) - - def _encode_search_query(self, req, params): - if 'rows' in params: - req.rows = params['rows'] - if 'start' in params: - req.start = params['start'] - if 'sort' in params: - req.sort = str_to_bytes(params['sort']) - if 'filter' in params: - req.filter = str_to_bytes(params['filter']) - if 'df' in params: - req.df = str_to_bytes(params['df']) - if 'op' in params: - req.op = str_to_bytes(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'] - - def _decode_search_doc(self, doc): - resultdoc = MultiDict() - for pair in doc.fields: - if PY2: - ukey = unicode(pair.key, 'utf-8') # noqa - uval = unicode(pair.value, 'utf-8') # noqa - else: - ukey = bytes_to_str(pair.key) - uval = bytes_to_str(pair.value) - resultdoc.add(ukey, uval) - return resultdoc.mixed() - - def _decode_dt_fetch(self, resp): - dtype = DT_FETCH_TYPES.get(resp.type) - if dtype is None: - raise ValueError("Unknown datatype on wire: {}".format(resp.type)) - - value = self._decode_dt_value(dtype, resp.value) - - if resp.HasField('context'): - context = resp.context[:] - else: - context = None - - return dtype, value, context - - def _decode_dt_value(self, dtype, msg): - if dtype == 'counter': - return msg.counter_value - elif dtype == 'set': - return self._decode_set_value(msg.set_value) - elif dtype == 'map': - return self._decode_map_value(msg.map_value) - - def _encode_dt_options(self, req, params): - for q in ['r', 'pr', 'w', 'dw', 'pw']: - if q in params and params[q] is not None: - setattr(req, q, self._encode_quorum(params[q])) - - for o in ['basic_quorum', 'notfound_ok', 'timeout', 'return_body', - 'include_context']: - if o in params and params[o] is not None: - setattr(req, o, params[o]) - - def _decode_map_value(self, entries): - out = {} - for entry in entries: - name = bytes_to_str(entry.field.name[:]) - dtype = MAP_FIELD_TYPES[entry.field.type] - if dtype == 'counter': - value = entry.counter_value - elif dtype == 'set': - value = self._decode_set_value(entry.set_value) - elif dtype == 'register': - value = bytes_to_str(entry.register_value[:]) - elif dtype == 'flag': - value = entry.flag_value - elif dtype == 'map': - value = self._decode_map_value(entry.map_value) - out[(name, dtype)] = value - return out - - def _decode_set_value(self, set_value): - return [bytes_to_str(string[:]) for string in set_value] - - def _encode_dt_op(self, dtype, req, op): - if dtype == 'counter': - req.op.counter_op.increment = op[1] - elif dtype == 'set': - self._encode_set_op(req.op, op) - elif dtype == 'map': - self._encode_map_op(req.op.map_op, op) - else: - raise TypeError("Cannot send operation on datatype {!r}". - format(dtype)) - - def _encode_set_op(self, msg, op): - if 'adds' in op: - msg.set_op.adds.extend(str_to_bytes(op['adds'])) - if 'removes' in op: - msg.set_op.removes.extend(str_to_bytes(op['removes'])) - - def _encode_map_op(self, msg, ops): - for op in ops: - name, dtype = op[1] - ftype = MAP_FIELD_TYPES[dtype] - if op[0] == 'add': - add = msg.adds.add() - add.name = str_to_bytes(name) - add.type = ftype - elif op[0] == 'remove': - remove = msg.removes.add() - remove.name = str_to_bytes(name) - remove.type = ftype - elif op[0] == 'update': - update = msg.updates.add() - update.field.name = str_to_bytes(name) - update.field.type = ftype - self._encode_map_update(dtype, update, op[2]) - - def _encode_map_update(self, dtype, msg, op): - if dtype == 'counter': - # ('increment', some_int) - msg.counter_op.increment = op[1] - elif dtype == 'set': - self._encode_set_op(msg, op) - elif dtype == 'map': - self._encode_map_op(msg.map_op, op) - elif dtype == 'register': - # ('assign', some_str) - msg.register_op = str_to_bytes(op[1]) - elif dtype == 'flag': - if op == 'enable': - msg.flag_op = riak.pb.riak_dt_pb2.MapUpdate.ENABLE - else: - msg.flag_op = riak.pb.riak_dt_pb2.MapUpdate.DISABLE - - def _encode_to_ts_cell(self, cell, ts_cell): - if cell is not None: - if isinstance(cell, datetime.datetime): - ts_cell.timestamp_value = self._unix_time_millis(cell) - elif isinstance(cell, bool): - ts_cell.boolean_value = cell - elif isinstance(cell, string_types): - logging.debug("cell -> str: '%s'", cell) - ts_cell.varchar_value = str_to_bytes(cell) - elif (isinstance(cell, int) or - (PY2 and isinstance(cell, long))): # noqa - logging.debug("cell -> int/long: '%s'", cell) - ts_cell.sint64_value = cell - elif isinstance(cell, float): - ts_cell.double_value = cell - else: - t = type(cell) - raise RiakError("can't serialize type '{}', value '{}'" - .format(t, cell)) - - def _encode_timeseries_keyreq(self, table, key, req): - key_vals = None - if isinstance(key, list): - key_vals = key - else: - raise ValueError("key must be a list") - - req.table = str_to_bytes(table.name) - for cell in key_vals: - ts_cell = req.key.add() - self._encode_to_ts_cell(cell, ts_cell) - - def _encode_timeseries_listkeysreq(self, table, req, timeout=None): - req.table = str_to_bytes(table.name) - if timeout: - req.timeout = timeout - - def _encode_timeseries_put(self, tsobj, req): - """ - Fills an TsPutReq message with the appropriate data and - metadata from a TsObject. - - :param tsobj: a TsObject - :type tsobj: TsObject - :param req: the protobuf message to fill - :type req: riak.pb.riak_ts_pb2.TsPutReq - """ - req.table = str_to_bytes(tsobj.table.name) - - if tsobj.columns: - raise NotImplementedError("columns are not implemented yet") - - if tsobj.rows and isinstance(tsobj.rows, list): - for row in tsobj.rows: - tsr = req.rows.add() # NB: type TsRow - if not isinstance(row, list): - raise ValueError("TsObject row must be a list of values") - for cell in row: - tsc = tsr.cells.add() # NB: type TsCell - self._encode_to_ts_cell(cell, tsc) - else: - raise RiakError("TsObject requires a list of rows") - - def _decode_timeseries(self, resp, tsobj): - """ - Fills an TsObject with the appropriate data and - metadata from a TsQueryResp. - - :param resp: the protobuf message from which to process data - :type resp: riak.pb.TsQueryRsp or riak.pb.riak_ts_pb2.TsGetResp - :param tsobj: a TsObject - :type tsobj: TsObject - """ - if tsobj.columns is not None: - for col in resp.columns: - col_name = bytes_to_str(col.name) - col_type = col.type - col = (col_name, col_type) - tsobj.columns.append(col) - - for row in resp.rows: - tsobj.rows.append( - self._decode_timeseries_row(row, resp.columns)) - - def _decode_timeseries_row(self, tsrow, tscols=None): - """ - Decodes a TsRow into a list - - :param tsrow: the protobuf TsRow to decode. - :type tsrow: riak.pb.riak_ts_pb2.TsRow - :param tscols: the protobuf TsColumn data to help decode. - :type tscols: list - :rtype list - """ - row = [] - for i, cell in enumerate(tsrow.cells): - col = None - if tscols is not None: - col = tscols[i] - if cell.HasField('varchar_value'): - if col and col.type != TsColumnType.Value('VARCHAR'): - raise TypeError('expected VARCHAR column') - else: - row.append(bytes_to_str(cell.varchar_value)) - elif cell.HasField('sint64_value'): - if col and col.type != TsColumnType.Value('SINT64'): - raise TypeError('expected SINT64 column') - else: - row.append(cell.sint64_value) - elif cell.HasField('double_value'): - if col and col.type != TsColumnType.Value('DOUBLE'): - raise TypeError('expected DOUBLE column') - else: - row.append(cell.double_value) - elif cell.HasField('timestamp_value'): - if col and col.type != TsColumnType.Value('TIMESTAMP'): - raise TypeError('expected TIMESTAMP column') - else: - dt = self._datetime_from_unix_time_millis( - cell.timestamp_value) - row.append(dt) - elif cell.HasField('boolean_value'): - if col and col.type != TsColumnType.Value('BOOLEAN'): - raise TypeError('expected BOOLEAN column') - else: - row.append(cell.boolean_value) - else: - row.append(None) - return row - - def _decode_preflist(self, item): - """ - Decodes a preflist response - - :param preflist: a bucket/key preflist - :type preflist: list of - riak.pb.riak_kv_pb2.RpbBucketKeyPreflistItem - :rtype dict - """ - result = {'partition': item.partition, - 'node': bytes_to_str(item.node), - 'primary': item. primary} - return result diff --git a/riak/transports/pbc/transport.py b/riak/transports/pbc/transport.py deleted file mode 100644 index 53df2181..00000000 --- a/riak/transports/pbc/transport.py +++ /dev/null @@ -1,770 +0,0 @@ -import riak.pb.messages -import riak.pb.riak_pb2 -import riak.pb.riak_kv_pb2 -import riak.pb.riak_ts_pb2 - -from riak import RiakError -from riak.transports.transport import RiakTransport -from riak.riak_object import VClock -from riak.ts_object import TsObject -from riak.util import decode_index_value, str_to_bytes, bytes_to_str -from riak.transports.pbc.connection import RiakPbcConnection -from riak.transports.pbc.stream import (RiakPbcKeyStream, - RiakPbcMapredStream, - RiakPbcBucketStream, - RiakPbcIndexStream, - RiakPbcTsKeyStream) -from riak.transports.pbc.codec import RiakPbcCodec -from six import PY2, PY3 - - -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, - 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._timeout = timeout - self._socket = None - - # FeatureDetection API - def _server_version(self): - return bytes_to_str(self.get_server_info()['server_version']) - - def ping(self): - """ - Ping the remote server - """ - - msg_code, msg = self._request(riak.pb.messages.MSG_CODE_PING_REQ) - if msg_code == riak.pb.messages.MSG_CODE_PING_RESP: - return True - else: - return False - - def get_server_info(self): - """ - Get information about the server - """ - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_SERVER_INFO_REQ, - expect=riak.pb.messages.MSG_CODE_GET_SERVER_INFO_RESP) - return {'node': bytes_to_str(resp.node), - 'server_version': bytes_to_str(resp.server_version)} - - def _get_client_id(self): - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_CLIENT_ID_REQ, - expect=riak.pb.messages.MSG_CODE_GET_CLIENT_ID_RESP) - return bytes_to_str(resp.client_id) - - def _set_client_id(self, client_id): - req = riak.pb.riak_kv_pb2.RpbSetClientIdReq() - req.client_id = str_to_bytes(client_id) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_SET_CLIENT_ID_REQ, req, - riak.pb.messages.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, timeout=None, basic_quorum=None, - notfound_ok=None): - """ - Serialize get request and deserialize response - """ - bucket = robj.bucket - - req = riak.pb.riak_kv_pb2.RpbGetReq() - if r: - req.r = self._encode_quorum(r) - if self.quorum_controls(): - if pr: - req.pr = self._encode_quorum(pr) - if basic_quorum is not None: - req.basic_quorum = basic_quorum - if notfound_ok is not None: - req.notfound_ok = notfound_ok - if self.client_timeouts() and timeout: - req.timeout = timeout - if self.tombstone_vclocks(): - req.deletedvclock = True - - req.bucket = str_to_bytes(bucket.name) - self._add_bucket_type(req, bucket.bucket_type) - - req.key = str_to_bytes(robj.key) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_REQ, req, - riak.pb.messages.MSG_CODE_GET_RESP) - - if resp is not None: - if resp.HasField('vclock'): - robj.vclock = VClock(resp.vclock, 'binary') - # We should do this even if there are no contents, i.e. - # the object is tombstoned - self._decode_contents(resp.content, robj) - else: - # "not found" returns an empty message, - # so let's make sure to clear the siblings - robj.siblings = [] - - return robj - - def put(self, robj, w=None, dw=None, pw=None, return_body=True, - if_none_match=False, timeout=None): - bucket = robj.bucket - - req = riak.pb.riak_kv_pb2.RpbPutReq() - if w: - req.w = self._encode_quorum(w) - if dw: - req.dw = self._encode_quorum(dw) - if self.quorum_controls() and pw: - req.pw = self._encode_quorum(pw) - - if return_body: - req.return_body = 1 - if if_none_match: - req.if_none_match = 1 - if self.client_timeouts() and timeout: - req.timeout = timeout - - req.bucket = str_to_bytes(bucket.name) - self._add_bucket_type(req, bucket.bucket_type) - - if robj.key: - req.key = str_to_bytes(robj.key) - if robj.vclock: - req.vclock = robj.vclock.encode('binary') - - self._encode_content(robj, req.content) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_PUT_REQ, req, - riak.pb.messages.MSG_CODE_PUT_RESP) - - if resp is not None: - if resp.HasField('key'): - robj.key = bytes_to_str(resp.key) - if resp.HasField("vclock"): - robj.vclock = VClock(resp.vclock, 'binary') - if resp.content: - self._decode_contents(resp.content, robj) - elif not robj.key: - raise RiakError("missing response object") - - return robj - - def ts_describe(self, table): - query = 'DESCRIBE {table}'.format(table=table.name) - return self.ts_query(table, query) - - def ts_get(self, table, key): - req = riak.pb.riak_ts_pb2.TsGetReq() - self._encode_timeseries_keyreq(table, key, req) - - msg_code, ts_get_resp = self._request( - riak.pb.messages.MSG_CODE_TS_GET_REQ, req, - riak.pb.messages.MSG_CODE_TS_GET_RESP) - - tsobj = TsObject(self._client, table, [], None) - self._decode_timeseries(ts_get_resp, tsobj) - return tsobj - - def ts_put(self, tsobj): - req = riak.pb.riak_ts_pb2.TsPutReq() - self._encode_timeseries_put(tsobj, req) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_TS_PUT_REQ, req, - riak.pb.messages.MSG_CODE_TS_PUT_RESP) - - if resp is not None: - return True - else: - raise RiakError("missing response object") - - def ts_delete(self, table, key): - req = riak.pb.riak_ts_pb2.TsDelReq() - self._encode_timeseries_keyreq(table, key, req) - - msg_code, ts_del_resp = self._request( - riak.pb.messages.MSG_CODE_TS_DEL_REQ, req, - riak.pb.messages.MSG_CODE_TS_DEL_RESP) - - if ts_del_resp is not None: - return True - else: - raise RiakError("missing response object") - - def ts_query(self, table, query, interpolations=None): - req = riak.pb.riak_ts_pb2.TsQueryReq() - - q = query - if '{table}' in q: - q = q.format(table=table.name) - - req.query.base = str_to_bytes(q) - - msg_code, ts_query_resp = self._request( - riak.pb.messages.MSG_CODE_TS_QUERY_REQ, req, - riak.pb.messages.MSG_CODE_TS_QUERY_RESP) - - tsobj = TsObject(self._client, table, [], []) - self._decode_timeseries(ts_query_resp, tsobj) - return tsobj - - def ts_stream_keys(self, table, timeout=None): - """ - Streams keys from a timeseries table, returning an iterator that - yields lists of keys. - """ - req = riak.pb.riak_ts_pb2.TsListKeysReq() - t = None - if self.client_timeouts() and timeout: - t = timeout - self._encode_timeseries_listkeysreq(table, req, t) - - self._send_msg(riak.pb.messages.MSG_CODE_TS_LIST_KEYS_REQ, req) - - return RiakPbcTsKeyStream(self) - - def delete(self, robj, rw=None, r=None, w=None, dw=None, pr=None, pw=None, - timeout=None): - req = riak.pb.riak_kv_pb2.RpbDelReq() - if rw: - req.rw = self._encode_quorum(rw) - if r: - req.r = self._encode_quorum(r) - if w: - req.w = self._encode_quorum(w) - if dw: - req.dw = self._encode_quorum(dw) - - if self.quorum_controls(): - if pr: - req.pr = self._encode_quorum(pr) - if pw: - req.pw = self._encode_quorum(pw) - - if self.client_timeouts() and timeout: - req.timeout = timeout - - use_vclocks = (self.tombstone_vclocks() and - hasattr(robj, 'vclock') and robj.vclock) - if use_vclocks: - req.vclock = robj.vclock.encode('binary') - - bucket = robj.bucket - req.bucket = str_to_bytes(bucket.name) - self._add_bucket_type(req, bucket.bucket_type) - req.key = str_to_bytes(robj.key) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_DEL_REQ, req, - riak.pb.messages.MSG_CODE_DEL_RESP) - return self - - def get_keys(self, bucket, timeout=None): - """ - Lists all keys within a bucket. - """ - keys = [] - for keylist in self.stream_keys(bucket, timeout=timeout): - for key in keylist: - keys.append(bytes_to_str(key)) - - return keys - - def stream_keys(self, bucket, timeout=None): - """ - Streams keys from a bucket, returning an iterator that yields - lists of keys. - """ - req = riak.pb.riak_kv_pb2.RpbListKeysReq() - req.bucket = str_to_bytes(bucket.name) - self._add_bucket_type(req, bucket.bucket_type) - if self.client_timeouts() and timeout: - req.timeout = timeout - - self._send_msg(riak.pb.messages.MSG_CODE_LIST_KEYS_REQ, req) - - return RiakPbcKeyStream(self) - - def get_buckets(self, bucket_type=None, timeout=None): - """ - Serialize bucket listing request and deserialize response - """ - req = riak.pb.riak_kv_pb2.RpbListBucketsReq() - self._add_bucket_type(req, bucket_type) - - if self.client_timeouts() and timeout: - req.timeout = timeout - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_LIST_BUCKETS_REQ, req, - riak.pb.messages.MSG_CODE_LIST_BUCKETS_RESP) - return resp.buckets - - def stream_buckets(self, bucket_type=None, timeout=None): - """ - Stream list of buckets through an iterator - """ - - if not self.bucket_stream(): - raise NotImplementedError('Streaming list-buckets is not ' - 'supported') - - req = riak.pb.riak_kv_pb2.RpbListBucketsReq() - req.stream = True - self._add_bucket_type(req, bucket_type) - # Bucket streaming landed in the same release as timeouts, so - # we don't need to check the capability. - if timeout: - req.timeout = timeout - - self._send_msg(riak.pb.messages.MSG_CODE_LIST_BUCKETS_REQ, req) - - return RiakPbcBucketStream(self) - - def get_bucket_props(self, bucket): - """ - Serialize bucket property request and deserialize response - """ - req = riak.pb.riak_pb2.RpbGetBucketReq() - req.bucket = str_to_bytes(bucket.name) - self._add_bucket_type(req, bucket.bucket_type) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_BUCKET_REQ, req, - riak.pb.messages.MSG_CODE_GET_BUCKET_RESP) - - return self._decode_bucket_props(resp.props) - - def set_bucket_props(self, bucket, props): - """ - Serialize set bucket property request and deserialize response - """ - req = riak.pb.riak_pb2.RpbSetBucketReq() - req.bucket = str_to_bytes(bucket.name) - self._add_bucket_type(req, bucket.bucket_type) - - if not self.pb_all_bucket_props(): - for key in props: - if key not in ('n_val', 'allow_mult'): - raise NotImplementedError('Server only supports n_val and ' - 'allow_mult properties over PBC') - - self._encode_bucket_props(props, req) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_SET_BUCKET_REQ, req, - riak.pb.messages.MSG_CODE_SET_BUCKET_RESP) - return True - - def clear_bucket_props(self, bucket): - """ - Clear bucket properties, resetting them to their defaults - """ - if not self.pb_clear_bucket_props(): - return False - - req = riak.pb.riak_pb2.RpbResetBucketReq() - req.bucket = str_to_bytes(bucket.name) - self._add_bucket_type(req, bucket.bucket_type) - self._request( - riak.pb.messages.MSG_CODE_RESET_BUCKET_REQ, req, - riak.pb.messages.MSG_CODE_RESET_BUCKET_RESP) - return True - - def get_bucket_type_props(self, bucket_type): - """ - Fetch bucket-type properties - """ - self._check_bucket_types(bucket_type) - - req = riak.pb.riak_pb2.RpbGetBucketTypeReq() - req.type = str_to_bytes(bucket_type.name) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_BUCKET_TYPE_REQ, req, - riak.pb.messages.MSG_CODE_GET_BUCKET_RESP) - - return self._decode_bucket_props(resp.props) - - def set_bucket_type_props(self, bucket_type, props): - """ - Set bucket-type properties - """ - self._check_bucket_types(bucket_type) - - req = riak.pb.riak_pb2.RpbSetBucketTypeReq() - req.type = str_to_bytes(bucket_type.name) - - self._encode_bucket_props(props, req) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_SET_BUCKET_TYPE_REQ, req, - riak.pb.messages.MSG_CODE_SET_BUCKET_RESP) - - return True - - 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.riak_kv_pb2.RpbMapRedReq() - req.request = str_to_bytes(content) - req.content_type = str_to_bytes("application/json") - - self._send_msg(riak.pb.messages.MSG_CODE_MAP_RED_REQ, req) - - return RiakPbcMapredStream(self) - - def get_index(self, bucket, index, startkey, endkey=None, - return_terms=None, max_results=None, continuation=None, - timeout=None, term_regex=None): - if not self.pb_indexes(): - return self._get_index_mapred_emu(bucket, index, startkey, endkey) - - if term_regex and not self.index_term_regex(): - raise NotImplementedError("Secondary index term_regex is not " - "supported") - - req = self._encode_index_req(bucket, index, startkey, endkey, - return_terms, max_results, continuation, - timeout, term_regex) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_INDEX_REQ, req, - riak.pb.messages.MSG_CODE_INDEX_RESP) - - if return_terms and resp.results: - results = [(decode_index_value(index, pair.key), - bytes_to_str(pair.value)) - for pair in resp.results] - else: - results = resp.keys[:] - if PY3: - results = [bytes_to_str(key) for key in resp.keys] - - if max_results is not None and resp.HasField('continuation'): - return (results, bytes_to_str(resp.continuation)) - else: - return (results, None) - - def stream_index(self, bucket, index, startkey, endkey=None, - return_terms=None, max_results=None, continuation=None, - timeout=None, term_regex=None): - if not self.stream_indexes(): - raise NotImplementedError("Secondary index streaming is not " - "supported") - - if term_regex and not self.index_term_regex(): - raise NotImplementedError("Secondary index term_regex is not " - "supported") - - req = self._encode_index_req(bucket, index, startkey, endkey, - return_terms, max_results, continuation, - timeout, term_regex) - req.stream = True - - self._send_msg(riak.pb.messages.MSG_CODE_INDEX_REQ, req) - - return RiakPbcIndexStream(self, index, return_terms) - - def create_search_index(self, index, schema=None, n_val=None, - timeout=None): - if not self.pb_search_admin(): - raise NotImplementedError("Search 2.0 administration is not " - "supported for this version") - index = str_to_bytes(index) - idx = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndex(name=index) - if schema: - idx.schema = str_to_bytes(schema) - if n_val: - idx.n_val = n_val - req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexPutReq(index=idx) - if timeout is not None: - req.timeout = timeout - - self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_PUT_REQ, req, - riak.pb.messages.MSG_CODE_PUT_RESP) - - return True - - def get_search_index(self, index): - if not self.pb_search_admin(): - raise NotImplementedError("Search 2.0 administration is not " - "supported for this version") - req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexGetReq( - name=str_to_bytes(index)) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_REQ, req, - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_RESP) - if len(resp.index) > 0: - return self._decode_search_index(resp.index[0]) - else: - raise RiakError('notfound') - - def list_search_indexes(self): - if not self.pb_search_admin(): - raise NotImplementedError("Search 2.0 administration is not " - "supported for this version") - req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexGetReq() - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_REQ, req, - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_RESP) - - return [self._decode_search_index(index) for index in resp.index] - - def delete_search_index(self, index): - if not self.pb_search_admin(): - raise NotImplementedError("Search 2.0 administration is not " - "supported for this version") - req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexDeleteReq( - name=str_to_bytes(index)) - - self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_DELETE_REQ, req, - riak.pb.messages.MSG_CODE_DEL_RESP) - - return True - - def create_search_schema(self, schema, content): - if not self.pb_search_admin(): - raise NotImplementedError("Search 2.0 administration is not " - "supported for this version") - scma = riak.pb.riak_yokozuna_pb2.RpbYokozunaSchema( - name=str_to_bytes(schema), - content=str_to_bytes(content)) - req = riak.pb.riak_yokozuna_pb2.RpbYokozunaSchemaPutReq( - schema=scma) - - self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_PUT_REQ, req, - riak.pb.messages.MSG_CODE_PUT_RESP) - - return True - - def get_search_schema(self, schema): - if not self.pb_search_admin(): - raise NotImplementedError("Search 2.0 administration is not " - "supported for this version") - req = riak.pb.riak_yokozuna_pb2.RpbYokozunaSchemaGetReq( - name=str_to_bytes(schema)) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_GET_REQ, req, - riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_GET_RESP) - - result = {} - result['name'] = bytes_to_str(resp.schema.name) - result['content'] = bytes_to_str(resp.schema.content) - return result - - def search(self, index, query, **params): - if not self.pb_search(): - return self._search_mapred_emu(index, query) - - if PY2 and isinstance(query, unicode): # noqa - query = query.encode('utf8') - - req = riak.pb.riak_search_pb2.RpbSearchQueryReq( - index=str_to_bytes(index), - q=str_to_bytes(query)) - self._encode_search_query(req, params) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_SEARCH_QUERY_REQ, req, - riak.pb.messages.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 - result['docs'] = [self._decode_search_doc(doc) for doc in resp.docs] - return result - - def get_counter(self, bucket, key, **params): - if not bucket.bucket_type.is_default(): - raise NotImplementedError("Counters are not " - "supported with bucket-types, " - "use datatypes instead.") - - if not self.counters(): - raise NotImplementedError("Counters are not supported") - - req = riak.pb.riak_kv_pb2.RpbCounterGetReq() - req.bucket = str_to_bytes(bucket.name) - req.key = str_to_bytes(key) - if params.get('r') is not None: - req.r = self._encode_quorum(params['r']) - if params.get('pr') is not None: - req.pr = self._encode_quorum(params['pr']) - if params.get('basic_quorum') is not None: - req.basic_quorum = params['basic_quorum'] - if params.get('notfound_ok') is not None: - req.notfound_ok = params['notfound_ok'] - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_COUNTER_GET_REQ, req, - riak.pb.messages.MSG_CODE_COUNTER_GET_RESP) - if resp.HasField('value'): - return resp.value - else: - return None - - def update_counter(self, bucket, key, value, **params): - if not bucket.bucket_type.is_default(): - raise NotImplementedError("Counters are not " - "supported with bucket-types, " - "use datatypes instead.") - - if not self.counters(): - raise NotImplementedError("Counters are not supported") - - req = riak.pb.riak_kv_pb2.RpbCounterUpdateReq() - req.bucket = str_to_bytes(bucket.name) - req.key = str_to_bytes(key) - req.amount = value - if params.get('w') is not None: - req.w = self._encode_quorum(params['w']) - if params.get('dw') is not None: - req.dw = self._encode_quorum(params['dw']) - if params.get('pw') is not None: - req.pw = self._encode_quorum(params['pw']) - if params.get('returnvalue') is not None: - req.returnvalue = params['returnvalue'] - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_COUNTER_UPDATE_REQ, req, - riak.pb.messages.MSG_CODE_COUNTER_UPDATE_RESP) - if resp.HasField('value'): - return resp.value - else: - return True - - def fetch_datatype(self, bucket, key, **options): - - if bucket.bucket_type.is_default(): - raise NotImplementedError("Datatypes cannot be used in the default" - " bucket-type.") - - if not self.datatypes(): - raise NotImplementedError("Datatypes are not supported.") - - req = riak.pb.riak_dt_pb2.DtFetchReq() - req.type = str_to_bytes(bucket.bucket_type.name) - req.bucket = str_to_bytes(bucket.name) - req.key = str_to_bytes(key) - self._encode_dt_options(req, options) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_DT_FETCH_REQ, req, - riak.pb.messages.MSG_CODE_DT_FETCH_RESP) - - return self._decode_dt_fetch(resp) - - def update_datatype(self, datatype, **options): - - if datatype.bucket.bucket_type.is_default(): - raise NotImplementedError("Datatypes cannot be used in the default" - " bucket-type.") - - if not self.datatypes(): - raise NotImplementedError("Datatypes are not supported.") - - op = datatype.to_op() - type_name = datatype.type_name - if not op: - raise ValueError("No operation to send on datatype {!r}". - format(datatype)) - - req = riak.pb.riak_dt_pb2.DtUpdateReq() - req.bucket = str_to_bytes(datatype.bucket.name) - req.type = str_to_bytes(datatype.bucket.bucket_type.name) - - if datatype.key: - req.key = str_to_bytes(datatype.key) - if datatype._context: - req.context = datatype._context - - self._encode_dt_options(req, options) - - self._encode_dt_op(type_name, req, op) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_DT_UPDATE_REQ, req, - riak.pb.messages.MSG_CODE_DT_UPDATE_RESP) - if resp.HasField('key'): - datatype.key = resp.key[:] - if resp.HasField('context'): - datatype._context = resp.context[:] - - if options.get('return_body'): - datatype._set_value(self._decode_dt_value(type_name, resp)) - - return True - - def get_preflist(self, bucket, key): - """ - Get the preflist for a bucket/key - - :param bucket: Riak Bucket - :type bucket: :class:`~riak.bucket.RiakBucket` - :param key: Riak Key - :type key: string - :rtype: list of dicts - """ - req = riak.pb.riak_kv_pb2.RpbGetBucketKeyPreflistReq() - req.bucket = str_to_bytes(bucket.name) - req.key = str_to_bytes(key) - req.type = str_to_bytes(bucket.bucket_type.name) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_BUCKET_KEY_PREFLIST_REQ, req, - riak.pb.messages.MSG_CODE_GET_BUCKET_KEY_PREFLIST_RESP) - - return [self._decode_preflist(item) for item in resp.preflist] diff --git a/riak/transports/pool.py b/riak/transports/pool.py index 4b21fd8e..d0a9ee7f 100644 --- a/riak/transports/pool.py +++ b/riak/transports/pool.py @@ -1,24 +1,8 @@ -""" -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 +from __future__ import print_function -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 -from __future__ import print_function from contextlib import contextmanager -import threading # This file is a rough port of the Innertube Ruby library diff --git a/riak/transports/tcp/__init__.py b/riak/transports/tcp/__init__.py new file mode 100644 index 00000000..312f9194 --- /dev/null +++ b/riak/transports/tcp/__init__.py @@ -0,0 +1,54 @@ +import errno +import socket + +from riak.transports.pool import Pool +from riak.transports.tcp.transport import TcpTransport + + +class TcpPool(Pool): + """ + A resource pool of TCP transports. + """ + def __init__(self, client, **options): + super(TcpPool, self).__init__() + self._client = client + self._options = options + + def create_resource(self): + node = self._client._choose_node() + return TcpTransport(node=node, + client=self._client, + **self._options) + + def destroy_resource(self, tcp): + tcp.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/connection.py b/riak/transports/tcp/connection.py similarity index 60% rename from riak/transports/pbc/connection.py rename to riak/transports/tcp/connection.py index 60f264c1..a9d75603 100644 --- a/riak/transports/pbc/connection.py +++ b/riak/transports/tcp/connection.py @@ -1,13 +1,13 @@ import socket import struct + import riak.pb.riak_pb2 import riak.pb.messages -from riak.security import SecurityError, USE_STDLIB_SSL from riak import RiakError -from riak.util import bytes_to_str, str_to_bytes +from riak.codecs.pbuf import PbufCodec +from riak.security import SecurityError, USE_STDLIB_SSL -from six import PY2 if not USE_STDLIB_SSL: from OpenSSL.SSL import Connection from riak.transports.security import configure_pyopenssl_context @@ -16,41 +16,42 @@ from riak.transports.security import configure_ssl_context -class RiakPbcConnection(object): +class TcpConnection(object): """ - Connection-related methods for RiakPbcTransport. + Connection-related methods for TcpTransport. """ - - def _encode_msg(self, msg_code, msg=None): - if msg is None: + def _encode_msg(self, msg_code, data=None): + if data 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 + hdr = struct.pack("!iB", 1 + len(data), msg_code) + return hdr + data - def _request(self, msg_code, msg=None, expect=None): - self._send_msg(msg_code, msg) - return self._recv_msg(expect) + def _send_recv(self, msg_code, data=None): + self._send_msg(msg_code, data) + return self._recv_msg() - def _non_connect_request(self, msg_code, msg=None, expect=None): + def _non_connect_send_recv(self, msg_code, data=None): """ - Similar to self._request, but doesn't try to initiate a connection, + Similar to self._send_recv, but doesn't try to initiate a connection, thus preventing an infinite loop. """ - self._non_connect_send_msg(msg_code, msg) - return self._recv_msg(expect) + self._non_connect_send_msg(msg_code, data) + return self._recv_msg() - def _non_connect_send_msg(self, msg_code, msg): + def _non_connect_send_recv_msg(self, msg): + self._non_connect_send_msg(msg.msg_code, msg.data) + return self._recv_msg() + + def _non_connect_send_msg(self, msg_code, data): """ Similar to self._send, but doesn't try to initiate a connection, thus preventing an infinite loop. """ - self._socket.sendall(self._encode_msg(msg_code, msg)) + self._socket.sendall(self._encode_msg(msg_code, data)) - def _send_msg(self, msg_code, msg): + def _send_msg(self, msg_code, data): self._connect() - self._non_connect_send_msg(msg_code, msg) + self._non_connect_send_msg(msg_code, data) def _init_security(self): """ @@ -68,9 +69,9 @@ def _starttls(self): Exchange a STARTTLS message with Riak to initiate secure communications return True is Riak responds with a STARTTLS response, False otherwise """ - msg_code, _ = self._non_connect_request( + resp_code, _ = self._non_connect_send_recv( riak.pb.messages.MSG_CODE_START_TLS) - if msg_code == riak.pb.messages.MSG_CODE_START_TLS: + if resp_code == riak.pb.messages.MSG_CODE_START_TLS: return True else: return False @@ -82,17 +83,14 @@ def _auth(self): Note: Riak will sleep for a short period of time upon a failed auth request/response to prevent denial of service attacks """ - req = riak.pb.riak_pb2.RpbAuthReq() - req.user = str_to_bytes(self._client._credentials.username) + codec = PbufCodec() + username = self._client._credentials.username password = self._client._credentials.password if not password: password = '' - req.password = str_to_bytes(password) - msg_code, _ = self._non_connect_request( - riak.pb.messages.MSG_CODE_AUTH_REQ, - req, - riak.pb.messages.MSG_CODE_AUTH_RESP) - if msg_code == riak.pb.messages.MSG_CODE_AUTH_RESP: + msg = codec.encode_auth(username, password) + resp_code, _ = self._non_connect_send_recv_msg(msg) + if resp_code == riak.pb.messages.MSG_CODE_AUTH_RESP: return True else: return False @@ -146,7 +144,6 @@ def _ssl_handshake(self): # ssl handshake successful ssl_socket.do_handshake() self._socket = ssl_socket - return True except ssl.SSLError as e: raise SecurityError(e) @@ -154,51 +151,36 @@ def _ssl_handshake(self): # fail if *any* exceptions are thrown during SSL handshake raise SecurityError(e) - def _recv_msg(self, expect=None): - self._recv_pkt() - msg_code, = struct.unpack("B", self._inbuf[:1]) - if msg_code is riak.pb.messages.MSG_CODE_ERROR_RESP: - err = self._parse_msg(msg_code, self._inbuf[1:]) - if err is None: - raise RiakError('no error provided!') - else: - raise RiakError(bytes_to_str(err.errmsg)) - elif msg_code in riak.pb.messages.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_msg(self): + msgbuf = self._recv_pkt() + mv = memoryview(msgbuf) + msg_code, = struct.unpack("B", mv[0:1]) + data = mv[1:].tobytes() + return (msg_code, data) def _recv_pkt(self): - nmsglen = self._socket.recv(4) - while len(nmsglen) < 4: - x = self._socket.recv(4 - len(nmsglen)) - if not x: - break - nmsglen += x - 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 - if PY2: - self._inbuf = '' - else: - self._inbuf = bytes() - 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: + # TODO FUTURE re-use buffer + msglen_buf = self._recv(4) + # NB: msg length is an unsigned int + msglen, = struct.unpack('!I', msglen_buf) + return self._recv(msglen) + + def _recv(self, msglen): + # TODO FUTURE re-use buffer + # http://stackoverflow.com/a/15964489 + msgbuf = bytearray(msglen) + view = memoryview(msgbuf) + nread = 0 + toread = msglen + while toread: + nbytes = self._socket.recv_into(view, toread) + view = view[nbytes:] # slicing views is cheap + toread -= nbytes + nread += nbytes + if nread != msglen: raise RiakError("Socket returned short packet %d - expected %d" - % (len(self._inbuf), self._inbuf_len)) + % (nread, msglen)) + return msgbuf def _connect(self): if not self._socket: @@ -218,19 +200,6 @@ def close(self): self._socket.close() del self._socket - def _parse_msg(self, code, packet): - try: - pbclass = riak.pb.messages.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 + # These are set in the TcpTransport initializer _address = None _timeout = None diff --git a/riak/transports/pbc/stream.py b/riak/transports/tcp/stream.py similarity index 70% rename from riak/transports/pbc/stream.py rename to riak/transports/tcp/stream.py index ed649279..3cf0e974 100644 --- a/riak/transports/pbc/stream.py +++ b/riak/transports/tcp/stream.py @@ -1,23 +1,25 @@ import json + import riak.pb.messages from riak.util import decode_index_value, bytes_to_str from riak.client.index_page import CONTINUATION -from riak.transports.pbc.codec import RiakPbcCodec +from riak.codecs.ttb import TtbCodec from six import PY2 -class RiakPbcStream(object): +class PbufStream(object): """ - Used internally by RiakPbcTransport to implement streaming + Used internally by TcpTransport to implement streaming operations. Implements the iterator interface. """ _expect = None - def __init__(self, transport): + def __init__(self, transport, codec): self.finished = False self.transport = transport + self.codec = codec self.resource = None def __iter__(self): @@ -28,7 +30,11 @@ def next(self): raise StopIteration try: - msg_code, resp = self.transport._recv_msg(expect=self._expect) + resp_code, data = self.transport._recv_msg() + self.codec.maybe_riak_error(resp_code, data) + expect = self._expect + self.codec.maybe_incorrect_code(resp_code, expect) + resp = self.codec.parse_msg(expect, data) except: self.finished = True raise @@ -62,15 +68,15 @@ def close(self): self.resource.release() -class RiakPbcKeyStream(RiakPbcStream): +class PbufKeyStream(PbufStream): """ - Used internally by RiakPbcTransport to implement key-list streams. + Used internally by TcpTransport to implement key-list streams. """ _expect = riak.pb.messages.MSG_CODE_LIST_KEYS_RESP def next(self): - response = super(RiakPbcKeyStream, self).next() + response = super(PbufKeyStream, self).next() if response.done and len(response.keys) is 0: raise StopIteration @@ -82,16 +88,16 @@ def __next__(self): return self.next() -class RiakPbcMapredStream(RiakPbcStream): +class PbufMapredStream(PbufStream): """ - Used internally by RiakPbcTransport to implement MapReduce + Used internally by TcpTransport to implement MapReduce streams. """ _expect = riak.pb.messages.MSG_CODE_MAP_RED_RESP def next(self): - response = super(RiakPbcMapredStream, self).next() + response = super(PbufMapredStream, self).next() if response.done and not response.HasField('response'): raise StopIteration @@ -103,15 +109,15 @@ def __next__(self): return self.next() -class RiakPbcBucketStream(RiakPbcStream): +class PbufBucketStream(PbufStream): """ - Used internally by RiakPbcTransport to implement key-list streams. + Used internally by TcpTransport to implement key-list streams. """ _expect = riak.pb.messages.MSG_CODE_LIST_BUCKETS_RESP def next(self): - response = super(RiakPbcBucketStream, self).next() + response = super(PbufBucketStream, self).next() if response.done and len(response.buckets) is 0: raise StopIteration @@ -123,21 +129,21 @@ def __next__(self): return self.next() -class RiakPbcIndexStream(RiakPbcStream): +class PbufIndexStream(PbufStream): """ - Used internally by RiakPbcTransport to implement Secondary Index + Used internally by TcpTransport to implement Secondary Index streams. """ _expect = riak.pb.messages.MSG_CODE_INDEX_RESP - def __init__(self, transport, index, return_terms=False): - super(RiakPbcIndexStream, self).__init__(transport) + def __init__(self, transport, codec, index, return_terms=False): + super(PbufIndexStream, self).__init__(transport, codec) self.index = index self.return_terms = return_terms def next(self): - response = super(RiakPbcIndexStream, self).next() + response = super(PbufIndexStream, self).next() if response.done and not (response.keys or response.results or @@ -161,22 +167,22 @@ def __next__(self): return self.next() -class RiakPbcTsKeyStream(RiakPbcStream, RiakPbcCodec): +class PbufTsKeyStream(PbufStream, TtbCodec): """ - Used internally by RiakPbcTransport to implement key-list streams. + Used internally by TcpTransport to implement TS key-list streams. """ _expect = riak.pb.messages.MSG_CODE_TS_LIST_KEYS_RESP def next(self): - response = super(RiakPbcTsKeyStream, self).next() + response = super(PbufTsKeyStream, self).next() if response.done and len(response.keys) is 0: raise StopIteration keys = [] for tsrow in response.keys: - keys.append(self._decode_timeseries_row(tsrow)) + keys.append(self.codec.decode_timeseries_row(tsrow)) return keys diff --git a/riak/transports/tcp/transport.py b/riak/transports/tcp/transport.py new file mode 100644 index 00000000..58420767 --- /dev/null +++ b/riak/transports/tcp/transport.py @@ -0,0 +1,535 @@ +import six +import riak.pb.messages + +from riak import RiakError +from riak.codecs import Codec, Msg +from riak.codecs.pbuf import PbufCodec +from riak.codecs.ttb import TtbCodec, MSG_CODE_TS_TTB +from riak.transports.transport import Transport +from riak.ts_object import TsObject + +from riak.transports.tcp.connection import TcpConnection +from riak.transports.tcp.stream import (PbufKeyStream, + PbufMapredStream, + PbufBucketStream, + PbufIndexStream, + PbufTsKeyStream) + + +class TcpTransport(Transport, TcpConnection): + """ + The TcpTransport object holds a connection to the TCP + socket on the Riak server. + """ + def __init__(self, + node=None, + client=None, + timeout=None, + **kwargs): + super(TcpTransport, self).__init__() + + self._client = client + self._node = node + self._address = (node.host, node.pb_port) + self._timeout = timeout + self._socket = None + self._pbuf_c = None + self._ttb_c = None + self._use_ttb = kwargs.get('use_ttb', True) + + def _get_pbuf_codec(self): + if not self._pbuf_c: + self._pbuf_c = PbufCodec( + self.client_timeouts(), self.quorum_controls(), + self.tombstone_vclocks(), self.bucket_types()) + return self._pbuf_c + + def _get_ttb_codec(self): + if self._use_ttb: + if not self._ttb_c: + self._ttb_c = TtbCodec() + codec = self._ttb_c + else: + codec = self._get_pbuf_codec() + return codec + + def _get_codec(self, msg_code): + if msg_code == MSG_CODE_TS_TTB: + codec = self._get_ttb_codec() + elif msg_code == riak.pb.messages.MSG_CODE_TS_GET_REQ: + codec = self._get_ttb_codec() + elif msg_code == riak.pb.messages.MSG_CODE_TS_PUT_REQ: + codec = self._get_ttb_codec() + elif msg_code == riak.pb.messages.MSG_CODE_TS_QUERY_REQ: + codec = self._get_ttb_codec() + else: + codec = self._get_pbuf_codec() + return codec + + # FeatureDetection API + def _server_version(self): + server_info = self.get_server_info() + return server_info['server_version'] + + def ping(self): + """ + Ping the remote server + """ + msg_code = riak.pb.messages.MSG_CODE_PING_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_ping() + resp_code, _ = self._request(msg, codec) + if resp_code == riak.pb.messages.MSG_CODE_PING_RESP: + return True + else: + return False + + def get_server_info(self): + """ + Get information about the server + """ + # NB: can't do it this way due to recursion + # codec = self._get_codec(ttb_supported=False) + codec = PbufCodec() + msg = Msg(riak.pb.messages.MSG_CODE_GET_SERVER_INFO_REQ, None, + riak.pb.messages.MSG_CODE_GET_SERVER_INFO_RESP) + resp_code, resp = self._request(msg, codec) + return codec.decode_get_server_info(resp) + + def _get_client_id(self): + msg_code = riak.pb.messages.MSG_CODE_GET_CLIENT_ID_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_get_client_id() + resp_code, resp = self._request(msg, codec) + return codec.decode_get_client_id(resp) + + def _set_client_id(self, client_id): + msg_code = riak.pb.messages.MSG_CODE_SET_CLIENT_ID_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_set_client_id(client_id) + resp_code, resp = self._request(msg, codec) + 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, timeout=None, basic_quorum=None, + notfound_ok=None): + """ + Serialize get request and deserialize response + """ + msg_code = riak.pb.messages.MSG_CODE_GET_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_get(robj, r, pr, + timeout, basic_quorum, + notfound_ok) + resp_code, resp = self._request(msg, codec) + return codec.decode_get(robj, resp) + + def put(self, robj, w=None, dw=None, pw=None, return_body=True, + if_none_match=False, timeout=None): + msg_code = riak.pb.messages.MSG_CODE_PUT_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_put(robj, w, dw, pw, return_body, + if_none_match, timeout) + resp_code, resp = self._request(msg, codec) + return codec.decode_put(robj, resp) + + def ts_describe(self, table): + query = 'DESCRIBE {table}'.format(table=table.name) + return self.ts_query(table, query) + + def ts_get(self, table, key): + msg_code = MSG_CODE_TS_TTB + codec = self._get_codec(msg_code) + msg = codec.encode_timeseries_keyreq(table, key) + resp_code, resp = self._request(msg, codec) + tsobj = TsObject(self._client, table) + codec.decode_timeseries(resp, tsobj) + return tsobj + + def ts_put(self, tsobj): + msg_code = MSG_CODE_TS_TTB + codec = self._get_codec(msg_code) + msg = codec.encode_timeseries_put(tsobj) + resp_code, resp = self._request(msg, codec) + return codec.validate_timeseries_put_resp(resp_code, resp) + + def ts_delete(self, table, key): + msg_code = riak.pb.messages.MSG_CODE_TS_DEL_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_timeseries_keyreq(table, key, is_delete=True) + resp_code, resp = self._request(msg, codec) + if resp is not None: + return True + else: + raise RiakError("missing response object") + + def ts_query(self, table, query, interpolations=None): + msg_code = riak.pb.messages.MSG_CODE_TS_QUERY_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_timeseries_query(table, query, interpolations) + resp_code, resp = self._request(msg, codec) + tsobj = TsObject(self._client, table) + codec.decode_timeseries(resp, tsobj) + return tsobj + + def ts_stream_keys(self, table, timeout=None): + """ + Streams keys from a timeseries table, returning an iterator that + yields lists of keys. + """ + msg_code = riak.pb.messages.MSG_CODE_TS_LIST_KEYS_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_timeseries_listkeysreq(table, timeout) + self._send_msg(msg.msg_code, msg.data) + return PbufTsKeyStream(self, codec) + + def delete(self, robj, rw=None, r=None, w=None, dw=None, + pr=None, pw=None, timeout=None): + msg_code = riak.pb.messages.MSG_CODE_DEL_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_delete(robj, rw, r, w, dw, pr, pw, timeout) + resp_code, resp = self._request(msg, codec) + return self + + def get_keys(self, bucket, timeout=None): + """ + Lists all keys within a bucket. + """ + msg_code = riak.pb.messages.MSG_CODE_LIST_KEYS_REQ + codec = self._get_codec(msg_code) + stream = self.stream_keys(bucket, timeout=timeout) + return codec.decode_get_keys(stream) + + def stream_keys(self, bucket, timeout=None): + """ + Streams keys from a bucket, returning an iterator that yields + lists of keys. + """ + msg_code = riak.pb.messages.MSG_CODE_LIST_KEYS_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_stream_keys(bucket, timeout) + self._send_msg(msg.msg_code, msg.data) + return PbufKeyStream(self, codec) + + def get_buckets(self, bucket_type=None, timeout=None): + """ + Serialize bucket listing request and deserialize response + """ + msg_code = riak.pb.messages.MSG_CODE_LIST_BUCKETS_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_get_buckets(bucket_type, + timeout, streaming=False) + resp_code, resp = self._request(msg, codec) + return resp.buckets + + def stream_buckets(self, bucket_type=None, timeout=None): + """ + Stream list of buckets through an iterator + """ + if not self.bucket_stream(): + raise NotImplementedError('Streaming list-buckets is not ' + 'supported') + msg_code = riak.pb.messages.MSG_CODE_LIST_BUCKETS_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_get_buckets(bucket_type, + timeout, streaming=True) + self._send_msg(msg.msg_code, msg.data) + return PbufBucketStream(self, codec) + + def get_bucket_props(self, bucket): + """ + Serialize bucket property request and deserialize response + """ + msg_code = riak.pb.messages.MSG_CODE_GET_BUCKET_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_get_bucket_props(bucket) + resp_code, resp = self._request(msg, codec) + return codec.decode_bucket_props(resp.props) + + def set_bucket_props(self, bucket, props): + """ + Serialize set bucket property request and deserialize response + """ + if not self.pb_all_bucket_props(): + for key in props: + if key not in ('n_val', 'allow_mult'): + raise NotImplementedError('Server only supports n_val and ' + 'allow_mult properties over PBC') + msg_code = riak.pb.messages.MSG_CODE_SET_BUCKET_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_set_bucket_props(bucket, props) + resp_code, resp = self._request(msg, codec) + return True + + def clear_bucket_props(self, bucket): + """ + Clear bucket properties, resetting them to their defaults + """ + if not self.pb_clear_bucket_props(): + return False + msg_code = riak.pb.messages.MSG_CODE_RESET_BUCKET_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_clear_bucket_props(bucket) + self._request(msg, codec) + return True + + def get_bucket_type_props(self, bucket_type): + """ + Fetch bucket-type properties + """ + self._check_bucket_types(bucket_type) + msg_code = riak.pb.messages.MSG_CODE_GET_BUCKET_TYPE_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_get_bucket_type_props(bucket_type) + resp_code, resp = self._request(msg, codec) + return codec.decode_bucket_props(resp.props) + + def set_bucket_type_props(self, bucket_type, props): + """ + Set bucket-type properties + """ + self._check_bucket_types(bucket_type) + msg_code = riak.pb.messages.MSG_CODE_SET_BUCKET_TYPE_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_set_bucket_type_props(bucket_type, props) + resp_code, resp = self._request(msg, codec) + return True + + 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... + msg_code = riak.pb.messages.MSG_CODE_MAP_RED_REQ + codec = self._get_codec(msg_code) + content = self._construct_mapred_json(inputs, query, timeout) + msg = codec.encode_stream_mapred(content) + self._send_msg(msg.msg_code, msg.data) + return PbufMapredStream(self, codec) + + def get_index(self, bucket, index, startkey, endkey=None, + return_terms=None, max_results=None, continuation=None, + timeout=None, term_regex=None): + # TODO RTS-842 NUKE THIS + if not self.pb_indexes(): + return self._get_index_mapred_emu(bucket, index, startkey, endkey) + + if term_regex and not self.index_term_regex(): + raise NotImplementedError("Secondary index term_regex is not " + "supported") + + msg_code = riak.pb.messages.MSG_CODE_INDEX_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_index_req(bucket, index, startkey, endkey, + return_terms, max_results, + continuation, timeout, + term_regex, streaming=False) + resp_code, resp = self._request(msg, codec) + return codec.decode_index_req(resp, index, + return_terms, max_results) + + def stream_index(self, bucket, index, startkey, endkey=None, + return_terms=None, max_results=None, continuation=None, + timeout=None, term_regex=None): + if not self.stream_indexes(): + raise NotImplementedError("Secondary index streaming is not " + "supported") + if term_regex and not self.index_term_regex(): + raise NotImplementedError("Secondary index term_regex is not " + "supported") + msg_code = riak.pb.messages.MSG_CODE_INDEX_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_index_req(bucket, index, startkey, endkey, + return_terms, max_results, + continuation, timeout, + term_regex, streaming=True) + self._send_msg(msg.msg_code, msg.data) + return PbufIndexStream(self, codec, index, return_terms) + + def create_search_index(self, index, schema=None, n_val=None, + timeout=None): + if not self.pb_search_admin(): + raise NotImplementedError("Search 2.0 administration is not " + "supported for this version") + msg_code = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_PUT_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_create_search_index(index, schema, n_val, timeout) + self._request(msg, codec) + return True + + def get_search_index(self, index): + if not self.pb_search_admin(): + raise NotImplementedError("Search 2.0 administration is not " + "supported for this version") + msg_code = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_get_search_index(index) + resp_code, resp = self._request(msg, codec) + if len(resp.index) > 0: + return codec.decode_search_index(resp.index[0]) + else: + raise RiakError('notfound') + + def list_search_indexes(self): + if not self.pb_search_admin(): + raise NotImplementedError("Search 2.0 administration is not " + "supported for this version") + msg_code = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_list_search_indexes() + resp_code, resp = self._request(msg, codec) + return [codec.decode_search_index(index) for index in resp.index] + + def delete_search_index(self, index): + if not self.pb_search_admin(): + raise NotImplementedError("Search 2.0 administration is not " + "supported for this version") + msg_code = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_DELETE_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_delete_search_index(index) + self._request(msg, codec) + return True + + def create_search_schema(self, schema, content): + if not self.pb_search_admin(): + raise NotImplementedError("Search 2.0 administration is not " + "supported for this version") + msg_code = riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_PUT_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_create_search_schema(schema, content) + self._request(msg, codec) + return True + + def get_search_schema(self, schema): + if not self.pb_search_admin(): + raise NotImplementedError("Search 2.0 administration is not " + "supported for this version") + msg_code = riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_GET_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_get_search_schema(schema) + resp_code, resp = self._request(msg, codec) + return codec.decode_get_search_schema(resp) + + def search(self, index, query, **kwargs): + # TODO RTS-842 NUKE THIS + if not self.pb_search(): + return self._search_mapred_emu(index, query) + # TODO RTS-842 six.u() instead? + if six.PY2 and isinstance(query, unicode): # noqa + query = query.encode('utf8') + msg_code = riak.pb.messages.MSG_CODE_SEARCH_QUERY_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_search(index, query, **kwargs) + resp_code, resp = self._request(msg, codec) + return codec.decode_search(resp) + + def get_counter(self, bucket, key, **kwargs): + if not bucket.bucket_type.is_default(): + raise NotImplementedError("Counters are not " + "supported with bucket-types, " + "use datatypes instead.") + if not self.counters(): + raise NotImplementedError("Counters are not supported") + msg_code = riak.pb.messages.MSG_CODE_COUNTER_GET_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_get_counter(bucket, key, **kwargs) + resp_code, resp = self._request(msg, codec) + if resp.HasField('value'): + return resp.value + else: + return None + + def update_counter(self, bucket, key, value, **kwargs): + if not bucket.bucket_type.is_default(): + raise NotImplementedError("Counters are not " + "supported with bucket-types, " + "use datatypes instead.") + if not self.counters(): + raise NotImplementedError("Counters are not supported") + msg_code = riak.pb.messages.MSG_CODE_COUNTER_UPDATE_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_update_counter(bucket, key, value, **kwargs) + resp_code, resp = self._request(msg, codec) + if resp.HasField('value'): + return resp.value + else: + return True + + def fetch_datatype(self, bucket, key, **kwargs): + if bucket.bucket_type.is_default(): + raise NotImplementedError("Datatypes cannot be used in the default" + " bucket-type.") + if not self.datatypes(): + raise NotImplementedError("Datatypes are not supported.") + msg_code = riak.pb.messages.MSG_CODE_DT_FETCH_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_fetch_datatype(bucket, key, **kwargs) + resp_code, resp = self._request(msg, codec) + return codec.decode_dt_fetch(resp) + + def update_datatype(self, datatype, **kwargs): + if datatype.bucket.bucket_type.is_default(): + raise NotImplementedError("Datatypes cannot be used in the default" + " bucket-type.") + if not self.datatypes(): + raise NotImplementedError("Datatypes are not supported.") + msg_code = riak.pb.messages.MSG_CODE_DT_UPDATE_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_update_datatype(datatype, **kwargs) + resp_code, resp = self._request(msg, codec) + codec.decode_update_datatype(datatype, resp, **kwargs) + return True + + def get_preflist(self, bucket, key): + """ + Get the preflist for a bucket/key + + :param bucket: Riak Bucket + :type bucket: :class:`~riak.bucket.RiakBucket` + :param key: Riak Key + :type key: string + :rtype: list of dicts + """ + msg_code = riak.pb.messages.MSG_CODE_GET_BUCKET_KEY_PREFLIST_REQ + codec = self._get_codec(msg_code) + msg = codec.encode_get_preflist(bucket, key) + resp_code, resp = self._request(msg, codec) + return [codec.decode_preflist(item) for item in resp.preflist] + + def _request(self, msg, codec=None): + if isinstance(msg, Msg): + msg_code = msg.msg_code + data = msg.data + expect = msg.resp_code + else: + raise ValueError('expected a Msg argument') + + if not isinstance(codec, Codec): + raise ValueError('expected a Codec argument') + + resp_code, data = self._send_recv(msg_code, data) + codec.maybe_riak_error(resp_code, data) + codec.maybe_incorrect_code(resp_code, expect) + if resp_code == MSG_CODE_TS_TTB or \ + resp_code in riak.pb.messages.MESSAGE_CLASSES: + msg = codec.parse_msg(resp_code, data) + else: + raise Exception("unknown msg code %s" % resp_code) + return resp_code, msg diff --git a/riak/transports/transport.py b/riak/transports/transport.py index 4f33168c..6e5fee2c 100644 --- a/riak/transports/transport.py +++ b/riak/transports/transport.py @@ -1,33 +1,15 @@ -""" -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 base64 import random import threading import os import json import platform + from six import PY2 from riak.transports.feature_detect import FeatureDetection -class RiakTransport(FeatureDetection): +class Transport(FeatureDetection): """ Class to encapsulate transport details and methods. All protocol transports are subclasses of this class. @@ -313,6 +295,7 @@ def get_preflist(self, bucket, key): """ raise NotImplementedError + # TODO RTS-842 NUKE THIS def _search_mapred_emu(self, index, query): """ Emulates a search request via MapReduce. Used in the case @@ -338,6 +321,7 @@ def _search_mapred_emu(self, index, query): result['docs'].append({u'id': key}) return result + # TODO RTS-842 NUKE THIS def _get_index_mapred_emu(self, bucket, index, startkey, endkey=None): """ Emulates a secondary index request via MapReduce. Used in the @@ -378,6 +362,5 @@ def _construct_mapred_json(self, inputs, query, timeout=None): def _check_bucket_types(self, bucket_type): if not self.bucket_types(): raise NotImplementedError('Server does not support bucket-types') - if bucket_type.is_default(): raise ValueError('Cannot manipulate the default bucket-type') diff --git a/riak/ts_object.py b/riak/ts_object.py index ef01baff..24eccbe1 100644 --- a/riak/ts_object.py +++ b/riak/ts_object.py @@ -1,13 +1,17 @@ +import collections + from riak import RiakError from riak.table import Table +TsColumns = collections.namedtuple('TsColumns', ['names', 'types']) + class TsObject(object): """ The TsObject holds information about Timeseries data, plus the data itself. """ - def __init__(self, client, table, rows=[], columns=[]): + def __init__(self, client, table, rows=None, columns=None): """ Construct a new TsObject. @@ -17,8 +21,8 @@ def __init__(self, client, table, rows=[], columns=[]): :type table: :class:`Table` :param rows: An list of lists with timeseries data :type rows: list - :param columns: An list of Column names and types. Optional. - :type columns: list + :param columns: A TsColumns tuple. Optional + :type columns: :class:`TsColumns` """ if not isinstance(table, Table): @@ -27,13 +31,17 @@ def __init__(self, client, table, rows=[], columns=[]): self.client = client self.table = table - self.rows = rows - if not isinstance(self.rows, list): - raise RiakError("TsObject requires a list of rows") - - self.columns = columns - if self.columns is not None and not isinstance(self.columns, list): - raise RiakError("TsObject columns must be a list") + if rows is not None and not isinstance(rows, list): + raise RiakError("TsObject rows parameter must be a list.") + else: + self.rows = rows + + if columns is not None and \ + not isinstance(columns, TsColumns): + raise RiakError( + "TsObject columns parameter must be a TsColumns instance") + else: + self.columns = columns def store(self): """ diff --git a/riak/util.py b/riak/util.py index 3932be99..4cbe6c0f 100644 --- a/riak/util.py +++ b/riak/util.py @@ -5,6 +5,18 @@ from collections import Mapping from six import string_types, PY2 +import datetime + +epoch = datetime.datetime.utcfromtimestamp(0) + + +def unix_time_millis(dt): + td = dt - epoch + return int(td.total_seconds() * 1000.0) + + +def datetime_from_unix_time_millis(ut): + return datetime.datetime.utcfromtimestamp(ut / 1000.0) def is_timeseries_supported(v=None): @@ -60,7 +72,6 @@ class lazy_property(object): memoization of an object attribute. The property should represent immutable data, as it replaces itself on first access. ''' - def __init__(self, fget): self.fget = fget self.func_name = fget.__name__ diff --git a/riak_pb b/riak_pb index 7fffa81b..341269c1 160000 --- a/riak_pb +++ b/riak_pb @@ -1 +1 @@ -Subproject commit 7fffa81b38804c18fffbec8d1677966c37d49d55 +Subproject commit 341269c19c75fa0557d5aa5fd5ac1f0dfe18cfae diff --git a/setup.py b/setup.py index faa91886..6338efd8 100755 --- a/setup.py +++ b/setup.py @@ -1,32 +1,35 @@ #!/usr/bin/env python -import platform +import codecs +import six +import sys + from setuptools import setup, find_packages from version import get_version from commands import setup_timeseries, build_messages -install_requires = ['six >= 1.8.0'] -requires = ['six(>=1.8.0)'] -if platform.python_version() < '2.7.9': +install_requires = ['six >= 1.8.0', 'basho_erlastic >= 2.1.0'] +requires = ['six(>=1.8.0)', 'basho_erlastic(>= 2.1.0)'] + +if sys.version_info[0:3] <= (2, 7, 9): install_requires.append("pyOpenSSL >= 0.14") requires.append("pyOpenSSL(>=0.14)") -if platform.python_version() < '3.0': +if six.PY2: install_requires.append('protobuf >=2.4.1, <2.7.0') requires.append('protobuf(>=2.4.1, <2.7.0)') else: install_requires.append('python3_protobuf >=2.4.1, <2.6.0') requires.append('python3_protobuf(>=2.4.1, <2.6.0)') -tests_require = [] -if platform.python_version() < '2.7.0': - tests_require.append("unittest2") - try: import pypandoc long_description = pypandoc.convert('README.md', 'rst') + with codecs.open('README.rst', 'w', 'utf-8') as f: + f.write(long_description) except(IOError, ImportError): - long_description = open('README.md').read() + with open('README.md') as f: + long_description = f.read() setup( name='riak', @@ -34,7 +37,6 @@ packages=find_packages(), requires=requires, install_requires=install_requires, - tests_require=tests_require, package_data={'riak': ['erl_src/*']}, description='Python client for Riak', long_description=long_description, diff --git a/tools b/tools index 5ff5850e..4dae68dd 160000 --- a/tools +++ b/tools @@ -1 +1 @@ -Subproject commit 5ff5850e1d7164f4f64f45a31d9b257e01a19e58 +Subproject commit 4dae68ddca2d405090d64a97c7e99b4607263892