From 3b20639cbeaf443cfb51ff2c15c39d005fe26513 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Thu, 28 Jan 2016 10:54:42 -0800 Subject: [PATCH 01/44] Get benchmarks in their own spot to add some for timeseries. --- riak/benchmarks/multiget.py | 40 ++++++++++++++++++++++++ riak/client/multiget.py | 61 ++----------------------------------- 2 files changed, 42 insertions(+), 59 deletions(-) create mode 100644 riak/benchmarks/multiget.py diff --git a/riak/benchmarks/multiget.py b/riak/benchmarks/multiget.py new file mode 100644 index 00000000..6baaba71 --- /dev/null +++ b/riak/benchmarks/multiget.py @@ -0,0 +1,40 @@ +from riak import RiakClient +from multiprocessing import cpu_count +import binascii +import os +import riak.benchmark as benchmark +import riak.client.multiget as mget + +client = RiakClient(protocol='pbc') +bkeys = [('default', 'multiget', str(key)) for key in range(10000)] + +data = binascii.b2a_hex(os.urandom(1024)) + +print("Benchmarking multiget:") +print(" CPUs: {0}".format(cpu_count())) +print(" Threads: {0}".format(mget.POOL_SIZE)) +print(" Keys: {0}".format(len(bkeys))) +print() + +with benchmark.measure() as b: + with b.report('populate'): + for _, bucket, key in bkeys: + client.bucket(bucket).new(key, encoded_data=data, + content_type='text/plain' + ).store() +for b in benchmark.measure_with_rehearsal(): + client.protocol = 'http' + with b.report('http seq'): + for _, bucket, key in bkeys: + client.bucket(bucket).get(key) + + with b.report('http multi'): + mget.multiget(client, bkeys) + + client.protocol = 'pbc' + with b.report('pbc seq'): + for _, bucket, key in bkeys: + client.bucket(bucket).get(key) + + with b.report('pbc multi'): + mget.multiget(client, bkeys) diff --git a/riak/client/multiget.py b/riak/client/multiget.py index 20d02801..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: @@ -201,43 +184,3 @@ def multiget(client, keys, **options): outq.task_done() return results - -if __name__ == '__main__': - # Run a benchmark! - from riak import RiakClient - import riak.benchmark as benchmark - client = RiakClient(protocol='pbc') - bkeys = [('default', 'multiget', str(key)) for key in range(10000)] - - data = None - with open(__file__) as f: - data = f.read() - - print("Benchmarking multiget:") - print(" CPUs: {0}".format(cpu_count())) - print(" Threads: {0}".format(POOL_SIZE)) - print(" Keys: {0}".format(len(bkeys))) - print() - - with benchmark.measure() as b: - with b.report('populate'): - for _, bucket, key in bkeys: - client.bucket(bucket).new(key, encoded_data=data, - content_type='text/plain' - ).store() - for b in benchmark.measure_with_rehearsal(): - client.protocol = 'http' - with b.report('http seq'): - for _, bucket, key in bkeys: - client.bucket(bucket).get(key) - - with b.report('http multi'): - multiget(client, bkeys) - - client.protocol = 'pbc' - with b.report('pbc seq'): - for _, bucket, key in bkeys: - client.bucket(bucket).get(key) - - with b.report('pbc multi'): - multiget(client, bkeys) From 02f509ed1c840ccd3777c6ee80c26aeee799d134 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Thu, 28 Jan 2016 13:22:49 -0800 Subject: [PATCH 02/44] Begin timeseries benchmarks --- riak/benchmarks/timeseries.py | 39 +++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 riak/benchmarks/timeseries.py diff --git a/riak/benchmarks/timeseries.py b/riak/benchmarks/timeseries.py new file mode 100644 index 00000000..ef3d789c --- /dev/null +++ b/riak/benchmarks/timeseries.py @@ -0,0 +1,39 @@ +from riak import RiakClient +from multiprocessing import cpu_count +import binascii +import os +import riak.benchmark as benchmark + +client = RiakClient(protocol='pbc') +bkeys = [('default', 'multiget', str(key)) for key in range(10000)] + +data = binascii.b2a_hex(os.urandom(1024)) + +print("Benchmarking timeseries:") +print(" CPUs: {0}".format(cpu_count())) +print(" Keys: {0}".format(len(bkeys))) +print() + +with benchmark.measure() as b: + with b.report('populate'): + for _, bucket, key in bkeys: + client.bucket(bucket).new(key, encoded_data=data, + content_type='text/plain' + ).store() + +for b in benchmark.measure_with_rehearsal(): + client.protocol = 'http' + with b.report('http seq'): + for _, bucket, key in bkeys: + client.bucket(bucket).get(key) + + with b.report('http multi'): + mget.multiget(client, bkeys) + + client.protocol = 'pbc' + with b.report('pbc seq'): + for _, bucket, key in bkeys: + client.bucket(bucket).get(key) + + with b.report('pbc multi'): + mget.multiget(client, bkeys) From 2baadfb8cde23bb8ec521c7e7bf7ca80d58dd478 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Thu, 28 Jan 2016 15:51:18 -0800 Subject: [PATCH 03/44] No need to limit connections to localhost. --- commands.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/commands.py b/commands.py index c12b9191..73a83ce2 100644 --- a/commands.py +++ b/commands.py @@ -266,9 +266,9 @@ class setup_security(Command, security_commands): _commands = [ "add-user $USERNAME password=$PASSWORD", - "add-source $USERNAME 127.0.0.1/32 password", + "add-source $USERNAME 0.0.0.0/32 password", "add-user $CERTUSER password=$CERTPASS", - "add-source $CERTUSER 127.0.0.1/32 certificate" + "add-source $CERTUSER 0.0.0.0/32 certificate" ] _grants = { @@ -392,9 +392,9 @@ class preconfigure(Command): * Update these lines in riak.conf * storage_backend = leveldb * search = on - * listener.protobuf.internal = 127.0.0.1:8087 - * listener.http.internal = 127.0.0.1:8098 - * listener.https.internal = 127.0.0.1:18098 + * listener.protobuf.internal = 0.0.0.0:8087 + * listener.http.internal = 0.0.0.0:8098 + * listener.https.internal = 0.0.0.0:18098 * ssl.certfile = $pwd/tests/resources/server.crt * ssl.keyfile = $pwd/tests/resources/server.key * ssl.cacertfile = $pwd/tests/resources/ca.crt @@ -412,7 +412,7 @@ class preconfigure(Command): def initialize_options(self): self.riak_conf = None - self.host = "127.0.0.1" + self.host = "0.0.0.0" self.pb_port = "8087" self.http_port = "8098" self.https_port = "18098" From 270a856d13230a25f37f7dbfaa3eeafa3f3ae557 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Mon, 1 Feb 2016 07:19:04 -0800 Subject: [PATCH 04/44] Put / Get TS benchmark complete. --- riak/benchmark.py | 18 -------- riak/benchmarks/multiget.py | 4 +- riak/benchmarks/timeseries.py | 84 +++++++++++++++++++++++------------ riak/tests/test_timeseries.py | 13 +++--- riak/transports/pbc/codec.py | 20 +++------ riak/util.py | 19 +++++++- 6 files changed, 88 insertions(+), 70 deletions(-) diff --git a/riak/benchmark.py b/riak/benchmark.py index 13286100..c26a0a49 100644 --- a/riak/benchmark.py +++ b/riak/benchmark.py @@ -1,21 +1,3 @@ -""" -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 diff --git a/riak/benchmarks/multiget.py b/riak/benchmarks/multiget.py index 6baaba71..505069a4 100644 --- a/riak/benchmarks/multiget.py +++ b/riak/benchmarks/multiget.py @@ -20,8 +20,8 @@ with b.report('populate'): for _, bucket, key in bkeys: client.bucket(bucket).new(key, encoded_data=data, - content_type='text/plain' - ).store() + content_type='text/plain' + ).store() for b in benchmark.measure_with_rehearsal(): client.protocol = 'http' with b.report('http seq'): diff --git a/riak/benchmarks/timeseries.py b/riak/benchmarks/timeseries.py index ef3d789c..57ed6f5c 100644 --- a/riak/benchmarks/timeseries.py +++ b/riak/benchmarks/timeseries.py @@ -1,39 +1,67 @@ -from riak import RiakClient from multiprocessing import cpu_count -import binascii -import os +from riak import RiakClient import riak.benchmark as benchmark +import datetime +import random -client = RiakClient(protocol='pbc') -bkeys = [('default', 'multiget', str(key)) for key in range(10000)] +epoch = datetime.datetime.utcfromtimestamp(0) +onesec = datetime.timedelta(0, 1) -data = binascii.b2a_hex(os.urandom(1024)) +rowcount = 32768 +batchsz = 32 +if rowcount % batchsz != 0: + raise AssertionError('rowcount must be divisible by batchsz') + +weather = ['typhoon', 'hurricane', 'rain', 'wind', 'snow'] +rows = [] +keys = [] +for i in range(rowcount): + ts = datetime.datetime(2016, 1, 1, 12, 0, 0) + \ + datetime.timedelta(seconds=i) + family_idx = i % 4 + series_idx = i % 4 + 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) + keys.append(key) print("Benchmarking timeseries:") print(" CPUs: {0}".format(cpu_count())) -print(" Keys: {0}".format(len(bkeys))) +print(" Rows: {0}".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} +] +client = RiakClient(nodes=n, protocol='pbc') +table = client.table(tbl) + with benchmark.measure() as b: with b.report('populate'): - for _, bucket, key in bkeys: - client.bucket(bucket).new(key, encoded_data=data, - content_type='text/plain' - ).store() - -for b in benchmark.measure_with_rehearsal(): - client.protocol = 'http' - with b.report('http seq'): - for _, bucket, key in bkeys: - client.bucket(bucket).get(key) - - with b.report('http multi'): - mget.multiget(client, bkeys) - - client.protocol = 'pbc' - with b.report('pbc seq'): - for _, bucket, key in bkeys: - client.bucket(bucket).get(key) - - with b.report('pbc multi'): - mget.multiget(client, bkeys) + 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") + with b.report('get'): + for k in keys: + ts_obj = client.ts_get(tbl, k) + if ts_obj is None: + raise AssertionError("expected obj") + if len(ts_obj.rows) != 1: + raise AssertionError("expected one row") + row = ts_obj.rows[0] + if len(row) != 5: + raise AssertionError("expected row to have five items") diff --git a/riak/tests/test_timeseries.py b/riak/tests/test_timeseries.py index e36cd95e..d1b64642 100644 --- a/riak/tests/test_timeseries.py +++ b/riak/tests/test_timeseries.py @@ -9,7 +9,7 @@ from riak.table import Table from riak.ts_object import TsObject from riak.transports.pbc.codec import RiakPbcCodec -from riak.util import str_to_bytes, bytes_to_str +from riak.util import str_to_bytes, bytes_to_str, unix_time_millis from riak.tests import RUN_TIMESERIES from riak.tests.base import IntegrationTestBase from riak.pb.riak_ts_pb2 import TsColumnType @@ -32,8 +32,8 @@ class TimeseriesUnitTests(unittest.TestCase): def setUp(self): self.c = RiakPbcCodec() - self.ts0ms = self.c._unix_time_millis(ts0) - self.ts1ms = self.c._unix_time_millis(ts1) + self.ts0ms = unix_time_millis(ts0) + self.ts1ms = unix_time_millis(ts1) self.rows = [ [bd0, 0, 1.2, ts0, True], [bd1, 3, 4.5, ts1, False] @@ -193,13 +193,12 @@ 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 diff --git a/riak/transports/pbc/codec.py b/riak/transports/pbc/codec.py index a7be2a89..e5da11c2 100644 --- a/riak/transports/pbc/codec.py +++ b/riak/transports/pbc/codec.py @@ -8,14 +8,13 @@ from riak import RiakError from riak.content import RiakContent -from riak.util import decode_index_value, str_to_bytes, bytes_to_str +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 from riak.pb.riak_ts_pb2 import TsColumnType from six import string_types, PY2 -epoch = datetime.datetime.utcfromtimestamp(0) - def _invert(d): out = {} @@ -84,17 +83,10 @@ def __init__(self, **unused_args): super(RiakPbcCodec, self).__init__(**unused_args) def _unix_time_millis(self, dt): - td = dt - epoch - try: - return int(dt.total_seconds() * 1000.0) - except AttributeError: - # NB: python 2.6 must use this method - return int(((td.microseconds + - (td.seconds + td.days * 24 * 3600) * 10**6) / - 10**6) * 1000.0) + return unix_time_millis(dt) def _datetime_from_unix_time_millis(self, ut): - return datetime.datetime.utcfromtimestamp(ut / 1000.0) + return datetime_from_unix_time_millis(ut) def _encode_quorum(self, rw): """ @@ -637,7 +629,7 @@ def _encode_map_update(self, dtype, msg, op): 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) + ts_cell.timestamp_value = unix_time_millis(cell) elif isinstance(cell, bool): ts_cell.boolean_value = cell elif isinstance(cell, string_types): @@ -752,7 +744,7 @@ def _decode_timeseries_row(self, tsrow, tscols=None): if col and col.type != TsColumnType.Value('TIMESTAMP'): raise TypeError('expected TIMESTAMP column') else: - dt = self._datetime_from_unix_time_millis( + dt = datetime_from_unix_time_millis( cell.timestamp_value) row.append(dt) elif cell.HasField('boolean_value'): diff --git a/riak/util.py b/riak/util.py index 5dc3e61a..4ea50389 100644 --- a/riak/util.py +++ b/riak/util.py @@ -2,6 +2,24 @@ import warnings from collections import Mapping from six import string_types, PY2 +import datetime + +epoch = datetime.datetime.utcfromtimestamp(0) + + +def unix_time_millis(dt): + try: + return int(dt.total_seconds() * 1000.0) + except AttributeError: + # NB: python 2.6 must use this method + td = dt - epoch + return int(((td.microseconds + + (td.seconds + td.days * 24 * 3600) * 10**6) / + 10**6) * 1000.0) + + +def datetime_from_unix_time_millis(ut): + return datetime.datetime.utcfromtimestamp(ut / 1000.0) def quacks_like_dict(object): @@ -51,7 +69,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__ From e99c6965ecb6edfdfbd8393a7c1274f8b557147e Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Mon, 1 Feb 2016 15:42:01 -0800 Subject: [PATCH 05/44] Adding TTB encoding and tests. --- riak/tests/test_timeseries.py | 2 +- riak/tests/test_timeseries_ttb.py | 40 +++++++++++++++++++++++++++ riak/transports/ttb/codec.py | 46 +++++++++++++++++++++++++++++++ setup.py | 4 +-- 4 files changed, 89 insertions(+), 3 deletions(-) create mode 100644 riak/tests/test_timeseries_ttb.py create mode 100644 riak/transports/ttb/codec.py diff --git a/riak/tests/test_timeseries.py b/riak/tests/test_timeseries.py index d1b64642..21a53497 100644 --- a/riak/tests/test_timeseries.py +++ b/riak/tests/test_timeseries.py @@ -39,7 +39,7 @@ def setUp(self): [bd1, 3, 4.5, ts1, False] ] self.test_key = ['hash1', 'user2', ts0] - self.table = Table(None, 'test-table') + self.table = Table(None, table_name) def validate_keyreq(self, req): self.assertEqual(self.table.name, bytes_to_str(req.table)) diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py new file mode 100644 index 00000000..914758b2 --- /dev/null +++ b/riak/tests/test_timeseries_ttb.py @@ -0,0 +1,40 @@ +# -*- coding: utf-8 -*- +import datetime +import platform +import random +import string + +from riak.table import Table +from riak.transports.ttb.codec import RiakTtbCodec +from riak.util import str_to_bytes, bytes_to_str, unix_time_millis + +if platform.python_version() < '2.7': + unittest = __import__('unittest2') +else: + import unittest + +table_name = 'GeoCheckin' + +bd0 = '时间序列' +bd1 = 'временные ряды' + +fiveMins = datetime.timedelta(0, 300) +ts0 = datetime.datetime(2015, 1, 1, 12, 0, 0) +ts1 = ts0 + fiveMins + + +class TimeseriesTtbUnitTests(unittest.TestCase): + def setUp(self): + self.c = RiakTtbCodec() + self.ts0ms = unix_time_millis(ts0) + self.ts1ms = unix_time_millis(ts1) + self.rows = [ + [bd0, 0, 1.2, ts0, True], + [bd1, 3, 4.5, ts1, False] + ] + self.test_key = ['hash1', 'user2', ts0] + self.table = Table(None, table_name) + + def test_encode_data_for_get(self): + req = self.c._encode_timeseries_keyreq(self.table, self.test_key) + self.assertIsNotNone(req) diff --git a/riak/transports/ttb/codec.py b/riak/transports/ttb/codec.py new file mode 100644 index 00000000..bdeb7513 --- /dev/null +++ b/riak/transports/ttb/codec.py @@ -0,0 +1,46 @@ +import erlastic +import datetime +import logging + +from riak import RiakError +from riak.content import RiakContent +from riak.util import decode_index_value, str_to_bytes, bytes_to_str, \ + unix_time_millis, datetime_from_unix_time_millis +from six import string_types, PY2 + + +class RiakTtbCodec(object): + ''' + Erlang term-to-binary Encoding and decoding methods for RiakTtbTransport + ''' + + def __init__(self, **unused_args): + super(RiakTtbCodec, self).__init__(**unused_args) + + 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, 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): + key_vals = None + if isinstance(key, list): + key_vals = key + else: + raise ValueError("key must be a list") + return None diff --git a/setup.py b/setup.py index 5219948d..b2c9d39a 100755 --- a/setup.py +++ b/setup.py @@ -7,8 +7,8 @@ setup_security, enable_security, disable_security, setup_timeseries, \ build_messages -install_requires = ['six >= 1.8.0'] -requires = ['six(>=1.8.0)'] +install_requires = ['six >= 1.8.0', 'erlastic >= 2.0.0'] +requires = ['six(>=1.8.0)', 'erlastic(>= 2.0.0)'] if platform.python_version() < '2.7.9': install_requires.append("pyOpenSSL >= 0.14") requires.append("pyOpenSSL(>=0.14)") From 2dd27d6d800de0022fb765b4914483e701f78d72 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Mon, 1 Feb 2016 17:56:14 -0800 Subject: [PATCH 06/44] Unit test for tsgetreq ttb encoding complete --- riak/tests/test_timeseries_ttb.py | 18 +++++++++++--- riak/transports/ttb/codec.py | 40 +++++++++++++++++++++---------- 2 files changed, 43 insertions(+), 15 deletions(-) diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index 914758b2..dfeabcb8 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -4,15 +4,20 @@ import random import string +from erlastic import decode, encode +from erlastic.types import Atom + from riak.table import Table from riak.transports.ttb.codec import RiakTtbCodec -from riak.util import str_to_bytes, bytes_to_str, unix_time_millis +from riak.util import str_to_bytes, unix_time_millis if platform.python_version() < '2.7': unittest = __import__('unittest2') else: import unittest +udef_a = Atom('undefined') +tsc_a = Atom('tscell') table_name = 'GeoCheckin' bd0 = '时间序列' @@ -36,5 +41,12 @@ def setUp(self): self.table = Table(None, table_name) def test_encode_data_for_get(self): - req = self.c._encode_timeseries_keyreq(self.table, self.test_key) - self.assertIsNotNone(req) + keylist = [ + (tsc_a, str_to_bytes('hash1'), udef_a, udef_a, udef_a, udef_a), + (tsc_a, str_to_bytes('user2'), udef_a, udef_a, udef_a, udef_a), + (tsc_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a) + ] + req = Atom('tsgetreq'), str_to_bytes(table_name), keylist, udef_a + req_test = encode(req) + req_encoded = self.c._encode_timeseries_keyreq(self.table, self.test_key) + self.assertEqual(req_test, req_encoded) diff --git a/riak/transports/ttb/codec.py b/riak/transports/ttb/codec.py index bdeb7513..90d82048 100644 --- a/riak/transports/ttb/codec.py +++ b/riak/transports/ttb/codec.py @@ -1,13 +1,20 @@ -import erlastic import datetime import logging -from riak import RiakError -from riak.content import RiakContent -from riak.util import decode_index_value, str_to_bytes, bytes_to_str, \ +from erlastic import decode, encode +from erlastic.types import Atom + +from riak.util import str_to_bytes, bytes_to_str, \ unix_time_millis, datetime_from_unix_time_millis from six import string_types, PY2 +udef_a = Atom('undefined') + +tsgetreq_a = Atom('tsgetreq') +tsputreq_a = Atom('tsputreq') +tscell_a = Atom('tscell') + +tscell_empty = (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a) class RiakTtbCodec(object): ''' @@ -17,21 +24,28 @@ class RiakTtbCodec(object): def __init__(self, **unused_args): super(RiakTtbCodec, self).__init__(**unused_args) - def _encode_to_ts_cell(self, cell, ts_cell): - if cell is not None: + def _encode_to_ts_cell(self, cell): + if cell is None: + return tscell_empty + else: if isinstance(cell, datetime.datetime): - ts_cell.timestamp_value = unix_time_millis(cell) + ts = unix_time_millis(cell) + logging.debug("cell -> timestamp: '%s'", ts) + return (tscell_a, udef_a, udef_a, ts, udef_a, udef_a) elif isinstance(cell, bool): - ts_cell.boolean_value = cell + logging.debug("cell -> bool: '%s'", cell) + return (tscell_a, udef_a, udef_a, udef_a, cell, udef_a) elif isinstance(cell, string_types): logging.debug("cell -> str: '%s'", cell) - ts_cell.varchar_value = str_to_bytes(cell) + return (tscell_a, str_to_bytes(cell), + udef_a, udef_a, udef_a, udef_a) elif (isinstance(cell, int) or (PY2 and isinstance(cell, long))): # noqa logging.debug("cell -> int/long: '%s'", cell) - ts_cell.sint64_value = cell + return (tscell_a, udef_a, cell, udef_a, udef_a, udef_a) elif isinstance(cell, float): - ts_cell.double_value = cell + logging.debug("cell -> float: '%s'", cell) + return (tscell_a, udef_a, udef_a, udef_a, udef_a, cell) else: t = type(cell) raise RiakError("can't serialize type '{}', value '{}'" @@ -43,4 +57,6 @@ def _encode_timeseries_keyreq(self, table, key): key_vals = key else: raise ValueError("key must be a list") - return None + req = tsgetreq_a, str_to_bytes(table.name), \ + [self._encode_to_ts_cell(k) for k in key_vals], udef_a + return encode(req) From b251fa9bddbbd90d20e2e05d9b805e4afbd6a84c Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Mon, 1 Feb 2016 18:35:42 -0800 Subject: [PATCH 07/44] Unit test for tsputreq ttb encoding complete --- riak/tests/test_timeseries_ttb.py | 31 +++++++++++++++++++++++++++++-- riak/transports/ttb/codec.py | 25 +++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index dfeabcb8..46a5797a 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -8,6 +8,7 @@ from erlastic.types import Atom from riak.table import Table +from riak.ts_object import TsObject from riak.transports.ttb.codec import RiakTtbCodec from riak.util import str_to_bytes, unix_time_millis @@ -34,8 +35,8 @@ def setUp(self): self.ts0ms = unix_time_millis(ts0) self.ts1ms = unix_time_millis(ts1) self.rows = [ - [bd0, 0, 1.2, ts0, True], - [bd1, 3, 4.5, ts1, False] + [bd0, 0, 1.2, ts0, True, None], + [bd1, 3, 4.5, ts1, False, None] ] self.test_key = ['hash1', 'user2', ts0] self.table = Table(None, table_name) @@ -48,5 +49,31 @@ def test_encode_data_for_get(self): ] req = Atom('tsgetreq'), str_to_bytes(table_name), keylist, udef_a req_test = encode(req) + req_encoded = self.c._encode_timeseries_keyreq(self.table, self.test_key) self.assertEqual(req_test, req_encoded) + + def test_encode_data_for_put(self): + r0 = [ + (tsc_a, bd0, udef_a, udef_a, udef_a, udef_a), + (tsc_a, udef_a, 0, udef_a, udef_a, udef_a), + (tsc_a, udef_a, udef_a, udef_a, udef_a, 1.2), + (tsc_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a), + (tsc_a, udef_a, udef_a, udef_a, True, udef_a), + (tsc_a, udef_a, udef_a, udef_a, udef_a, udef_a) + ] + r1 = [ + (tsc_a, bd1, udef_a, udef_a, udef_a, udef_a), + (tsc_a, udef_a, 3, udef_a, udef_a, udef_a), + (tsc_a, udef_a, udef_a, udef_a, udef_a, 4.5), + (tsc_a, udef_a, udef_a, unix_time_millis(ts1), udef_a, udef_a), + (tsc_a, udef_a, udef_a, udef_a, False, udef_a), + (tsc_a, udef_a, udef_a, udef_a, udef_a, udef_a) + ] + rows = [r0, r1] + req = Atom('tsputreq'), str_to_bytes(table_name), udef_a, rows + req_test = encode(req) + + tsobj = TsObject(None, self.table, self.rows, None) + req_encoded = self.c._encode_timeseries_put(tsobj) + self.assertEqual(req_test, req_encoded) diff --git a/riak/transports/ttb/codec.py b/riak/transports/ttb/codec.py index 90d82048..9eb5240d 100644 --- a/riak/transports/ttb/codec.py +++ b/riak/transports/ttb/codec.py @@ -60,3 +60,28 @@ def _encode_timeseries_keyreq(self, table, key): req = tsgetreq_a, str_to_bytes(table.name), \ [self._encode_to_ts_cell(k) for k in key_vals], udef_a return encode(req) + + 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 implemented yet") + + 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(req_r) + req = tsputreq_a, str_to_bytes(tsobj.table.name), \ + udef_a, req_rows + return encode(req) + else: + raise RiakError("TsObject requires a list of rows") From aabd6009d135de9a2f402d32d17747dda66008e2 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Mon, 22 Feb 2016 10:59:06 -0800 Subject: [PATCH 08/44] Merge fixes. --- riak/benchmarks/timeseries.py | 14 ++++-- riak/tests/test_timeseries.py | 16 +++--- riak/tests/test_timeseries_ttb.py | 56 +++++++++++++++++++-- riak/transports/pbc/codec.py | 8 +-- riak/transports/pbc/connection.py | 83 +++++++++++++++++++++---------- riak/transports/pbc/transport.py | 38 ++++++++++---- riak/transports/ttb/codec.py | 70 +++++++++++++++++++++++++- riak/util.py | 10 +--- 8 files changed, 232 insertions(+), 63 deletions(-) diff --git a/riak/benchmarks/timeseries.py b/riak/benchmarks/timeseries.py index 57ed6f5c..1ddbac83 100644 --- a/riak/benchmarks/timeseries.py +++ b/riak/benchmarks/timeseries.py @@ -1,8 +1,16 @@ -from multiprocessing import cpu_count -from riak import RiakClient -import riak.benchmark as benchmark import datetime +import logging 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)) epoch = datetime.datetime.utcfromtimestamp(0) onesec = datetime.timedelta(0, 1) diff --git a/riak/tests/test_timeseries.py b/riak/tests/test_timeseries.py index c1d44a7d..cad20ad3 100644 --- a/riak/tests/test_timeseries.py +++ b/riak/tests/test_timeseries.py @@ -10,7 +10,8 @@ from riak.table import Table 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 +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 @@ -39,14 +40,13 @@ class TimeseriesUnitTests(unittest.TestCase): @classmethod def setUpClass(cls): - self.ts0ms = unix_time_millis(ts0) - self.ts1ms = unix_time_millis(ts1) - cls.ts0ms = cls.c._unix_time_millis(ts0) + cls.c = RiakPbcCodec() + 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 +56,7 @@ def setUpClass(cls): [bd1, 3, 4.5, ts1, False] ] cls.test_key = ['hash1', 'user2', ts0] - self.table = Table(None, table_name) + cls.table = Table(None, table_name) def validate_keyreq(self, req): self.assertEqual(self.table.name, bytes_to_str(req.table)) @@ -66,9 +66,9 @@ 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): diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index 46a5797a..16e2c5fb 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -10,13 +10,19 @@ from riak.table import Table from riak.ts_object import TsObject from riak.transports.ttb.codec import RiakTtbCodec -from riak.util import str_to_bytes, unix_time_millis +from riak.util import str_to_bytes, \ + unix_time_millis, datetime_from_unix_time_millis if platform.python_version() < '2.7': unittest = __import__('unittest2') else: import unittest +rpberrorresp_a = Atom('rpberrorresp') +tsgetreq_a = Atom('tsgetreq') +tsgetresp_a = Atom('tsgetresp') +tsputreq_a = Atom('tsputreq') + udef_a = Atom('undefined') tsc_a = Atom('tscell') table_name = 'GeoCheckin' @@ -47,12 +53,54 @@ def test_encode_data_for_get(self): (tsc_a, str_to_bytes('user2'), udef_a, udef_a, udef_a, udef_a), (tsc_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a) ] - req = Atom('tsgetreq'), str_to_bytes(table_name), keylist, udef_a + req = tsgetreq_a, str_to_bytes(table_name), keylist, udef_a req_test = encode(req) - req_encoded = self.c._encode_timeseries_keyreq(self.table, self.test_key) + req_encoded = self.c._encode_timeseries_keyreq_ttb(self.table, self.test_key) self.assertEqual(req_test, req_encoded) + # def test_decode_riak_error(self): + + def test_decode_data_from_get(self): + cols = [] + r0 = [ + (tsc_a, bd0, udef_a, udef_a, udef_a, udef_a), + (tsc_a, udef_a, 0, udef_a, udef_a, udef_a), + (tsc_a, udef_a, udef_a, udef_a, udef_a, 1.2), + (tsc_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a), + (tsc_a, udef_a, udef_a, udef_a, True, udef_a), + (tsc_a, udef_a, udef_a, udef_a, udef_a, udef_a) + ] + r1 = [ + (tsc_a, bd1, udef_a, udef_a, udef_a, udef_a), + (tsc_a, udef_a, 3, udef_a, udef_a, udef_a), + (tsc_a, udef_a, udef_a, udef_a, udef_a, 4.5), + (tsc_a, udef_a, udef_a, unix_time_millis(ts1), udef_a, udef_a), + (tsc_a, udef_a, udef_a, udef_a, False, udef_a), + (tsc_a, udef_a, udef_a, udef_a, udef_a, udef_a) + ] + rows = [r0, r1] + # { tsgetresp, [cols], [rows] } + rsp_data = tsgetresp_a, cols, rows # NB: Python tuple notation + rsp_ttb = encode(rsp_data) + + tsobj = TsObject(None, self.table, [], []) + self.c._decode_timeseries_ttb(rsp_ttb, tsobj) + + for i in range(0, 1): + dr = rows[i] + r = tsobj.rows[i] + self.assertEqual(r[0], dr[0][1]) + self.assertEqual(r[1], dr[1][2]) + self.assertEqual(r[2], dr[2][5]) + self.assertEqual(r[3], + datetime_from_unix_time_millis(dr[3][3])) + if i == 0: + self.assertEqual(r[4], True) + else: + self.assertEqual(r[4], False) + self.assertEqual(r[5], None) + def test_encode_data_for_put(self): r0 = [ (tsc_a, bd0, udef_a, udef_a, udef_a, udef_a), @@ -71,7 +119,7 @@ def test_encode_data_for_put(self): (tsc_a, udef_a, udef_a, udef_a, udef_a, udef_a) ] rows = [r0, r1] - req = Atom('tsputreq'), str_to_bytes(table_name), udef_a, rows + req = tsputreq_a, str_to_bytes(table_name), udef_a, rows req_test = encode(req) tsobj = TsObject(None, self.table, self.rows, None) diff --git a/riak/transports/pbc/codec.py b/riak/transports/pbc/codec.py index 9321356a..e44ab42b 100644 --- a/riak/transports/pbc/codec.py +++ b/riak/transports/pbc/codec.py @@ -633,11 +633,11 @@ def _encode_to_ts_cell(self, cell, ts_cell): elif isinstance(cell, bool): ts_cell.boolean_value = cell elif isinstance(cell, string_types): - logging.debug("cell -> str: '%s'", cell) + # 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) + # logging.debug("cell -> int/long: '%s'", cell) ts_cell.sint64_value = cell elif isinstance(cell, float): ts_cell.double_value = cell @@ -692,10 +692,10 @@ def _encode_timeseries_put(self, tsobj, req): def _decode_timeseries(self, resp, tsobj): """ Fills an TsObject with the appropriate data and - metadata from a TsQueryResp. + metadata from a TsGetResp / TsQueryResp. :param resp: the protobuf message from which to process data - :type resp: riak.pb.TsQueryRsp or riak.pb.riak_ts_pb2.TsGetResp + :type resp: riak.pb.riak_ts_pb2.TsQueryRsp or riak.pb.riak_ts_pb2.TsGetResp :param tsobj: a TsObject :type tsobj: TsObject """ diff --git a/riak/transports/pbc/connection.py b/riak/transports/pbc/connection.py index 60f264c1..c268d4f4 100644 --- a/riak/transports/pbc/connection.py +++ b/riak/transports/pbc/connection.py @@ -1,7 +1,9 @@ +import logging import socket import struct import riak.pb.riak_pb2 import riak.pb.messages +import erlastic from riak.security import SecurityError, USE_STDLIB_SSL from riak import RiakError @@ -21,17 +23,25 @@ class RiakPbcConnection(object): Connection-related methods for RiakPbcTransport. """ - def _encode_msg(self, msg_code, msg=None): + def __init__(self): + self._ttb_enabled = False + + def _encode_msg(self, msg_code, msg=None, is_ttb=False): if msg is None: return struct.pack("!iB", 1, msg_code) - msgstr = msg.SerializeToString() - slen = len(msgstr) - hdr = struct.pack("!iB", 1 + slen, msg_code) - return hdr + msgstr - def _request(self, msg_code, msg=None, expect=None): - self._send_msg(msg_code, msg) - return self._recv_msg(expect) + if is_ttb: + data = msg + else: + data = msg.SerializeToString() + + datalen = len(data) + hdr = struct.pack("!iB", 1 + datalen, msg_code) + return hdr + data + + def _request(self, msg_code, msg=None, expect=None, is_ttb=False): + self._send_msg(msg_code, msg, is_ttb) + return self._recv_msg(expect, is_ttb) def _non_connect_request(self, msg_code, msg=None, expect=None): """ @@ -41,16 +51,18 @@ def _non_connect_request(self, msg_code, msg=None, expect=None): self._non_connect_send_msg(msg_code, msg) return self._recv_msg(expect) - def _non_connect_send_msg(self, msg_code, msg): + def _non_connect_send_msg(self, msg_code, msg, is_ttb=False): """ 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, msg, is_ttb)) - def _send_msg(self, msg_code, msg): + def _send_msg(self, msg_code, msg, is_ttb=False): self._connect() - self._non_connect_send_msg(msg_code, msg) + if is_ttb and not self._enable_ttb(): + raise RiakError('could not switch to TTB encoding!') + self._non_connect_send_msg(msg_code, msg, is_ttb) def _init_security(self): """ @@ -75,6 +87,20 @@ def _starttls(self): else: return False + def _enable_ttb(self): + if self._ttb_enabled: + return True + else: + logging.debug("pbc/connection enabling TTB") + msg_code, _ = self._non_connect_request( + riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_REQ) + if msg_code == riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_RESP: + self._ttb_enabled = True + logging.debug("pbc/connection TTB IS ENABLED") + return True + else: + return False + def _auth(self): """ Perform an authorization request against Riak @@ -154,23 +180,24 @@ def _ssl_handshake(self): # fail if *any* exceptions are thrown during SSL handshake raise SecurityError(e) - def _recv_msg(self, expect=None): + def _recv_msg(self, expect=None, is_ttb=False): 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:]) + err = self._parse_msg(msg_code, self._inbuf[1:], is_ttb) 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:]) + msg = self._parse_msg(msg_code, self._inbuf[1:], is_ttb) 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)) + logging.debug("pbc/connection received msg_code %d msg %s", msg_code, msg) return msg_code, msg def _recv_pkt(self): @@ -218,18 +245,24 @@ 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 + def _parse_msg(self, code, packet, is_ttb=False): + if is_ttb: + if code != riak.pb.messages.MSG_CODE_TS_GET_RESP and \ + code != riak.pb.messages.MSG_CODE_TS_PUT_RESP: + raise RiakError("TTB can't parse code: %d" % code) + return erlastic.decode(packet) + else: + try: + pbclass = riak.pb.messages.MESSAGE_CLASSES[code] + except KeyError: + pbclass = None - if pbclass is None: - return None + if pbclass is None: + return None - pbo = pbclass() - pbo.ParseFromString(packet) - return pbo + pbo = pbclass() + pbo.ParseFromString(packet) + return pbo # These are set in the RiakPbcTransport initializer _address = None diff --git a/riak/transports/pbc/transport.py b/riak/transports/pbc/transport.py index 53df2181..c3e4f749 100644 --- a/riak/transports/pbc/transport.py +++ b/riak/transports/pbc/transport.py @@ -1,3 +1,4 @@ +import logging import riak.pb.messages import riak.pb.riak_pb2 import riak.pb.riak_kv_pb2 @@ -15,10 +16,13 @@ RiakPbcIndexStream, RiakPbcTsKeyStream) from riak.transports.pbc.codec import RiakPbcCodec +from riak.transports.ttb.codec import RiakTtbCodec + from six import PY2, PY3 -class RiakPbcTransport(RiakTransport, RiakPbcConnection, RiakPbcCodec): +class RiakPbcTransport(RiakTransport, RiakPbcConnection, + RiakPbcCodec, RiakTtbCodec): """ The RiakPbcTransport object holds a connection to the protocol buffers interface on the riak server. @@ -28,7 +32,7 @@ def __init__(self, node=None, client=None, timeout=None, - *unused_options): + **transport_options): """ Construct a new RiakPbcTransport object. """ @@ -39,6 +43,7 @@ def __init__(self, self._address = (node.host, node.pb_port) self._timeout = timeout self._socket = None + self._use_ttb = transport_options.get('use_ttb', False) # FeatureDetection API def _server_version(self): @@ -178,24 +183,39 @@ def ts_describe(self, table): 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) + ts_get_resp = None + if self._use_ttb: + encoded = self._encode_timeseries_keyreq_ttb(table, key) + else: + 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) + riak.pb.messages.MSG_CODE_TS_GET_RESP, + self._use_ttb) tsobj = TsObject(self._client, table, [], None) - self._decode_timeseries(ts_get_resp, tsobj) + if self._use_ttb: + self._decode_timeseries_ttb(ts_get_resp, tsobj) + else: + 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) + if self._use_ttb: + req = self._encode_timeseries_put_ttb(tsobj) + else: + req = riak.pb.riak_ts_pb2.TsPutReq() + self._encode_timeseries_put(tsobj, req) + + logging.debug("pbc/transport ts_put _use_ttb: '%s'", + self._use_ttb) msg_code, resp = self._request( riak.pb.messages.MSG_CODE_TS_PUT_REQ, req, - riak.pb.messages.MSG_CODE_TS_PUT_RESP) + riak.pb.messages.MSG_CODE_TS_PUT_RESP, + self._use_ttb) if resp is not None: return True diff --git a/riak/transports/ttb/codec.py b/riak/transports/ttb/codec.py index 9eb5240d..127483c7 100644 --- a/riak/transports/ttb/codec.py +++ b/riak/transports/ttb/codec.py @@ -3,14 +3,17 @@ from erlastic import decode, encode from erlastic.types import Atom +from six import string_types, PY2 +from riak import RiakError from riak.util import str_to_bytes, bytes_to_str, \ unix_time_millis, datetime_from_unix_time_millis -from six import string_types, PY2 udef_a = Atom('undefined') +rpberrorresp_a = Atom('rpberrorresp') tsgetreq_a = Atom('tsgetreq') +tsgetresp_a = Atom('tsgetresp') tsputreq_a = Atom('tsputreq') tscell_a = Atom('tscell') @@ -51,7 +54,7 @@ def _encode_to_ts_cell(self, cell): raise RiakError("can't serialize type '{}', value '{}'" .format(t, cell)) - def _encode_timeseries_keyreq(self, table, key): + def _encode_timeseries_keyreq_ttb(self, table, key): key_vals = None if isinstance(key, list): key_vals = key @@ -85,3 +88,66 @@ def _encode_timeseries_put(self, tsobj): return encode(req) else: raise RiakError("TsObject requires a list of rows") + + def _decode_timeseries_ttb(self, resp_ttb, tsobj): + """ + Fills an TsObject with the appropriate data and + metadata from a TTB-encoded TsGetResp / TsQueryResp. + + :param resp_ttb: the protobuf message from which to process data + :type resp_ttb: TTB-encoded tsqueryrsp or 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) + resp = decode(resp_ttb) + resp_a = resp[0] + if resp_a == tsgetresp_a: + resp_cols = resp[1] + resp_rows = resp[2] + for row_ttb in resp_rows: + tsobj.rows.append( + self._decode_timeseries_row(row_ttb, None)) # TODO cols + # elif resp_a == rpberrorresp_a: + else: + raise RiakError("Unknown TTB response type: {}".format(resp_a)) + + def _decode_timeseries_row(self, tsrow_ttb, tscols=None): + """ + Decodes a TTB-encoded TsRow into a list + + :param tsrow: the TTB-encoded TsRow to decode. + :type tsrow: TTB encoded row + :param tscols: the TTB-encoded TsColumn data to help decode. + :type tscols: list + :rtype list + """ + row = [] + for tsc_ttb in tsrow_ttb: + if tsc_ttb == tscell_empty: + row.append(None) + else: + val = None + if tsc_ttb[0] == tscell_a: + if tsc_ttb[1] != udef_a: + row.append(bytes_to_str(tsc_ttb[1])) + elif tsc_ttb[2] != udef_a: + row.append(tsc_ttb[2]) + elif tsc_ttb[3] != udef_a: + row.append( + datetime_from_unix_time_millis(tsc_ttb[3])) + elif tsc_ttb[4] != udef_a: + row.append(tsc_ttb[4]) + elif tsc_ttb[5] != udef_a: + row.append(tsc_ttb[5]) + else: + row.append(None) + else: + raise RiakError( + "Expected tscell atom, got: {}".format(tsc_ttb)) + return row diff --git a/riak/util.py b/riak/util.py index c54e7c4c..4cbe6c0f 100644 --- a/riak/util.py +++ b/riak/util.py @@ -11,14 +11,8 @@ def unix_time_millis(dt): - try: - return int(dt.total_seconds() * 1000.0) - except AttributeError: - # NB: python 2.6 must use this method - td = dt - epoch - return int(((td.microseconds + - (td.seconds + td.days * 24 * 3600) * 10**6) / - 10**6) * 1000.0) + td = dt - epoch + return int(td.total_seconds() * 1000.0) def datetime_from_unix_time_millis(ut): From c8e7341b653b3809ab5a4aaef0564e0bc826a7b6 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Tue, 23 Feb 2016 20:56:38 -0800 Subject: [PATCH 09/44] Update riak_pb to origin/riak_ts-develop-1.2 --- riak/pb/messages.py | 10 ++ riak/pb/riak_ts_pb2.py | 338 +++++++++++++++++++++++++++++++++++++---- riak_pb | 2 +- 3 files changed, 319 insertions(+), 31 deletions(-) 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_pb b/riak_pb index 7fffa81b..e0986adb 160000 --- a/riak_pb +++ b/riak_pb @@ -1 +1 @@ -Subproject commit 7fffa81b38804c18fffbec8d1677966c37d49d55 +Subproject commit e0986adb3b0b79765b91c04bfee5b16cfedb1165 From 714e7f4185dfaf789dfd2c695bf493f70abeb4b7 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 24 Feb 2016 10:44:25 -0800 Subject: [PATCH 10/44] TTB fix --- riak/benchmarks/timeseries.py | 2 +- riak/transports/ttb/codec.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/riak/benchmarks/timeseries.py b/riak/benchmarks/timeseries.py index 1ddbac83..c222189c 100644 --- a/riak/benchmarks/timeseries.py +++ b/riak/benchmarks/timeseries.py @@ -50,7 +50,7 @@ {'host': h, 'pb_port': 10037}, {'host': h, 'pb_port': 10047} ] -client = RiakClient(nodes=n, protocol='pbc') +client = RiakClient(nodes=n, protocol='pbc', transport_options={'use_ttb': True}) table = client.table(tbl) with benchmark.measure() as b: diff --git a/riak/transports/ttb/codec.py b/riak/transports/ttb/codec.py index 127483c7..ac669863 100644 --- a/riak/transports/ttb/codec.py +++ b/riak/transports/ttb/codec.py @@ -27,7 +27,7 @@ class RiakTtbCodec(object): def __init__(self, **unused_args): super(RiakTtbCodec, self).__init__(**unused_args) - def _encode_to_ts_cell(self, cell): + def _encode_to_ts_cell_ttb(self, cell): if cell is None: return tscell_empty else: @@ -61,10 +61,10 @@ def _encode_timeseries_keyreq_ttb(self, table, key): else: raise ValueError("key must be a list") req = tsgetreq_a, str_to_bytes(table.name), \ - [self._encode_to_ts_cell(k) for k in key_vals], udef_a + [self._encode_to_ts_cell_ttb(k) for k in key_vals], udef_a return encode(req) - def _encode_timeseries_put(self, tsobj): + def _encode_timeseries_put_ttb(self, tsobj): ''' Returns an Erlang-TTB encoded tuple with the appropriate data and metadata from a TsObject. @@ -81,7 +81,7 @@ def _encode_timeseries_put(self, tsobj): for row in tsobj.rows: req_r = [] for cell in row: - req_r.append(self._encode_to_ts_cell(cell)) + req_r.append(self._encode_to_ts_cell_ttb(cell)) req_rows.append(req_r) req = tsputreq_a, str_to_bytes(tsobj.table.name), \ udef_a, req_rows From 4cddfe5304bd1dbf35c48a768c1905a7d08171ce Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Tue, 23 Feb 2016 21:44:40 -0800 Subject: [PATCH 11/44] Fix toggle encoding request --- riak/transports/pbc/connection.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/riak/transports/pbc/connection.py b/riak/transports/pbc/connection.py index c268d4f4..a1718e82 100644 --- a/riak/transports/pbc/connection.py +++ b/riak/transports/pbc/connection.py @@ -92,8 +92,12 @@ def _enable_ttb(self): return True else: logging.debug("pbc/connection enabling TTB") + req = riak.pb.riak_pb2.RpbToggleEncodingReq() + req.use_native = True msg_code, _ = self._non_connect_request( - riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_REQ) + riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_REQ, + req, + riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_RESP) if msg_code == riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_RESP: self._ttb_enabled = True logging.debug("pbc/connection TTB IS ENABLED") From 246ad6ba91c3c33a0953f544d9f5abddc13471b6 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 24 Feb 2016 11:43:44 -0800 Subject: [PATCH 12/44] Add ttb integration test. --- riak/tests/test_timeseries_ttb.py | 43 +++++++++++++++++++++++++++++-- riak/transports/pbc/connection.py | 5 +++- riak/transports/pbc/transport.py | 3 +++ 3 files changed, 48 insertions(+), 3 deletions(-) diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index 16e2c5fb..8e010bec 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -7,11 +7,15 @@ from erlastic import decode, encode from erlastic.types import Atom +from riak.client import RiakClient from riak.table import Table from riak.ts_object import TsObject from riak.transports.ttb.codec import RiakTtbCodec from riak.util import str_to_bytes, \ - unix_time_millis, datetime_from_unix_time_millis + unix_time_millis, datetime_from_unix_time_millis, \ + is_timeseries_supported +from riak.tests import RUN_TIMESERIES +from riak.tests.base import IntegrationTestBase if platform.python_version() < '2.7': unittest = __import__('unittest2') @@ -35,6 +39,7 @@ ts1 = ts0 + fiveMins +@unittest.skipUnless(is_timeseries_supported(), "Timeseries not supported") class TimeseriesTtbUnitTests(unittest.TestCase): def setUp(self): self.c = RiakTtbCodec() @@ -123,5 +128,39 @@ def test_encode_data_for_put(self): req_test = encode(req) tsobj = TsObject(None, self.table, self.rows, None) - req_encoded = self.c._encode_timeseries_put(tsobj) + req_encoded = self.c._encode_timeseries_put_ttb(tsobj) self.assertEqual(req_test, req_encoded) + + +@unittest.skipUnless(is_timeseries_supported() and RUN_TIMESERIES, + 'Timeseries not supported or RUN_TIMESERIES is 0') +class TimeseriesTtbTests(IntegrationTestBase, unittest.TestCase): + @classmethod + def setUpClass(cls): + super(TimeseriesTtbTests, cls).setUpClass() + + def test_store_data_ttb(self): + now = datetime.datetime.utcfromtimestamp(144379690.987000) + fiveMinsAgo = now - fiveMins + tenMinsAgo = fiveMinsAgo - fiveMins + fifteenMinsAgo = tenMinsAgo - fiveMins + twentyMinsAgo = fifteenMinsAgo - fiveMins + twentyFiveMinsAgo = twentyMinsAgo - fiveMins + + client = RiakClient(protocol='pbc', + host='riak-test', + pb_port=10017, + transport_options={'use_ttb': True}) + + table = 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] + ] + ts_obj = table.new(rows) + result = ts_obj.store() + self.assertTrue(result) + client.close() diff --git a/riak/transports/pbc/connection.py b/riak/transports/pbc/connection.py index a1718e82..eec9de23 100644 --- a/riak/transports/pbc/connection.py +++ b/riak/transports/pbc/connection.py @@ -254,7 +254,10 @@ def _parse_msg(self, code, packet, is_ttb=False): if code != riak.pb.messages.MSG_CODE_TS_GET_RESP and \ code != riak.pb.messages.MSG_CODE_TS_PUT_RESP: raise RiakError("TTB can't parse code: %d" % code) - return erlastic.decode(packet) + if len(packet) > 0: + return erlastic.decode(packet) + else: + return None else: try: pbclass = riak.pb.messages.MESSAGE_CLASSES[code] diff --git a/riak/transports/pbc/transport.py b/riak/transports/pbc/transport.py index c3e4f749..595aa7c3 100644 --- a/riak/transports/pbc/transport.py +++ b/riak/transports/pbc/transport.py @@ -217,6 +217,9 @@ def ts_put(self, tsobj): riak.pb.messages.MSG_CODE_TS_PUT_RESP, self._use_ttb) + if self._use_ttb and resp is None: + return True + if resp is not None: return True else: From c89872d2e90c5428c2d1025547bf1ad50cdd9e1c Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 24 Feb 2016 12:44:35 -0800 Subject: [PATCH 13/44] TTB is ALIVE --- riak/benchmark.py | 9 +++++++-- riak/tests/test_timeseries_ttb.py | 2 +- riak/transports/pbc/connection.py | 2 +- riak/transports/pbc/transport.py | 6 +++--- riak/transports/ttb/codec.py | 33 +++++++++++++------------------ 5 files changed, 26 insertions(+), 26 deletions(-) diff --git a/riak/benchmark.py b/riak/benchmark.py index c26a0a49..cfb220c1 100644 --- a/riak/benchmark.py +++ b/riak/benchmark.py @@ -1,6 +1,9 @@ from __future__ import print_function + import os import gc +import sys +import traceback __all__ = ['measure', 'measure_with_rehearsal'] @@ -154,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/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index 8e010bec..6fffd161 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -90,7 +90,7 @@ def test_decode_data_from_get(self): rsp_ttb = encode(rsp_data) tsobj = TsObject(None, self.table, [], []) - self.c._decode_timeseries_ttb(rsp_ttb, tsobj) + self.c._decode_timeseries_ttb(decode(rsp_ttb), tsobj) for i in range(0, 1): dr = rows[i] diff --git a/riak/transports/pbc/connection.py b/riak/transports/pbc/connection.py index eec9de23..f13c75a4 100644 --- a/riak/transports/pbc/connection.py +++ b/riak/transports/pbc/connection.py @@ -201,7 +201,7 @@ def _recv_msg(self, expect=None, is_ttb=False): if expect and msg_code != expect: raise RiakError("unexpected protocol buffer message code: %d, %r" % (msg_code, msg)) - logging.debug("pbc/connection received msg_code %d msg %s", msg_code, msg) + # logging.debug("pbc/connection received msg_code %d msg %s", msg_code, msg) return msg_code, msg def _recv_pkt(self): diff --git a/riak/transports/pbc/transport.py b/riak/transports/pbc/transport.py index 595aa7c3..837508e8 100644 --- a/riak/transports/pbc/transport.py +++ b/riak/transports/pbc/transport.py @@ -185,7 +185,7 @@ def ts_describe(self, table): def ts_get(self, table, key): ts_get_resp = None if self._use_ttb: - encoded = self._encode_timeseries_keyreq_ttb(table, key) + req = self._encode_timeseries_keyreq_ttb(table, key) else: req = riak.pb.riak_ts_pb2.TsGetReq() self._encode_timeseries_keyreq(table, key, req) @@ -209,8 +209,8 @@ def ts_put(self, tsobj): req = riak.pb.riak_ts_pb2.TsPutReq() self._encode_timeseries_put(tsobj, req) - logging.debug("pbc/transport ts_put _use_ttb: '%s'", - self._use_ttb) + # logging.debug("pbc/transport ts_put _use_ttb: '%s'", + # self._use_ttb) msg_code, resp = self._request( riak.pb.messages.MSG_CODE_TS_PUT_REQ, req, diff --git a/riak/transports/ttb/codec.py b/riak/transports/ttb/codec.py index ac669863..e5afa5f2 100644 --- a/riak/transports/ttb/codec.py +++ b/riak/transports/ttb/codec.py @@ -1,5 +1,4 @@ import datetime -import logging from erlastic import decode, encode from erlastic.types import Atom @@ -15,6 +14,7 @@ tsgetreq_a = Atom('tsgetreq') tsgetresp_a = Atom('tsgetresp') tsputreq_a = Atom('tsputreq') +tsrow_a = Atom('tsrow') tscell_a = Atom('tscell') tscell_empty = (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a) @@ -33,21 +33,16 @@ def _encode_to_ts_cell_ttb(self, cell): else: if isinstance(cell, datetime.datetime): ts = unix_time_millis(cell) - logging.debug("cell -> timestamp: '%s'", ts) return (tscell_a, udef_a, udef_a, ts, udef_a, udef_a) elif isinstance(cell, bool): - logging.debug("cell -> bool: '%s'", cell) return (tscell_a, udef_a, udef_a, udef_a, cell, udef_a) elif isinstance(cell, string_types): - logging.debug("cell -> str: '%s'", cell) return (tscell_a, str_to_bytes(cell), udef_a, udef_a, udef_a, udef_a) elif (isinstance(cell, int) or (PY2 and isinstance(cell, long))): # noqa - logging.debug("cell -> int/long: '%s'", cell) return (tscell_a, udef_a, cell, udef_a, udef_a, udef_a) elif isinstance(cell, float): - logging.debug("cell -> float: '%s'", cell) return (tscell_a, udef_a, udef_a, udef_a, udef_a, cell) else: t = type(cell) @@ -94,7 +89,7 @@ def _decode_timeseries_ttb(self, resp_ttb, tsobj): Fills an TsObject with the appropriate data and metadata from a TTB-encoded TsGetResp / TsQueryResp. - :param resp_ttb: the protobuf message from which to process data + :param resp_ttb: the decoded TTB data :type resp_ttb: TTB-encoded tsqueryrsp or tsgetresp :param tsobj: a TsObject :type tsobj: TsObject @@ -105,19 +100,18 @@ def _decode_timeseries_ttb(self, resp_ttb, tsobj): # col_type = col.type # col = (col_name, col_type) # tsobj.columns.append(col) - resp = decode(resp_ttb) - resp_a = resp[0] + resp_a = resp_ttb[0] if resp_a == tsgetresp_a: - resp_cols = resp[1] - resp_rows = resp[2] + resp_cols = resp_ttb[1] + resp_rows = resp_ttb[2] for row_ttb in resp_rows: tsobj.rows.append( - self._decode_timeseries_row(row_ttb, None)) # TODO cols + self._decode_timeseries_row_ttb(row_ttb, None)) # TODO cols # elif resp_a == rpberrorresp_a: else: raise RiakError("Unknown TTB response type: {}".format(resp_a)) - def _decode_timeseries_row(self, tsrow_ttb, tscols=None): + def _decode_timeseries_row_ttb(self, tsrow_ttb, tscols=None): """ Decodes a TTB-encoded TsRow into a list @@ -127,11 +121,9 @@ def _decode_timeseries_row(self, tsrow_ttb, tscols=None): :type tscols: list :rtype list """ - row = [] - for tsc_ttb in tsrow_ttb: - if tsc_ttb == tscell_empty: - row.append(None) - else: + if tsrow_ttb[0] == tsrow_a: + row = [] + for tsc_ttb in tsrow_ttb[1]: val = None if tsc_ttb[0] == tscell_a: if tsc_ttb[1] != udef_a: @@ -149,5 +141,8 @@ def _decode_timeseries_row(self, tsrow_ttb, tscols=None): row.append(None) else: raise RiakError( - "Expected tscell atom, got: {}".format(tsc_ttb)) + "Expected tscell atom, got: {}".format(tsc_ttb[0])) + else: + raise RiakError( + "Expected tsrow atom, got: {}".format(tsrow_ttb[0])) return row From ac1afcea2a4bff1ff043b6b798a30982a13ab438 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 24 Feb 2016 17:18:07 -0800 Subject: [PATCH 14/44] TTB encode rows correctly --- riak/benchmarks/timeseries.py | 4 ++-- riak/tests/test_timeseries_ttb.py | 29 +++++++++++++++++++---------- riak/transports/pbc/transport.py | 3 ++- riak/transports/ttb/codec.py | 3 ++- 4 files changed, 25 insertions(+), 14 deletions(-) diff --git a/riak/benchmarks/timeseries.py b/riak/benchmarks/timeseries.py index c222189c..12734719 100644 --- a/riak/benchmarks/timeseries.py +++ b/riak/benchmarks/timeseries.py @@ -50,7 +50,7 @@ {'host': h, 'pb_port': 10037}, {'host': h, 'pb_port': 10047} ] -client = RiakClient(nodes=n, protocol='pbc', transport_options={'use_ttb': True}) +client = RiakClient(nodes=n, protocol='pbc', transport_options={'use_ttb': False}) table = client.table(tbl) with benchmark.measure() as b: @@ -69,7 +69,7 @@ if ts_obj is None: raise AssertionError("expected obj") if len(ts_obj.rows) != 1: - raise AssertionError("expected one row") + raise AssertionError("expected one row, got: %d" % len(tsobj.rows)) row = ts_obj.rows[0] if len(row) != 5: raise AssertionError("expected row to have five items") diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index 6fffd161..23d2d787 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -28,6 +28,7 @@ tsputreq_a = Atom('tsputreq') udef_a = Atom('undefined') +tsr_a = Atom('tsrow') tsc_a = Atom('tscell') table_name = 'GeoCheckin' @@ -68,22 +69,22 @@ def test_encode_data_for_get(self): def test_decode_data_from_get(self): cols = [] - r0 = [ + r0 = (tsr_a, [ (tsc_a, bd0, udef_a, udef_a, udef_a, udef_a), (tsc_a, udef_a, 0, udef_a, udef_a, udef_a), (tsc_a, udef_a, udef_a, udef_a, udef_a, 1.2), (tsc_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a), (tsc_a, udef_a, udef_a, udef_a, True, udef_a), (tsc_a, udef_a, udef_a, udef_a, udef_a, udef_a) - ] - r1 = [ + ]) + r1 = (tsr_a, [ (tsc_a, bd1, udef_a, udef_a, udef_a, udef_a), (tsc_a, udef_a, 3, udef_a, udef_a, udef_a), (tsc_a, udef_a, udef_a, udef_a, udef_a, 4.5), (tsc_a, udef_a, udef_a, unix_time_millis(ts1), udef_a, udef_a), (tsc_a, udef_a, udef_a, udef_a, False, udef_a), (tsc_a, udef_a, udef_a, udef_a, udef_a, udef_a) - ] + ]) rows = [r0, r1] # { tsgetresp, [cols], [rows] } rsp_data = tsgetresp_a, cols, rows # NB: Python tuple notation @@ -93,7 +94,7 @@ def test_decode_data_from_get(self): self.c._decode_timeseries_ttb(decode(rsp_ttb), tsobj) for i in range(0, 1): - dr = rows[i] + dr = rows[i][1] r = tsobj.rows[i] self.assertEqual(r[0], dr[0][1]) self.assertEqual(r[1], dr[1][2]) @@ -107,22 +108,22 @@ def test_decode_data_from_get(self): self.assertEqual(r[5], None) def test_encode_data_for_put(self): - r0 = [ + r0 = (tsr_a, [ (tsc_a, bd0, udef_a, udef_a, udef_a, udef_a), (tsc_a, udef_a, 0, udef_a, udef_a, udef_a), (tsc_a, udef_a, udef_a, udef_a, udef_a, 1.2), (tsc_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a), (tsc_a, udef_a, udef_a, udef_a, True, udef_a), (tsc_a, udef_a, udef_a, udef_a, udef_a, udef_a) - ] - r1 = [ + ]) + r1 = (tsr_a, [ (tsc_a, bd1, udef_a, udef_a, udef_a, udef_a), (tsc_a, udef_a, 3, udef_a, udef_a, udef_a), (tsc_a, udef_a, udef_a, udef_a, udef_a, 4.5), (tsc_a, udef_a, udef_a, unix_time_millis(ts1), udef_a, udef_a), (tsc_a, udef_a, udef_a, udef_a, False, udef_a), (tsc_a, udef_a, udef_a, udef_a, udef_a, udef_a) - ] + ]) rows = [r0, r1] req = tsputreq_a, str_to_bytes(table_name), udef_a, rows req_test = encode(req) @@ -139,7 +140,7 @@ class TimeseriesTtbTests(IntegrationTestBase, unittest.TestCase): def setUpClass(cls): super(TimeseriesTtbTests, cls).setUpClass() - def test_store_data_ttb(self): + def test_store_and_fetch_ttb(self): now = datetime.datetime.utcfromtimestamp(144379690.987000) fiveMinsAgo = now - fiveMins tenMinsAgo = fiveMinsAgo - fiveMins @@ -163,4 +164,12 @@ def test_store_data_ttb(self): ts_obj = table.new(rows) result = ts_obj.store() self.assertTrue(result) + + for r in rows: + k = r[0:3] + ts_obj = client.ts_get(table_name, k) + self.assertIsNotNone(ts_obj) + self.assertEqual(len(ts_obj.rows), 1) + self.assertEqual(len(ts_obj.rows[0]), 5) + client.close() diff --git a/riak/transports/pbc/transport.py b/riak/transports/pbc/transport.py index 837508e8..2e822adf 100644 --- a/riak/transports/pbc/transport.py +++ b/riak/transports/pbc/transport.py @@ -217,7 +217,8 @@ def ts_put(self, tsobj): riak.pb.messages.MSG_CODE_TS_PUT_RESP, self._use_ttb) - if self._use_ttb and resp is None: + if self._use_ttb and resp is None and \ + msg_code == riak.pb.messages.MSG_CODE_TS_PUT_RESP: return True if resp is not None: diff --git a/riak/transports/ttb/codec.py b/riak/transports/ttb/codec.py index e5afa5f2..fb7a14ef 100644 --- a/riak/transports/ttb/codec.py +++ b/riak/transports/ttb/codec.py @@ -77,7 +77,8 @@ def _encode_timeseries_put_ttb(self, tsobj): req_r = [] for cell in row: req_r.append(self._encode_to_ts_cell_ttb(cell)) - req_rows.append(req_r) + req_t = (tsrow_a, req_r) + req_rows.append(req_t) req = tsputreq_a, str_to_bytes(tsobj.table.name), \ udef_a, req_rows return encode(req) From 2db41577d31a4f9bb9f28b17674d41723014cbc0 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Thu, 25 Feb 2016 12:47:58 -0800 Subject: [PATCH 15/44] Add args to benchmark program. --- riak/benchmarks/timeseries.py | 57 +++++++++++++++++------------------ 1 file changed, 27 insertions(+), 30 deletions(-) diff --git a/riak/benchmarks/timeseries.py b/riak/benchmarks/timeseries.py index 12734719..4bd8a676 100644 --- a/riak/benchmarks/timeseries.py +++ b/riak/benchmarks/timeseries.py @@ -12,22 +12,26 @@ # logger.level = logging.DEBUG # logger.addHandler(logging.StreamHandler(sys.stdout)) -epoch = datetime.datetime.utcfromtimestamp(0) -onesec = datetime.timedelta(0, 1) +# 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 = 32 +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 = [] -keys = [] for i in range(rowcount): ts = datetime.datetime(2016, 1, 1, 12, 0, 0) + \ datetime.timedelta(seconds=i) - family_idx = i % 4 - series_idx = i % 4 + 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)] @@ -35,11 +39,12 @@ row = [family, series, ts, w, temp] key = [family, series, ts] rows.append(row) - keys.append(key) print("Benchmarking timeseries:") -print(" CPUs: {0}".format(cpu_count())) -print(" Rows: {0}".format(len(rows))) +print(" Use TTB: {}".format(use_ttb)) +print("Batch Size: {}".format(batchsz)) +print(" CPUs: {}".format(cpu_count())) +print(" Rows: {}".format(len(rows))) print() tbl = 'GeoCheckin' @@ -48,28 +53,20 @@ {'host': h, 'pb_port': 10017}, {'host': h, 'pb_port': 10027}, {'host': h, 'pb_port': 10037}, - {'host': h, 'pb_port': 10047} + {'host': h, 'pb_port': 10047}, + {'host': h, 'pb_port': 10057} ] -client = RiakClient(nodes=n, protocol='pbc', transport_options={'use_ttb': False}) +client = RiakClient(nodes=n, protocol='pbc', transport_options={'use_ttb': use_ttb}) table = client.table(tbl) with benchmark.measure() as b: - with b.report('populate'): - 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") - with b.report('get'): - for k in keys: - ts_obj = client.ts_get(tbl, k) - if ts_obj is None: - raise AssertionError("expected obj") - if len(ts_obj.rows) != 1: - raise AssertionError("expected one row, got: %d" % len(tsobj.rows)) - row = ts_obj.rows[0] - if len(row) != 5: - raise AssertionError("expected row to have five items") + 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") From c340dd5b8acc31901d0886d689aa2ab451cfa019 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Mon, 21 Mar 2016 15:13:27 -0700 Subject: [PATCH 16/44] bump submodules --- riak_pb | 2 +- tools | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/riak_pb b/riak_pb index e0986adb..341269c1 160000 --- a/riak_pb +++ b/riak_pb @@ -1 +1 @@ -Subproject commit e0986adb3b0b79765b91c04bfee5b16cfedb1165 +Subproject commit 341269c19c75fa0557d5aa5fd5ac1f0dfe18cfae diff --git a/tools b/tools index 5ff5850e..4dae68dd 160000 --- a/tools +++ b/tools @@ -1 +1 @@ -Subproject commit 5ff5850e1d7164f4f64f45a31d9b257e01a19e58 +Subproject commit 4dae68ddca2d405090d64a97c7e99b4607263892 From fd157fbbf4cf501ed64f6ec5a8cac7c5705817bf Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Mon, 28 Mar 2016 14:41:11 -0700 Subject: [PATCH 17/44] Python 3 + TTB --- riak/tests/suite.py | 7 +- riak/tests/test_2i.py | 7 +- riak/tests/test_btypes.py | 8 +- riak/tests/test_client.py | 8 +- riak/tests/test_comparison.py | 8 +- riak/tests/test_datatypes.py | 10 +-- riak/tests/test_feature_detection.py | 8 +- riak/tests/test_filters.py | 8 +- riak/tests/test_kv.py | 11 +-- riak/tests/test_mapreduce.py | 7 +- riak/tests/test_pool.py | 8 +- riak/tests/test_search.py | 9 +- riak/tests/test_security.py | 7 +- riak/tests/test_timeseries.py | 7 +- riak/tests/test_timeseries_ttb.py | 126 ++++++++++++++++----------- riak/tests/test_util.py | 7 +- riak/tests/test_yokozuna.py | 8 +- riak/transports/ttb/codec.py | 21 +++-- setup.py | 11 +-- 19 files changed, 120 insertions(+), 166 deletions(-) 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..f6e6d54a 100644 --- a/riak/tests/test_2i.py +++ b/riak/tests/test_2i.py @@ -1,13 +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): 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..ffcd6ea0 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): 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..abcf4e95 100644 --- a/riak/tests/test_kv.py +++ b/riak/tests/test_kv.py @@ -1,9 +1,9 @@ # -*- coding: utf-8 -*- +import copy import os -import platform -from six import string_types, PY2, PY3 +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 +16,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 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_timeseries.py b/riak/tests/test_timeseries.py index cad20ad3..bace0543 100644 --- a/riak/tests/test_timeseries.py +++ b/riak/tests/test_timeseries.py @@ -1,8 +1,8 @@ # -*- coding: utf-8 -*- import datetime -import platform import random import string +import unittest import riak.pb.riak_ts_pb2 @@ -17,11 +17,6 @@ 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' bd0 = '时间序列' diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index 23d2d787..70ae6d88 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -1,8 +1,9 @@ # -*- coding: utf-8 -*- import datetime -import platform import random +import six import string +import unittest from erlastic import decode, encode from erlastic.types import Atom @@ -17,23 +18,26 @@ from riak.tests import RUN_TIMESERIES from riak.tests.base import IntegrationTestBase -if platform.python_version() < '2.7': - unittest = __import__('unittest2') -else: - import unittest - rpberrorresp_a = Atom('rpberrorresp') tsgetreq_a = Atom('tsgetreq') tsgetresp_a = Atom('tsgetresp') tsputreq_a = Atom('tsputreq') udef_a = Atom('undefined') -tsr_a = Atom('tsrow') -tsc_a = Atom('tscell') +tsrow_a = Atom('tsrow') +tscell_a = Atom('tscell') table_name = 'GeoCheckin' -bd0 = '时间序列' -bd1 = 'временные ряды' +str0 = 'ascii-0' +str1 = 'ascii-1' + +if six.PY2: + # https://docs.python.org/2/library/functions.html#unicode + bd0 = unicode('时间序列', 'utf-8') + bd1 = unicode('временные ряды', 'utf-8') +else: + bd0 = u'时间序列' + bd1 = u'временные ряды' fiveMins = datetime.timedelta(0, 300) ts0 = datetime.datetime(2015, 1, 1, 12, 0, 0) @@ -44,46 +48,42 @@ class TimeseriesTtbUnitTests(unittest.TestCase): def setUp(self): self.c = RiakTtbCodec() - self.ts0ms = unix_time_millis(ts0) - self.ts1ms = unix_time_millis(ts1) - self.rows = [ - [bd0, 0, 1.2, ts0, True, None], - [bd1, 3, 4.5, ts1, False, None] - ] - self.test_key = ['hash1', 'user2', ts0] self.table = Table(None, table_name) def test_encode_data_for_get(self): keylist = [ - (tsc_a, str_to_bytes('hash1'), udef_a, udef_a, udef_a, udef_a), - (tsc_a, str_to_bytes('user2'), udef_a, udef_a, udef_a, udef_a), - (tsc_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a) + (tscell_a, str_to_bytes('hash1'), udef_a, udef_a, udef_a, udef_a), + (tscell_a, str_to_bytes('user2'), udef_a, udef_a, udef_a, udef_a), + (tscell_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a) ] req = tsgetreq_a, str_to_bytes(table_name), keylist, udef_a req_test = encode(req) - req_encoded = self.c._encode_timeseries_keyreq_ttb(self.table, self.test_key) + test_key = ['hash1', 'user2', ts0] + req_encoded = self.c._encode_timeseries_keyreq_ttb(self.table, test_key) self.assertEqual(req_test, req_encoded) # def test_decode_riak_error(self): def test_decode_data_from_get(self): cols = [] - r0 = (tsr_a, [ - (tsc_a, bd0, udef_a, udef_a, udef_a, udef_a), - (tsc_a, udef_a, 0, udef_a, udef_a, udef_a), - (tsc_a, udef_a, udef_a, udef_a, udef_a, 1.2), - (tsc_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a), - (tsc_a, udef_a, udef_a, udef_a, True, udef_a), - (tsc_a, udef_a, udef_a, udef_a, udef_a, udef_a) + r0 = (tsrow_a, [ + (tscell_a, bd0, udef_a, udef_a, udef_a, udef_a), + (tscell_a, udef_a, 0, udef_a, udef_a, udef_a), + (tscell_a, udef_a, udef_a, udef_a, udef_a, 1.2), + (tscell_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a), + (tscell_a, udef_a, udef_a, udef_a, True, udef_a), + (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a), + (tscell_a, str1, udef_a, udef_a, udef_a, udef_a) ]) - r1 = (tsr_a, [ - (tsc_a, bd1, udef_a, udef_a, udef_a, udef_a), - (tsc_a, udef_a, 3, udef_a, udef_a, udef_a), - (tsc_a, udef_a, udef_a, udef_a, udef_a, 4.5), - (tsc_a, udef_a, udef_a, unix_time_millis(ts1), udef_a, udef_a), - (tsc_a, udef_a, udef_a, udef_a, False, udef_a), - (tsc_a, udef_a, udef_a, udef_a, udef_a, udef_a) + r1 = (tsrow_a, [ + (tscell_a, bd1, udef_a, udef_a, udef_a, udef_a), + (tscell_a, udef_a, 3, udef_a, udef_a, udef_a), + (tscell_a, udef_a, udef_a, udef_a, udef_a, 4.5), + (tscell_a, udef_a, udef_a, unix_time_millis(ts1), udef_a, udef_a), + (tscell_a, udef_a, udef_a, udef_a, False, udef_a), + (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a), + (tscell_a, str1, udef_a, udef_a, udef_a, udef_a) ]) rows = [r0, r1] # { tsgetresp, [cols], [rows] } @@ -94,41 +94,63 @@ def test_decode_data_from_get(self): self.c._decode_timeseries_ttb(decode(rsp_ttb), tsobj) for i in range(0, 1): + self.assertEqual(tsrow_a, rows[i][0]) dr = rows[i][1] - r = tsobj.rows[i] - self.assertEqual(r[0], dr[0][1]) + r = tsobj.rows[i] # encoded + + # cells + self.assertEqual(tscell_a, dr[0][0]) + self.assertEqual(r[0], dr[0][1].encode('utf-8')) + + self.assertEqual(tscell_a, dr[1][0]) self.assertEqual(r[1], dr[1][2]) + + self.assertEqual(tscell_a, dr[2][0]) self.assertEqual(r[2], dr[2][5]) + + self.assertEqual(tscell_a, dr[3][0]) self.assertEqual(r[3], datetime_from_unix_time_millis(dr[3][3])) + + self.assertEqual(tscell_a, dr[4][0]) if i == 0: self.assertEqual(r[4], True) else: self.assertEqual(r[4], False) + + self.assertEqual(tscell_a, dr[5][0]) self.assertEqual(r[5], None) + self.assertEqual(tscell_a, dr[6][0]) + self.assertEqual(r[6], dr[6][1].encode('ascii')) + def test_encode_data_for_put(self): - r0 = (tsr_a, [ - (tsc_a, bd0, udef_a, udef_a, udef_a, udef_a), - (tsc_a, udef_a, 0, udef_a, udef_a, udef_a), - (tsc_a, udef_a, udef_a, udef_a, udef_a, 1.2), - (tsc_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a), - (tsc_a, udef_a, udef_a, udef_a, True, udef_a), - (tsc_a, udef_a, udef_a, udef_a, udef_a, udef_a) + r0 = (tsrow_a, [ + (tscell_a, bd0, udef_a, udef_a, udef_a, udef_a), + (tscell_a, udef_a, 0, udef_a, udef_a, udef_a), + (tscell_a, udef_a, udef_a, udef_a, udef_a, 1.2), + (tscell_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a), + (tscell_a, udef_a, udef_a, udef_a, True, udef_a), + (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a) ]) - r1 = (tsr_a, [ - (tsc_a, bd1, udef_a, udef_a, udef_a, udef_a), - (tsc_a, udef_a, 3, udef_a, udef_a, udef_a), - (tsc_a, udef_a, udef_a, udef_a, udef_a, 4.5), - (tsc_a, udef_a, udef_a, unix_time_millis(ts1), udef_a, udef_a), - (tsc_a, udef_a, udef_a, udef_a, False, udef_a), - (tsc_a, udef_a, udef_a, udef_a, udef_a, udef_a) + r1 = (tsrow_a, [ + (tscell_a, bd1, udef_a, udef_a, udef_a, udef_a), + (tscell_a, udef_a, 3, udef_a, udef_a, udef_a), + (tscell_a, udef_a, udef_a, udef_a, udef_a, 4.5), + (tscell_a, udef_a, udef_a, unix_time_millis(ts1), udef_a, udef_a), + (tscell_a, udef_a, udef_a, udef_a, False, udef_a), + (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a) ]) rows = [r0, r1] req = tsputreq_a, str_to_bytes(table_name), udef_a, rows req_test = encode(req) - tsobj = TsObject(None, self.table, self.rows, None) + 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) req_encoded = self.c._encode_timeseries_put_ttb(tsobj) self.assertEqual(req_test, req_encoded) 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/ttb/codec.py b/riak/transports/ttb/codec.py index fb7a14ef..93995902 100644 --- a/riak/transports/ttb/codec.py +++ b/riak/transports/ttb/codec.py @@ -2,11 +2,12 @@ from erlastic import decode, encode from erlastic.types import Atom -from six import string_types, PY2 +from six import text_type, binary_type, \ + string_types, PY2 from riak import RiakError -from riak.util import str_to_bytes, bytes_to_str, \ - unix_time_millis, datetime_from_unix_time_millis +from riak.util import unix_time_millis, \ + datetime_from_unix_time_millis udef_a = Atom('undefined') @@ -36,8 +37,10 @@ def _encode_to_ts_cell_ttb(self, cell): return (tscell_a, udef_a, udef_a, ts, udef_a, udef_a) elif isinstance(cell, bool): return (tscell_a, udef_a, udef_a, udef_a, cell, udef_a) - elif isinstance(cell, string_types): - return (tscell_a, str_to_bytes(cell), + elif isinstance(cell, text_type) or \ + isinstance(cell, binary_type) or \ + isinstance(cell, string_types): + return (tscell_a, cell, udef_a, udef_a, udef_a, udef_a) elif (isinstance(cell, int) or (PY2 and isinstance(cell, long))): # noqa @@ -55,7 +58,7 @@ def _encode_timeseries_keyreq_ttb(self, table, key): key_vals = key else: raise ValueError("key must be a list") - req = tsgetreq_a, str_to_bytes(table.name), \ + req = tsgetreq_a, table.name, \ [self._encode_to_ts_cell_ttb(k) for k in key_vals], udef_a return encode(req) @@ -79,7 +82,7 @@ def _encode_timeseries_put_ttb(self, tsobj): req_r.append(self._encode_to_ts_cell_ttb(cell)) req_t = (tsrow_a, req_r) req_rows.append(req_t) - req = tsputreq_a, str_to_bytes(tsobj.table.name), \ + req = tsputreq_a, tsobj.table.name, \ udef_a, req_rows return encode(req) else: @@ -95,6 +98,7 @@ def _decode_timeseries_ttb(self, resp_ttb, tsobj): :param tsobj: a TsObject :type tsobj: TsObject """ + # TODO TODO RTS-842 CLIENTS-814 GH-445 # if tsobj.columns is not None: # for col in resp.columns: # col_name = bytes_to_str(col.name) @@ -125,10 +129,9 @@ def _decode_timeseries_row_ttb(self, tsrow_ttb, tscols=None): if tsrow_ttb[0] == tsrow_a: row = [] for tsc_ttb in tsrow_ttb[1]: - val = None if tsc_ttb[0] == tscell_a: if tsc_ttb[1] != udef_a: - row.append(bytes_to_str(tsc_ttb[1])) + row.append(tsc_ttb[1]) elif tsc_ttb[2] != udef_a: row.append(tsc_ttb[2]) elif tsc_ttb[3] != udef_a: diff --git a/setup.py b/setup.py index 9fd8f156..e7109118 100755 --- a/setup.py +++ b/setup.py @@ -1,17 +1,19 @@ #!/usr/bin/env python import platform +import six + from setuptools import setup, find_packages from version import get_version from commands import setup_timeseries, build_messages -install_requires = ['six >= 1.8.0', 'erlastic >= 2.0.0'] +install_requires = ['six >= 1.8.0', 'erlastic >= 2.1.0'] requires = ['six(>=1.8.0)', 'erlastic(>= 2.0.0)'] if platform.python_version() < '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: @@ -19,17 +21,12 @@ requires.append('python3_protobuf(>=2.4.1, <2.6.0)') -tests_require = [] -if platform.python_version() < '2.7.0': - tests_require.append("unittest2") - setup( name='riak', version=get_version(), packages=find_packages(), requires=requires, install_requires=install_requires, - tests_require=tests_require, package_data={'riak': ['erl_src/*']}, description='Python client for Riak', zip_safe=True, From f9176969894eaa872a5f5377809e82f9b5fe2a03 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Mon, 28 Mar 2016 14:51:32 -0700 Subject: [PATCH 18/44] Fix setup version check, TTB for Python 2 --- riak/transports/ttb/__init__.py | 0 setup.py | 3 ++- 2 files changed, 2 insertions(+), 1 deletion(-) create mode 100644 riak/transports/ttb/__init__.py diff --git a/riak/transports/ttb/__init__.py b/riak/transports/ttb/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/setup.py b/setup.py index e7109118..f20bf524 100755 --- a/setup.py +++ b/setup.py @@ -9,7 +9,8 @@ install_requires = ['six >= 1.8.0', 'erlastic >= 2.1.0'] requires = ['six(>=1.8.0)', 'erlastic(>= 2.0.0)'] -if platform.python_version() < '2.7.9': + +if platform.python_version_tuple() <= (2, 7, 9): install_requires.append("pyOpenSSL >= 0.14") requires.append("pyOpenSSL(>=0.14)") From 77f09de7fa9da117f48f41ebd74f56c2c8a46027 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Mon, 28 Mar 2016 15:21:03 -0700 Subject: [PATCH 19/44] make lint happy --- riak/benchmarks/timeseries.py | 10 ++++++---- riak/tests/test_2i.py | 1 - riak/tests/test_timeseries_ttb.py | 20 ++++++++++---------- riak/transports/pbc/codec.py | 4 ++-- riak/transports/pbc/connection.py | 3 ++- riak/transports/pbc/transport.py | 8 ++++---- riak/transports/ttb/codec.py | 17 ++++++++++------- 7 files changed, 34 insertions(+), 29 deletions(-) diff --git a/riak/benchmarks/timeseries.py b/riak/benchmarks/timeseries.py index 4bd8a676..84960962 100644 --- a/riak/benchmarks/timeseries.py +++ b/riak/benchmarks/timeseries.py @@ -1,5 +1,4 @@ import datetime -import logging import random import sys @@ -14,13 +13,15 @@ # 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') + 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' +use_ttb = sys.argv[2].lower() == 'true' epoch = datetime.datetime.utcfromtimestamp(0) onesec = datetime.timedelta(0, 1) @@ -56,7 +57,8 @@ {'host': h, 'pb_port': 10047}, {'host': h, 'pb_port': 10057} ] -client = RiakClient(nodes=n, protocol='pbc', transport_options={'use_ttb': use_ttb}) +client = RiakClient(nodes=n, protocol='pbc', + transport_options={'use_ttb': use_ttb}) table = client.table(tbl) with benchmark.measure() as b: diff --git a/riak/tests/test_2i.py b/riak/tests/test_2i.py index f6e6d54a..6db10602 100644 --- a/riak/tests/test_2i.py +++ b/riak/tests/test_2i.py @@ -6,7 +6,6 @@ from riak.tests.base import IntegrationTestBase - class TwoITests(IntegrationTestBase, unittest.TestCase): def is_2i_supported(self): # Immediate test to see if 2i is even supported w/ the backend diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index 70ae6d88..47c5c39a 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -1,8 +1,6 @@ # -*- coding: utf-8 -*- import datetime -import random import six -import string import unittest from erlastic import decode, encode @@ -60,7 +58,8 @@ def test_encode_data_for_get(self): req_test = encode(req) test_key = ['hash1', 'user2', ts0] - req_encoded = self.c._encode_timeseries_keyreq_ttb(self.table, test_key) + req_encoded = self.c._encode_timeseries_keyreq_ttb( + self.table, test_key) self.assertEqual(req_test, req_encoded) # def test_decode_riak_error(self): @@ -87,7 +86,7 @@ def test_decode_data_from_get(self): ]) rows = [r0, r1] # { tsgetresp, [cols], [rows] } - rsp_data = tsgetresp_a, cols, rows # NB: Python tuple notation + rsp_data = tsgetresp_a, cols, rows # NB: Python tuple notation rsp_ttb = encode(rsp_data) tsobj = TsObject(None, self.table, [], []) @@ -96,7 +95,7 @@ def test_decode_data_from_get(self): for i in range(0, 1): self.assertEqual(tsrow_a, rows[i][0]) dr = rows[i][1] - r = tsobj.rows[i] # encoded + r = tsobj.rows[i] # encoded # cells self.assertEqual(tscell_a, dr[0][0]) @@ -109,8 +108,8 @@ def test_decode_data_from_get(self): self.assertEqual(r[2], dr[2][5]) self.assertEqual(tscell_a, dr[3][0]) - self.assertEqual(r[3], - datetime_from_unix_time_millis(dr[3][3])) + dt = datetime_from_unix_time_millis(dr[3][3]) + self.assertEqual(r[3], dt) self.assertEqual(tscell_a, dr[4][0]) if i == 0: @@ -170,10 +169,11 @@ def test_store_and_fetch_ttb(self): twentyMinsAgo = fifteenMinsAgo - fiveMins twentyFiveMinsAgo = twentyMinsAgo - fiveMins + opts = {'use_ttb': True} client = RiakClient(protocol='pbc', - host='riak-test', - pb_port=10017, - transport_options={'use_ttb': True}) + host='riak-test', + pb_port=10017, + transport_options=opts) table = client.table(table_name) rows = [ diff --git a/riak/transports/pbc/codec.py b/riak/transports/pbc/codec.py index e44ab42b..47932512 100644 --- a/riak/transports/pbc/codec.py +++ b/riak/transports/pbc/codec.py @@ -1,5 +1,4 @@ import datetime -import logging import riak.pb import riak.pb.riak_pb2 import riak.pb.riak_dt_pb2 @@ -695,7 +694,8 @@ def _decode_timeseries(self, resp, tsobj): 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 + :type resp: riak.pb.riak_ts_pb2.TsQueryRsp or + riak.pb.riak_ts_pb2.TsGetResp :param tsobj: a TsObject :type tsobj: TsObject """ diff --git a/riak/transports/pbc/connection.py b/riak/transports/pbc/connection.py index f13c75a4..dfaee06b 100644 --- a/riak/transports/pbc/connection.py +++ b/riak/transports/pbc/connection.py @@ -201,7 +201,8 @@ def _recv_msg(self, expect=None, is_ttb=False): if expect and msg_code != expect: raise RiakError("unexpected protocol buffer message code: %d, %r" % (msg_code, msg)) - # logging.debug("pbc/connection received msg_code %d msg %s", msg_code, msg) + # logging.debug("pbc/connection received msg_code %d msg %s", + # msg_code, msg) return msg_code, msg def _recv_pkt(self): diff --git a/riak/transports/pbc/transport.py b/riak/transports/pbc/transport.py index 2e822adf..c7c4c387 100644 --- a/riak/transports/pbc/transport.py +++ b/riak/transports/pbc/transport.py @@ -1,4 +1,3 @@ -import logging import riak.pb.messages import riak.pb.riak_pb2 import riak.pb.riak_kv_pb2 @@ -22,7 +21,7 @@ class RiakPbcTransport(RiakTransport, RiakPbcConnection, - RiakPbcCodec, RiakTtbCodec): + RiakPbcCodec, RiakTtbCodec): """ The RiakPbcTransport object holds a connection to the protocol buffers interface on the riak server. @@ -217,8 +216,9 @@ def ts_put(self, tsobj): riak.pb.messages.MSG_CODE_TS_PUT_RESP, self._use_ttb) - if self._use_ttb and resp is None and \ - msg_code == riak.pb.messages.MSG_CODE_TS_PUT_RESP: + if self._use_ttb and \ + resp is None and \ + msg_code == riak.pb.messages.MSG_CODE_TS_PUT_RESP: return True if resp is not None: diff --git a/riak/transports/ttb/codec.py b/riak/transports/ttb/codec.py index 93995902..e0e7afdd 100644 --- a/riak/transports/ttb/codec.py +++ b/riak/transports/ttb/codec.py @@ -1,6 +1,6 @@ import datetime -from erlastic import decode, encode +from erlastic import encode from erlastic.types import Atom from six import text_type, binary_type, \ string_types, PY2 @@ -20,6 +20,7 @@ tscell_empty = (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a) + class RiakTtbCodec(object): ''' Erlang term-to-binary Encoding and decoding methods for RiakTtbTransport @@ -38,8 +39,8 @@ def _encode_to_ts_cell_ttb(self, cell): elif isinstance(cell, bool): return (tscell_a, udef_a, udef_a, udef_a, cell, udef_a) elif isinstance(cell, text_type) or \ - isinstance(cell, binary_type) or \ - isinstance(cell, string_types): + isinstance(cell, binary_type) or \ + isinstance(cell, string_types): return (tscell_a, cell, udef_a, udef_a, udef_a, udef_a) elif (isinstance(cell, int) or @@ -82,8 +83,7 @@ def _encode_timeseries_put_ttb(self, tsobj): req_r.append(self._encode_to_ts_cell_ttb(cell)) req_t = (tsrow_a, req_r) req_rows.append(req_t) - req = tsputreq_a, tsobj.table.name, \ - udef_a, req_rows + req = tsputreq_a, tsobj.table.name, udef_a, req_rows return encode(req) else: raise RiakError("TsObject requires a list of rows") @@ -99,6 +99,8 @@ def _decode_timeseries_ttb(self, resp_ttb, tsobj): :type tsobj: TsObject """ # TODO TODO RTS-842 CLIENTS-814 GH-445 + # TODO COLUMNS + # TODO TODO RTS-842 CLIENTS-814 GH-445 # if tsobj.columns is not None: # for col in resp.columns: # col_name = bytes_to_str(col.name) @@ -107,11 +109,12 @@ def _decode_timeseries_ttb(self, resp_ttb, tsobj): # tsobj.columns.append(col) resp_a = resp_ttb[0] if resp_a == tsgetresp_a: - resp_cols = resp_ttb[1] + # TODO resp_cols = resp_ttb[1] resp_rows = resp_ttb[2] for row_ttb in resp_rows: tsobj.rows.append( - self._decode_timeseries_row_ttb(row_ttb, None)) # TODO cols + self._decode_timeseries_row_ttb(row_ttb, None)) + # TODO # elif resp_a == rpberrorresp_a: else: raise RiakError("Unknown TTB response type: {}".format(resp_a)) From 80aac21258af45ecb5405254f55a419d6ff6922a Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Mon, 28 Mar 2016 15:52:13 -0700 Subject: [PATCH 20/44] Add basho-erlastic as a dependency --- setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index f20bf524..bdd0ff1d 100755 --- a/setup.py +++ b/setup.py @@ -7,8 +7,8 @@ from version import get_version from commands import setup_timeseries, build_messages -install_requires = ['six >= 1.8.0', 'erlastic >= 2.1.0'] -requires = ['six(>=1.8.0)', 'erlastic(>= 2.0.0)'] +install_requires = ['six >= 1.8.0', 'basho_erlastic >= 2.1.0'] +requires = ['six(>=1.8.0)', 'basho_erlastic(>= 2.1.0)'] if platform.python_version_tuple() <= (2, 7, 9): install_requires.append("pyOpenSSL >= 0.14") From 5cf97fa97179dba046febe9ddbad453e132cea1a Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Mon, 28 Mar 2016 16:10:31 -0700 Subject: [PATCH 21/44] Use Python version comparison that uses ints instead of strings --- setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index bdd0ff1d..a408c1b7 100755 --- a/setup.py +++ b/setup.py @@ -1,7 +1,7 @@ #!/usr/bin/env python -import platform import six +import sys from setuptools import setup, find_packages from version import get_version @@ -10,7 +10,7 @@ install_requires = ['six >= 1.8.0', 'basho_erlastic >= 2.1.0'] requires = ['six(>=1.8.0)', 'basho_erlastic(>= 2.1.0)'] -if platform.python_version_tuple() <= (2, 7, 9): +if sys.version_info[0:3] <= (2, 7, 9): install_requires.append("pyOpenSSL >= 0.14") requires.append("pyOpenSSL(>=0.14)") From 432513040c293b2154ac5197967516ecfc65dbad Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Tue, 29 Mar 2016 10:03:00 -0700 Subject: [PATCH 22/44] Improve README generation --- .gitignore | 3 +++ MANIFEST.in | 1 + setup.py | 6 +++++- 3 files changed, 9 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 81fd6c28..68a831bf 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ +README.rst + *.pyc .python-version __pycache__/ @@ -11,6 +13,7 @@ docs/_build .*.swp .coverage +riak-*/ py-build/ dist/ 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/setup.py b/setup.py index 65dd9995..6338efd8 100755 --- a/setup.py +++ b/setup.py @@ -1,5 +1,6 @@ #!/usr/bin/env python +import codecs import six import sys @@ -24,8 +25,11 @@ 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', From 85158938c832372e63d367773b43170a940645d0 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Tue, 29 Mar 2016 10:49:13 -0700 Subject: [PATCH 23/44] Move transports/pbc dir to transports/tcp --- docs/advanced.rst | 4 ++-- riak/client/__init__.py | 2 +- riak/client/transport.py | 19 +---------------- riak/tests/test_timeseries.py | 2 +- riak/transports/{pbc => tcp}/__init__.py | 24 ++-------------------- riak/transports/{pbc => tcp}/codec.py | 0 riak/transports/{pbc => tcp}/connection.py | 0 riak/transports/{pbc => tcp}/stream.py | 2 +- riak/transports/{pbc => tcp}/transport.py | 6 +++--- 9 files changed, 11 insertions(+), 48 deletions(-) rename riak/transports/{pbc => tcp}/__init__.py (63%) rename riak/transports/{pbc => tcp}/codec.py (100%) rename riak/transports/{pbc => tcp}/connection.py (100%) rename riak/transports/{pbc => tcp}/stream.py (98%) rename riak/transports/{pbc => tcp}/transport.py (99%) diff --git a/docs/advanced.rst b/docs/advanced.rst index 523b465d..b9e209c9 100644 --- a/docs/advanced.rst +++ b/docs/advanced.rst @@ -132,10 +132,10 @@ HTTP Transport :members: ^^^^^^^^^^^^^^^^^^^^^^^^^^ -Protocol Buffers Transport +TCP Transport ^^^^^^^^^^^^^^^^^^^^^^^^^^ -.. currentmodule:: riak.transports.pbc +.. currentmodule:: riak.transports.tcp .. autoclass:: RiakPbcTransport :members: diff --git a/riak/client/__init__.py b/riak/client/__init__.py index d623f970..548275da 100644 --- a/riak/client/__init__.py +++ b/riak/client/__init__.py @@ -12,7 +12,7 @@ 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.tcp import RiakPbcPool from riak.security import SecurityCreds from riak.util import lazy_property, bytes_to_str, str_to_bytes from six import string_types, PY2 diff --git a/riak/client/transport.py b/riak/client/transport.py index 027951d6..8c3eb92e 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 diff --git a/riak/tests/test_timeseries.py b/riak/tests/test_timeseries.py index bace0543..d2f76a9a 100644 --- a/riak/tests/test_timeseries.py +++ b/riak/tests/test_timeseries.py @@ -9,7 +9,7 @@ from riak import RiakError from riak.table import Table from riak.ts_object import TsObject -from riak.transports.pbc.codec import RiakPbcCodec +from riak.transports.tcp.codec import RiakPbcCodec from riak.util import str_to_bytes, bytes_to_str, \ unix_time_millis, datetime_from_unix_time_millis, \ is_timeseries_supported diff --git a/riak/transports/pbc/__init__.py b/riak/transports/tcp/__init__.py similarity index 63% rename from riak/transports/pbc/__init__.py rename to riak/transports/tcp/__init__.py index fc8914b6..8ccdaa10 100644 --- a/riak/transports/pbc/__init__.py +++ b/riak/transports/tcp/__init__.py @@ -1,28 +1,8 @@ -""" -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 +from riak.transports.tcp.transport import RiakPbcTransport class RiakPbcPool(Pool): diff --git a/riak/transports/pbc/codec.py b/riak/transports/tcp/codec.py similarity index 100% rename from riak/transports/pbc/codec.py rename to riak/transports/tcp/codec.py diff --git a/riak/transports/pbc/connection.py b/riak/transports/tcp/connection.py similarity index 100% rename from riak/transports/pbc/connection.py rename to riak/transports/tcp/connection.py diff --git a/riak/transports/pbc/stream.py b/riak/transports/tcp/stream.py similarity index 98% rename from riak/transports/pbc/stream.py rename to riak/transports/tcp/stream.py index ed649279..38357bbd 100644 --- a/riak/transports/pbc/stream.py +++ b/riak/transports/tcp/stream.py @@ -3,7 +3,7 @@ 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.transports.tcp.codec import RiakPbcCodec from six import PY2 diff --git a/riak/transports/pbc/transport.py b/riak/transports/tcp/transport.py similarity index 99% rename from riak/transports/pbc/transport.py rename to riak/transports/tcp/transport.py index c7c4c387..01eddfa8 100644 --- a/riak/transports/pbc/transport.py +++ b/riak/transports/tcp/transport.py @@ -8,13 +8,13 @@ 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, +from riak.transports.tcp.connection import RiakPbcConnection +from riak.transports.tcp.stream import (RiakPbcKeyStream, RiakPbcMapredStream, RiakPbcBucketStream, RiakPbcIndexStream, RiakPbcTsKeyStream) -from riak.transports.pbc.codec import RiakPbcCodec +from riak.transports.tcp.codec import RiakPbcCodec from riak.transports.ttb.codec import RiakTtbCodec from six import PY2, PY3 From cb1cffd40b4ac27cd1a0bd43feaa08d6860c35b4 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Tue, 29 Mar 2016 11:27:32 -0700 Subject: [PATCH 24/44] Test suite fixes for Windows --- riak/tests/test_kv.py | 3 +++ riak/tests/test_server_test.py | 5 ++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/riak/tests/test_kv.py b/riak/tests/test_kv.py index abcf4e95..5513c603 100644 --- a/riak/tests/test_kv.py +++ b/riak/tests/test_kv.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- import copy import os +import sys import unittest from six import string_types, PY2, PY3 @@ -690,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_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() From ac43d1be947d554dd9590f55014d8408b396711e Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Tue, 29 Mar 2016 13:04:24 -0700 Subject: [PATCH 25/44] Moving code around to separate codecs from transports. Make lint happy. --- docs/advanced.rst | 16 ++++-- riak/client/__init__.py | 14 ++--- riak/client/transport.py | 6 +- riak/{transports/ttb => codecs}/__init__.py | 0 .../http/codec.py => codecs/http.py} | 25 +------- .../tcp/codec.py => codecs/pbuf.py} | 8 +-- .../ttb/codec.py => codecs/ttb.py} | 6 +- riak/node.py | 18 +----- riak/tests/test_client.py | 4 +- riak/tests/test_timeseries.py | 12 ++-- riak/tests/test_timeseries_ttb.py | 13 ++--- riak/transports/feature_detect.py | 2 +- riak/transports/http/__init__.py | 35 ++++-------- riak/transports/http/connection.py | 26 ++------- riak/transports/http/resources.py | 26 ++------- riak/transports/http/stream.py | 41 ++++--------- riak/transports/http/transport.py | 57 +++++++------------ riak/transports/pool.py | 20 +------ riak/transports/tcp/__init__.py | 18 +++--- riak/transports/tcp/connection.py | 12 ++-- riak/transports/tcp/stream.py | 38 ++++++------- riak/transports/tcp/transport.py | 41 +++++++------ riak/transports/transport.py | 22 +------ 23 files changed, 152 insertions(+), 308 deletions(-) rename riak/{transports/ttb => codecs}/__init__.py (100%) rename riak/{transports/http/codec.py => codecs/http.py} (93%) rename riak/{transports/tcp/codec.py => codecs/pbuf.py} (99%) rename riak/{transports/ttb/codec.py => codecs/ttb.py} (97%) diff --git a/docs/advanced.rst b/docs/advanced.rst index b9e209c9..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: -^^^^^^^^^^^^^^^^^^^^^^^^^^ +^^^^^^^^^^^^^ TCP Transport -^^^^^^^^^^^^^^^^^^^^^^^^^^ +^^^^^^^^^^^^^ .. currentmodule:: riak.transports.tcp -.. autoclass:: RiakPbcTransport +.. autoclass:: TcpPool + +.. autofunction:: is_retryable + +.. autoclass:: TcpTransport :members: --------- diff --git a/riak/client/__init__.py b/riak/client/__init__.py index 548275da..3a3caad5 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.tcp 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 @@ -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/transport.py b/riak/client/transport.py index 8c3eb92e..6aca7f24 100644 --- a/riak/client/transport.py +++ b/riak/client/transport.py @@ -32,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): @@ -146,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/transports/ttb/__init__.py b/riak/codecs/__init__.py similarity index 100% rename from riak/transports/ttb/__init__.py rename to riak/codecs/__init__.py 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/transports/tcp/codec.py b/riak/codecs/pbuf.py similarity index 99% rename from riak/transports/tcp/codec.py rename to riak/codecs/pbuf.py index 47932512..2d5231d8 100644 --- a/riak/transports/tcp/codec.py +++ b/riak/codecs/pbuf.py @@ -71,15 +71,15 @@ def _invert(d): } -class RiakPbcCodec(object): +class PbufCodec(object): """ - Protobuffs Encoding and decoding methods for RiakPbcTransport. + Protobuffs Encoding and decoding methods for TcpTransport. """ def __init__(self, **unused_args): if riak.pb is None: - raise NotImplementedError("this transport is not available") - super(RiakPbcCodec, self).__init__(**unused_args) + raise NotImplementedError("this codec is not available") + super(PbufCodec, self).__init__(**unused_args) def _unix_time_millis(self, dt): return unix_time_millis(dt) diff --git a/riak/transports/ttb/codec.py b/riak/codecs/ttb.py similarity index 97% rename from riak/transports/ttb/codec.py rename to riak/codecs/ttb.py index e0e7afdd..2ecb846d 100644 --- a/riak/transports/ttb/codec.py +++ b/riak/codecs/ttb.py @@ -21,13 +21,13 @@ tscell_empty = (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a) -class RiakTtbCodec(object): +class TtbCodec(object): ''' - Erlang term-to-binary Encoding and decoding methods for RiakTtbTransport + Erlang term-to-binary Encoding and decoding methods for TcpTransport ''' def __init__(self, **unused_args): - super(RiakTtbCodec, self).__init__(**unused_args) + super(TtbCodec, self).__init__(**unused_args) def _encode_to_ts_cell_ttb(self, cell): if cell is None: 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/tests/test_client.py b/riak/tests/test_client.py index ffcd6ea0..19379d06 100644 --- a/riak/tests/test_client.py +++ b/riak/tests/test_client.py @@ -197,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_timeseries.py b/riak/tests/test_timeseries.py index d2f76a9a..6f906726 100644 --- a/riak/tests/test_timeseries.py +++ b/riak/tests/test_timeseries.py @@ -5,17 +5,17 @@ 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.tcp.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 table_name = 'GeoCheckin' @@ -35,7 +35,7 @@ class TimeseriesUnitTests(unittest.TestCase): @classmethod def setUpClass(cls): - cls.c = RiakPbcCodec() + cls.c = PbufCodec() cls.ts0ms = unix_time_millis(ts0) if cls.ts0ms != ex0ms: raise AssertionError( @@ -151,7 +151,7 @@ def test_decode_data_from_query(self): r1c4.boolean_value = self.rows[1][4] tsobj = TsObject(None, self.table, [], []) - c = RiakPbcCodec() + c = PbufCodec() c._decode_timeseries(tqr, tsobj) self.assertEqual(len(self.rows), len(tsobj.rows)) diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index 47c5c39a..aae5f73e 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -9,7 +9,7 @@ from riak.client import RiakClient from riak.table import Table from riak.ts_object import TsObject -from riak.transports.ttb.codec import RiakTtbCodec +from riak.codecs.ttb import TtbCodec from riak.util import str_to_bytes, \ unix_time_millis, datetime_from_unix_time_millis, \ is_timeseries_supported @@ -29,13 +29,8 @@ str0 = 'ascii-0' str1 = 'ascii-1' -if six.PY2: - # https://docs.python.org/2/library/functions.html#unicode - bd0 = unicode('时间序列', 'utf-8') - bd1 = unicode('временные ряды', 'utf-8') -else: - bd0 = u'时间序列' - bd1 = u'временные ряды' +bd0 = six.text_type('时间序列') +bd1 = six.text_type('временные ряды') fiveMins = datetime.timedelta(0, 300) ts0 = datetime.datetime(2015, 1, 1, 12, 0, 0) @@ -45,7 +40,7 @@ @unittest.skipUnless(is_timeseries_supported(), "Timeseries not supported") class TimeseriesTtbUnitTests(unittest.TestCase): def setUp(self): - self.c = RiakTtbCodec() + self.c = TtbCodec() self.table = Table(None, table_name) def test_encode_data_for_get(self): 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/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 index 8ccdaa10..312f9194 100644 --- a/riak/transports/tcp/__init__.py +++ b/riak/transports/tcp/__init__.py @@ -2,26 +2,26 @@ import socket from riak.transports.pool import Pool -from riak.transports.tcp.transport import RiakPbcTransport +from riak.transports.tcp.transport import TcpTransport -class RiakPbcPool(Pool): +class TcpPool(Pool): """ - A resource pool of PBC transports. + A resource pool of TCP transports. """ def __init__(self, client, **options): - super(RiakPbcPool, self).__init__() + super(TcpPool, 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) + return TcpTransport(node=node, + client=self._client, + **self._options) - def destroy_resource(self, pbc): - pbc.close() + def destroy_resource(self, tcp): + tcp.close() # These are a specific set of socket errors # that could be raised on send/recv that indicate diff --git a/riak/transports/tcp/connection.py b/riak/transports/tcp/connection.py index dfaee06b..8468a72d 100644 --- a/riak/transports/tcp/connection.py +++ b/riak/transports/tcp/connection.py @@ -18,9 +18,9 @@ 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 __init__(self): @@ -91,7 +91,7 @@ def _enable_ttb(self): if self._ttb_enabled: return True else: - logging.debug("pbc/connection enabling TTB") + logging.debug("tcp/connection enabling TTB") req = riak.pb.riak_pb2.RpbToggleEncodingReq() req.use_native = True msg_code, _ = self._non_connect_request( @@ -100,7 +100,7 @@ def _enable_ttb(self): riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_RESP) if msg_code == riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_RESP: self._ttb_enabled = True - logging.debug("pbc/connection TTB IS ENABLED") + logging.debug("tcp/connection TTB IS ENABLED") return True else: return False @@ -201,7 +201,7 @@ def _recv_msg(self, expect=None, is_ttb=False): if expect and msg_code != expect: raise RiakError("unexpected protocol buffer message code: %d, %r" % (msg_code, msg)) - # logging.debug("pbc/connection received msg_code %d msg %s", + # logging.debug("tcp/connection received msg_code %d msg %s", # msg_code, msg) return msg_code, msg @@ -272,6 +272,6 @@ def _parse_msg(self, code, packet, is_ttb=False): 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/tcp/stream.py b/riak/transports/tcp/stream.py index 38357bbd..986059d8 100644 --- a/riak/transports/tcp/stream.py +++ b/riak/transports/tcp/stream.py @@ -3,13 +3,13 @@ from riak.util import decode_index_value, bytes_to_str from riak.client.index_page import CONTINUATION -from riak.transports.tcp.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. """ @@ -62,15 +62,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 +82,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 +103,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 +123,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) + super(PbufIndexStream, self).__init__(transport) 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,15 +161,15 @@ 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 diff --git a/riak/transports/tcp/transport.py b/riak/transports/tcp/transport.py index 01eddfa8..f07ed0b2 100644 --- a/riak/transports/tcp/transport.py +++ b/riak/transports/tcp/transport.py @@ -4,27 +4,27 @@ import riak.pb.riak_ts_pb2 from riak import RiakError -from riak.transports.transport import RiakTransport +from riak.codecs.pbuf import PbufCodec +from riak.codecs.ttb import TtbCodec +from riak.transports.transport import Transport 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.tcp.connection import RiakPbcConnection -from riak.transports.tcp.stream import (RiakPbcKeyStream, - RiakPbcMapredStream, - RiakPbcBucketStream, - RiakPbcIndexStream, - RiakPbcTsKeyStream) -from riak.transports.tcp.codec import RiakPbcCodec -from riak.transports.ttb.codec import RiakTtbCodec +from riak.transports.tcp.connection import TcpConnection +from riak.transports.tcp.stream import (PbufKeyStream, + PbufMapredStream, + PbufBucketStream, + PbufIndexStream, + PbufTsKeyStream) from six import PY2, PY3 -class RiakPbcTransport(RiakTransport, RiakPbcConnection, - RiakPbcCodec, RiakTtbCodec): +class TcpTransport(Transport, TcpConnection, + PbufCodec, TtbCodec): """ - The RiakPbcTransport object holds a connection to the protocol - buffers interface on the riak server. + The TcpTransport object holds a connection to the TCP + socket on the Riak server. """ def __init__(self, @@ -32,10 +32,7 @@ def __init__(self, client=None, timeout=None, **transport_options): - """ - Construct a new RiakPbcTransport object. - """ - super(RiakPbcTransport, self).__init__() + super(TcpTransport, self).__init__() self._client = client self._node = node @@ -269,7 +266,7 @@ def ts_stream_keys(self, table, timeout=None): self._send_msg(riak.pb.messages.MSG_CODE_TS_LIST_KEYS_REQ, req) - return RiakPbcTsKeyStream(self) + return PbufTsKeyStream(self) def delete(self, robj, rw=None, r=None, w=None, dw=None, pr=None, pw=None, timeout=None): @@ -331,7 +328,7 @@ def stream_keys(self, bucket, timeout=None): self._send_msg(riak.pb.messages.MSG_CODE_LIST_KEYS_REQ, req) - return RiakPbcKeyStream(self) + return PbufKeyStream(self) def get_buckets(self, bucket_type=None, timeout=None): """ @@ -367,7 +364,7 @@ def stream_buckets(self, bucket_type=None, timeout=None): self._send_msg(riak.pb.messages.MSG_CODE_LIST_BUCKETS_REQ, req) - return RiakPbcBucketStream(self) + return PbufBucketStream(self) def get_bucket_props(self, bucket): """ @@ -480,7 +477,7 @@ def stream_mapred(self, inputs, query, timeout=None): self._send_msg(riak.pb.messages.MSG_CODE_MAP_RED_REQ, req) - return RiakPbcMapredStream(self) + return PbufMapredStream(self) def get_index(self, bucket, index, startkey, endkey=None, return_terms=None, max_results=None, continuation=None, @@ -532,7 +529,7 @@ def stream_index(self, bucket, index, startkey, endkey=None, self._send_msg(riak.pb.messages.MSG_CODE_INDEX_REQ, req) - return RiakPbcIndexStream(self, index, return_terms) + return PbufIndexStream(self, index, return_terms) def create_search_index(self, index, schema=None, n_val=None, timeout=None): diff --git a/riak/transports/transport.py b/riak/transports/transport.py index 4f33168c..e30f2d5e 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. From 0bcea331160d9ffb8eafcd63c1cde4ce86d7faff Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Tue, 29 Mar 2016 13:06:47 -0700 Subject: [PATCH 26/44] fix test on PY2 --- riak/tests/test_timeseries_ttb.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index aae5f73e..64ff0888 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -29,8 +29,8 @@ str0 = 'ascii-0' str1 = 'ascii-1' -bd0 = six.text_type('时间序列') -bd1 = six.text_type('временные ряды') +bd0 = six.u('时间序列') +bd1 = six.u('временные ряды') fiveMins = datetime.timedelta(0, 300) ts0 = datetime.datetime(2015, 1, 1, 12, 0, 0) From dd52c945094b3f98d990468538efd4411812d40f Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 30 Mar 2016 08:29:32 -0700 Subject: [PATCH 27/44] Rewrite TCP connection class to be much more efficient. --- .gitignore | 10 +----- riak/benchmarks/multiget.py | 39 ++++++++++++++--------- riak/transports/tcp/connection.py | 52 +++++++++++++++---------------- 3 files changed, 52 insertions(+), 49 deletions(-) diff --git a/.gitignore b/.gitignore index 68a831bf..e20b3037 100644 --- a/.gitignore +++ b/.gitignore @@ -1,25 +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/riak/benchmarks/multiget.py b/riak/benchmarks/multiget.py index 505069a4..22c01c8b 100644 --- a/riak/benchmarks/multiget.py +++ b/riak/benchmarks/multiget.py @@ -1,18 +1,31 @@ -from riak import RiakClient -from multiprocessing import cpu_count import binascii import os + import riak.benchmark as benchmark -import riak.client.multiget as mget -client = RiakClient(protocol='pbc') +from riak import RiakClient +from multiprocessing import cpu_count + +nodes = [ + ('riak-test', 8098, 8087), + # ('riak-test', 10018, 10017), + # ('riak-test', 10028, 10027), + # ('riak-test', 10038, 10037), + # ('riak-test', 10048, 10047), + # ('riak-test', 10058, 10057), +] +client = RiakClient( + nodes=nodes, + protocol='pbc', + multiget_pool_size=128) + bkeys = [('default', 'multiget', str(key)) for key in range(10000)] data = binascii.b2a_hex(os.urandom(1024)) print("Benchmarking multiget:") print(" CPUs: {0}".format(cpu_count())) -print(" Threads: {0}".format(mget.POOL_SIZE)) +print(" Threads: {0}".format(client._multiget_pool._size)) print(" Keys: {0}".format(len(bkeys))) print() @@ -23,18 +36,16 @@ content_type='text/plain' ).store() for b in benchmark.measure_with_rehearsal(): - client.protocol = 'http' - with b.report('http seq'): - for _, bucket, key in bkeys: - client.bucket(bucket).get(key) - - with b.report('http multi'): - mget.multiget(client, bkeys) + # client.protocol = 'http' + # with b.report('http seq'): + # for _, bucket, key in bkeys: + # client.bucket(bucket).get(key) + # with b.report('http multi'): + # client.multiget(bkeys) client.protocol = 'pbc' with b.report('pbc seq'): for _, bucket, key in bkeys: client.bucket(bucket).get(key) - with b.report('pbc multi'): - mget.multiget(client, bkeys) + client.multiget(bkeys) diff --git a/riak/transports/tcp/connection.py b/riak/transports/tcp/connection.py index 8468a72d..dfaf3e55 100644 --- a/riak/transports/tcp/connection.py +++ b/riak/transports/tcp/connection.py @@ -185,16 +185,17 @@ def _ssl_handshake(self): raise SecurityError(e) def _recv_msg(self, expect=None, is_ttb=False): - self._recv_pkt() - msg_code, = struct.unpack("B", self._inbuf[:1]) + msgbuf = self._recv_pkt() + mv = memoryview(msgbuf) + msg_code, = struct.unpack("B", mv[0:1]) if msg_code is riak.pb.messages.MSG_CODE_ERROR_RESP: - err = self._parse_msg(msg_code, self._inbuf[1:], is_ttb) + err = self._parse_msg(msg_code, mv[1:].tobytes(), is_ttb) 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:], is_ttb) + msg = self._parse_msg(msg_code, mv[1:].tobytes(), is_ttb) else: raise Exception("unknown msg code %s" % msg_code) @@ -206,31 +207,30 @@ def _recv_msg(self, expect=None, is_ttb=False): return msg_code, msg 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: + # TODO FUTURE re-use buffer + msglen_buf = bytearray(4) + recv_len = self._socket.recv_into(msglen_buf) + if recv_len != 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: + % recv_len) + # NB: msg length is an unsigned int + msglen, = struct.unpack('!I', msglen_buf) + # 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: From 460b0151ffea6625510ed43baacae0952f9ebd00 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 30 Mar 2016 12:49:58 -0700 Subject: [PATCH 28/44] Moving code to PbufCodec --- riak/codecs/pbuf.py | 354 ++++++++++++++++- riak/codecs/ttb.py | 16 +- riak/transports/tcp/connection.py | 101 ++--- riak/transports/tcp/stream.py | 6 +- riak/transports/tcp/transport.py | 611 ++++++++++-------------------- riak/transports/transport.py | 3 +- 6 files changed, 595 insertions(+), 496 deletions(-) diff --git a/riak/codecs/pbuf.py b/riak/codecs/pbuf.py index 2d5231d8..4c746f08 100644 --- a/riak/codecs/pbuf.py +++ b/riak/codecs/pbuf.py @@ -7,6 +7,7 @@ from riak import RiakError from riak.content import RiakContent +from riak.riak_object import VClock 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 @@ -76,10 +77,15 @@ class PbufCodec(object): Protobuffs Encoding and decoding methods for TcpTransport. """ - def __init__(self, **unused_args): + 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") - super(PbufCodec, self).__init__(**unused_args) + self._client_timeouts = client_timeouts + self._quorum_controls = quorum_controls + self._tombstone_vclocks = tombstone_vclocks + self._bucket_types = bucket_types def _unix_time_millis(self, dt): return unix_time_millis(dt) @@ -305,7 +311,6 @@ def _decode_bucket_props(self, msg): :rtype dict """ props = {} - for prop in NORMAL_PROPS: if msg.HasField(prop): props[prop] = getattr(msg, prop) @@ -322,7 +327,6 @@ def _decode_bucket_props(self, msg): 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): @@ -411,7 +415,8 @@ def _encode_hook(self, hook, msg): def _encode_index_req(self, bucket, index, startkey, endkey=None, return_terms=None, max_results=None, - continuation=None, timeout=None, term_regex=None): + continuation=None, timeout=None, term_regex=None, + streaming=False): """ Encodes a secondary index request into the protobuf message. @@ -434,6 +439,8 @@ def _encode_index_req(self, bucket, index, startkey, endkey=None, :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( @@ -460,7 +467,8 @@ def _encode_index_req(self, bucket, index, startkey, endkey=None, req.timeout = timeout if term_regex: req.term_regex = str_to_bytes(term_regex) - return req + req.stream = streaming + return req.SerializeToString() def _decode_search_index(self, index): """ @@ -480,7 +488,7 @@ def _decode_search_index(self, index): def _add_bucket_type(self, req, bucket_type): if bucket_type and not bucket_type.is_default(): - if not self.bucket_types(): + if not self._bucket_types: raise NotImplementedError( 'Server does not support bucket-types') req.type = str_to_bytes(bucket_type.name) @@ -645,22 +653,30 @@ def _encode_to_ts_cell(self, cell, ts_cell): raise RiakError("can't serialize type '{}', value '{}'" .format(t, cell)) - def _encode_timeseries_keyreq(self, table, key, req): + 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") + if is_delete: + req = riak.pb.riak_ts_pb2.TsDelReq() + else: + req = riak.pb.riak_ts_pb2.TsGetReq() + 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 req.SerializeToString() - def _encode_timeseries_listkeysreq(self, table, req, timeout=None): + def _encode_timeseries_listkeysreq(self, table, timeout=None): + req = riak.pb.riak_ts_pb2.TsListKeysReq() req.table = str_to_bytes(table.name) - if timeout: + if self._client_timeouts and timeout: req.timeout = timeout + return req.SerializeToString() def _encode_timeseries_put(self, tsobj, req): """ @@ -672,6 +688,7 @@ def _encode_timeseries_put(self, tsobj, req): :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: @@ -687,6 +704,15 @@ def _encode_timeseries_put(self, tsobj, req): self._encode_to_ts_cell(cell, tsc) else: raise RiakError("TsObject requires a list of rows") + return req.SerializeToString() + + 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) + return req.SerializeToString() def _decode_timeseries(self, resp, tsobj): """ @@ -769,3 +795,311 @@ def _decode_preflist(self, item): '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) + return req.SerializeToString() + + 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) + return req.SerializeToString() + + 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) + return req.SerializeToString() + + 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) + return req.SerializeToString() + + 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 _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) + return req.SerializeToString() + + def _encode_get_buckets(self, bucket_type, timeout): + req = riak.pb.riak_kv_pb2.RpbListBucketsReq() + self._add_bucket_type(req, bucket_type) + if self._client_timeouts and timeout: + req.timeout = timeout + return req.SerializeToString() + + def _encode_stream_buckets(self, bucket_type, timeout): + 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 + return req.SerializeToString() + + 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) + return req.SerializeToString() + + 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) + return req.SerializeToString() + + 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) + return req.SerializeToString() + + def _encode_get_bucket_type_props(self, bucket_type): + req = riak.pb.riak_pb2.RpbGetBucketTypeReq() + req.type = str_to_bytes(bucket_type.name) + return req.SerializeToString() + + 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) + return req.SerializeToString() + + def _encode_stream_mapred(content): + req = riak.pb.riak_kv_pb2.RpbMapRedReq() + req.request = str_to_bytes(content) + req.content_type = str_to_bytes("application/json") + return req.SerializeToString() + + 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 + return req.SerializeToString() + + def _encode_get_search_index(self, index): + req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexGetReq( + name=str_to_bytes(index)) + return req.SerializeToString() + + def _encode_list_search_indexes(self): + req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexGetReq() + return req.SerializeToString() + + def _encode_delete_search_indexes(self): + req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexDeleteReq( + name=str_to_bytes(index)) + return req.SerializeToString() + + 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) + return req.SerializeToString() + + def _encode_get_search_schema(self, schema): + req = riak.pb.riak_yokozuna_pb2.RpbYokozunaSchemaGetReq( + name=str_to_bytes(schema)) + return req.SerializeToString() + + 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, **params): + req = riak.pb.riak_search_pb2.RpbSearchQueryReq( + index=str_to_bytes(index), + q=str_to_bytes(query)) + self._encode_search_query(req, params) + return req.SerializeToString() + + def _decode_search(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, **params): + 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'] + return req.SerializeToString() + + def _encode_update_counter(self, bucket, key, value, **params): + 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'] + return req.SerializeToString() + + def _encode_fetch_datatype(self, bucket, key, **options): + 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) + return req.SerializeToString() + + def _encode_update_datatype(self, datatype, **options): + 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) + return req.SerializeToString() + + 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) + return req.SerializeToString() diff --git a/riak/codecs/ttb.py b/riak/codecs/ttb.py index 2ecb846d..8d795675 100644 --- a/riak/codecs/ttb.py +++ b/riak/codecs/ttb.py @@ -29,7 +29,7 @@ class TtbCodec(object): def __init__(self, **unused_args): super(TtbCodec, self).__init__(**unused_args) - def _encode_to_ts_cell_ttb(self, cell): + def _encode_to_ts_cell(self, cell): if cell is None: return tscell_empty else: @@ -53,17 +53,17 @@ def _encode_to_ts_cell_ttb(self, cell): raise RiakError("can't serialize type '{}', value '{}'" .format(t, cell)) - def _encode_timeseries_keyreq_ttb(self, table, key): + def _encode_timeseries_keyreq(self, table, key): key_vals = None if isinstance(key, list): key_vals = key else: raise ValueError("key must be a list") req = tsgetreq_a, table.name, \ - [self._encode_to_ts_cell_ttb(k) for k in key_vals], udef_a + [self._encode_to_ts_cell(k) for k in key_vals], udef_a return encode(req) - def _encode_timeseries_put_ttb(self, tsobj): + def _encode_timeseries_put(self, tsobj): ''' Returns an Erlang-TTB encoded tuple with the appropriate data and metadata from a TsObject. @@ -80,7 +80,7 @@ def _encode_timeseries_put_ttb(self, tsobj): for row in tsobj.rows: req_r = [] for cell in row: - req_r.append(self._encode_to_ts_cell_ttb(cell)) + req_r.append(self._encode_to_ts_cell(cell)) req_t = (tsrow_a, req_r) req_rows.append(req_t) req = tsputreq_a, tsobj.table.name, udef_a, req_rows @@ -88,7 +88,7 @@ def _encode_timeseries_put_ttb(self, tsobj): else: raise RiakError("TsObject requires a list of rows") - def _decode_timeseries_ttb(self, resp_ttb, tsobj): + def _decode_timeseries(self, resp_ttb, tsobj): """ Fills an TsObject with the appropriate data and metadata from a TTB-encoded TsGetResp / TsQueryResp. @@ -113,13 +113,13 @@ def _decode_timeseries_ttb(self, resp_ttb, tsobj): resp_rows = resp_ttb[2] for row_ttb in resp_rows: tsobj.rows.append( - self._decode_timeseries_row_ttb(row_ttb, None)) + self._decode_timeseries_row(row_ttb, None)) # TODO # elif resp_a == rpberrorresp_a: else: raise RiakError("Unknown TTB response type: {}".format(resp_a)) - def _decode_timeseries_row_ttb(self, tsrow_ttb, tscols=None): + def _decode_timeseries_row(self, tsrow_ttb, tscols=None): """ Decodes a TTB-encoded TsRow into a list diff --git a/riak/transports/tcp/connection.py b/riak/transports/tcp/connection.py index dfaf3e55..6ba885a1 100644 --- a/riak/transports/tcp/connection.py +++ b/riak/transports/tcp/connection.py @@ -5,11 +5,12 @@ import riak.pb.messages import erlastic -from riak.security import SecurityError, USE_STDLIB_SSL from riak import RiakError -from riak.util import bytes_to_str, str_to_bytes +from riak.security import SecurityError, USE_STDLIB_SSL +from riak.util import str_to_bytes from six import PY2 + if not USE_STDLIB_SSL: from OpenSSL.SSL import Connection from riak.transports.security import configure_pyopenssl_context @@ -22,47 +23,37 @@ class TcpConnection(object): """ Connection-related methods for TcpTransport. """ - def __init__(self): self._ttb_enabled = False - def _encode_msg(self, msg_code, msg=None, is_ttb=False): - if msg is None: + def _encode_msg(self, msg_code, data=None): + if data is None: return struct.pack("!iB", 1, msg_code) - - if is_ttb: - data = msg - else: - data = msg.SerializeToString() - - datalen = len(data) - hdr = struct.pack("!iB", 1 + datalen, msg_code) + hdr = struct.pack("!iB", 1 + len(data), msg_code) return hdr + data - def _request(self, msg_code, msg=None, expect=None, is_ttb=False): - self._send_msg(msg_code, msg, is_ttb) - return self._recv_msg(expect, is_ttb) + def _send_recv(self, msg_code, data=None, expect=None): + self._send_msg(msg_code, data) + return self._recv_msg(expect) - 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) + self._non_connect_send_msg(msg_code, data) return self._recv_msg(expect) - def _non_connect_send_msg(self, msg_code, msg, is_ttb=False): + 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, is_ttb)) + self._socket.sendall(self._encode_msg(msg_code, data)) - def _send_msg(self, msg_code, msg, is_ttb=False): + def _send_msg(self, msg_code, data): self._connect() - if is_ttb and not self._enable_ttb(): - raise RiakError('could not switch to TTB encoding!') - self._non_connect_send_msg(msg_code, msg, is_ttb) + self._non_connect_send_msg(msg_code, data) def _init_security(self): """ @@ -80,7 +71,7 @@ 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( + msg_code, _ = self._non_connect_send_recv( riak.pb.messages.MSG_CODE_START_TLS) if msg_code == riak.pb.messages.MSG_CODE_START_TLS: return True @@ -94,7 +85,7 @@ def _enable_ttb(self): logging.debug("tcp/connection enabling TTB") req = riak.pb.riak_pb2.RpbToggleEncodingReq() req.use_native = True - msg_code, _ = self._non_connect_request( + msg_code, _ = self._non_connect_send_recv( riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_REQ, req, riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_RESP) @@ -118,7 +109,7 @@ def _auth(self): if not password: password = '' req.password = str_to_bytes(password) - msg_code, _ = self._non_connect_request( + msg_code, _ = self._non_connect_send_recv( riak.pb.messages.MSG_CODE_AUTH_REQ, req, riak.pb.messages.MSG_CODE_AUTH_RESP) @@ -176,7 +167,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) @@ -184,38 +174,21 @@ def _ssl_handshake(self): # fail if *any* exceptions are thrown during SSL handshake raise SecurityError(e) - def _recv_msg(self, expect=None, is_ttb=False): + def _recv_msg(self, expect=None): msgbuf = self._recv_pkt() mv = memoryview(msgbuf) msg_code, = struct.unpack("B", mv[0:1]) - if msg_code is riak.pb.messages.MSG_CODE_ERROR_RESP: - err = self._parse_msg(msg_code, mv[1:].tobytes(), is_ttb) - 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, mv[1:].tobytes(), is_ttb) - 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)) - # logging.debug("tcp/connection received msg_code %d msg %s", - # msg_code, msg) - return msg_code, msg + data = mv[1:].tobytes() + return (msg_code, data) def _recv_pkt(self): # TODO FUTURE re-use buffer - msglen_buf = bytearray(4) - recv_len = self._socket.recv_into(msglen_buf) - if recv_len != 4: - raise RiakError( - "Socket returned short packet length %d - expected 4" - % recv_len) + 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) @@ -250,28 +223,6 @@ def close(self): self._socket.close() del self._socket - def _parse_msg(self, code, packet, is_ttb=False): - if is_ttb: - if code != riak.pb.messages.MSG_CODE_TS_GET_RESP and \ - code != riak.pb.messages.MSG_CODE_TS_PUT_RESP: - raise RiakError("TTB can't parse code: %d" % code) - if len(packet) > 0: - return erlastic.decode(packet) - else: - return None - else: - 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 TcpTransport initializer _address = None _timeout = None diff --git a/riak/transports/tcp/stream.py b/riak/transports/tcp/stream.py index 986059d8..2288c607 100644 --- a/riak/transports/tcp/stream.py +++ b/riak/transports/tcp/stream.py @@ -1,4 +1,5 @@ import json + import riak.pb.messages from riak.util import decode_index_value, bytes_to_str @@ -28,7 +29,10 @@ def next(self): raise StopIteration try: - msg_code, resp = self.transport._recv_msg(expect=self._expect) + expected_code = self._expect + msg_code, data = self.transport._recv_msg(expect=expected_code) + self.transport._maybe_riak_error(msg_code, data) + resp = self.transport._parse_msg(expected_code, data, is_ttb=False) except: self.finished = True raise diff --git a/riak/transports/tcp/transport.py b/riak/transports/tcp/transport.py index f07ed0b2..e73e82ac 100644 --- a/riak/transports/tcp/transport.py +++ b/riak/transports/tcp/transport.py @@ -1,15 +1,16 @@ +# TODO RTS-842 codecs should return msg codes too +import six 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.codecs.pbuf import PbufCodec from riak.codecs.ttb import TtbCodec from riak.transports.transport import Transport -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 + +# TODO RTS-842 ideally these would not be needed +from riak.util import decode_index_value, bytes_to_str + from riak.transports.tcp.connection import TcpConnection from riak.transports.tcp.stream import (PbufKeyStream, PbufMapredStream, @@ -17,16 +18,12 @@ PbufIndexStream, PbufTsKeyStream) -from six import PY2, PY3 - -class TcpTransport(Transport, TcpConnection, - PbufCodec, TtbCodec): +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, @@ -39,18 +36,41 @@ def __init__(self, self._address = (node.host, node.pb_port) self._timeout = timeout self._socket = None - self._use_ttb = transport_options.get('use_ttb', False) + self._pbuf_c = None + self._ttb_c = None + self._use_ttb = transport_options.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_codec(self, ttb_supported=False): + if ttb_supported: + if self._use_ttb: + if not self._enable_ttb(): + raise RiakError('could not switch to TTB encoding!') + if not self._ttb_c: + self._ttb_c = TtbCodec() + codec = self._ttb_c + else: + codec = self._get_pbuf_codec() + else: + codec = self._get_pbuf_codec() + return codec # FeatureDetection API def _server_version(self): - return bytes_to_str(self.get_server_info()['server_version']) + server_info = self.get_server_info() + return server_info['server_version'] def ping(self): """ Ping the remote server """ - - msg_code, msg = self._request(riak.pb.messages.MSG_CODE_PING_REQ) + msg_code, _ = self._request(riak.pb.messages.MSG_CODE_PING_REQ) if msg_code == riak.pb.messages.MSG_CODE_PING_RESP: return True else: @@ -60,26 +80,27 @@ 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_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)} + return codec._decode_get_server_info(resp) def _get_client_id(self): + codec = self._get_codec(ttb_supported=False) 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) + return codec._decode_get_client_id(resp) def _set_client_id(self, client_id): - req = riak.pb.riak_kv_pb2.RpbSetClientIdReq() - req.client_id = str_to_bytes(client_id) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_set_client_id(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_REQ, data, riak.pb.messages.MSG_CODE_SET_CLIENT_ID_RESP) - self._client_id = client_id client_id = property(_get_client_id, _set_client_id, @@ -90,165 +111,74 @@ def get(self, robj, r=None, pr=None, timeout=None, basic_quorum=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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_get(robj, r, pr, + timeout, basic_quorum, notfound_ok) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_REQ, req, + riak.pb.messages.MSG_CODE_GET_REQ, data, 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 + 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): - 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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_put(robj, w, dw, pw, return_body, + if_none_match, timeout) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_PUT_REQ, req, + riak.pb.messages.MSG_CODE_PUT_REQ, data, 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 + 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): - ts_get_resp = None - if self._use_ttb: - req = self._encode_timeseries_keyreq_ttb(table, key) - else: - req = riak.pb.riak_ts_pb2.TsGetReq() - self._encode_timeseries_keyreq(table, key, req) - + codec = self._get_codec(ttb_supported=True) + data = codec._encode_timeseries_keyreq(table, key) msg_code, ts_get_resp = self._request( - riak.pb.messages.MSG_CODE_TS_GET_REQ, req, - riak.pb.messages.MSG_CODE_TS_GET_RESP, - self._use_ttb) - + riak.pb.messages.MSG_CODE_TS_GET_REQ, data, + riak.pb.messages.MSG_CODE_TS_GET_RESP) tsobj = TsObject(self._client, table, [], None) - if self._use_ttb: - self._decode_timeseries_ttb(ts_get_resp, tsobj) - else: - self._decode_timeseries(ts_get_resp, tsobj) + codec._decode_timeseries(ts_get_resp, tsobj) return tsobj def ts_put(self, tsobj): - if self._use_ttb: - req = self._encode_timeseries_put_ttb(tsobj) - else: - req = riak.pb.riak_ts_pb2.TsPutReq() - self._encode_timeseries_put(tsobj, req) - + codec = self._get_codec(ttb_supported=True) + # TODO RTS-842 codecs should return msg codes too + data = codec._encode_timeseries_put(tsobj) # logging.debug("pbc/transport ts_put _use_ttb: '%s'", # self._use_ttb) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_TS_PUT_REQ, req, + riak.pb.messages.MSG_CODE_TS_PUT_REQ, data, riak.pb.messages.MSG_CODE_TS_PUT_RESP, self._use_ttb) - if self._use_ttb and \ resp is None and \ msg_code == riak.pb.messages.MSG_CODE_TS_PUT_RESP: return True - 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) - + codec = self._get_codec(ttb_supported=True) + data = codec._encode_timeseries_keyreq(table, key, is_delete=True) msg_code, ts_del_resp = self._request( - riak.pb.messages.MSG_CODE_TS_DEL_REQ, req, + riak.pb.messages.MSG_CODE_TS_DEL_REQ, data, 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) - + codec = self._get_codec(ttb_supported=True) + data = codec._encode_timeseries_query(table, query, interpolations) msg_code, ts_query_resp = self._request( - riak.pb.messages.MSG_CODE_TS_QUERY_REQ, req, + riak.pb.messages.MSG_CODE_TS_QUERY_REQ, data, riak.pb.messages.MSG_CODE_TS_QUERY_RESP) - tsobj = TsObject(self._client, table, [], []) self._decode_timeseries(ts_query_resp, tsobj) return tsobj @@ -258,49 +188,17 @@ 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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_timeseries_listkeysreq(table, t) + self._send_msg(riak.pb.messages.MSG_CODE_TS_LIST_KEYS_REQ, data) return PbufTsKeyStream(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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_delete(robj, rw, r, w, dw, pr, pw, timeout) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_DEL_REQ, req, + riak.pb.messages.MSG_CODE_DEL_REQ, data, riak.pb.messages.MSG_CODE_DEL_RESP) return self @@ -308,40 +206,28 @@ 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 + codec = self._get_codec(ttb_supported=False) + 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. """ - 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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_stream_keys(bucket, timeout) + self._send_msg(riak.pb.messages.MSG_CODE_LIST_KEYS_REQ, data) return PbufKeyStream(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 - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_get_buckets(bucket_type, timeout) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_LIST_BUCKETS_REQ, req, + riak.pb.messages.MSG_CODE_LIST_BUCKETS_REQ, data, riak.pb.messages.MSG_CODE_LIST_BUCKETS_RESP) return resp.buckets @@ -349,55 +235,38 @@ 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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_stream_buckets(bucket_type, timeout) + self._send_msg(riak.pb.messages.MSG_CODE_LIST_BUCKETS_REQ, data) return PbufBucketStream(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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_get_bucket_props(bucket) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_BUCKET_REQ, req, + riak.pb.messages.MSG_CODE_GET_BUCKET_REQ, data, riak.pb.messages.MSG_CODE_GET_BUCKET_RESP) - - return self._decode_bucket_props(resp.props) + return codec._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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_set_bucket_props(bucket, props) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_SET_BUCKET_REQ, req, + riak.pb.messages.MSG_CODE_SET_BUCKET_REQ, data, riak.pb.messages.MSG_CODE_SET_BUCKET_RESP) return True @@ -407,12 +276,10 @@ def clear_bucket_props(self, bucket): """ 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) + codec = self._get_codec(ttb_supported=False) + data = codec._encode_clear_bucket_props(bucket) self._request( - riak.pb.messages.MSG_CODE_RESET_BUCKET_REQ, req, + riak.pb.messages.MSG_CODE_RESET_BUCKET_REQ, data, riak.pb.messages.MSG_CODE_RESET_BUCKET_RESP) return True @@ -421,31 +288,23 @@ 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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_get_bucket_type_props(bucket_type) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_BUCKET_TYPE_REQ, req, + riak.pb.messages.MSG_CODE_GET_BUCKET_TYPE_REQ, data, riak.pb.messages.MSG_CODE_GET_BUCKET_RESP) - - return self._decode_bucket_props(resp.props) + 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) - - req = riak.pb.riak_pb2.RpbSetBucketTypeReq() - req.type = str_to_bytes(bucket_type.name) - - self._encode_bucket_props(props, req) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_set_bucket_type_props(bucket_type, props) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_SET_BUCKET_TYPE_REQ, req, + riak.pb.messages.MSG_CODE_SET_BUCKET_TYPE_REQ, data, riak.pb.messages.MSG_CODE_SET_BUCKET_RESP) - return True def mapred(self, inputs, query, timeout=None): @@ -457,7 +316,6 @@ def mapred(self, inputs, query, timeout=None): 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): @@ -470,18 +328,15 @@ def mapred(self, inputs, query, timeout=None): 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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_stream_mapred(content) + self._send_msg(riak.pb.messages.MSG_CODE_MAP_RED_REQ, data) return PbufMapredStream(self) 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) @@ -489,12 +344,13 @@ def get_index(self, bucket, index, startkey, endkey=None, raise NotImplementedError("Secondary index term_regex is not " "supported") - req = self._encode_index_req(bucket, index, startkey, endkey, + codec = self._get_codec(ttb_supported=False) + data = codec._encode_index_req(bucket, index, startkey, endkey, return_terms, max_results, continuation, - timeout, term_regex) + timeout, term_regex, streaming=False) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_INDEX_REQ, req, + riak.pb.messages.MSG_CODE_INDEX_REQ, data, riak.pb.messages.MSG_CODE_INDEX_RESP) if return_terms and resp.results: @@ -503,7 +359,7 @@ def get_index(self, bucket, index, startkey, endkey=None, for pair in resp.results] else: results = resp.keys[:] - if PY3: + if six.PY3: results = [bytes_to_str(key) for key in resp.keys] if max_results is not None and resp.HasField('continuation'): @@ -517,18 +373,14 @@ def stream_index(self, bucket, index, startkey, endkey=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, + codec = self._get_codec(ttb_supported=False) + data = codec._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) - + timeout, term_regex, streaming=True) + self._send_msg(riak.pb.messages.MSG_CODE_INDEX_REQ, data) return PbufIndexStream(self, index, return_terms) def create_search_index(self, index, schema=None, n_val=None, @@ -536,34 +388,24 @@ def create_search_index(self, index, schema=None, n_val=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 - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_create_search_index(index, schema, n_val, timeout) self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_PUT_REQ, req, + riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_PUT_REQ, data, 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)) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_get_search_index(index) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_REQ, req, + riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_REQ, data, riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_RESP) if len(resp.index) > 0: - return self._decode_search_index(resp.index[0]) + return codec._decode_search_index(resp.index[0]) else: raise RiakError('notfound') @@ -571,106 +413,71 @@ 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() - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_list_search_indexes() msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_REQ, req, + riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_REQ, data, riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_RESP) - - return [self._decode_search_index(index) for index in resp.index] + 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") - req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexDeleteReq( - name=str_to_bytes(index)) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_delete_search_index(index) self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_DELETE_REQ, req, + riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_DELETE_REQ, data, 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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_create_search_schema(schema, content) self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_PUT_REQ, req, + riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_PUT_REQ, data, 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)) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_get_search_schema(schema) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_GET_REQ, req, + riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_GET_REQ, data, 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 + return codec._decode_get_search_schema(resp) def search(self, index, query, **params): + # TODO RTS-842 NUKE THIS if not self.pb_search(): return self._search_mapred_emu(index, query) - - if PY2 and isinstance(query, unicode): # noqa + # TODO RTS-842 six.u() instead? + if six.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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_search(index, query, params) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_SEARCH_QUERY_REQ, req, + riak.pb.messages.MSG_CODE_SEARCH_QUERY_REQ, data, 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 + return codec._decode_search(resp) 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'] - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_get_counter(bucket, key, params) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_COUNTER_GET_REQ, req, + riak.pb.messages.MSG_CODE_COUNTER_GET_REQ, data, riak.pb.messages.MSG_CODE_COUNTER_GET_RESP) if resp.HasField('value'): return resp.value @@ -682,25 +489,12 @@ def update_counter(self, bucket, key, value, **params): 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'] - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_update_counter(bucket, key, value, params) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_COUNTER_UPDATE_REQ, req, + riak.pb.messages.MSG_CODE_COUNTER_UPDATE_REQ, data, riak.pb.messages.MSG_CODE_COUNTER_UPDATE_RESP) if resp.HasField('value'): return resp.value @@ -708,65 +502,35 @@ def update_counter(self, bucket, key, value, **params): 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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_fetch_datatype(bucket, key, options) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_DT_FETCH_REQ, req, + riak.pb.messages.MSG_CODE_DT_FETCH_REQ, data, riak.pb.messages.MSG_CODE_DT_FETCH_RESP) - - return self._decode_dt_fetch(resp) + return codec._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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_update_datatype(datatype, options) msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_DT_UPDATE_REQ, req, + riak.pb.messages.MSG_CODE_DT_UPDATE_REQ, data, 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): @@ -779,13 +543,58 @@ def get_preflist(self, bucket, 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) - + codec = self._get_codec(ttb_supported=False) + data = codec._encode_get_preflist(bucket, key) 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_REQ, data, riak.pb.messages.MSG_CODE_GET_BUCKET_KEY_PREFLIST_RESP) - - return [self._decode_preflist(item) for item in resp.preflist] + return [codec._decode_preflist(item) for item in resp.preflist] + + # TODO RTS-842 is_ttb + def _parse_msg(self, code, packet, is_ttb=False): + if is_ttb: + if code != riak.pb.messages.MSG_CODE_TS_GET_RESP and \ + code != riak.pb.messages.MSG_CODE_TS_PUT_RESP: + raise RiakError("TTB can't parse code: %d" % code) + if len(packet) > 0: + return erlastic.decode(packet) + else: + return None + else: + 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 + + def _maybe_riak_error(self, msg_code, data=None, is_ttb=False): + if msg_code is riak.pb.messages.MSG_CODE_ERROR_RESP: + if data is None: + raise RiakError('no error provided!') + err = self._parse_msg(msg_code, data, is_ttb) + if err is None: + raise RiakError('no error provided!') + else: + raise RiakError(bytes_to_str(err.errmsg)) + + # TODO RTS-842 is_ttb + def _request(self, msg_code, data=None, expect=None, is_ttb=False): + msg_code, data = self._send_recv(msg_code, data, expect) + self._maybe_riak_error(msg_code, data, is_ttb) + if msg_code in riak.pb.messages.MESSAGE_CLASSES: + msg = self._parse_msg(msg_code, data, is_ttb) + 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)) + # logging.debug("tcp/connection received msg_code %d msg %s", + # msg_code, msg) + return msg_code, msg diff --git a/riak/transports/transport.py b/riak/transports/transport.py index e30f2d5e..6e5fee2c 100644 --- a/riak/transports/transport.py +++ b/riak/transports/transport.py @@ -295,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 @@ -320,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 @@ -360,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') From 22ec2a46156ba5218e6a1f4c21df4f0312278c15 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 30 Mar 2016 15:15:38 -0700 Subject: [PATCH 29/44] Standardize on kwargs for double-star parameter name. --- riak/codecs/pbuf.py | 117 ++++++++++++++++-------------- riak/tests/test_timeseries.py | 29 +++++--- riak/tests/test_timeseries_ttb.py | 11 +-- riak/transports/tcp/transport.py | 31 ++++---- 4 files changed, 101 insertions(+), 87 deletions(-) diff --git a/riak/codecs/pbuf.py b/riak/codecs/pbuf.py index 4c746f08..eb59317e 100644 --- a/riak/codecs/pbuf.py +++ b/riak/codecs/pbuf.py @@ -493,28 +493,28 @@ def _add_bucket_type(self, req, bucket_type): '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']) + 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(params['fl']) - if 'presort' in params: - req.presort = params['presort'] + req.fl.append(kwargs['fl']) + if 'presort' in kwargs: + req.presort = kwargs['presort'] def _decode_search_doc(self, doc): resultdoc = MultiDict() @@ -550,15 +550,15 @@ def _decode_dt_value(self, dtype, msg): elif dtype == 'map': return self._decode_map_value(msg.map_value) - def _encode_dt_options(self, req, params): + def _encode_dt_options(self, req, **kwargs): 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])) + 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 params and params[o] is not None: - setattr(req, o, params[o]) + if o in kwargs and kwargs[o] is not None: + setattr(req, o, kwargs[o]) def _decode_map_value(self, entries): out = {} @@ -678,7 +678,7 @@ def _encode_timeseries_listkeysreq(self, table, timeout=None): req.timeout = timeout return req.SerializeToString() - def _encode_timeseries_put(self, tsobj, req): + def _encode_timeseries_put(self, tsobj): """ Fills an TsPutReq message with the appropriate data and metadata from a TsObject. @@ -975,7 +975,7 @@ def _encode_set_bucket_type_props(self, bucket_type, props): self._encode_bucket_props(props, req) return req.SerializeToString() - def _encode_stream_mapred(content): + 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") @@ -1003,7 +1003,7 @@ def _encode_list_search_indexes(self): req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexGetReq() return req.SerializeToString() - def _encode_delete_search_indexes(self): + def _encode_delete_search_index(self, index): req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexDeleteReq( name=str_to_bytes(index)) return req.SerializeToString() @@ -1027,14 +1027,14 @@ def _decode_get_search_schema(self, resp): result['content'] = bytes_to_str(resp.schema.content) return result - def _encode_search(self, index, query, **params): + 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, params) + self._encode_search_query(req, **kwargs) return req.SerializeToString() - def _decode_search(resp): + def _decode_search(self, resp): result = {} if resp.HasField('max_score'): result['max_score'] = resp.max_score @@ -1043,44 +1043,44 @@ def _decode_search(resp): result['docs'] = [self._decode_search_doc(doc) for doc in resp.docs] return result - def _encode_get_counter(self, bucket, key, **params): + 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 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'] + 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'] return req.SerializeToString() - def _encode_update_counter(self, bucket, key, value, **params): + 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 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'] + 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'] return req.SerializeToString() - def _encode_fetch_datatype(self, bucket, key, **options): + 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, options) + self._encode_dt_options(req, **kwargs) return req.SerializeToString() - def _encode_update_datatype(self, datatype, **options): + def _encode_update_datatype(self, datatype, **kwargs): op = datatype.to_op() type_name = datatype.type_name if not op: @@ -1093,10 +1093,19 @@ def _encode_update_datatype(self, datatype, **options): req.key = str_to_bytes(datatype.key) if datatype._context: req.context = datatype._context - self._encode_dt_options(req, options) + self._encode_dt_options(req, **kwargs) self._encode_dt_op(type_name, req, op) return req.SerializeToString() + 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) diff --git a/riak/tests/test_timeseries.py b/riak/tests/test_timeseries.py index 6f906726..30210099 100644 --- a/riak/tests/test_timeseries.py +++ b/riak/tests/test_timeseries.py @@ -35,7 +35,6 @@ class TimeseriesUnitTests(unittest.TestCase): @classmethod def setUpClass(cls): - cls.c = PbufCodec() cls.ts0ms = unix_time_millis(ts0) if cls.ts0ms != ex0ms: raise AssertionError( @@ -67,25 +66,33 @@ def test_encode_decode_timestamp(self): self.assertEqual(ts0, ts0_d) def test_encode_data_for_get(self): + c = PbufCodec() + data = 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(data) self.validate_keyreq(req) def test_encode_data_for_delete(self): + c = PbufCodec() + data = 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(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) + data = c._encode_timeseries_put(tsobj) + req = riak.pb.riak_ts_pb2.TsPutReq() + req.ParseFromString(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]) @@ -93,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]) @@ -102,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) + data = 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(data) self.assertEqual(self.table.name, bytes_to_str(req.table)) self.assertEqual(1234, req.timeout) diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index 64ff0888..959b0467 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -40,7 +40,6 @@ @unittest.skipUnless(is_timeseries_supported(), "Timeseries not supported") class TimeseriesTtbUnitTests(unittest.TestCase): def setUp(self): - self.c = TtbCodec() self.table = Table(None, table_name) def test_encode_data_for_get(self): @@ -53,8 +52,8 @@ def test_encode_data_for_get(self): req_test = encode(req) test_key = ['hash1', 'user2', ts0] - req_encoded = self.c._encode_timeseries_keyreq_ttb( - self.table, test_key) + c = TtbCodec() + req_encoded = c._encode_timeseries_keyreq(self.table, test_key) self.assertEqual(req_test, req_encoded) # def test_decode_riak_error(self): @@ -85,7 +84,8 @@ def test_decode_data_from_get(self): rsp_ttb = encode(rsp_data) tsobj = TsObject(None, self.table, [], []) - self.c._decode_timeseries_ttb(decode(rsp_ttb), tsobj) + c = TtbCodec() + c._decode_timeseries(decode(rsp_ttb), tsobj) for i in range(0, 1): self.assertEqual(tsrow_a, rows[i][0]) @@ -145,7 +145,8 @@ def test_encode_data_for_put(self): ] tsobj = TsObject(None, self.table, rows_to_encode, None) - req_encoded = self.c._encode_timeseries_put_ttb(tsobj) + c = TtbCodec() + req_encoded = c._encode_timeseries_put(tsobj) self.assertEqual(req_test, req_encoded) diff --git a/riak/transports/tcp/transport.py b/riak/transports/tcp/transport.py index e73e82ac..9679cdd7 100644 --- a/riak/transports/tcp/transport.py +++ b/riak/transports/tcp/transport.py @@ -28,7 +28,7 @@ def __init__(self, node=None, client=None, timeout=None, - **transport_options): + **kwargs): super(TcpTransport, self).__init__() self._client = client @@ -38,7 +38,7 @@ def __init__(self, self._socket = None self._pbuf_c = None self._ttb_c = None - self._use_ttb = transport_options.get('use_ttb', True) + self._use_ttb = kwargs.get('use_ttb', True) def _get_pbuf_codec(self): if not self._pbuf_c: @@ -453,7 +453,7 @@ def get_search_schema(self, schema): riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_GET_RESP) return codec._decode_get_search_schema(resp) - def search(self, index, query, **params): + def search(self, index, query, **kwargs): # TODO RTS-842 NUKE THIS if not self.pb_search(): return self._search_mapred_emu(index, query) @@ -461,13 +461,13 @@ def search(self, index, query, **params): if six.PY2 and isinstance(query, unicode): # noqa query = query.encode('utf8') codec = self._get_codec(ttb_supported=False) - data = codec._encode_search(index, query, params) + data = codec._encode_search(index, query, **kwargs) msg_code, resp = self._request( riak.pb.messages.MSG_CODE_SEARCH_QUERY_REQ, data, riak.pb.messages.MSG_CODE_SEARCH_QUERY_RESP) return codec._decode_search(resp) - def get_counter(self, bucket, key, **params): + def get_counter(self, bucket, key, **kwargs): if not bucket.bucket_type.is_default(): raise NotImplementedError("Counters are not " "supported with bucket-types, " @@ -475,7 +475,7 @@ def get_counter(self, bucket, key, **params): if not self.counters(): raise NotImplementedError("Counters are not supported") codec = self._get_codec(ttb_supported=False) - data = codec._encode_get_counter(bucket, key, params) + data = codec._encode_get_counter(bucket, key, **kwargs) msg_code, resp = self._request( riak.pb.messages.MSG_CODE_COUNTER_GET_REQ, data, riak.pb.messages.MSG_CODE_COUNTER_GET_RESP) @@ -484,7 +484,7 @@ def get_counter(self, bucket, key, **params): else: return None - def update_counter(self, bucket, key, value, **params): + def update_counter(self, bucket, key, value, **kwargs): if not bucket.bucket_type.is_default(): raise NotImplementedError("Counters are not " "supported with bucket-types, " @@ -492,7 +492,7 @@ def update_counter(self, bucket, key, value, **params): if not self.counters(): raise NotImplementedError("Counters are not supported") codec = self._get_codec(ttb_supported=False) - data = codec._encode_update_counter(bucket, key, value, params) + data = codec._encode_update_counter(bucket, key, value, **kwargs) msg_code, resp = self._request( riak.pb.messages.MSG_CODE_COUNTER_UPDATE_REQ, data, riak.pb.messages.MSG_CODE_COUNTER_UPDATE_RESP) @@ -501,36 +501,31 @@ def update_counter(self, bucket, key, value, **params): else: return True - def fetch_datatype(self, bucket, key, **options): + 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.") codec = self._get_codec(ttb_supported=False) - data = codec._encode_fetch_datatype(bucket, key, options) + data = codec._encode_fetch_datatype(bucket, key, **kwargs) msg_code, resp = self._request( riak.pb.messages.MSG_CODE_DT_FETCH_REQ, data, riak.pb.messages.MSG_CODE_DT_FETCH_RESP) return codec._decode_dt_fetch(resp) - def update_datatype(self, datatype, **options): + 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.") codec = self._get_codec(ttb_supported=False) - data = codec._encode_update_datatype(datatype, options) + data = codec._encode_update_datatype(datatype, **kwargs) msg_code, resp = self._request( riak.pb.messages.MSG_CODE_DT_UPDATE_REQ, data, 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)) + codec._decode_update_datatype(datatype, resp, **kwargs) return True def get_preflist(self, bucket, key): From 38bd2862e0014bc0a96a774fc69621557475a1cb Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 30 Mar 2016 15:22:55 -0700 Subject: [PATCH 30/44] Make linters happy. --- riak/codecs/pbuf.py | 17 +++++++++-------- riak/transports/tcp/connection.py | 8 +++----- riak/transports/tcp/transport.py | 17 ++++++++++------- 3 files changed, 22 insertions(+), 20 deletions(-) diff --git a/riak/codecs/pbuf.py b/riak/codecs/pbuf.py index eb59317e..44f5889b 100644 --- a/riak/codecs/pbuf.py +++ b/riak/codecs/pbuf.py @@ -78,8 +78,8 @@ class PbufCodec(object): """ def __init__(self, - client_timeouts=False, quorum_controls=False, - tombstone_vclocks=False, bucket_types=False): + 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 @@ -797,7 +797,7 @@ def _decode_preflist(self, item): return result def _encode_get(self, robj, r=None, pr=None, timeout=None, - basic_quorum=None, notfound_ok=None): + basic_quorum=None, notfound_ok=None): bucket = robj.bucket req = riak.pb.riak_kv_pb2.RpbGetReq() if r: @@ -819,7 +819,8 @@ def _encode_get(self, robj, r=None, pr=None, timeout=None, return req.SerializeToString() def _encode_put(self, robj, w=None, dw=None, pw=None, - return_body=True, if_none_match=False, timeout=None): + return_body=True, if_none_match=False, + timeout=None): bucket = robj.bucket req = riak.pb.riak_kv_pb2.RpbPutReq() if w: @@ -869,8 +870,8 @@ def _decode_put(self, robj, resp): return robj def _encode_delete(self, robj, rw=None, r=None, - w=None, dw=None, pr=None, pw=None, - timeout=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) @@ -982,7 +983,7 @@ def _encode_stream_mapred(self, content): return req.SerializeToString() def _encode_create_search_index(self, index, schema=None, - n_val=None, timeout=None): + n_val=None, timeout=None): index = str_to_bytes(index) idx = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndex(name=index) if schema: @@ -1007,7 +1008,7 @@ def _encode_delete_search_index(self, index): req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexDeleteReq( name=str_to_bytes(index)) return req.SerializeToString() - + def _encode_create_search_schema(self, schema, content): scma = riak.pb.riak_yokozuna_pb2.RpbYokozunaSchema( name=str_to_bytes(schema), diff --git a/riak/transports/tcp/connection.py b/riak/transports/tcp/connection.py index 6ba885a1..776ef4a6 100644 --- a/riak/transports/tcp/connection.py +++ b/riak/transports/tcp/connection.py @@ -1,16 +1,14 @@ import logging import socket import struct + import riak.pb.riak_pb2 import riak.pb.messages -import erlastic from riak import RiakError from riak.security import SecurityError, USE_STDLIB_SSL from riak.util import str_to_bytes -from six import PY2 - if not USE_STDLIB_SSL: from OpenSSL.SSL import Connection from riak.transports.security import configure_pyopenssl_context @@ -36,7 +34,7 @@ def _send_recv(self, msg_code, data=None, expect=None): self._send_msg(msg_code, data) return self._recv_msg(expect) - def _non_connect_send_recv(self, msg_code, data=None): + def _non_connect_send_recv(self, msg_code, data=None, expect=None): """ Similar to self._send_recv, but doesn't try to initiate a connection, thus preventing an infinite loop. @@ -197,7 +195,7 @@ def _recv(self, msglen): toread = msglen while toread: nbytes = self._socket.recv_into(view, toread) - view = view[nbytes:] # slicing views is cheap + view = view[nbytes:] # slicing views is cheap toread -= nbytes nread += nbytes if nread != msglen: diff --git a/riak/transports/tcp/transport.py b/riak/transports/tcp/transport.py index 9679cdd7..c6072f82 100644 --- a/riak/transports/tcp/transport.py +++ b/riak/transports/tcp/transport.py @@ -1,4 +1,5 @@ # TODO RTS-842 codecs should return msg codes too +import erlastic import six import riak.pb.messages @@ -113,7 +114,7 @@ def get(self, robj, r=None, pr=None, timeout=None, basic_quorum=None, """ codec = self._get_codec(ttb_supported=False) data = codec._encode_get(robj, r, pr, - timeout, basic_quorum, notfound_ok) + timeout, basic_quorum, notfound_ok) msg_code, resp = self._request( riak.pb.messages.MSG_CODE_GET_REQ, data, riak.pb.messages.MSG_CODE_GET_RESP) @@ -123,7 +124,7 @@ def put(self, robj, w=None, dw=None, pw=None, return_body=True, if_none_match=False, timeout=None): codec = self._get_codec(ttb_supported=False) data = codec._encode_put(robj, w, dw, pw, return_body, - if_none_match, timeout) + if_none_match, timeout) msg_code, resp = self._request( riak.pb.messages.MSG_CODE_PUT_REQ, data, riak.pb.messages.MSG_CODE_PUT_RESP) @@ -189,7 +190,7 @@ def ts_stream_keys(self, table, timeout=None): yields lists of keys. """ codec = self._get_codec(ttb_supported=False) - data = codec._encode_timeseries_listkeysreq(table, t) + data = codec._encode_timeseries_listkeysreq(table, timeout) self._send_msg(riak.pb.messages.MSG_CODE_TS_LIST_KEYS_REQ, data) return PbufTsKeyStream(self) @@ -346,8 +347,9 @@ def get_index(self, bucket, index, startkey, endkey=None, codec = self._get_codec(ttb_supported=False) data = codec._encode_index_req(bucket, index, startkey, endkey, - return_terms, max_results, continuation, - timeout, term_regex, streaming=False) + return_terms, max_results, + continuation, timeout, term_regex, + streaming=False) msg_code, resp = self._request( riak.pb.messages.MSG_CODE_INDEX_REQ, data, @@ -378,8 +380,9 @@ def stream_index(self, bucket, index, startkey, endkey=None, "supported") codec = self._get_codec(ttb_supported=False) data = codec._encode_index_req(bucket, index, startkey, endkey, - return_terms, max_results, continuation, - timeout, term_regex, streaming=True) + return_terms, max_results, + continuation, timeout, + term_regex, streaming=True) self._send_msg(riak.pb.messages.MSG_CODE_INDEX_REQ, data) return PbufIndexStream(self, index, return_terms) From 5dff949754424a85e6fbc646b99534ef9a286482 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 30 Mar 2016 15:52:09 -0700 Subject: [PATCH 31/44] Fix message encoding for Auth and ToggleEncoding PB requests --- riak/transports/tcp/connection.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/riak/transports/tcp/connection.py b/riak/transports/tcp/connection.py index 776ef4a6..0f414bee 100644 --- a/riak/transports/tcp/connection.py +++ b/riak/transports/tcp/connection.py @@ -83,9 +83,10 @@ def _enable_ttb(self): logging.debug("tcp/connection enabling TTB") req = riak.pb.riak_pb2.RpbToggleEncodingReq() req.use_native = True + data = req.SerializeToString() msg_code, _ = self._non_connect_send_recv( riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_REQ, - req, + data, riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_RESP) if msg_code == riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_RESP: self._ttb_enabled = True @@ -107,9 +108,10 @@ def _auth(self): if not password: password = '' req.password = str_to_bytes(password) + data = req.SerializeToString() msg_code, _ = self._non_connect_send_recv( riak.pb.messages.MSG_CODE_AUTH_REQ, - req, + data, riak.pb.messages.MSG_CODE_AUTH_RESP) if msg_code == riak.pb.messages.MSG_CODE_AUTH_RESP: return True From 3ac26e86db859d70d9e97a24332794798ee5f6f2 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 30 Mar 2016 16:29:00 -0700 Subject: [PATCH 32/44] Added Msg named tuple to contain data between codecs and transport layer. --- riak/codecs/__init__.py | 4 ++++ riak/codecs/pbuf.py | 11 ++++++++- riak/transports/tcp/connection.py | 40 +++++++++++++++---------------- riak/transports/tcp/stream.py | 10 ++++---- riak/transports/tcp/transport.py | 34 +++++++++++++++----------- 5 files changed, 60 insertions(+), 39 deletions(-) diff --git a/riak/codecs/__init__.py b/riak/codecs/__init__.py index e69de29b..479675d4 100644 --- a/riak/codecs/__init__.py +++ b/riak/codecs/__init__.py @@ -0,0 +1,4 @@ +import collections + +Msg = collections.namedtuple('Msg', + ['msg_code', 'data', 'resp_code'], verbose=False) diff --git a/riak/codecs/pbuf.py b/riak/codecs/pbuf.py index 44f5889b..7e860e28 100644 --- a/riak/codecs/pbuf.py +++ b/riak/codecs/pbuf.py @@ -1,11 +1,12 @@ import datetime -import riak.pb +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 Msg from riak.content import RiakContent from riak.riak_object import VClock from riak.util import decode_index_value, str_to_bytes, bytes_to_str, \ @@ -93,6 +94,14 @@ def _unix_time_millis(self, dt): def _datetime_from_unix_time_millis(self, ut): return datetime_from_unix_time_millis(ut) + def _encode_auth(self, username, password): + req = riak.pb.riak_pb2.RpbAuthReq() + req.user = str_to_bytes(username) + req.password = str_to_bytes(password) + return Msg(riak.pb.messages.MSG_CODE_AUTH_REQ, + req.SerializeToString(), + riak.pb.messages.MSG_CODE_AUTH_RESP) + def _encode_quorum(self, rw): """ Converts a symbolic quorum value into its on-the-wire diff --git a/riak/transports/tcp/connection.py b/riak/transports/tcp/connection.py index 0f414bee..d888cd6a 100644 --- a/riak/transports/tcp/connection.py +++ b/riak/transports/tcp/connection.py @@ -6,6 +6,7 @@ import riak.pb.messages from riak import RiakError +from riak.codecs.pbuf import PbufCodec from riak.security import SecurityError, USE_STDLIB_SSL from riak.util import str_to_bytes @@ -30,17 +31,21 @@ def _encode_msg(self, msg_code, data=None): hdr = struct.pack("!iB", 1 + len(data), msg_code) return hdr + data - def _send_recv(self, msg_code, data=None, expect=None): + def _send_recv(self, msg_code, data=None): self._send_msg(msg_code, data) - return self._recv_msg(expect) + return self._recv_msg() - def _non_connect_send_recv(self, msg_code, data=None, expect=None): + def _non_connect_send_recv(self, msg_code, data=None): """ 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, data) - return self._recv_msg(expect) + return self._recv_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): """ @@ -69,9 +74,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_send_recv( + 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 @@ -84,11 +89,10 @@ def _enable_ttb(self): req = riak.pb.riak_pb2.RpbToggleEncodingReq() req.use_native = True data = req.SerializeToString() - msg_code, _ = self._non_connect_send_recv( + resp_code, _ = self._non_connect_send_recv( riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_REQ, - data, - riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_RESP) - if msg_code == riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_RESP: + data) + if resp_code == riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_RESP: self._ttb_enabled = True logging.debug("tcp/connection TTB IS ENABLED") return True @@ -102,18 +106,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) + c = PbufCodec() + username = self._client._credentials.username password = self._client._credentials.password if not password: password = '' - req.password = str_to_bytes(password) - data = req.SerializeToString() - msg_code, _ = self._non_connect_send_recv( - riak.pb.messages.MSG_CODE_AUTH_REQ, - data, - riak.pb.messages.MSG_CODE_AUTH_RESP) - if msg_code == riak.pb.messages.MSG_CODE_AUTH_RESP: + msg = c._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 @@ -174,7 +174,7 @@ def _ssl_handshake(self): # fail if *any* exceptions are thrown during SSL handshake raise SecurityError(e) - def _recv_msg(self, expect=None): + def _recv_msg(self): msgbuf = self._recv_pkt() mv = memoryview(msgbuf) msg_code, = struct.unpack("B", mv[0:1]) diff --git a/riak/transports/tcp/stream.py b/riak/transports/tcp/stream.py index 2288c607..8c290e7c 100644 --- a/riak/transports/tcp/stream.py +++ b/riak/transports/tcp/stream.py @@ -29,10 +29,12 @@ def next(self): raise StopIteration try: - expected_code = self._expect - msg_code, data = self.transport._recv_msg(expect=expected_code) - self.transport._maybe_riak_error(msg_code, data) - resp = self.transport._parse_msg(expected_code, data, is_ttb=False) + # TODO RTS-842 - should be part of passed-in codec + resp_code, data = self.transport._recv_msg() + self.transport._maybe_riak_error(resp_code, data) + expect = self._expect + self.transport._maybe_incorrect_code(resp_code, expect) + resp = self.transport._parse_msg(expect, data, is_ttb=False) except: self.finished = True raise diff --git a/riak/transports/tcp/transport.py b/riak/transports/tcp/transport.py index c6072f82..e688da41 100644 --- a/riak/transports/tcp/transport.py +++ b/riak/transports/tcp/transport.py @@ -71,8 +71,11 @@ def ping(self): """ Ping the remote server """ - msg_code, _ = self._request(riak.pb.messages.MSG_CODE_PING_REQ) - if msg_code == riak.pb.messages.MSG_CODE_PING_RESP: + resp_code, _ = self._request( + riak.pb.messages.MSG_CODE_PING_REQ, + None, + riak.pb.messages.MSG_CODE_PING_RESP) + if resp_code == riak.pb.messages.MSG_CODE_PING_RESP: return True else: return False @@ -575,24 +578,27 @@ def _maybe_riak_error(self, msg_code, data=None, is_ttb=False): if msg_code is riak.pb.messages.MSG_CODE_ERROR_RESP: if data is None: raise RiakError('no error provided!') + # TODO RTS-842 TTB-specific version err = self._parse_msg(msg_code, data, is_ttb) if err is None: raise RiakError('no error provided!') else: raise RiakError(bytes_to_str(err.errmsg)) + 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)) + # TODO RTS-842 is_ttb def _request(self, msg_code, data=None, expect=None, is_ttb=False): - msg_code, data = self._send_recv(msg_code, data, expect) - self._maybe_riak_error(msg_code, data, is_ttb) - if msg_code in riak.pb.messages.MESSAGE_CLASSES: - msg = self._parse_msg(msg_code, data, is_ttb) + resp_code, data = self._send_recv(msg_code, data) + self._maybe_riak_error(resp_code, data, is_ttb) + self._maybe_incorrect_code(resp_code, expect) + if resp_code in riak.pb.messages.MESSAGE_CLASSES: + msg = self._parse_msg(resp_code, data, is_ttb) 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)) - # logging.debug("tcp/connection received msg_code %d msg %s", - # msg_code, msg) - return msg_code, msg + raise Exception("unknown msg code %s" % resp_code) + # logging.debug("tcp/connection received resp_code %d msg %s", + # resp_code, msg) + return resp_code, msg From 86511b442ab4064a80a6077adeb17fb65b5d0605 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 30 Mar 2016 16:32:25 -0700 Subject: [PATCH 33/44] make linters happy --- riak/codecs/__init__.py | 3 ++- riak/codecs/pbuf.py | 4 ++-- riak/transports/tcp/connection.py | 1 - 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/riak/codecs/__init__.py b/riak/codecs/__init__.py index 479675d4..8aab82f0 100644 --- a/riak/codecs/__init__.py +++ b/riak/codecs/__init__.py @@ -1,4 +1,5 @@ import collections Msg = collections.namedtuple('Msg', - ['msg_code', 'data', 'resp_code'], verbose=False) + ['msg_code', 'data', 'resp_code'], + verbose=False) diff --git a/riak/codecs/pbuf.py b/riak/codecs/pbuf.py index 7e860e28..1ecb35d1 100644 --- a/riak/codecs/pbuf.py +++ b/riak/codecs/pbuf.py @@ -99,8 +99,8 @@ def _encode_auth(self, username, password): req.user = str_to_bytes(username) req.password = str_to_bytes(password) return Msg(riak.pb.messages.MSG_CODE_AUTH_REQ, - req.SerializeToString(), - riak.pb.messages.MSG_CODE_AUTH_RESP) + req.SerializeToString(), + riak.pb.messages.MSG_CODE_AUTH_RESP) def _encode_quorum(self, rw): """ diff --git a/riak/transports/tcp/connection.py b/riak/transports/tcp/connection.py index d888cd6a..92802009 100644 --- a/riak/transports/tcp/connection.py +++ b/riak/transports/tcp/connection.py @@ -8,7 +8,6 @@ from riak import RiakError from riak.codecs.pbuf import PbufCodec from riak.security import SecurityError, USE_STDLIB_SSL -from riak.util import str_to_bytes if not USE_STDLIB_SSL: from OpenSSL.SSL import Connection From 47eaee13766dea9fd7b0b5d6b81b39b95093c625 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Thu, 31 Mar 2016 09:52:41 -0700 Subject: [PATCH 34/44] Move more code into the codecs. --- buildbot/Makefile | 1 - riak/codecs/pbuf.py | 185 ++++++++++----- riak/tests/test_timeseries.py | 16 +- riak/transports/tcp/transport.py | 388 +++++++++++++++---------------- 4 files changed, 321 insertions(+), 269 deletions(-) diff --git a/buildbot/Makefile b/buildbot/Makefile index 8272fdf1..abf4863f 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 diff --git a/riak/codecs/pbuf.py b/riak/codecs/pbuf.py index 1ecb35d1..9432d3c9 100644 --- a/riak/codecs/pbuf.py +++ b/riak/codecs/pbuf.py @@ -1,4 +1,6 @@ import datetime +import six + import riak.pb.messages import riak.pb.riak_pb2 import riak.pb.riak_dt_pb2 @@ -14,8 +16,6 @@ from riak.multidict import MultiDict from riak.pb.riak_ts_pb2 import TsColumnType -from six import string_types, PY2 - def _invert(d): out = {} @@ -88,19 +88,13 @@ def __init__(self, self._tombstone_vclocks = tombstone_vclocks self._bucket_types = bucket_types - def _unix_time_millis(self, dt): - return unix_time_millis(dt) - - def _datetime_from_unix_time_millis(self, ut): - return datetime_from_unix_time_millis(ut) - def _encode_auth(self, username, password): req = riak.pb.riak_pb2.RpbAuthReq() req.user = str_to_bytes(username) req.password = str_to_bytes(password) - return Msg(riak.pb.messages.MSG_CODE_AUTH_REQ, - req.SerializeToString(), - riak.pb.messages.MSG_CODE_AUTH_RESP) + mc = riak.pb.messages.MSG_CODE_AUTH_REQ + rc = riak.pb.messages.MSG_CODE_AUTH_RESP + return Msg(mc, req.SerializeToString(), rc) def _encode_quorum(self, rw): """ @@ -233,7 +227,7 @@ def _encode_content(self, robj, rpb_content): pair.value = str_to_bytes(str(value)) # Python 2.x data is stored in a string - if PY2: + if six.PY2: rpb_content.value = str(robj.encoded_data) else: rpb_content.value = robj.encoded_data @@ -287,7 +281,7 @@ def _encode_bucket_props(self, props, msg): """ for prop in NORMAL_PROPS: if prop in props and props[prop] is not None: - if isinstance(props[prop], string_types): + if isinstance(props[prop], six.string_types): setattr(msg.props, prop, str_to_bytes(props[prop])) else: setattr(msg.props, prop, props[prop]) @@ -302,7 +296,7 @@ def _encode_bucket_props(self, props, msg): 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): + if isinstance(value, six.string_types): setattr(msg.props, prop, str_to_bytes(value)) else: setattr(msg.props, prop, value) @@ -477,7 +471,25 @@ def _encode_index_req(self, bucket, index, startkey, endkey=None, if term_regex: req.term_regex = str_to_bytes(term_regex) req.stream = streaming - return req.SerializeToString() + 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): """ @@ -528,7 +540,7 @@ def _encode_search_query(self, req, **kwargs): def _decode_search_doc(self, doc): resultdoc = MultiDict() for pair in doc.fields: - if PY2: + if six.PY2: ukey = unicode(pair.key, 'utf-8') # noqa uval = unicode(pair.value, 'utf-8') # noqa else: @@ -648,11 +660,11 @@ def _encode_to_ts_cell(self, cell, ts_cell): ts_cell.timestamp_value = unix_time_millis(cell) elif isinstance(cell, bool): ts_cell.boolean_value = cell - elif isinstance(cell, string_types): + elif isinstance(cell, six.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 + (six.PY2 and isinstance(cell, long))): # noqa # logging.debug("cell -> int/long: '%s'", cell) ts_cell.sint64_value = cell elif isinstance(cell, float): @@ -678,14 +690,21 @@ def _encode_timeseries_keyreq(self, table, key, is_delete=False): for cell in key_vals: ts_cell = req.key.add() self._encode_to_ts_cell(cell, ts_cell) - return req.SerializeToString() + mc = riak.pb.messages.MSG_CODE_TS_GET_REQ + rc = riak.pb.messages.MSG_CODE_TS_GET_RESP + if is_delete: + mc = riak.pb.messages.MSG_CODE_TS_DEL_REQ + rc = riak.pb.messages.MSG_CODE_TS_DEL_RESP + 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 - return req.SerializeToString() + 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 _encode_timeseries_put(self, tsobj): """ @@ -713,7 +732,10 @@ def _encode_timeseries_put(self, tsobj): self._encode_to_ts_cell(cell, tsc) else: raise RiakError("TsObject requires a list of rows") - return req.SerializeToString() + + 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() @@ -721,7 +743,9 @@ def _encode_timeseries_query(self, table, query, interpolations=None): if '{table}' in q: q = q.format(table=table.name) req.query.base = str_to_bytes(q) - return req.SerializeToString() + 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): """ @@ -825,7 +849,9 @@ def _encode_get(self, robj, r=None, pr=None, timeout=None, req.bucket = str_to_bytes(bucket.name) self._add_bucket_type(req, bucket.bucket_type) req.key = str_to_bytes(robj.key) - return req.SerializeToString() + 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, @@ -851,7 +877,9 @@ def _encode_put(self, robj, w=None, dw=None, pw=None, if robj.vclock: req.vclock = robj.vclock.encode('binary') self._encode_content(robj, req.content) - return req.SerializeToString() + 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: @@ -909,7 +937,9 @@ def _encode_delete(self, robj, rw=None, r=None, req.bucket = str_to_bytes(bucket.name) self._add_bucket_type(req, bucket.bucket_type) req.key = str_to_bytes(robj.key) - return req.SerializeToString() + 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() @@ -917,7 +947,9 @@ def _encode_stream_keys(self, bucket, timeout=None): if self._client_timeouts and timeout: req.timeout = timeout self._add_bucket_type(req, bucket.bucket_type) - return req.SerializeToString() + 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 = [] @@ -930,66 +962,81 @@ 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) - return req.SerializeToString() + 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): + 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 - return req.SerializeToString() - - def _encode_stream_buckets(self, bucket_type, timeout): - 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 - return req.SerializeToString() + 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) - return req.SerializeToString() + 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) - return req.SerializeToString() + 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) - return req.SerializeToString() + 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) - return req.SerializeToString() + 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) - return req.SerializeToString() + 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") - return req.SerializeToString() + 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): @@ -1002,21 +1049,29 @@ def _encode_create_search_index(self, index, schema=None, req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexPutReq(index=idx) if timeout is not None: req.timeout = timeout - return req.SerializeToString() + 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)) - return req.SerializeToString() + 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() - return req.SerializeToString() + 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)) - return req.SerializeToString() + 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( @@ -1024,12 +1079,16 @@ def _encode_create_search_schema(self, schema, content): content=str_to_bytes(content)) req = riak.pb.riak_yokozuna_pb2.RpbYokozunaSchemaPutReq( schema=scma) - return req.SerializeToString() + 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)) - return req.SerializeToString() + 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 = {} @@ -1042,7 +1101,9 @@ def _encode_search(self, index, query, **kwargs): index=str_to_bytes(index), q=str_to_bytes(query)) self._encode_search_query(req, **kwargs) - return req.SerializeToString() + 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 = {} @@ -1065,7 +1126,9 @@ def _encode_get_counter(self, bucket, key, **kwargs): req.basic_quorum = kwargs['basic_quorum'] if kwargs.get('notfound_ok') is not None: req.notfound_ok = kwargs['notfound_ok'] - return req.SerializeToString() + 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() @@ -1080,7 +1143,9 @@ def _encode_update_counter(self, bucket, key, value, **kwargs): req.pw = self._encode_quorum(kwargs['pw']) if kwargs.get('returnvalue') is not None: req.returnvalue = kwargs['returnvalue'] - return req.SerializeToString() + 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() @@ -1088,7 +1153,9 @@ def _encode_fetch_datatype(self, bucket, key, **kwargs): req.bucket = str_to_bytes(bucket.name) req.key = str_to_bytes(key) self._encode_dt_options(req, **kwargs) - return req.SerializeToString() + 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() @@ -1105,7 +1172,9 @@ def _encode_update_datatype(self, datatype, **kwargs): req.context = datatype._context self._encode_dt_options(req, **kwargs) self._encode_dt_op(type_name, req, op) - return req.SerializeToString() + 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 @@ -1121,4 +1190,6 @@ def _encode_get_preflist(self, bucket, key): req.bucket = str_to_bytes(bucket.name) req.key = str_to_bytes(key) req.type = str_to_bytes(bucket.bucket_type.name) - return req.SerializeToString() + 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/tests/test_timeseries.py b/riak/tests/test_timeseries.py index 30210099..7c5d68fb 100644 --- a/riak/tests/test_timeseries.py +++ b/riak/tests/test_timeseries.py @@ -67,26 +67,26 @@ def test_encode_decode_timestamp(self): def test_encode_data_for_get(self): c = PbufCodec() - data = c._encode_timeseries_keyreq( + msg = c._encode_timeseries_keyreq( self.table, self.test_key, is_delete=False) req = riak.pb.riak_ts_pb2.TsGetReq() - req.ParseFromString(data) + req.ParseFromString(msg.data) self.validate_keyreq(req) def test_encode_data_for_delete(self): c = PbufCodec() - data = c._encode_timeseries_keyreq( + msg = c._encode_timeseries_keyreq( self.table, self.test_key, is_delete=True) req = riak.pb.riak_ts_pb2.TsDelReq() - req.ParseFromString(data) + 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) - data = c._encode_timeseries_put(tsobj) + msg = c._encode_timeseries_put(tsobj) req = riak.pb.riak_ts_pb2.TsPutReq() - req.ParseFromString(data) + req.ParseFromString(msg.data) # NB: expected, actual self.assertEqual(self.table.name, bytes_to_str(req.table)) @@ -110,9 +110,9 @@ def test_encode_data_for_put(self): def test_encode_data_for_listkeys(self): c = PbufCodec(client_timeouts=True) - data = c._encode_timeseries_listkeysreq(self.table, 1234) + msg = c._encode_timeseries_listkeysreq(self.table, 1234) req = riak.pb.riak_ts_pb2.TsListKeysReq() - req.ParseFromString(data) + req.ParseFromString(msg.data) self.assertEqual(self.table.name, bytes_to_str(req.table)) self.assertEqual(1234, req.timeout) diff --git a/riak/transports/tcp/transport.py b/riak/transports/tcp/transport.py index e688da41..c2a9c053 100644 --- a/riak/transports/tcp/transport.py +++ b/riak/transports/tcp/transport.py @@ -4,13 +4,14 @@ import riak.pb.messages from riak import RiakError +from riak.codecs import Msg from riak.codecs.pbuf import PbufCodec from riak.codecs.ttb import TtbCodec from riak.transports.transport import Transport from riak.ts_object import TsObject # TODO RTS-842 ideally these would not be needed -from riak.util import decode_index_value, bytes_to_str +from riak.util import bytes_to_str from riak.transports.tcp.connection import TcpConnection from riak.transports.tcp.stream import (PbufKeyStream, @@ -48,16 +49,24 @@ def _get_pbuf_codec(self): self.tombstone_vclocks(), self.bucket_types()) return self._pbuf_c - def _get_codec(self, ttb_supported=False): - if ttb_supported: - if self._use_ttb: - if not self._enable_ttb(): - raise RiakError('could not switch to TTB encoding!') - if not self._ttb_c: - self._ttb_c = TtbCodec() - codec = self._ttb_c - else: - codec = self._get_pbuf_codec() + def _get_ttb_codec(self): + if self._use_ttb: + if not self._enable_ttb(): + raise RiakError('could not switch to TTB encoding!') + 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 == 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_DEL_REQ: + codec = self._get_ttb_codec() else: codec = self._get_pbuf_codec() return codec @@ -71,10 +80,9 @@ def ping(self): """ Ping the remote server """ - resp_code, _ = self._request( - riak.pb.messages.MSG_CODE_PING_REQ, - None, - riak.pb.messages.MSG_CODE_PING_RESP) + msg = Msg(riak.pb.messages.MSG_CODE_PING_REQ, None, + riak.pb.messages.MSG_CODE_PING_RESP) + resp_code, _ = self._request(msg) if resp_code == riak.pb.messages.MSG_CODE_PING_RESP: return True else: @@ -87,24 +95,23 @@ def get_server_info(self): # NB: can't do it this way due to recursion # codec = self._get_codec(ttb_supported=False) codec = PbufCodec() - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_SERVER_INFO_REQ, - expect=riak.pb.messages.MSG_CODE_GET_SERVER_INFO_RESP) + 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) return codec._decode_get_server_info(resp) def _get_client_id(self): - codec = self._get_codec(ttb_supported=False) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_CLIENT_ID_REQ, - expect=riak.pb.messages.MSG_CODE_GET_CLIENT_ID_RESP) + 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) return codec._decode_get_client_id(resp) def _set_client_id(self, client_id): - codec = self._get_codec(ttb_supported=False) - data = codec._encode_set_client_id(client_id) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_SET_CLIENT_ID_REQ, data, - riak.pb.messages.MSG_CODE_SET_CLIENT_ID_RESP) + 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) self._client_id = client_id client_id = property(_get_client_id, _set_client_id, @@ -115,22 +122,21 @@ def get(self, robj, r=None, pr=None, timeout=None, basic_quorum=None, """ Serialize get request and deserialize response """ - codec = self._get_codec(ttb_supported=False) - data = codec._encode_get(robj, r, pr, - timeout, basic_quorum, notfound_ok) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_REQ, data, - riak.pb.messages.MSG_CODE_GET_RESP) + 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) 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): - codec = self._get_codec(ttb_supported=False) - data = codec._encode_put(robj, w, dw, pw, return_body, - if_none_match, timeout) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_PUT_REQ, data, - riak.pb.messages.MSG_CODE_PUT_RESP) + 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) return codec._decode_put(robj, resp) def ts_describe(self, table): @@ -138,28 +144,24 @@ def ts_describe(self, table): return self.ts_query(table, query) def ts_get(self, table, key): - codec = self._get_codec(ttb_supported=True) - data = codec._encode_timeseries_keyreq(table, key) - msg_code, ts_get_resp = self._request( - riak.pb.messages.MSG_CODE_TS_GET_REQ, data, - riak.pb.messages.MSG_CODE_TS_GET_RESP) + msg_code = riak.pb.messages.MSG_CODE_TS_GET_REQ + codec = self._get_codec(msg_code) + msg = codec._encode_timeseries_keyreq(table, key) + resp_code, resp = self._request(msg) tsobj = TsObject(self._client, table, [], None) - codec._decode_timeseries(ts_get_resp, tsobj) + codec._decode_timeseries(resp, tsobj) return tsobj def ts_put(self, tsobj): - codec = self._get_codec(ttb_supported=True) - # TODO RTS-842 codecs should return msg codes too - data = codec._encode_timeseries_put(tsobj) + msg_code = riak.pb.messages.MSG_CODE_TS_PUT_REQ + codec = self._get_codec(msg_code) + msg = codec._encode_timeseries_put(tsobj) # logging.debug("pbc/transport ts_put _use_ttb: '%s'", # self._use_ttb) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_TS_PUT_REQ, data, - riak.pb.messages.MSG_CODE_TS_PUT_RESP, - self._use_ttb) + resp_code, resp = self._request(msg, self._use_ttb) if self._use_ttb and \ resp is None and \ - msg_code == riak.pb.messages.MSG_CODE_TS_PUT_RESP: + resp_code == riak.pb.messages.MSG_CODE_TS_PUT_RESP: return True if resp is not None: return True @@ -167,24 +169,22 @@ def ts_put(self, tsobj): raise RiakError("missing response object") def ts_delete(self, table, key): - codec = self._get_codec(ttb_supported=True) - data = codec._encode_timeseries_keyreq(table, key, is_delete=True) - msg_code, ts_del_resp = self._request( - riak.pb.messages.MSG_CODE_TS_DEL_REQ, data, - riak.pb.messages.MSG_CODE_TS_DEL_RESP) - if ts_del_resp is not None: + 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) + if resp is not None: return True else: raise RiakError("missing response object") def ts_query(self, table, query, interpolations=None): - codec = self._get_codec(ttb_supported=True) - data = codec._encode_timeseries_query(table, query, interpolations) - msg_code, ts_query_resp = self._request( - riak.pb.messages.MSG_CODE_TS_QUERY_REQ, data, - riak.pb.messages.MSG_CODE_TS_QUERY_RESP) + 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) tsobj = TsObject(self._client, table, [], []) - self._decode_timeseries(ts_query_resp, tsobj) + self._decode_timeseries(resp, tsobj) return tsobj def ts_stream_keys(self, table, timeout=None): @@ -192,25 +192,26 @@ def ts_stream_keys(self, table, timeout=None): Streams keys from a timeseries table, returning an iterator that yields lists of keys. """ - codec = self._get_codec(ttb_supported=False) - data = codec._encode_timeseries_listkeysreq(table, timeout) - self._send_msg(riak.pb.messages.MSG_CODE_TS_LIST_KEYS_REQ, data) + 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) def delete(self, robj, rw=None, r=None, w=None, dw=None, pr=None, pw=None, timeout=None): - codec = self._get_codec(ttb_supported=False) - data = codec._encode_delete(robj, rw, r, w, dw, pr, pw, timeout) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_DEL_REQ, data, - riak.pb.messages.MSG_CODE_DEL_RESP) + 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) return self def get_keys(self, bucket, timeout=None): """ Lists all keys within a bucket. """ - codec = self._get_codec(ttb_supported=False) + 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) @@ -219,20 +220,21 @@ def stream_keys(self, bucket, timeout=None): Streams keys from a bucket, returning an iterator that yields lists of keys. """ - codec = self._get_codec(ttb_supported=False) - data = codec._encode_stream_keys(bucket, timeout) - self._send_msg(riak.pb.messages.MSG_CODE_LIST_KEYS_REQ, data) + 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) def get_buckets(self, bucket_type=None, timeout=None): """ Serialize bucket listing request and deserialize response """ - codec = self._get_codec(ttb_supported=False) - data = codec._encode_get_buckets(bucket_type, timeout) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_LIST_BUCKETS_REQ, data, - riak.pb.messages.MSG_CODE_LIST_BUCKETS_RESP) + 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) return resp.buckets def stream_buckets(self, bucket_type=None, timeout=None): @@ -242,20 +244,21 @@ def stream_buckets(self, bucket_type=None, timeout=None): if not self.bucket_stream(): raise NotImplementedError('Streaming list-buckets is not ' 'supported') - codec = self._get_codec(ttb_supported=False) - data = codec._encode_stream_buckets(bucket_type, timeout) - self._send_msg(riak.pb.messages.MSG_CODE_LIST_BUCKETS_REQ, data) + 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) def get_bucket_props(self, bucket): """ Serialize bucket property request and deserialize response """ - codec = self._get_codec(ttb_supported=False) - data = codec._encode_get_bucket_props(bucket) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_BUCKET_REQ, data, - riak.pb.messages.MSG_CODE_GET_BUCKET_RESP) + 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) return codec._decode_bucket_props(resp.props) def set_bucket_props(self, bucket, props): @@ -267,11 +270,10 @@ def set_bucket_props(self, bucket, props): if key not in ('n_val', 'allow_mult'): raise NotImplementedError('Server only supports n_val and ' 'allow_mult properties over PBC') - codec = self._get_codec(ttb_supported=False) - data = codec._encode_set_bucket_props(bucket, props) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_SET_BUCKET_REQ, data, - riak.pb.messages.MSG_CODE_SET_BUCKET_RESP) + 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) return True def clear_bucket_props(self, bucket): @@ -280,11 +282,10 @@ def clear_bucket_props(self, bucket): """ if not self.pb_clear_bucket_props(): return False - codec = self._get_codec(ttb_supported=False) - data = codec._encode_clear_bucket_props(bucket) - self._request( - riak.pb.messages.MSG_CODE_RESET_BUCKET_REQ, data, - riak.pb.messages.MSG_CODE_RESET_BUCKET_RESP) + 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) return True def get_bucket_type_props(self, bucket_type): @@ -292,11 +293,10 @@ def get_bucket_type_props(self, bucket_type): Fetch bucket-type properties """ self._check_bucket_types(bucket_type) - codec = self._get_codec(ttb_supported=False) - data = codec._encode_get_bucket_type_props(bucket_type) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_BUCKET_TYPE_REQ, data, - riak.pb.messages.MSG_CODE_GET_BUCKET_RESP) + 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) return codec._decode_bucket_props(resp.props) def set_bucket_type_props(self, bucket_type, props): @@ -304,11 +304,10 @@ def set_bucket_type_props(self, bucket_type, props): Set bucket-type properties """ self._check_bucket_types(bucket_type) - codec = self._get_codec(ttb_supported=False) - data = codec._encode_set_bucket_type_props(bucket_type, props) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_SET_BUCKET_TYPE_REQ, data, - riak.pb.messages.MSG_CODE_SET_BUCKET_RESP) + 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) return True def mapred(self, inputs, query, timeout=None): @@ -331,10 +330,11 @@ def mapred(self, inputs, query, timeout=None): 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) - codec = self._get_codec(ttb_supported=False) - data = codec._encode_stream_mapred(content) - self._send_msg(riak.pb.messages.MSG_CODE_MAP_RED_REQ, data) + msg = codec._encode_stream_mapred(content) + self._send_msg(msg.msg_code, msg.data) return PbufMapredStream(self) def get_index(self, bucket, index, startkey, endkey=None, @@ -348,29 +348,15 @@ def get_index(self, bucket, index, startkey, endkey=None, raise NotImplementedError("Secondary index term_regex is not " "supported") - codec = self._get_codec(ttb_supported=False) - data = codec._encode_index_req(bucket, index, startkey, endkey, - return_terms, max_results, - continuation, timeout, term_regex, - streaming=False) - - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_INDEX_REQ, data, - 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 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) + 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) + 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, @@ -381,12 +367,13 @@ def stream_index(self, bucket, index, startkey, endkey=None, if term_regex and not self.index_term_regex(): raise NotImplementedError("Secondary index term_regex is not " "supported") - codec = self._get_codec(ttb_supported=False) - data = codec._encode_index_req(bucket, index, startkey, endkey, - return_terms, max_results, - continuation, timeout, - term_regex, streaming=True) - self._send_msg(riak.pb.messages.MSG_CODE_INDEX_REQ, data) + 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, index, return_terms) def create_search_index(self, index, schema=None, n_val=None, @@ -394,22 +381,20 @@ def create_search_index(self, index, schema=None, n_val=None, if not self.pb_search_admin(): raise NotImplementedError("Search 2.0 administration is not " "supported for this version") - codec = self._get_codec(ttb_supported=False) - data = codec._encode_create_search_index(index, schema, n_val, timeout) - self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_PUT_REQ, data, - riak.pb.messages.MSG_CODE_PUT_RESP) + 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) 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") - codec = self._get_codec(ttb_supported=False) - data = codec._encode_get_search_index(index) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_REQ, data, - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_RESP) + 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) if len(resp.index) > 0: return codec._decode_search_index(resp.index[0]) else: @@ -419,44 +404,40 @@ def list_search_indexes(self): if not self.pb_search_admin(): raise NotImplementedError("Search 2.0 administration is not " "supported for this version") - codec = self._get_codec(ttb_supported=False) - data = codec._encode_list_search_indexes() - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_REQ, data, - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_RESP) + 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) 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") - codec = self._get_codec(ttb_supported=False) - data = codec._encode_delete_search_index(index) - self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_DELETE_REQ, data, - riak.pb.messages.MSG_CODE_DEL_RESP) + 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) 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") - codec = self._get_codec(ttb_supported=False) - data = codec._encode_create_search_schema(schema, content) - self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_PUT_REQ, data, - riak.pb.messages.MSG_CODE_PUT_RESP) + 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) 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") - codec = self._get_codec(ttb_supported=False) - data = codec._encode_get_search_schema(schema) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_GET_REQ, data, - riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_GET_RESP) + 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) return codec._decode_get_search_schema(resp) def search(self, index, query, **kwargs): @@ -466,11 +447,10 @@ def search(self, index, query, **kwargs): # TODO RTS-842 six.u() instead? if six.PY2 and isinstance(query, unicode): # noqa query = query.encode('utf8') - codec = self._get_codec(ttb_supported=False) - data = codec._encode_search(index, query, **kwargs) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_SEARCH_QUERY_REQ, data, - riak.pb.messages.MSG_CODE_SEARCH_QUERY_RESP) + 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) return codec._decode_search(resp) def get_counter(self, bucket, key, **kwargs): @@ -480,11 +460,10 @@ def get_counter(self, bucket, key, **kwargs): "use datatypes instead.") if not self.counters(): raise NotImplementedError("Counters are not supported") - codec = self._get_codec(ttb_supported=False) - data = codec._encode_get_counter(bucket, key, **kwargs) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_COUNTER_GET_REQ, data, - riak.pb.messages.MSG_CODE_COUNTER_GET_RESP) + 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) if resp.HasField('value'): return resp.value else: @@ -497,11 +476,10 @@ def update_counter(self, bucket, key, value, **kwargs): "use datatypes instead.") if not self.counters(): raise NotImplementedError("Counters are not supported") - codec = self._get_codec(ttb_supported=False) - data = codec._encode_update_counter(bucket, key, value, **kwargs) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_COUNTER_UPDATE_REQ, data, - riak.pb.messages.MSG_CODE_COUNTER_UPDATE_RESP) + 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) if resp.HasField('value'): return resp.value else: @@ -513,11 +491,10 @@ def fetch_datatype(self, bucket, key, **kwargs): " bucket-type.") if not self.datatypes(): raise NotImplementedError("Datatypes are not supported.") - codec = self._get_codec(ttb_supported=False) - data = codec._encode_fetch_datatype(bucket, key, **kwargs) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_DT_FETCH_REQ, data, - riak.pb.messages.MSG_CODE_DT_FETCH_RESP) + 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) return codec._decode_dt_fetch(resp) def update_datatype(self, datatype, **kwargs): @@ -526,11 +503,10 @@ def update_datatype(self, datatype, **kwargs): " bucket-type.") if not self.datatypes(): raise NotImplementedError("Datatypes are not supported.") - codec = self._get_codec(ttb_supported=False) - data = codec._encode_update_datatype(datatype, **kwargs) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_DT_UPDATE_REQ, data, - riak.pb.messages.MSG_CODE_DT_UPDATE_RESP) + 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._decode_update_datatype(datatype, resp, **kwargs) return True @@ -544,11 +520,10 @@ def get_preflist(self, bucket, key): :type key: string :rtype: list of dicts """ - codec = self._get_codec(ttb_supported=False) - data = codec._encode_get_preflist(bucket, key) - msg_code, resp = self._request( - riak.pb.messages.MSG_CODE_GET_BUCKET_KEY_PREFLIST_REQ, data, - riak.pb.messages.MSG_CODE_GET_BUCKET_KEY_PREFLIST_RESP) + 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) return [codec._decode_preflist(item) for item in resp.preflist] # TODO RTS-842 is_ttb @@ -574,6 +549,7 @@ def _parse_msg(self, code, packet, is_ttb=False): pbo.ParseFromString(packet) return pbo + # TODO RTS-842 move to base Codec object def _maybe_riak_error(self, msg_code, data=None, is_ttb=False): if msg_code is riak.pb.messages.MSG_CODE_ERROR_RESP: if data is None: @@ -591,7 +567,13 @@ def _maybe_incorrect_code(self, resp_code, expect=None): % (resp_code, expect)) # TODO RTS-842 is_ttb - def _request(self, msg_code, data=None, expect=None, is_ttb=False): + def _request(self, msg, is_ttb=False): + if isinstance(msg, Msg): + msg_code = msg.msg_code + data = msg.data + expect = msg.resp_code + else: + raise ValueError('expected a Msg argument') resp_code, data = self._send_recv(msg_code, data) self._maybe_riak_error(resp_code, data, is_ttb) self._maybe_incorrect_code(resp_code, expect) From da5c13a30f6e2c5b6c1ce118896737e6cb6b5250 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Thu, 31 Mar 2016 16:50:03 -0700 Subject: [PATCH 35/44] make timeseries tests happy --- buildbot/Makefile | 2 +- riak/codecs/pbuf.py | 24 +++++++++---------- riak/codecs/ttb.py | 40 ++++++++++++++++++++++--------- riak/tests/test_timeseries.py | 29 +++++++++++++++------- riak/tests/test_timeseries_ttb.py | 14 ++++------- riak/transports/tcp/connection.py | 2 +- riak/transports/tcp/stream.py | 6 ++++- riak/transports/tcp/transport.py | 10 ++++---- 8 files changed, 79 insertions(+), 48 deletions(-) diff --git a/buildbot/Makefile b/buildbot/Makefile index abf4863f..ecc933b8 100644 --- a/buildbot/Makefile +++ b/buildbot/Makefile @@ -45,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/riak/codecs/pbuf.py b/riak/codecs/pbuf.py index 9432d3c9..b55a4253 100644 --- a/riak/codecs/pbuf.py +++ b/riak/codecs/pbuf.py @@ -660,12 +660,13 @@ def _encode_to_ts_cell(self, cell, ts_cell): 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): - # logging.debug("cell -> str: '%s'", cell) ts_cell.varchar_value = str_to_bytes(cell) - elif (isinstance(cell, int) or - (six.PY2 and isinstance(cell, long))): # noqa - # logging.debug("cell -> int/long: '%s'", cell) + elif (isinstance(cell, six.integer_types)): ts_cell.sint64_value = cell elif isinstance(cell, float): ts_cell.double_value = cell @@ -681,20 +682,18 @@ def _encode_timeseries_keyreq(self, table, key, is_delete=False): 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() - else: - req = riak.pb.riak_ts_pb2.TsGetReq() + 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) - mc = riak.pb.messages.MSG_CODE_TS_GET_REQ - rc = riak.pb.messages.MSG_CODE_TS_GET_RESP - if is_delete: - mc = riak.pb.messages.MSG_CODE_TS_DEL_REQ - rc = riak.pb.messages.MSG_CODE_TS_DEL_RESP return Msg(mc, req.SerializeToString(), rc) def _encode_timeseries_listkeysreq(self, table, timeout=None): @@ -788,7 +787,8 @@ def _decode_timeseries_row(self, tsrow, tscols=None): if col and col.type != TsColumnType.Value('VARCHAR'): raise TypeError('expected VARCHAR column') else: - row.append(bytes_to_str(cell.varchar_value)) + # TODO RTS-842 - keep as bytes? + row.append(cell.varchar_value) elif cell.HasField('sint64_value'): if col and col.type != TsColumnType.Value('SINT64'): raise TypeError('expected SINT64 column') diff --git a/riak/codecs/ttb.py b/riak/codecs/ttb.py index 8d795675..eef65fdc 100644 --- a/riak/codecs/ttb.py +++ b/riak/codecs/ttb.py @@ -1,11 +1,13 @@ import datetime +import six + +import riak.pb.messages from erlastic import encode from erlastic.types import Atom -from six import text_type, binary_type, \ - string_types, PY2 from riak import RiakError +from riak.codecs import Msg from riak.util import unix_time_millis, \ datetime_from_unix_time_millis @@ -15,6 +17,7 @@ tsgetreq_a = Atom('tsgetreq') tsgetresp_a = Atom('tsgetresp') tsputreq_a = Atom('tsputreq') +tsdelreq_a = Atom('tsdelreq') tsrow_a = Atom('tsrow') tscell_a = Atom('tscell') @@ -38,13 +41,12 @@ def _encode_to_ts_cell(self, cell): return (tscell_a, udef_a, udef_a, ts, udef_a, udef_a) elif isinstance(cell, bool): return (tscell_a, udef_a, udef_a, udef_a, cell, udef_a) - elif isinstance(cell, text_type) or \ - isinstance(cell, binary_type) or \ - isinstance(cell, string_types): + elif isinstance(cell, six.text_type) or \ + isinstance(cell, six.binary_type) or \ + isinstance(cell, six.string_types): return (tscell_a, cell, udef_a, udef_a, udef_a, udef_a) - elif (isinstance(cell, int) or - (PY2 and isinstance(cell, long))): # noqa + elif (isinstance(cell, six.integer_types)): return (tscell_a, udef_a, cell, udef_a, udef_a, udef_a) elif isinstance(cell, float): return (tscell_a, udef_a, udef_a, udef_a, udef_a, cell) @@ -53,15 +55,24 @@ def _encode_to_ts_cell(self, cell): raise RiakError("can't serialize type '{}', value '{}'" .format(t, cell)) - def _encode_timeseries_keyreq(self, table, key): + 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 = tsgetreq_a, table.name, \ + + mc = riak.pb.messages.MSG_CODE_TS_GET_REQ + rc = riak.pb.messages.MSG_CODE_TS_GET_RESP + req_atom = tsgetreq_a + if is_delete: + mc = riak.pb.messages.MSG_CODE_TS_DEL_REQ + rc = riak.pb.messages.MSG_CODE_TS_DEL_RESP + req_atom = tsdelreq_a + + req = req_atom, table.name, \ [self._encode_to_ts_cell(k) for k in key_vals], udef_a - return encode(req) + return Msg(mc, encode(req), rc) def _encode_timeseries_put(self, tsobj): ''' @@ -84,7 +95,9 @@ def _encode_timeseries_put(self, tsobj): req_t = (tsrow_a, req_r) req_rows.append(req_t) req = tsputreq_a, tsobj.table.name, udef_a, req_rows - return encode(req) + mc = riak.pb.messages.MSG_CODE_TS_PUT_REQ + rc = riak.pb.messages.MSG_CODE_TS_PUT_RESP + return Msg(mc, encode(req), rc) else: raise RiakError("TsObject requires a list of rows") @@ -107,6 +120,11 @@ def _decode_timeseries(self, resp_ttb, tsobj): # col_type = col.type # col = (col_name, col_type) # tsobj.columns.append(col) + # + # TODO RTS-842 is this correct? + if resp_ttb is None: + return tsobj + resp_a = resp_ttb[0] if resp_a == tsgetresp_a: # TODO resp_cols = resp_ttb[1] diff --git a/riak/tests/test_timeseries.py b/riak/tests/test_timeseries.py index 7c5d68fb..e50d2613 100644 --- a/riak/tests/test_timeseries.py +++ b/riak/tests/test_timeseries.py @@ -179,14 +179,14 @@ def test_decode_data_from_query(self): self.assertEqual(c[4][1], TsColumnType.Value('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) @@ -229,17 +229,30 @@ def setUpClass(cls): 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_data(self, ts_obj): if ts_obj.columns is not None: self.assertEqual(len(ts_obj.columns), 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): @@ -359,7 +372,7 @@ def test_query_that_matches_more_data(self): t2=self.nowMsec) ts_obj = self.client.ts_query('GeoCheckin', 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] @@ -393,8 +406,8 @@ 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) diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index 959b0467..81c0554c 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -6,7 +6,6 @@ from erlastic import decode, encode from erlastic.types import Atom -from riak.client import RiakClient from riak.table import Table from riak.ts_object import TsObject from riak.codecs.ttb import TtbCodec @@ -53,8 +52,8 @@ def test_encode_data_for_get(self): test_key = ['hash1', 'user2', ts0] c = TtbCodec() - req_encoded = c._encode_timeseries_keyreq(self.table, test_key) - self.assertEqual(req_test, req_encoded) + msg = c._encode_timeseries_keyreq(self.table, test_key) + self.assertEqual(req_test, msg.data) # def test_decode_riak_error(self): @@ -146,8 +145,8 @@ def test_encode_data_for_put(self): tsobj = TsObject(None, self.table, rows_to_encode, None) c = TtbCodec() - req_encoded = c._encode_timeseries_put(tsobj) - self.assertEqual(req_test, req_encoded) + msg = c._encode_timeseries_put(tsobj) + self.assertEqual(req_test, msg.data) @unittest.skipUnless(is_timeseries_supported() and RUN_TIMESERIES, @@ -166,10 +165,7 @@ def test_store_and_fetch_ttb(self): twentyFiveMinsAgo = twentyMinsAgo - fiveMins opts = {'use_ttb': True} - client = RiakClient(protocol='pbc', - host='riak-test', - pb_port=10017, - transport_options=opts) + client = self.create_client(transport_options=opts) table = client.table(table_name) rows = [ diff --git a/riak/transports/tcp/connection.py b/riak/transports/tcp/connection.py index 92802009..8972e480 100644 --- a/riak/transports/tcp/connection.py +++ b/riak/transports/tcp/connection.py @@ -88,7 +88,7 @@ def _enable_ttb(self): req = riak.pb.riak_pb2.RpbToggleEncodingReq() req.use_native = True data = req.SerializeToString() - resp_code, _ = self._non_connect_send_recv( + resp_code, _ = self._send_recv( riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_REQ, data) if resp_code == riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_RESP: diff --git a/riak/transports/tcp/stream.py b/riak/transports/tcp/stream.py index 8c290e7c..46fb5941 100644 --- a/riak/transports/tcp/stream.py +++ b/riak/transports/tcp/stream.py @@ -174,6 +174,10 @@ class PbufTsKeyStream(PbufStream, TtbCodec): _expect = riak.pb.messages.MSG_CODE_TS_LIST_KEYS_RESP + def __init__(self, transport, codec): + super(PbufTsKeyStream, self).__init__(transport) + self._codec = codec + def next(self): response = super(PbufTsKeyStream, self).next() @@ -182,7 +186,7 @@ def next(self): 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 index c2a9c053..ee8898f2 100644 --- a/riak/transports/tcp/transport.py +++ b/riak/transports/tcp/transport.py @@ -65,8 +65,6 @@ def _get_codec(self, msg_code): 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_DEL_REQ: - codec = self._get_ttb_codec() else: codec = self._get_pbuf_codec() return codec @@ -147,7 +145,8 @@ def ts_get(self, table, key): msg_code = riak.pb.messages.MSG_CODE_TS_GET_REQ codec = self._get_codec(msg_code) msg = codec._encode_timeseries_keyreq(table, key) - resp_code, resp = self._request(msg) + # TODO RTS-842 is_ttb + resp_code, resp = self._request(msg, self._use_ttb) tsobj = TsObject(self._client, table, [], None) codec._decode_timeseries(resp, tsobj) return tsobj @@ -158,6 +157,7 @@ def ts_put(self, tsobj): msg = codec._encode_timeseries_put(tsobj) # logging.debug("pbc/transport ts_put _use_ttb: '%s'", # self._use_ttb) + # TODO RTS-842 use_ttb resp_code, resp = self._request(msg, self._use_ttb) if self._use_ttb and \ resp is None and \ @@ -184,7 +184,7 @@ def ts_query(self, table, query, interpolations=None): msg = codec._encode_timeseries_query(table, query, interpolations) resp_code, resp = self._request(msg) tsobj = TsObject(self._client, table, [], []) - self._decode_timeseries(resp, tsobj) + codec._decode_timeseries(resp, tsobj) return tsobj def ts_stream_keys(self, table, timeout=None): @@ -196,7 +196,7 @@ def ts_stream_keys(self, table, timeout=None): codec = self._get_codec(msg_code) msg = codec._encode_timeseries_listkeysreq(table, timeout) self._send_msg(msg.msg_code, msg.data) - return PbufTsKeyStream(self) + return PbufTsKeyStream(self, codec) def delete(self, robj, rw=None, r=None, w=None, dw=None, pr=None, pw=None, timeout=None): From 642e4b77e783902c334b5ecb7d1a4dcbce3f1f8e Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Fri, 1 Apr 2016 12:09:21 -0700 Subject: [PATCH 36/44] All code for PB and TTB encoding / decoding now in dedicated classes. --- riak/client/__init__.py | 6 +- riak/codecs/__init__.py | 20 ++ riak/codecs/pbuf.py | 302 ++++++++++++++++-------------- riak/codecs/ttb.py | 64 +++++-- riak/tests/base.py | 13 +- riak/tests/test_timeseries.py | 50 ++--- riak/tests/test_timeseries_ttb.py | 26 ++- riak/transports/tcp/connection.py | 4 +- riak/transports/tcp/stream.py | 20 +- riak/transports/tcp/transport.py | 273 +++++++++++---------------- 10 files changed, 410 insertions(+), 368 deletions(-) diff --git a/riak/client/__init__.py b/riak/client/__init__.py index 3a3caad5..ea9abaca 100644 --- a/riak/client/__init__.py +++ b/riak/client/__init__.py @@ -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] diff --git a/riak/codecs/__init__.py b/riak/codecs/__init__.py index 8aab82f0..0e221d3b 100644 --- a/riak/codecs/__init__.py +++ b/riak/codecs/__init__.py @@ -1,5 +1,25 @@ import collections +import riak.pb.messages +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, msg_code, data=None): + if msg_code is riak.pb.messages.MSG_CODE_ERROR_RESP: + if data is None: + raise RiakError('no error provided!') + return data + else: + return None diff --git a/riak/codecs/pbuf.py b/riak/codecs/pbuf.py index b55a4253..b4da1e16 100644 --- a/riak/codecs/pbuf.py +++ b/riak/codecs/pbuf.py @@ -8,7 +8,7 @@ import riak.pb.riak_ts_pb2 from riak import RiakError -from riak.codecs import Msg +from riak.codecs import Codec, Msg from riak.content import RiakContent from riak.riak_object import VClock from riak.util import decode_index_value, str_to_bytes, bytes_to_str, \ @@ -73,10 +73,10 @@ def _invert(d): } -class PbufCodec(object): - """ +class PbufCodec(Codec): + ''' Protobuffs Encoding and decoding methods for TcpTransport. - """ + ''' def __init__(self, client_timeouts=False, quorum_controls=False, @@ -88,7 +88,21 @@ def __init__(self, self._tombstone_vclocks = tombstone_vclocks self._bucket_types = bucket_types - def _encode_auth(self, username, password): + 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_data = super(PbufCodec, self).maybe_riak_error(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) @@ -96,7 +110,11 @@ def _encode_auth(self, username, password): rc = riak.pb.messages.MSG_CODE_AUTH_RESP return Msg(mc, req.SerializeToString(), rc) - def _encode_quorum(self, rw): + 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. @@ -112,7 +130,7 @@ def _encode_quorum(self, rw): else: return None - def _decode_quorum(self, rw): + def decode_quorum(self, rw): """ Converts a protobuf quorum value to a symbolic value if necessary. @@ -126,7 +144,7 @@ def _decode_quorum(self, rw): else: return rw - def _decode_contents(self, contents, obj): + def decode_contents(self, contents, obj): """ Decodes the list of siblings from the protobuf representation into the object. @@ -137,14 +155,14 @@ def _decode_contents(self, contents, obj): :type obj: RiakObject :rtype RiakObject """ - obj.siblings = [self._decode_content(c, RiakContent(obj)) + 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): + def decode_content(self, rpb_content, sibling): """ Decodes a single sibling from the protobuf representation into a RiakObject. @@ -170,7 +188,7 @@ def _decode_content(self, rpb_content, sibling): if rpb_content.HasField("vtag"): sibling.etag = bytes_to_str(rpb_content.vtag) - sibling.links = [self._decode_link(link) + 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) @@ -187,7 +205,7 @@ def _decode_content(self, rpb_content, sibling): return sibling - def _encode_content(self, robj, rpb_content): + def encode_content(self, robj, rpb_content): """ Fills an RpbContent message with the appropriate data and metadata from a RiakObject. @@ -232,7 +250,7 @@ def _encode_content(self, robj, rpb_content): else: rpb_content.value = robj.encoded_data - def _decode_link(self, link): + def decode_link(self, link): """ Decodes an RpbLink message into a tuple @@ -256,7 +274,7 @@ def _decode_link(self, link): return (bucket, key, tag) - def _decode_index_value(self, index, value): + def decode_index_value(self, index, value): """ Decodes a secondary index value into the correct Python type. :param index: the name of the index @@ -270,7 +288,7 @@ def _decode_index_value(self, index, value): else: return bytes_to_str(value) - def _encode_bucket_props(self, props, msg): + def encode_bucket_props(self, props, msg): """ Encodes a dict of bucket properties into the protobuf message. @@ -288,13 +306,13 @@ def _encode_bucket_props(self, props, msg): 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)) + 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)) + 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]) + 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)) @@ -305,7 +323,7 @@ def _encode_bucket_props(self, props, msg): return msg - def _decode_bucket_props(self, msg): + def decode_bucket_props(self, msg): """ Decodes the protobuf bucket properties message into a dict. @@ -321,18 +339,18 @@ def _decode_bucket_props(self, msg): 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)) + 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)) + 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)) + 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): + def decode_modfun(self, modfun): """ Decodes a protobuf modfun pair into a dict with 'mod' and 'fun' keys. Used in bucket properties. @@ -344,7 +362,7 @@ def _decode_modfun(self, modfun): return {'mod': bytes_to_str(modfun.module), 'fun': bytes_to_str(modfun.function)} - def _encode_modfun(self, props, msg=None): + def encode_modfun(self, props, msg=None): """ Encodes a dict with 'mod' and 'fun' keys into a protobuf modfun pair. Used in bucket properties. @@ -361,7 +379,7 @@ def _encode_modfun(self, props, msg=None): msg.function = str_to_bytes(props['fun']) return msg - def _decode_hooklist(self, hooklist): + def decode_hooklist(self, hooklist): """ Decodes a list of protobuf commit hooks into their python equivalents. Used in bucket properties. @@ -370,9 +388,9 @@ def _decode_hooklist(self, hooklist): :type hooklist: list :rtype list """ - return [self._decode_hook(hook) for hook in hooklist] + return [self.decode_hook(hook) for hook in hooklist] - def _encode_hooklist(self, hooklist, msg): + def encode_hooklist(self, hooklist, msg): """ Encodes a list of commit hooks into their protobuf equivalent. Used in bucket properties. @@ -383,9 +401,9 @@ def _encode_hooklist(self, hooklist, msg): """ for hook in hooklist: pbhook = msg.add() - self._encode_hook(hook, pbhook) + self.encode_hook(hook, pbhook) - def _decode_hook(self, hook): + def decode_hook(self, hook): """ Decodes a protobuf commit hook message into a dict. Used in bucket properties. @@ -395,11 +413,11 @@ def _decode_hook(self, hook): :rtype dict """ if hook.HasField('modfun'): - return self._decode_modfun(hook.modfun) + return self.decode_modfun(hook.modfun) else: return {'name': bytes_to_str(hook.name)} - def _encode_hook(self, hook, msg): + def encode_hook(self, hook, msg): """ Encodes a commit hook dict into the protobuf message. Used in bucket properties. @@ -413,13 +431,13 @@ def _encode_hook(self, hook, msg): if 'name' in hook: msg.name = str_to_bytes(hook['name']) else: - self._encode_modfun(hook, msg.modfun) + 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): + 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. @@ -475,8 +493,8 @@ def _encode_index_req(self, bucket, index, startkey, endkey=None, 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): + 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)) @@ -491,7 +509,7 @@ def _decode_index_req(self, resp, index, else: return (results, None) - def _decode_search_index(self, index): + def decode_search_index(self, index): """ Fills an RpbYokozunaIndex message with the appropriate data. @@ -514,7 +532,7 @@ def _add_bucket_type(self, req, bucket_type): 'Server does not support bucket-types') req.type = str_to_bytes(bucket_type.name) - def _encode_search_query(self, req, **kwargs): + def encode_search_query(self, req, **kwargs): if 'rows' in kwargs: req.rows = kwargs['rows'] if 'start' in kwargs: @@ -537,7 +555,7 @@ def _encode_search_query(self, req, **kwargs): if 'presort' in kwargs: req.presort = kwargs['presort'] - def _decode_search_doc(self, doc): + def decode_search_doc(self, doc): resultdoc = MultiDict() for pair in doc.fields: if six.PY2: @@ -549,12 +567,12 @@ def _decode_search_doc(self, doc): resultdoc.add(ukey, uval) return resultdoc.mixed() - def _decode_dt_fetch(self, resp): + 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) + value = self.decode_dt_value(dtype, resp.value) if resp.HasField('context'): context = resp.context[:] @@ -563,25 +581,25 @@ def _decode_dt_fetch(self, resp): return dtype, value, context - def _decode_dt_value(self, dtype, msg): + 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) + return self.decode_set_value(msg.set_value) elif dtype == 'map': - return self._decode_map_value(msg.map_value) + return self.decode_map_value(msg.map_value) - def _encode_dt_options(self, req, **kwargs): + 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])) + 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): + def decode_map_value(self, entries): out = {} for entry in entries: name = bytes_to_str(entry.field.name[:]) @@ -589,37 +607,37 @@ def _decode_map_value(self, entries): if dtype == 'counter': value = entry.counter_value elif dtype == 'set': - value = self._decode_set_value(entry.set_value) + 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) + value = self.decode_map_value(entry.map_value) out[(name, dtype)] = value return out - def _decode_set_value(self, set_value): + 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): + 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) + self.encode_set_op(req.op, op) elif dtype == 'map': - self._encode_map_op(req.op.map_op, op) + 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): + 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): + def encode_map_op(self, msg, ops): for op in ops: name, dtype = op[1] ftype = MAP_FIELD_TYPES[dtype] @@ -635,16 +653,16 @@ def _encode_map_op(self, msg, ops): update = msg.updates.add() update.field.name = str_to_bytes(name) update.field.type = ftype - self._encode_map_update(dtype, update, op[2]) + self.encode_map_update(dtype, update, op[2]) - def _encode_map_update(self, dtype, msg, op): + 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) + self.encode_set_op(msg, op) elif dtype == 'map': - self._encode_map_op(msg.map_op, op) + self.encode_map_op(msg.map_op, op) elif dtype == 'register': # ('assign', some_str) msg.register_op = str_to_bytes(op[1]) @@ -654,7 +672,7 @@ def _encode_map_update(self, dtype, msg, op): else: msg.flag_op = riak.pb.riak_dt_pb2.MapUpdate.DISABLE - def _encode_to_ts_cell(self, cell, ts_cell): + 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) @@ -675,7 +693,7 @@ def _encode_to_ts_cell(self, cell, ts_cell): raise RiakError("can't serialize type '{}', value '{}'" .format(t, cell)) - def _encode_timeseries_keyreq(self, table, key, is_delete=False): + def encode_timeseries_keyreq(self, table, key, is_delete=False): key_vals = None if isinstance(key, list): key_vals = key @@ -693,10 +711,10 @@ def _encode_timeseries_keyreq(self, table, key, is_delete=False): 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) + self.encode_to_ts_cell(cell, ts_cell) return Msg(mc, req.SerializeToString(), rc) - def _encode_timeseries_listkeysreq(self, table, timeout=None): + 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: @@ -705,7 +723,13 @@ def _encode_timeseries_listkeysreq(self, table, timeout=None): rc = riak.pb.messages.MSG_CODE_TS_LIST_KEYS_RESP return Msg(mc, req.SerializeToString(), rc) - def _encode_timeseries_put(self, tsobj): + 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. @@ -728,7 +752,7 @@ def _encode_timeseries_put(self, tsobj): 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) + self.encode_to_ts_cell(cell, tsc) else: raise RiakError("TsObject requires a list of rows") @@ -736,7 +760,7 @@ def _encode_timeseries_put(self, tsobj): rc = riak.pb.messages.MSG_CODE_TS_PUT_RESP return Msg(mc, req.SerializeToString(), rc) - def _encode_timeseries_query(self, table, query, interpolations=None): + def encode_timeseries_query(self, table, query, interpolations=None): req = riak.pb.riak_ts_pb2.TsQueryReq() q = query if '{table}' in q: @@ -746,7 +770,7 @@ def _encode_timeseries_query(self, table, query, interpolations=None): rc = riak.pb.messages.MSG_CODE_TS_QUERY_RESP return Msg(mc, req.SerializeToString(), rc) - def _decode_timeseries(self, resp, tsobj): + def decode_timeseries(self, resp, tsobj): """ Fills an TsObject with the appropriate data and metadata from a TsGetResp / TsQueryResp. @@ -766,9 +790,9 @@ def _decode_timeseries(self, resp, tsobj): for row in resp.rows: tsobj.rows.append( - self._decode_timeseries_row(row, resp.columns)) + self.decode_timeseries_row(row, resp.columns)) - def _decode_timeseries_row(self, tsrow, tscols=None): + def decode_timeseries_row(self, tsrow, tscols=None): """ Decodes a TsRow into a list @@ -815,7 +839,7 @@ def _decode_timeseries_row(self, tsrow, tscols=None): row.append(None) return row - def _decode_preflist(self, item): + def decode_preflist(self, item): """ Decodes a preflist response @@ -829,15 +853,15 @@ def _decode_preflist(self, item): 'primary': item. primary} return result - def _encode_get(self, robj, r=None, pr=None, timeout=None, - basic_quorum=None, notfound_ok=None): + 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) + req.r = self.encode_quorum(r) if self._quorum_controls: if pr: - req.pr = self._encode_quorum(pr) + req.pr = self.encode_quorum(pr) if basic_quorum is not None: req.basic_quorum = basic_quorum if notfound_ok is not None: @@ -853,17 +877,17 @@ def _encode_get(self, robj, r=None, pr=None, timeout=None, 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): + 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) + req.w = self.encode_quorum(w) if dw: - req.dw = self._encode_quorum(dw) + req.dw = self.encode_quorum(dw) if self._quorum_controls and pw: - req.pw = self._encode_quorum(pw) + req.pw = self.encode_quorum(pw) if return_body: req.return_body = 1 if if_none_match: @@ -876,54 +900,54 @@ def _encode_put(self, robj, w=None, dw=None, pw=None, req.key = str_to_bytes(robj.key) if robj.vclock: req.vclock = robj.vclock.encode('binary') - self._encode_content(robj, req.content) + 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): + 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) + 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): + 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) + 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): + 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) + req.rw = self.encode_quorum(rw) if r: - req.r = self._encode_quorum(r) + req.r = self.encode_quorum(r) if w: - req.w = self._encode_quorum(w) + req.w = self.encode_quorum(w) if dw: - req.dw = self._encode_quorum(dw) + req.dw = self.encode_quorum(dw) if self._quorum_controls: if pr: - req.pr = self._encode_quorum(pr) + req.pr = self.encode_quorum(pr) if pw: - req.pw = self._encode_quorum(pw) + req.pw = self.encode_quorum(pw) if self._client_timeouts and timeout: req.timeout = timeout @@ -941,7 +965,7 @@ def _encode_delete(self, robj, rw=None, r=None, rc = riak.pb.messages.MSG_CODE_DEL_RESP return Msg(mc, req.SerializeToString(), rc) - def _encode_stream_keys(self, bucket, timeout=None): + 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: @@ -951,34 +975,34 @@ def _encode_stream_keys(self, bucket, timeout=None): rc = riak.pb.messages.MSG_CODE_LIST_KEYS_RESP return Msg(mc, req.SerializeToString(), rc) - def _decode_get_keys(self, stream): + 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): + 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): + 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): + def decode_get_client_id(self, resp): return bytes_to_str(resp.client_id) - def _encode_set_client_id(self, 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): + 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() @@ -990,7 +1014,7 @@ def _encode_get_buckets(self, bucket_type, rc = riak.pb.messages.MSG_CODE_LIST_BUCKETS_RESP return Msg(mc, req.SerializeToString(), rc) - def _encode_get_bucket_props(self, bucket): + 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) @@ -998,16 +1022,16 @@ def _encode_get_bucket_props(self, bucket): rc = riak.pb.messages.MSG_CODE_GET_BUCKET_RESP return Msg(mc, req.SerializeToString(), rc) - def _encode_set_bucket_props(self, bucket, props): + 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) + 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): + 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) @@ -1015,22 +1039,22 @@ def _encode_clear_bucket_props(self, bucket): rc = riak.pb.messages.MSG_CODE_RESET_BUCKET_RESP return Msg(mc, req.SerializeToString(), rc) - def _encode_get_bucket_type_props(self, bucket_type): + 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): + 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) + 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): + 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") @@ -1038,8 +1062,8 @@ def _encode_stream_mapred(self, content): 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): + 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: @@ -1053,27 +1077,27 @@ def _encode_create_search_index(self, index, schema=None, rc = riak.pb.messages.MSG_CODE_PUT_RESP return Msg(mc, req.SerializeToString(), rc) - def _encode_get_search_index(self, index): + 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): + 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): + 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): + 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)) @@ -1083,45 +1107,45 @@ def _encode_create_search_schema(self, schema, content): rc = riak.pb.messages.MSG_CODE_PUT_RESP return Msg(mc, req.SerializeToString(), rc) - def _encode_get_search_schema(self, schema): + 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): + 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): + 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) + 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): + 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] + result['docs'] = [self.decode_search_doc(doc) for doc in resp.docs] return result - def _encode_get_counter(self, bucket, key, **kwargs): + 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']) + req.r = self.encode_quorum(kwargs['r']) if kwargs.get('pr') is not None: - req.pr = self._encode_quorum(kwargs['pr']) + 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: @@ -1130,34 +1154,34 @@ def _encode_get_counter(self, bucket, key, **kwargs): rc = riak.pb.messages.MSG_CODE_COUNTER_GET_RESP return Msg(mc, req.SerializeToString(), rc) - def _encode_update_counter(self, bucket, key, value, **kwargs): + 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']) + req.w = self.encode_quorum(kwargs['w']) if kwargs.get('dw') is not None: - req.dw = self._encode_quorum(kwargs['dw']) + req.dw = self.encode_quorum(kwargs['dw']) if kwargs.get('pw') is not None: - req.pw = self._encode_quorum(kwargs['pw']) + 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): + 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) + 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): + def encode_update_datatype(self, datatype, **kwargs): op = datatype.to_op() type_name = datatype.type_name if not op: @@ -1170,22 +1194,22 @@ def _encode_update_datatype(self, datatype, **kwargs): 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) + 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): + 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)) + datatype._set_value(self.decode_dt_value(type_name, resp)) - def _encode_get_preflist(self, bucket, key): + 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) diff --git a/riak/codecs/ttb.py b/riak/codecs/ttb.py index eef65fdc..3ac7651a 100644 --- a/riak/codecs/ttb.py +++ b/riak/codecs/ttb.py @@ -3,12 +3,12 @@ import riak.pb.messages -from erlastic import encode +from erlastic import encode, decode from erlastic.types import Atom from riak import RiakError -from riak.codecs import Msg -from riak.util import unix_time_millis, \ +from riak.codecs import Codec, Msg +from riak.util import bytes_to_str, unix_time_millis, \ datetime_from_unix_time_millis udef_a = Atom('undefined') @@ -24,7 +24,7 @@ tscell_empty = (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a) -class TtbCodec(object): +class TtbCodec(Codec): ''' Erlang term-to-binary Encoding and decoding methods for TcpTransport ''' @@ -32,7 +32,32 @@ class TtbCodec(object): def __init__(self, **unused_args): super(TtbCodec, self).__init__(**unused_args) - def _encode_to_ts_cell(self, cell): + def parse_msg(self, msg_code, data): + if msg_code != riak.pb.messages.MSG_CODE_TS_GET_RESP and \ + msg_code != riak.pb.messages.MSG_CODE_TS_PUT_RESP and \ + msg_code != riak.pb.messages.MSG_CODE_ERROR_RESP: + raise RiakError("TTB can't parse code: {}".format(msg_code)) + if len(data) > 0: + return decode(data) + else: + return None + + def process_err_ttb(self, err_ttb): + resp_a = err_ttb[0] + if resp_a == rpberrorresp_a: + errmsg = err_ttb[1] + raise RiakError(bytes_to_str(errmsg)) + else: + raise RiakError( + "Unknown TTB error type: {}".format(resp_a)) + + def maybe_riak_error(self, msg_code, data=None): + err_data = super(TtbCodec, self).maybe_riak_error(msg_code, data) + if err_data: + err_ttb = decode(err_data) + self.process_err_ttb(err_ttb) + + def encode_to_ts_cell(self, cell): if cell is None: return tscell_empty else: @@ -55,7 +80,7 @@ def _encode_to_ts_cell(self, cell): raise RiakError("can't serialize type '{}', value '{}'" .format(t, cell)) - def _encode_timeseries_keyreq(self, table, key, is_delete=False): + def encode_timeseries_keyreq(self, table, key, is_delete=False): key_vals = None if isinstance(key, list): key_vals = key @@ -71,10 +96,19 @@ def _encode_timeseries_keyreq(self, table, key, is_delete=False): req_atom = tsdelreq_a req = req_atom, table.name, \ - [self._encode_to_ts_cell(k) for k in key_vals], udef_a + [self.encode_to_ts_cell(k) for k in key_vals], udef_a return Msg(mc, encode(req), rc) - def _encode_timeseries_put(self, tsobj): + def validate_timeseries_put_resp(self, resp_code, resp): + if resp is None and \ + resp_code == riak.pb.messages.MSG_CODE_TS_PUT_RESP: + 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. @@ -91,7 +125,7 @@ def _encode_timeseries_put(self, tsobj): for row in tsobj.rows: req_r = [] for cell in row: - req_r.append(self._encode_to_ts_cell(cell)) + req_r.append(self.encode_to_ts_cell(cell)) req_t = (tsrow_a, req_r) req_rows.append(req_t) req = tsputreq_a, tsobj.table.name, udef_a, req_rows @@ -101,7 +135,7 @@ def _encode_timeseries_put(self, tsobj): else: raise RiakError("TsObject requires a list of rows") - def _decode_timeseries(self, resp_ttb, tsobj): + def decode_timeseries(self, resp_ttb, tsobj): """ Fills an TsObject with the appropriate data and metadata from a TTB-encoded TsGetResp / TsQueryResp. @@ -126,18 +160,18 @@ def _decode_timeseries(self, resp_ttb, tsobj): return tsobj resp_a = resp_ttb[0] - if resp_a == tsgetresp_a: + if resp_a == rpberrorresp_a: + self.process_err_ttb(resp_ttb) + elif resp_a == tsgetresp_a: # TODO resp_cols = resp_ttb[1] resp_rows = resp_ttb[2] for row_ttb in resp_rows: tsobj.rows.append( - self._decode_timeseries_row(row_ttb, None)) - # TODO - # elif resp_a == rpberrorresp_a: + self.decode_timeseries_row(row_ttb, None)) else: raise RiakError("Unknown TTB response type: {}".format(resp_a)) - def _decode_timeseries_row(self, tsrow_ttb, tscols=None): + def decode_timeseries_row(self, tsrow_ttb, tscols=None): """ Decodes a TTB-encoded TsRow into a list 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/test_timeseries.py b/riak/tests/test_timeseries.py index e50d2613..0ebbe552 100644 --- a/riak/tests/test_timeseries.py +++ b/riak/tests/test_timeseries.py @@ -67,7 +67,7 @@ def test_encode_decode_timestamp(self): def test_encode_data_for_get(self): c = PbufCodec() - msg = c._encode_timeseries_keyreq( + msg = c.encode_timeseries_keyreq( self.table, self.test_key, is_delete=False) req = riak.pb.riak_ts_pb2.TsGetReq() req.ParseFromString(msg.data) @@ -75,7 +75,7 @@ def test_encode_data_for_get(self): def test_encode_data_for_delete(self): c = PbufCodec() - msg = c._encode_timeseries_keyreq( + msg = c.encode_timeseries_keyreq( self.table, self.test_key, is_delete=True) req = riak.pb.riak_ts_pb2.TsDelReq() req.ParseFromString(msg.data) @@ -84,7 +84,7 @@ def test_encode_data_for_delete(self): def test_encode_data_for_put(self): c = PbufCodec() tsobj = TsObject(None, self.table, self.rows, None) - msg = c._encode_timeseries_put(tsobj) + msg = c.encode_timeseries_put(tsobj) req = riak.pb.riak_ts_pb2.TsPutReq() req.ParseFromString(msg.data) @@ -110,7 +110,7 @@ def test_encode_data_for_put(self): def test_encode_data_for_listkeys(self): c = PbufCodec(client_timeouts=True) - msg = c._encode_timeseries_listkeysreq(self.table, 1234) + msg = c.encode_timeseries_listkeysreq(self.table, 1234) req = riak.pb.riak_ts_pb2.TsListKeysReq() req.ParseFromString(msg.data) self.assertEqual(self.table.name, bytes_to_str(req.table)) @@ -161,7 +161,7 @@ def test_decode_data_from_query(self): tsobj = TsObject(None, self.table, [], []) c = PbufCodec() - c._decode_timeseries(tqr, tsobj) + c.decode_timeseries(tqr, tsobj) self.assertEqual(len(self.rows), len(tsobj.rows)) self.assertEqual(len(tqr.columns), len(tsobj.columns)) @@ -196,6 +196,8 @@ 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): + client_options = {'transport_options': {'use_ttb': False}} + @classmethod def setUpClass(cls): super(TimeseriesTests, cls).setUpClass() @@ -276,34 +278,34 @@ def test_query_that_creates_table_using_interpolation(self): 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) 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) 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) 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) 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) @@ -317,7 +319,7 @@ def test_query_that_returns_no_data(self): user = 'user1' """ query = fmt.format(table=table_name) - ts_obj = self.client.ts_query('GeoCheckin', query) + ts_obj = self.client.ts_query(table_name, query) self.assertEqual(len(ts_obj.columns), 0) self.assertEqual(len(ts_obj.rows), 0) @@ -328,7 +330,7 @@ def test_query_that_returns_no_data_using_interpolation(self): geohash = 'hash1' and user = 'user1' """ - ts_obj = self.client.ts_query('GeoCheckin', query) + ts_obj = self.client.ts_query(table_name, query) self.assertEqual(len(ts_obj.columns), 0) self.assertEqual(len(ts_obj.rows), 0) @@ -343,7 +345,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): @@ -356,7 +358,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): @@ -370,7 +372,7 @@ 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.encoded_rows): if want[2] == self.twentyFiveMinsAgo: @@ -382,23 +384,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) @@ -413,7 +415,13 @@ def test_stream_keys(self): 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.assertEqual(len(ts_obj.rows), 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 index 81c0554c..ba368a1c 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -1,11 +1,13 @@ # -*- 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 @@ -52,7 +54,7 @@ def test_encode_data_for_get(self): test_key = ['hash1', 'user2', ts0] c = TtbCodec() - msg = c._encode_timeseries_keyreq(self.table, test_key) + msg = c.encode_timeseries_keyreq(self.table, test_key) self.assertEqual(req_test, msg.data) # def test_decode_riak_error(self): @@ -84,7 +86,7 @@ def test_decode_data_from_get(self): tsobj = TsObject(None, self.table, [], []) c = TtbCodec() - c._decode_timeseries(decode(rsp_ttb), tsobj) + c.decode_timeseries(decode(rsp_ttb), tsobj) for i in range(0, 1): self.assertEqual(tsrow_a, rows[i][0]) @@ -145,13 +147,15 @@ def test_encode_data_for_put(self): tsobj = TsObject(None, self.table, rows_to_encode, None) c = TtbCodec() - msg = c._encode_timeseries_put(tsobj) + 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() @@ -164,10 +168,7 @@ def test_store_and_fetch_ttb(self): twentyMinsAgo = fifteenMinsAgo - fiveMins twentyFiveMinsAgo = twentyMinsAgo - fiveMins - opts = {'use_ttb': True} - client = self.create_client(transport_options=opts) - - table = client.table(table_name) + table = self.client.table(table_name) rows = [ ['hash1', 'user2', twentyFiveMinsAgo, 'typhoon', 90.3], ['hash1', 'user2', twentyMinsAgo, 'hurricane', 82.3], @@ -181,9 +182,16 @@ def test_store_and_fetch_ttb(self): for r in rows: k = r[0:3] - ts_obj = client.ts_get(table_name, k) + ts_obj = self.client.ts_get(table_name, k) self.assertIsNotNone(ts_obj) self.assertEqual(len(ts_obj.rows), 1) self.assertEqual(len(ts_obj.rows[0]), 5) - client.close() + 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/transports/tcp/connection.py b/riak/transports/tcp/connection.py index 8972e480..a8433f30 100644 --- a/riak/transports/tcp/connection.py +++ b/riak/transports/tcp/connection.py @@ -105,12 +105,12 @@ 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 """ - c = PbufCodec() + codec = PbufCodec() username = self._client._credentials.username password = self._client._credentials.password if not password: password = '' - msg = c._encode_auth(username, password) + 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 diff --git a/riak/transports/tcp/stream.py b/riak/transports/tcp/stream.py index 46fb5941..3cf0e974 100644 --- a/riak/transports/tcp/stream.py +++ b/riak/transports/tcp/stream.py @@ -16,9 +16,10 @@ class PbufStream(object): _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): @@ -29,12 +30,11 @@ def next(self): raise StopIteration try: - # TODO RTS-842 - should be part of passed-in codec resp_code, data = self.transport._recv_msg() - self.transport._maybe_riak_error(resp_code, data) + self.codec.maybe_riak_error(resp_code, data) expect = self._expect - self.transport._maybe_incorrect_code(resp_code, expect) - resp = self.transport._parse_msg(expect, data, is_ttb=False) + self.codec.maybe_incorrect_code(resp_code, expect) + resp = self.codec.parse_msg(expect, data) except: self.finished = True raise @@ -137,8 +137,8 @@ class PbufIndexStream(PbufStream): _expect = riak.pb.messages.MSG_CODE_INDEX_RESP - def __init__(self, transport, index, return_terms=False): - super(PbufIndexStream, 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 @@ -174,10 +174,6 @@ class PbufTsKeyStream(PbufStream, TtbCodec): _expect = riak.pb.messages.MSG_CODE_TS_LIST_KEYS_RESP - def __init__(self, transport, codec): - super(PbufTsKeyStream, self).__init__(transport) - self._codec = codec - def next(self): response = super(PbufTsKeyStream, self).next() @@ -186,7 +182,7 @@ def next(self): keys = [] for tsrow in response.keys: - keys.append(self._codec._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 index ee8898f2..151f76c2 100644 --- a/riak/transports/tcp/transport.py +++ b/riak/transports/tcp/transport.py @@ -1,18 +1,13 @@ -# TODO RTS-842 codecs should return msg codes too -import erlastic import six import riak.pb.messages from riak import RiakError -from riak.codecs import Msg +from riak.codecs import Codec, Msg from riak.codecs.pbuf import PbufCodec from riak.codecs.ttb import TtbCodec from riak.transports.transport import Transport from riak.ts_object import TsObject -# TODO RTS-842 ideally these would not be needed -from riak.util import bytes_to_str - from riak.transports.tcp.connection import TcpConnection from riak.transports.tcp.stream import (PbufKeyStream, PbufMapredStream, @@ -78,9 +73,10 @@ def ping(self): """ Ping the remote server """ - msg = Msg(riak.pb.messages.MSG_CODE_PING_REQ, None, - riak.pb.messages.MSG_CODE_PING_RESP) - resp_code, _ = self._request(msg) + 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: @@ -95,21 +91,21 @@ def get_server_info(self): 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) - return codec._decode_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) - return codec._decode_get_client_id(resp) + 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) + 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, @@ -122,20 +118,20 @@ def get(self, robj, r=None, pr=None, timeout=None, basic_quorum=None, """ 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) - return codec._decode_get(robj, resp) + 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) - return codec._decode_put(robj, resp) + 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) @@ -144,35 +140,24 @@ def ts_describe(self, table): def ts_get(self, table, key): msg_code = riak.pb.messages.MSG_CODE_TS_GET_REQ codec = self._get_codec(msg_code) - msg = codec._encode_timeseries_keyreq(table, key) - # TODO RTS-842 is_ttb - resp_code, resp = self._request(msg, self._use_ttb) + msg = codec.encode_timeseries_keyreq(table, key) + resp_code, resp = self._request(msg, codec) tsobj = TsObject(self._client, table, [], None) - codec._decode_timeseries(resp, tsobj) + codec.decode_timeseries(resp, tsobj) return tsobj def ts_put(self, tsobj): msg_code = riak.pb.messages.MSG_CODE_TS_PUT_REQ codec = self._get_codec(msg_code) - msg = codec._encode_timeseries_put(tsobj) - # logging.debug("pbc/transport ts_put _use_ttb: '%s'", - # self._use_ttb) - # TODO RTS-842 use_ttb - resp_code, resp = self._request(msg, self._use_ttb) - if self._use_ttb and \ - resp is None and \ - resp_code == riak.pb.messages.MSG_CODE_TS_PUT_RESP: - return True - if resp is not None: - return True - else: - raise RiakError("missing response object") + 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) + 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: @@ -181,10 +166,10 @@ def ts_delete(self, table, key): 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) + 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) + codec.decode_timeseries(resp, tsobj) return tsobj def ts_stream_keys(self, table, timeout=None): @@ -194,16 +179,16 @@ def ts_stream_keys(self, table, timeout=None): """ msg_code = riak.pb.messages.MSG_CODE_TS_LIST_KEYS_REQ codec = self._get_codec(msg_code) - msg = codec._encode_timeseries_listkeysreq(table, timeout) + 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): + 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) + 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): @@ -213,7 +198,7 @@ def get_keys(self, bucket, timeout=None): 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) + return codec.decode_get_keys(stream) def stream_keys(self, bucket, timeout=None): """ @@ -222,9 +207,9 @@ def stream_keys(self, bucket, timeout=None): """ msg_code = riak.pb.messages.MSG_CODE_LIST_KEYS_REQ codec = self._get_codec(msg_code) - msg = codec._encode_stream_keys(bucket, timeout) + msg = codec.encode_stream_keys(bucket, timeout) self._send_msg(msg.msg_code, msg.data) - return PbufKeyStream(self) + return PbufKeyStream(self, codec) def get_buckets(self, bucket_type=None, timeout=None): """ @@ -232,9 +217,9 @@ def get_buckets(self, bucket_type=None, timeout=None): """ 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) + 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): @@ -246,10 +231,10 @@ def stream_buckets(self, bucket_type=None, timeout=None): '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) + msg = codec.encode_get_buckets(bucket_type, + timeout, streaming=True) self._send_msg(msg.msg_code, msg.data) - return PbufBucketStream(self) + return PbufBucketStream(self, codec) def get_bucket_props(self, bucket): """ @@ -257,9 +242,9 @@ def get_bucket_props(self, bucket): """ 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) - return codec._decode_bucket_props(resp.props) + 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): """ @@ -272,8 +257,8 @@ def set_bucket_props(self, bucket, props): '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) + msg = codec.encode_set_bucket_props(bucket, props) + resp_code, resp = self._request(msg, codec) return True def clear_bucket_props(self, bucket): @@ -284,8 +269,8 @@ def clear_bucket_props(self, bucket): 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) + msg = codec.encode_clear_bucket_props(bucket) + self._request(msg, codec) return True def get_bucket_type_props(self, bucket_type): @@ -295,9 +280,9 @@ def get_bucket_type_props(self, bucket_type): 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) - return codec._decode_bucket_props(resp.props) + 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): """ @@ -306,8 +291,8 @@ def set_bucket_type_props(self, bucket_type, props): 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) + 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): @@ -333,9 +318,9 @@ def stream_mapred(self, inputs, query, timeout=None): 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) + msg = codec.encode_stream_mapred(content) self._send_msg(msg.msg_code, msg.data) - return PbufMapredStream(self) + return PbufMapredStream(self, codec) def get_index(self, bucket, index, startkey, endkey=None, return_terms=None, max_results=None, continuation=None, @@ -350,13 +335,13 @@ def get_index(self, bucket, index, startkey, endkey=None, 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) - return codec._decode_index_req(resp, index, - return_terms, max_results) + 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, @@ -369,12 +354,12 @@ def stream_index(self, bucket, index, startkey, endkey=None, "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) + 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, index, return_terms) + return PbufIndexStream(self, codec, index, return_terms) def create_search_index(self, index, schema=None, n_val=None, timeout=None): @@ -383,8 +368,8 @@ def create_search_index(self, index, schema=None, n_val=None, "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) + msg = codec.encode_create_search_index(index, schema, n_val, timeout) + self._request(msg, codec) return True def get_search_index(self, index): @@ -393,10 +378,10 @@ def get_search_index(self, index): "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) + 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]) + return codec.decode_search_index(resp.index[0]) else: raise RiakError('notfound') @@ -406,9 +391,9 @@ def list_search_indexes(self): "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) - return [codec._decode_search_index(index) for index in resp.index] + 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(): @@ -416,8 +401,8 @@ def delete_search_index(self, index): "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) + msg = codec.encode_delete_search_index(index) + self._request(msg, codec) return True def create_search_schema(self, schema, content): @@ -426,8 +411,8 @@ def create_search_schema(self, schema, content): "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) + msg = codec.encode_create_search_schema(schema, content) + self._request(msg, codec) return True def get_search_schema(self, schema): @@ -436,9 +421,9 @@ def get_search_schema(self, schema): "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) - return codec._decode_get_search_schema(resp) + 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 @@ -449,9 +434,9 @@ def search(self, index, query, **kwargs): 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) - return codec._decode_search(resp) + 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(): @@ -462,8 +447,8 @@ def get_counter(self, bucket, key, **kwargs): 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) + msg = codec.encode_get_counter(bucket, key, **kwargs) + resp_code, resp = self._request(msg, codec) if resp.HasField('value'): return resp.value else: @@ -478,8 +463,8 @@ def update_counter(self, bucket, key, value, **kwargs): 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) + msg = codec.encode_update_counter(bucket, key, value, **kwargs) + resp_code, resp = self._request(msg, codec) if resp.HasField('value'): return resp.value else: @@ -493,9 +478,9 @@ def fetch_datatype(self, bucket, key, **kwargs): 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) - return codec._decode_dt_fetch(resp) + 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(): @@ -505,9 +490,9 @@ def update_datatype(self, datatype, **kwargs): 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._decode_update_datatype(datatype, resp, **kwargs) + 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): @@ -522,63 +507,27 @@ def get_preflist(self, bucket, key): """ 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) - return [codec._decode_preflist(item) for item in resp.preflist] + msg = codec.encode_get_preflist(bucket, key) + resp_code, resp = self._request(msg, codec) + return [codec.decode_preflist(item) for item in resp.preflist] # TODO RTS-842 is_ttb - def _parse_msg(self, code, packet, is_ttb=False): - if is_ttb: - if code != riak.pb.messages.MSG_CODE_TS_GET_RESP and \ - code != riak.pb.messages.MSG_CODE_TS_PUT_RESP: - raise RiakError("TTB can't parse code: %d" % code) - if len(packet) > 0: - return erlastic.decode(packet) - else: - return None - else: - 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 - - # TODO RTS-842 move to base Codec object - def _maybe_riak_error(self, msg_code, data=None, is_ttb=False): - if msg_code is riak.pb.messages.MSG_CODE_ERROR_RESP: - if data is None: - raise RiakError('no error provided!') - # TODO RTS-842 TTB-specific version - err = self._parse_msg(msg_code, data, is_ttb) - if err is None: - raise RiakError('no error provided!') - else: - raise RiakError(bytes_to_str(err.errmsg)) - - 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)) - - # TODO RTS-842 is_ttb - def _request(self, msg, is_ttb=False): + 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) - self._maybe_riak_error(resp_code, data, is_ttb) - self._maybe_incorrect_code(resp_code, expect) + codec.maybe_riak_error(resp_code, data) + codec.maybe_incorrect_code(resp_code, expect) if resp_code in riak.pb.messages.MESSAGE_CLASSES: - msg = self._parse_msg(resp_code, data, is_ttb) + msg = codec.parse_msg(resp_code, data) else: raise Exception("unknown msg code %s" % resp_code) # logging.debug("tcp/connection received resp_code %d msg %s", From a2f6b92a4a19b68244cc00a5a85e665f66378754 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 6 Apr 2016 09:09:38 -0700 Subject: [PATCH 37/44] Updating for TTB changes server-side. --- riak/codecs/__init__.py | 4 ++-- riak/codecs/pbuf.py | 3 ++- riak/codecs/ttb.py | 36 +++++++++++++++---------------- riak/tests/test_timeseries.py | 4 ++-- riak/tests/test_timeseries_ttb.py | 2 +- riak/transports/tcp/connection.py | 21 ------------------ riak/transports/tcp/transport.py | 18 +++++++--------- 7 files changed, 32 insertions(+), 56 deletions(-) diff --git a/riak/codecs/__init__.py b/riak/codecs/__init__.py index 0e221d3b..e2ef78de 100644 --- a/riak/codecs/__init__.py +++ b/riak/codecs/__init__.py @@ -16,8 +16,8 @@ def maybe_incorrect_code(self, resp_code, expect=None): raise RiakError("unexpected message code: %d, expected %d" % (resp_code, expect)) - def maybe_riak_error(self, msg_code, data=None): - if msg_code is riak.pb.messages.MSG_CODE_ERROR_RESP: + 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 diff --git a/riak/codecs/pbuf.py b/riak/codecs/pbuf.py index b4da1e16..2e08ec5c 100644 --- a/riak/codecs/pbuf.py +++ b/riak/codecs/pbuf.py @@ -97,7 +97,8 @@ def parse_msg(self, msg_code, data): return pbo def maybe_riak_error(self, msg_code, data=None): - err_data = super(PbufCodec, self).maybe_riak_error(msg_code, data) + 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)) diff --git a/riak/codecs/ttb.py b/riak/codecs/ttb.py index 3ac7651a..98e75da3 100644 --- a/riak/codecs/ttb.py +++ b/riak/codecs/ttb.py @@ -23,6 +23,9 @@ tscell_empty = (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a) +# TODO RTS-842 +MSG_CODE_TS_TTB = 104 + class TtbCodec(Codec): ''' @@ -33,29 +36,25 @@ def __init__(self, **unused_args): super(TtbCodec, self).__init__(**unused_args) def parse_msg(self, msg_code, data): - if msg_code != riak.pb.messages.MSG_CODE_TS_GET_RESP and \ - msg_code != riak.pb.messages.MSG_CODE_TS_PUT_RESP and \ - msg_code != riak.pb.messages.MSG_CODE_ERROR_RESP: + if msg_code != MSG_CODE_TS_TTB and \ + msg_code != riak.pb.messages.MSG_CODE_TS_GET_RESP and \ + msg_code != riak.pb.messages.MSG_CODE_TS_PUT_RESP: raise RiakError("TTB can't parse code: {}".format(msg_code)) if len(data) > 0: - return decode(data) + decoded = decode(data) + self.maybe_err_ttb(decoded) + return decoded else: return None - def process_err_ttb(self, err_ttb): + def maybe_err_ttb(self, err_ttb): resp_a = err_ttb[0] if resp_a == rpberrorresp_a: errmsg = err_ttb[1] raise RiakError(bytes_to_str(errmsg)) - else: - raise RiakError( - "Unknown TTB error type: {}".format(resp_a)) def maybe_riak_error(self, msg_code, data=None): - err_data = super(TtbCodec, self).maybe_riak_error(msg_code, data) - if err_data: - err_ttb = decode(err_data) - self.process_err_ttb(err_ttb) + pass def encode_to_ts_cell(self, cell): if cell is None: @@ -87,14 +86,13 @@ def encode_timeseries_keyreq(self, table, key, is_delete=False): else: raise ValueError("key must be a list") - mc = riak.pb.messages.MSG_CODE_TS_GET_REQ - rc = riak.pb.messages.MSG_CODE_TS_GET_RESP + mc = MSG_CODE_TS_TTB + rc = MSG_CODE_TS_TTB req_atom = tsgetreq_a if is_delete: - mc = riak.pb.messages.MSG_CODE_TS_DEL_REQ - rc = riak.pb.messages.MSG_CODE_TS_DEL_RESP 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) @@ -128,9 +126,9 @@ def encode_timeseries_put(self, tsobj): req_r.append(self.encode_to_ts_cell(cell)) req_t = (tsrow_a, req_r) req_rows.append(req_t) - req = tsputreq_a, tsobj.table.name, udef_a, req_rows - mc = riak.pb.messages.MSG_CODE_TS_PUT_REQ - rc = riak.pb.messages.MSG_CODE_TS_PUT_RESP + 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") diff --git a/riak/tests/test_timeseries.py b/riak/tests/test_timeseries.py index 0ebbe552..5d9e81cb 100644 --- a/riak/tests/test_timeseries.py +++ b/riak/tests/test_timeseries.py @@ -417,8 +417,8 @@ def test_delete_single_value(self): key = ['hash1', 'user2', self.twentyFiveMinsAgo] rslt = self.client.ts_delete(table_name, key) self.assertTrue(rslt) - ts_obj = self.client.ts_get(table_name, key) - self.assertEqual(len(ts_obj.rows), 0) + with self.assertRaises(RiakError): + self.client.ts_get(table_name, key) def test_create_error_via_put(self): table = Table(self.client, table_name) diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index ba368a1c..5da40aab 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -137,7 +137,7 @@ def test_encode_data_for_put(self): (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a) ]) rows = [r0, r1] - req = tsputreq_a, str_to_bytes(table_name), udef_a, rows + req = tsputreq_a, str_to_bytes(table_name), [], rows req_test = encode(req) rows_to_encode = [ diff --git a/riak/transports/tcp/connection.py b/riak/transports/tcp/connection.py index a8433f30..d0ffab13 100644 --- a/riak/transports/tcp/connection.py +++ b/riak/transports/tcp/connection.py @@ -21,9 +21,6 @@ class TcpConnection(object): """ Connection-related methods for TcpTransport. """ - def __init__(self): - self._ttb_enabled = False - def _encode_msg(self, msg_code, data=None): if data is None: return struct.pack("!iB", 1, msg_code) @@ -80,24 +77,6 @@ def _starttls(self): else: return False - def _enable_ttb(self): - if self._ttb_enabled: - return True - else: - logging.debug("tcp/connection enabling TTB") - req = riak.pb.riak_pb2.RpbToggleEncodingReq() - req.use_native = True - data = req.SerializeToString() - resp_code, _ = self._send_recv( - riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_REQ, - data) - if resp_code == riak.pb.messages.MSG_CODE_TOGGLE_ENCODING_RESP: - self._ttb_enabled = True - logging.debug("tcp/connection TTB IS ENABLED") - return True - else: - return False - def _auth(self): """ Perform an authorization request against Riak diff --git a/riak/transports/tcp/transport.py b/riak/transports/tcp/transport.py index 151f76c2..d7e7954b 100644 --- a/riak/transports/tcp/transport.py +++ b/riak/transports/tcp/transport.py @@ -4,7 +4,7 @@ from riak import RiakError from riak.codecs import Codec, Msg from riak.codecs.pbuf import PbufCodec -from riak.codecs.ttb import TtbCodec +from riak.codecs.ttb import TtbCodec, MSG_CODE_TS_TTB from riak.transports.transport import Transport from riak.ts_object import TsObject @@ -46,8 +46,6 @@ def _get_pbuf_codec(self): def _get_ttb_codec(self): if self._use_ttb: - if not self._enable_ttb(): - raise RiakError('could not switch to TTB encoding!') if not self._ttb_c: self._ttb_c = TtbCodec() codec = self._ttb_c @@ -56,7 +54,9 @@ def _get_ttb_codec(self): return codec def _get_codec(self, msg_code): - if msg_code == riak.pb.messages.MSG_CODE_TS_GET_REQ: + 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() @@ -138,7 +138,7 @@ def ts_describe(self, table): return self.ts_query(table, query) def ts_get(self, table, key): - msg_code = riak.pb.messages.MSG_CODE_TS_GET_REQ + 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) @@ -147,7 +147,7 @@ def ts_get(self, table, key): return tsobj def ts_put(self, tsobj): - msg_code = riak.pb.messages.MSG_CODE_TS_PUT_REQ + 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) @@ -511,7 +511,6 @@ def get_preflist(self, bucket, key): resp_code, resp = self._request(msg, codec) return [codec.decode_preflist(item) for item in resp.preflist] - # TODO RTS-842 is_ttb def _request(self, msg, codec=None): if isinstance(msg, Msg): msg_code = msg.msg_code @@ -526,10 +525,9 @@ def _request(self, msg, codec=None): 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 in riak.pb.messages.MESSAGE_CLASSES: + 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) - # logging.debug("tcp/connection received resp_code %d msg %s", - # resp_code, msg) return resp_code, msg From b946f0fe3eeff47106480ffac3a3c9a1ab301fe8 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 6 Apr 2016 17:30:52 -0700 Subject: [PATCH 38/44] Continuing to work on TTB changes --- riak/codecs/ttb.py | 20 ++++----- riak/tests/test_timeseries_ttb.py | 69 +++++++++++-------------------- 2 files changed, 33 insertions(+), 56 deletions(-) diff --git a/riak/codecs/ttb.py b/riak/codecs/ttb.py index 98e75da3..63479a55 100644 --- a/riak/codecs/ttb.py +++ b/riak/codecs/ttb.py @@ -21,8 +21,6 @@ tsrow_a = Atom('tsrow') tscell_a = Atom('tscell') -tscell_empty = (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a) - # TODO RTS-842 MSG_CODE_TS_TTB = 104 @@ -58,22 +56,21 @@ def maybe_riak_error(self, msg_code, data=None): def encode_to_ts_cell(self, cell): if cell is None: - return tscell_empty + return [] else: if isinstance(cell, datetime.datetime): ts = unix_time_millis(cell) - return (tscell_a, udef_a, udef_a, ts, udef_a, udef_a) + return ts elif isinstance(cell, bool): - return (tscell_a, udef_a, udef_a, udef_a, cell, udef_a) + return cell elif isinstance(cell, six.text_type) or \ isinstance(cell, six.binary_type) or \ isinstance(cell, six.string_types): - return (tscell_a, cell, - udef_a, udef_a, udef_a, udef_a) + return cell elif (isinstance(cell, six.integer_types)): - return (tscell_a, udef_a, cell, udef_a, udef_a, udef_a) + return cell elif isinstance(cell, float): - return (tscell_a, udef_a, udef_a, udef_a, udef_a, cell) + return cell else: t = type(cell) raise RiakError("can't serialize type '{}', value '{}'" @@ -116,7 +113,7 @@ def encode_timeseries_put(self, tsobj): :rtype: term-to-binary encoded object ''' if tsobj.columns: - raise NotImplementedError("columns are not implemented yet") + raise NotImplementedError('columns are not used') if tsobj.rows and isinstance(tsobj.rows, list): req_rows = [] @@ -124,8 +121,7 @@ def encode_timeseries_put(self, tsobj): req_r = [] for cell in row: req_r.append(self.encode_to_ts_cell(cell)) - req_t = (tsrow_a, req_r) - req_rows.append(req_t) + req_rows.append(tuple(req_r)) req = tsputreq_a, tsobj.table.name, [], req_rows mc = MSG_CODE_TS_TTB rc = MSG_CODE_TS_TTB diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index 5da40aab..811f90b3 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -23,8 +23,12 @@ tsputreq_a = Atom('tsputreq') udef_a = Atom('undefined') -tsrow_a = Atom('tsrow') -tscell_a = Atom('tscell') +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' @@ -45,9 +49,7 @@ def setUp(self): def test_encode_data_for_get(self): keylist = [ - (tscell_a, str_to_bytes('hash1'), udef_a, udef_a, udef_a, udef_a), - (tscell_a, str_to_bytes('user2'), udef_a, udef_a, udef_a, udef_a), - (tscell_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a) + 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) @@ -57,31 +59,24 @@ def test_encode_data_for_get(self): msg = c.encode_timeseries_keyreq(self.table, test_key) self.assertEqual(req_test, msg.data) - # def test_decode_riak_error(self): - + # {tsgetresp, + # { + # [<<"geohash">>, <<"user">>, <<"time">>, <<"weather">>, <<"temperature">>], + # [varchar, varchar, timestamp, varchar, double] + # }, + # [[<<"hash1">>, <<"user2">>, 144378190987, <<"typhoon">>, 90.3]] + # } def test_decode_data_from_get(self): - cols = [] - r0 = (tsrow_a, [ - (tscell_a, bd0, udef_a, udef_a, udef_a, udef_a), - (tscell_a, udef_a, 0, udef_a, udef_a, udef_a), - (tscell_a, udef_a, udef_a, udef_a, udef_a, 1.2), - (tscell_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a), - (tscell_a, udef_a, udef_a, udef_a, True, udef_a), - (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a), - (tscell_a, str1, udef_a, udef_a, udef_a, udef_a) - ]) - r1 = (tsrow_a, [ - (tscell_a, bd1, udef_a, udef_a, udef_a, udef_a), - (tscell_a, udef_a, 3, udef_a, udef_a, udef_a), - (tscell_a, udef_a, udef_a, udef_a, udef_a, 4.5), - (tscell_a, udef_a, udef_a, unix_time_millis(ts1), udef_a, udef_a), - (tscell_a, udef_a, udef_a, udef_a, False, udef_a), - (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a), - (tscell_a, str1, udef_a, udef_a, udef_a, udef_a) - ]) + 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) + r1 = (bd1, 3, 4.5, unix_time_millis(ts1), False, [], str1) rows = [r0, r1] - # { tsgetresp, [cols], [rows] } - rsp_data = tsgetresp_a, cols, rows # NB: Python tuple notation + # { tsgetresp, { [colnames], [coltypes] }, [rows] } + cols_t = colnames, coltypes + rsp_data = tsgetresp_a, cols_t, rows rsp_ttb = encode(rsp_data) tsobj = TsObject(None, self.table, [], []) @@ -120,22 +115,8 @@ def test_decode_data_from_get(self): self.assertEqual(r[6], dr[6][1].encode('ascii')) def test_encode_data_for_put(self): - r0 = (tsrow_a, [ - (tscell_a, bd0, udef_a, udef_a, udef_a, udef_a), - (tscell_a, udef_a, 0, udef_a, udef_a, udef_a), - (tscell_a, udef_a, udef_a, udef_a, udef_a, 1.2), - (tscell_a, udef_a, udef_a, unix_time_millis(ts0), udef_a, udef_a), - (tscell_a, udef_a, udef_a, udef_a, True, udef_a), - (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a) - ]) - r1 = (tsrow_a, [ - (tscell_a, bd1, udef_a, udef_a, udef_a, udef_a), - (tscell_a, udef_a, 3, udef_a, udef_a, udef_a), - (tscell_a, udef_a, udef_a, udef_a, udef_a, 4.5), - (tscell_a, udef_a, udef_a, unix_time_millis(ts1), udef_a, udef_a), - (tscell_a, udef_a, udef_a, udef_a, False, udef_a), - (tscell_a, udef_a, udef_a, udef_a, udef_a, udef_a) - ]) + 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) From 1eed9f253ba80f1c8945c5cc93091192a318ba88 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Thu, 7 Apr 2016 07:49:47 -0700 Subject: [PATCH 39/44] TTB and PBUF timeseries tests working correctly. --- riak/codecs/__init__.py | 2 +- riak/codecs/pbuf.py | 26 ++++--- riak/codecs/ttb.py | 70 ++++++++----------- ..._timeseries.py => test_timeseries_pbuf.py} | 69 +++++++++--------- riak/tests/test_timeseries_ttb.py | 42 ++++------- riak/transports/tcp/connection.py | 1 - riak/transports/tcp/transport.py | 4 +- riak/ts_object.py | 28 +++++--- 8 files changed, 114 insertions(+), 128 deletions(-) rename riak/tests/{test_timeseries.py => test_timeseries_pbuf.py} (89%) diff --git a/riak/codecs/__init__.py b/riak/codecs/__init__.py index e2ef78de..e356b5f9 100644 --- a/riak/codecs/__init__.py +++ b/riak/codecs/__init__.py @@ -1,5 +1,5 @@ import collections -import riak.pb.messages + from riak import RiakError Msg = collections.namedtuple('Msg', diff --git a/riak/codecs/pbuf.py b/riak/codecs/pbuf.py index 2e08ec5c..b17c5b6a 100644 --- a/riak/codecs/pbuf.py +++ b/riak/codecs/pbuf.py @@ -11,6 +11,7 @@ from riak.codecs import Codec, Msg from riak.content import RiakContent 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 @@ -98,7 +99,8 @@ def parse_msg(self, msg_code, data): 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) + 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)) @@ -782,16 +784,20 @@ def decode_timeseries(self, resp, tsobj): :param tsobj: a TsObject :type tsobj: TsObject """ - if tsobj.columns is not None: + if resp.columns is not None: + col_names = [] + col_types = [] 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)) + col_names.append(bytes_to_str(col.name)) + 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_row(self, tsrow, tscols=None): """ diff --git a/riak/codecs/ttb.py b/riak/codecs/ttb.py index 63479a55..f852dea7 100644 --- a/riak/codecs/ttb.py +++ b/riak/codecs/ttb.py @@ -8,6 +8,7 @@ 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 @@ -18,8 +19,7 @@ tsgetresp_a = Atom('tsgetresp') tsputreq_a = Atom('tsputreq') tsdelreq_a = Atom('tsdelreq') -tsrow_a = Atom('tsrow') -tscell_a = Atom('tscell') +timestamp_a = Atom('timestamp') # TODO RTS-842 MSG_CODE_TS_TTB = 104 @@ -139,17 +139,6 @@ def decode_timeseries(self, resp_ttb, tsobj): :param tsobj: a TsObject :type tsobj: TsObject """ - # TODO TODO RTS-842 CLIENTS-814 GH-445 - # TODO COLUMNS - # TODO TODO RTS-842 CLIENTS-814 GH-445 - # 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) - # - # TODO RTS-842 is this correct? if resp_ttb is None: return tsobj @@ -157,45 +146,42 @@ def decode_timeseries(self, resp_ttb, tsobj): if resp_a == rpberrorresp_a: self.process_err_ttb(resp_ttb) elif resp_a == tsgetresp_a: - # TODO resp_cols = resp_ttb[1] + resp_cols = resp_ttb[1] + tsobj.columns = self.decode_timeseries_cols(resp_cols) resp_rows = resp_ttb[2] - for row_ttb in resp_rows: + tsobj.rows = [] + for resp_row in resp_rows: tsobj.rows.append( - self.decode_timeseries_row(row_ttb, None)) + self.decode_timeseries_row(resp_row, resp_cols)) else: raise RiakError("Unknown TTB response type: {}".format(resp_a)) - def decode_timeseries_row(self, tsrow_ttb, tscols=None): + def decode_timeseries_cols(self, tscols): + cn, ct = tscols + cnames = [bytes_to_str(cname) for cname in cn] + ctypes = [str(ctype) for ctype in ct] + return TsColumns(cnames, ctypes) + + def decode_timeseries_row(self, tsrow, tscols): """ Decodes a TTB-encoded TsRow into a list - :param tsrow: the TTB-encoded TsRow to decode. - :type tsrow: TTB encoded row - :param tscols: the TTB-encoded TsColumn data to help decode. + :param tsrow: the TTB decoded TsRow to decode. + :type tsrow: TTB dncoded row + :param tscols: the TTB decoded TsColumn data to help decode rows. :type tscols: list :rtype list """ - if tsrow_ttb[0] == tsrow_a: - row = [] - for tsc_ttb in tsrow_ttb[1]: - if tsc_ttb[0] == tscell_a: - if tsc_ttb[1] != udef_a: - row.append(tsc_ttb[1]) - elif tsc_ttb[2] != udef_a: - row.append(tsc_ttb[2]) - elif tsc_ttb[3] != udef_a: - row.append( - datetime_from_unix_time_millis(tsc_ttb[3])) - elif tsc_ttb[4] != udef_a: - row.append(tsc_ttb[4]) - elif tsc_ttb[5] != udef_a: - row.append(tsc_ttb[5]) - else: - row.append(None) + cn, ct = tscols + row = [] + for i, cell in enumerate(tsrow): + if cell is None: + row.append(None) + elif cell is list and len(cell) == 0: + row.append(None) + else: + if ct[i] == timestamp_a: + row.append(datetime_from_unix_time_millis(cell)) else: - raise RiakError( - "Expected tscell atom, got: {}".format(tsc_ttb[0])) - else: - raise RiakError( - "Expected tsrow atom, got: {}".format(tsrow_ttb[0])) + row.append(cell) return row diff --git a/riak/tests/test_timeseries.py b/riak/tests/test_timeseries_pbuf.py similarity index 89% rename from riak/tests/test_timeseries.py rename to riak/tests/test_timeseries_pbuf.py index 5d9e81cb..93e088a2 100644 --- a/riak/tests/test_timeseries.py +++ b/riak/tests/test_timeseries_pbuf.py @@ -32,7 +32,7 @@ @unittest.skipUnless(is_timeseries_supported(), "Timeseries not supported") -class TimeseriesUnitTests(unittest.TestCase): +class TimeseriesPbufUnitTests(unittest.TestCase): @classmethod def setUpClass(cls): cls.ts0ms = unix_time_millis(ts0) @@ -159,24 +159,25 @@ def test_decode_data_from_query(self): r1c4 = r1.cells.add() r1c4.boolean_value = self.rows[1][4] - tsobj = TsObject(None, self.table, [], []) + tsobj = TsObject(None, self.table) c = PbufCodec() 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')) + 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], TsColumnType.Value('VARCHAR')) + self.assertEqual(cn[1], 'col_integer') + self.assertEqual(ct[1], TsColumnType.Value('SINT64')) + self.assertEqual(cn[2], 'col_double') + self.assertEqual(ct[2], TsColumnType.Value('DOUBLE')) + self.assertEqual(cn[3], 'col_timestamp') + self.assertEqual(ct[3], TsColumnType.Value('TIMESTAMP')) + self.assertEqual(cn[4], 'col_boolean') + self.assertEqual(ct[4], TsColumnType.Value('BOOLEAN')) r0 = tsobj.rows[0] self.assertEqual(bytes_to_str(r0[0]), self.rows[0][0]) @@ -195,12 +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 @@ -245,9 +246,15 @@ def setUpClass(cls): ] 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(bytes_to_str(row[0]), 'hash1') @@ -272,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(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(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, 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(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, 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 = """ @@ -320,8 +321,7 @@ def test_query_that_returns_no_data(self): """ query = fmt.format(table=table_name) ts_obj = self.client.ts_query(table_name, query) - 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_no_data_using_interpolation(self): query = """ @@ -331,8 +331,7 @@ def test_query_that_returns_no_data_using_interpolation(self): user = 'user1' """ ts_obj = self.client.ts_query(table_name, query) - self.assertEqual(len(ts_obj.columns), 0) - self.assertEqual(len(ts_obj.rows), 0) + self.validate_len(ts_obj, 0) def test_query_that_matches_some_data(self): fmt = """ diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index 811f90b3..8936567c 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -61,7 +61,8 @@ def test_encode_data_for_get(self): # {tsgetresp, # { - # [<<"geohash">>, <<"user">>, <<"time">>, <<"weather">>, <<"temperature">>], + # [<<"geohash">>, <<"user">>, <<"time">>, + # <<"weather">>, <<"temperature">>], # [varchar, varchar, timestamp, varchar, double] # }, # [[<<"hash1">>, <<"user2">>, 144378190987, <<"typhoon">>, 90.3]] @@ -71,48 +72,35 @@ def test_decode_data_from_get(self): "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) - r1 = (bd1, 3, 4.5, unix_time_millis(ts1), False, [], str1) + 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] } cols_t = colnames, coltypes rsp_data = tsgetresp_a, cols_t, rows rsp_ttb = encode(rsp_data) - tsobj = TsObject(None, self.table, [], []) + tsobj = TsObject(None, self.table, []) c = TtbCodec() c.decode_timeseries(decode(rsp_ttb), tsobj) for i in range(0, 1): - self.assertEqual(tsrow_a, rows[i][0]) - dr = rows[i][1] + dr = rows[i] r = tsobj.rows[i] # encoded - - # cells - self.assertEqual(tscell_a, dr[0][0]) - self.assertEqual(r[0], dr[0][1].encode('utf-8')) - - self.assertEqual(tscell_a, dr[1][0]) - self.assertEqual(r[1], dr[1][2]) - - self.assertEqual(tscell_a, dr[2][0]) - self.assertEqual(r[2], dr[2][5]) - - self.assertEqual(tscell_a, dr[3][0]) - dt = datetime_from_unix_time_millis(dr[3][3]) + 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) - - self.assertEqual(tscell_a, dr[4][0]) if i == 0: self.assertEqual(r[4], True) else: self.assertEqual(r[4], False) - - self.assertEqual(tscell_a, dr[5][0]) - self.assertEqual(r[5], None) - - self.assertEqual(tscell_a, dr[6][0]) - self.assertEqual(r[6], dr[6][1].encode('ascii')) + self.assertEqual(r[5], []) + 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, []) diff --git a/riak/transports/tcp/connection.py b/riak/transports/tcp/connection.py index d0ffab13..a9d75603 100644 --- a/riak/transports/tcp/connection.py +++ b/riak/transports/tcp/connection.py @@ -1,4 +1,3 @@ -import logging import socket import struct diff --git a/riak/transports/tcp/transport.py b/riak/transports/tcp/transport.py index d7e7954b..02d4b5e4 100644 --- a/riak/transports/tcp/transport.py +++ b/riak/transports/tcp/transport.py @@ -142,7 +142,7 @@ def ts_get(self, table, key): 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, [], None) + tsobj = TsObject(self._client, table) codec.decode_timeseries(resp, tsobj) return tsobj @@ -168,7 +168,7 @@ def ts_query(self, table, query, interpolations=None): 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, [], []) + tsobj = TsObject(self._client, table) codec.decode_timeseries(resp, tsobj) return tsobj 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): """ From 7736ae3ef78cf0221b4dec5aea9a137d40f42b55 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Thu, 7 Apr 2016 08:29:23 -0700 Subject: [PATCH 40/44] Decode column type the same way for ttb and pbuf --- riak/codecs/pbuf.py | 22 +++++++++++++++++++--- riak/tests/test_timeseries_pbuf.py | 10 +++++----- 2 files changed, 24 insertions(+), 8 deletions(-) diff --git a/riak/codecs/pbuf.py b/riak/codecs/pbuf.py index b17c5b6a..a976d1ff 100644 --- a/riak/codecs/pbuf.py +++ b/riak/codecs/pbuf.py @@ -10,12 +10,12 @@ 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 -from riak.pb.riak_ts_pb2 import TsColumnType def _invert(d): @@ -789,7 +789,8 @@ def decode_timeseries(self, resp, tsobj): col_types = [] for col in resp.columns: col_names.append(bytes_to_str(col.name)) - col_types.append(col.type) + col_type = self.decode_timeseries_col_type(col.type) + col_types.append(col_type) tsobj.columns = TsColumns(col_names, col_types) tsobj.rows = [] @@ -799,6 +800,22 @@ def decode_timeseries(self, resp, tsobj): 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 @@ -818,7 +835,6 @@ def decode_timeseries_row(self, tsrow, tscols=None): if col and col.type != TsColumnType.Value('VARCHAR'): raise TypeError('expected VARCHAR column') else: - # TODO RTS-842 - keep as bytes? row.append(cell.varchar_value) elif cell.HasField('sint64_value'): if col and col.type != TsColumnType.Value('SINT64'): diff --git a/riak/tests/test_timeseries_pbuf.py b/riak/tests/test_timeseries_pbuf.py index 93e088a2..9eddcc81 100644 --- a/riak/tests/test_timeseries_pbuf.py +++ b/riak/tests/test_timeseries_pbuf.py @@ -169,15 +169,15 @@ def test_decode_data_from_query(self): cn, ct = tsobj.columns self.assertEqual(cn[0], 'col_varchar') - self.assertEqual(ct[0], TsColumnType.Value('VARCHAR')) + self.assertEqual(ct[0], 'varchar') self.assertEqual(cn[1], 'col_integer') - self.assertEqual(ct[1], TsColumnType.Value('SINT64')) + self.assertEqual(ct[1], 'sint64') self.assertEqual(cn[2], 'col_double') - self.assertEqual(ct[2], TsColumnType.Value('DOUBLE')) + self.assertEqual(ct[2], 'double') self.assertEqual(cn[3], 'col_timestamp') - self.assertEqual(ct[3], TsColumnType.Value('TIMESTAMP')) + self.assertEqual(ct[3], 'timestamp') self.assertEqual(cn[4], 'col_boolean') - self.assertEqual(ct[4], TsColumnType.Value('BOOLEAN')) + self.assertEqual(ct[4], 'boolean') r0 = tsobj.rows[0] self.assertEqual(bytes_to_str(r0[0]), self.rows[0][0]) From 57daac0c451135a3a8964279603a0a2d190412c3 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Wed, 13 Apr 2016 16:19:48 -0700 Subject: [PATCH 41/44] More work on TTB encoding --- riak/codecs/ttb.py | 29 +++++++++++++++++++++-------- riak/tests/test_timeseries_ttb.py | 26 +++++++++++++++++++++++--- riak/transports/tcp/transport.py | 2 ++ 3 files changed, 46 insertions(+), 11 deletions(-) diff --git a/riak/codecs/ttb.py b/riak/codecs/ttb.py index f852dea7..353f955d 100644 --- a/riak/codecs/ttb.py +++ b/riak/codecs/ttb.py @@ -1,8 +1,6 @@ import datetime import six -import riak.pb.messages - from erlastic import encode, decode from erlastic.types import Atom @@ -17,6 +15,9 @@ 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') tsdelreq_a = Atom('tsdelreq') timestamp_a = Atom('timestamp') @@ -34,9 +35,7 @@ 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 and \ - msg_code != riak.pb.messages.MSG_CODE_TS_GET_RESP and \ - msg_code != riak.pb.messages.MSG_CODE_TS_PUT_RESP: + if msg_code != MSG_CODE_TS_TTB: raise RiakError("TTB can't parse code: {}".format(msg_code)) if len(data) > 0: decoded = decode(data) @@ -95,8 +94,7 @@ def encode_timeseries_keyreq(self, table, key, is_delete=False): return Msg(mc, encode(req), rc) def validate_timeseries_put_resp(self, resp_code, resp): - if resp is None and \ - resp_code == riak.pb.messages.MSG_CODE_TS_PUT_RESP: + if resp is None and resp_code == MSG_CODE_TS_TTB: return True if resp is not None: return True @@ -129,6 +127,16 @@ def encode_timeseries_put(self, tsobj): 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 @@ -142,10 +150,13 @@ def decode_timeseries(self, resp_ttb, tsobj): if resp_ttb is None: return tsobj + import sys resp_a = resp_ttb[0] + sys.stderr.write("resp_a: {}".format(resp_a)) if resp_a == rpberrorresp_a: self.process_err_ttb(resp_ttb) - elif resp_a == tsgetresp_a: + elif resp_a == tsgetresp_a or \ + resp_a == tsqueryresp_a: resp_cols = resp_ttb[1] tsobj.columns = self.decode_timeseries_cols(resp_cols) resp_rows = resp_ttb[2] @@ -172,9 +183,11 @@ def decode_timeseries_row(self, tsrow, tscols): :type tscols: list :rtype list """ + import sys cn, ct = tscols row = [] for i, cell in enumerate(tsrow): + sys.stderr.write("\ncell: {}\n".format(cell)) if cell is None: row.append(None) elif cell is list and len(cell) == 0: diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index 8936567c..b79993cb 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -129,7 +129,21 @@ class TimeseriesTtbTests(IntegrationTestBase, unittest.TestCase): def setUpClass(cls): super(TimeseriesTtbTests, cls).setUpClass() - def test_store_and_fetch_ttb(self): + def test_query_that_returns_table_description(self): + import sys + 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 + sys.stderr.write("\n\nts_cols: {}\n\n".format(ts_cols)) + sys.stderr.write("\n\nrows: {}\n\n".format(ts_obj.rows)) + 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 @@ -149,12 +163,18 @@ def test_store_and_fetch_ttb(self): result = ts_obj.store() self.assertTrue(result) - for r in rows: + 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) - self.assertEqual(len(ts_obj.rows[0]), 5) + row = ts_obj.rows[0] + 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) diff --git a/riak/transports/tcp/transport.py b/riak/transports/tcp/transport.py index 02d4b5e4..58420767 100644 --- a/riak/transports/tcp/transport.py +++ b/riak/transports/tcp/transport.py @@ -60,6 +60,8 @@ def _get_codec(self, msg_code): 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 From 4eec9379ca48ae6a3c40f36230073755b2901e24 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Thu, 14 Apr 2016 14:16:06 -0700 Subject: [PATCH 42/44] finish up Riak TS 1.3 changes --- riak/codecs/ttb.py | 49 ++++++++++++++++-------------- riak/tests/test_timeseries_pbuf.py | 7 +++-- riak/tests/test_timeseries_ttb.py | 34 +++++++++++++-------- 3 files changed, 53 insertions(+), 37 deletions(-) diff --git a/riak/codecs/ttb.py b/riak/codecs/ttb.py index 353f955d..7333943b 100644 --- a/riak/codecs/ttb.py +++ b/riak/codecs/ttb.py @@ -19,6 +19,7 @@ tsqueryresp_a = Atom('tsqueryresp') tsinterpolation_a = Atom('tsinterpolation') tsputreq_a = Atom('tsputreq') +tsputresp_a = Atom('tsputresp') tsdelreq_a = Atom('tsdelreq') timestamp_a = Atom('timestamp') @@ -150,50 +151,52 @@ def decode_timeseries(self, resp_ttb, tsobj): if resp_ttb is None: return tsobj - import sys resp_a = resp_ttb[0] - sys.stderr.write("resp_a: {}".format(resp_a)) if resp_a == rpberrorresp_a: self.process_err_ttb(resp_ttb) - elif resp_a == tsgetresp_a or \ - resp_a == tsqueryresp_a: - resp_cols = resp_ttb[1] - tsobj.columns = self.decode_timeseries_cols(resp_cols) - resp_rows = resp_ttb[2] - tsobj.rows = [] - for resp_row in resp_rows: - tsobj.rows.append( - self.decode_timeseries_row(resp_row, resp_cols)) + elif 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, tscols): - cn, ct = tscols - cnames = [bytes_to_str(cname) for cname in cn] - ctypes = [str(ctype) for ctype in ct] + 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, tscols): + 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 tscols: the TTB decoded TsColumn data to help decode rows. - :type tscols: list + :param tsct: the TTB decoded column types (atoms). + :type tsct: list :rtype list """ - import sys - cn, ct = tscols row = [] for i, cell in enumerate(tsrow): - sys.stderr.write("\ncell: {}\n".format(cell)) if cell is None: row.append(None) - elif cell is list and len(cell) == 0: + elif isinstance(cell, list) and len(cell) == 0: row.append(None) else: - if ct[i] == timestamp_a: + if tsct[i] == timestamp_a: row.append(datetime_from_unix_time_millis(cell)) else: row.append(cell) diff --git a/riak/tests/test_timeseries_pbuf.py b/riak/tests/test_timeseries_pbuf.py index 9eddcc81..50c0818c 100644 --- a/riak/tests/test_timeseries_pbuf.py +++ b/riak/tests/test_timeseries_pbuf.py @@ -416,8 +416,11 @@ def test_delete_single_value(self): key = ['hash1', 'user2', self.twentyFiveMinsAgo] rslt = self.client.ts_delete(table_name, key) self.assertTrue(rslt) - with self.assertRaises(RiakError): - self.client.ts_get(table_name, 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) diff --git a/riak/tests/test_timeseries_ttb.py b/riak/tests/test_timeseries_ttb.py index b79993cb..339dd040 100644 --- a/riak/tests/test_timeseries_ttb.py +++ b/riak/tests/test_timeseries_ttb.py @@ -63,9 +63,9 @@ def test_encode_data_for_get(self): # { # [<<"geohash">>, <<"user">>, <<"time">>, # <<"weather">>, <<"temperature">>], - # [varchar, varchar, timestamp, varchar, double] - # }, - # [[<<"hash1">>, <<"user2">>, 144378190987, <<"typhoon">>, 90.3]] + # [varchar, varchar, timestamp, varchar, double], + # [(<<"hash1">>, <<"user2">>, 144378190987, <<"typhoon">>, 90.3)] + # } # } def test_decode_data_from_get(self): colnames = ["varchar", "sint64", "double", "timestamp", @@ -77,12 +77,12 @@ def test_decode_data_from_get(self): r1 = (bd1, 3, 4.5, unix_time_millis(ts1), False, [], str1, None) rows = [r0, r1] - # { tsgetresp, { [colnames], [coltypes] }, [rows] } - cols_t = colnames, coltypes - rsp_data = tsgetresp_a, cols_t, rows + # { 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, []) + tsobj = TsObject(None, self.table) c = TtbCodec() c.decode_timeseries(decode(rsp_ttb), tsobj) @@ -98,7 +98,7 @@ def test_decode_data_from_get(self): self.assertEqual(r[4], True) else: self.assertEqual(r[4], False) - self.assertEqual(r[5], []) + self.assertEqual(r[5], None) self.assertEqual(r[6], dr[6].encode('ascii')) self.assertEqual(r[7], None) @@ -130,14 +130,11 @@ def setUpClass(cls): super(TimeseriesTtbTests, cls).setUpClass() def test_query_that_returns_table_description(self): - import sys 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 - sys.stderr.write("\n\nts_cols: {}\n\n".format(ts_cols)) - sys.stderr.write("\n\nrows: {}\n\n".format(ts_obj.rows)) self.assertEqual(len(ts_cols.names), 5) self.assertEqual(len(ts_cols.types), 5) row = ts_obj.rows[0] @@ -159,6 +156,19 @@ def test_store_and_fetch(self): ['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) @@ -172,7 +182,7 @@ def test_store_and_fetch(self): self.assertEqual(len(ts_cols.types), 5) self.assertEqual(len(ts_obj.rows), 1) row = ts_obj.rows[0] - exp = rows[i] + exp = exp_rows[i] self.assertEqual(len(row), 5) self.assertEqual(row, exp) From 2e5b0aab5846a919ceb52fc817ec9e6ca8cfba4f Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Tue, 19 Apr 2016 11:19:17 -0700 Subject: [PATCH 43/44] process rpberrorresp correctly --- riak/codecs/ttb.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/riak/codecs/ttb.py b/riak/codecs/ttb.py index 7333943b..c18912af 100644 --- a/riak/codecs/ttb.py +++ b/riak/codecs/ttb.py @@ -49,6 +49,7 @@ 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): @@ -151,10 +152,10 @@ def decode_timeseries(self, resp_ttb, tsobj): if resp_ttb is None: return tsobj + self.maybe_err_ttb(resp_ttb) + resp_a = resp_ttb[0] - if resp_a == rpberrorresp_a: - self.process_err_ttb(resp_ttb) - elif resp_a == tsputresp_a: + if resp_a == tsputresp_a: return elif resp_a == tsgetresp_a or resp_a == tsqueryresp_a: resp_data = resp_ttb[1] From c6ecbfb3dead3552c239e7c56661c9d67397d0b6 Mon Sep 17 00:00:00 2001 From: Luke Bakken Date: Fri, 22 Apr 2016 11:14:17 -0700 Subject: [PATCH 44/44] make lint happy --- riak/codecs/ttb.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/riak/codecs/ttb.py b/riak/codecs/ttb.py index c18912af..6bb81f0d 100644 --- a/riak/codecs/ttb.py +++ b/riak/codecs/ttb.py @@ -164,14 +164,16 @@ def decode_timeseries(self, resp_ttb, tsobj): 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) + 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)) + raise RiakError( + "Expected 3-tuple in response, got: {}".format(resp_data)) else: raise RiakError("Unknown TTB response type: {}".format(resp_a))