def encode_get(self,
robj,
r=None,
pr=None,
timeout=None,
basic_quorum=None,
notfound_ok=None,
head_only=False):
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)
req.head = head_only
mc = riak.pb.messages.MSG_CODE_GET_REQ
rc = riak.pb.messages.MSG_CODE_GET_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_timeseries_put(self, tsobj):
'''
Returns an Erlang-TTB encoded tuple with the appropriate data and
metadata from a TsObject.
:param tsobj: a TsObject
:type tsobj: TsObject
:rtype: term-to-binary encoded object
'''
if tsobj.columns:
raise NotImplementedError('columns are not used')
if tsobj.rows and isinstance(tsobj.rows, list):
req_rows = []
for row in tsobj.rows:
req_r = []
for cell in row:
req_r.append(self.encode_to_ts_cell(cell))
req_rows.append(tuple(req_r))
req = tsputreq_a, tsobj.table.name, [], req_rows
mc = MSG_CODE_TS_TTB_MSG
rc = MSG_CODE_TS_TTB_MSG
return Msg(mc, encode(req), rc)
else:
raise RiakError("TsObject requires a list of rows")
def encode_search(self, index, query, **kwargs):
req = riak.pb.riak_search_pb2.RpbSearchQueryReq(
index=str_to_bytes(index), q=str_to_bytes(query))
self.encode_search_query(req, **kwargs)
mc = riak.pb.messages.MSG_CODE_SEARCH_QUERY_REQ
rc = riak.pb.messages.MSG_CODE_SEARCH_QUERY_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_auth(self, username, password):
req = riak.pb.riak_pb2.RpbAuthReq()
req.user = str_to_bytes(username)
req.password = str_to_bytes(password)
mc = riak.pb.messages.MSG_CODE_AUTH_REQ
rc = riak.pb.messages.MSG_CODE_AUTH_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_clear_bucket_props(self, bucket):
req = riak.pb.riak_pb2.RpbResetBucketReq()
req.bucket = str_to_bytes(bucket.name)
self._add_bucket_type(req, bucket.bucket_type)
mc = riak.pb.messages.MSG_CODE_RESET_BUCKET_REQ
rc = riak.pb.messages.MSG_CODE_RESET_BUCKET_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_timeseries_put(self, tsobj):
"""
Fills an TsPutReq message with the appropriate data and
metadata from a TsObject.
:param tsobj: a TsObject
:type tsobj: TsObject
:param req: the protobuf message to fill
:type req: riak.pb.riak_ts_pb2.TsPutReq
"""
req = riak.pb.riak_ts_pb2.TsPutReq()
req.table = str_to_bytes(tsobj.table.name)
if tsobj.columns:
raise NotImplementedError("columns are not implemented yet")
if tsobj.rows and isinstance(tsobj.rows, list):
for row in tsobj.rows:
tsr = req.rows.add() # NB: type TsRow
if not isinstance(row, list):
raise ValueError("TsObject row must be a list of values")
for cell in row:
tsc = tsr.cells.add() # NB: type TsCell
self.encode_to_ts_cell(cell, tsc)
else:
raise RiakError("TsObject requires a list of rows")
mc = riak.pb.messages.MSG_CODE_TS_PUT_REQ
rc = riak.pb.messages.MSG_CODE_TS_PUT_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_create_search_schema(self, schema, content):
scma = riak.pb.riak_yokozuna_pb2.RpbYokozunaSchema(
name=str_to_bytes(schema), content=str_to_bytes(content))
req = riak.pb.riak_yokozuna_pb2.RpbYokozunaSchemaPutReq(schema=scma)
mc = riak.pb.messages.MSG_CODE_YOKOZUNA_SCHEMA_PUT_REQ
rc = riak.pb.messages.MSG_CODE_PUT_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_put(self, robj, w=None, dw=None, pw=None,
return_body=True, if_none_match=False,
timeout=None):
bucket = robj.bucket
req = riak.pb.riak_kv_pb2.RpbPutReq()
if w:
req.w = self.encode_quorum(w)
if dw:
req.dw = self.encode_quorum(dw)
if self._quorum_controls and pw:
req.pw = self.encode_quorum(pw)
if return_body:
req.return_body = 1
if if_none_match:
req.if_none_match = 1
if self._client_timeouts and timeout:
req.timeout = timeout
req.bucket = str_to_bytes(bucket.name)
self._add_bucket_type(req, bucket.bucket_type)
if robj.key:
req.key = str_to_bytes(robj.key)
if robj.vclock:
req.vclock = robj.vclock.encode('binary')
self.encode_content(robj, req.content)
mc = riak.pb.messages.MSG_CODE_PUT_REQ
rc = riak.pb.messages.MSG_CODE_PUT_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_set_bucket_type_props(self, bucket_type, props):
req = riak.pb.riak_pb2.RpbSetBucketTypeReq()
req.type = str_to_bytes(bucket_type.name)
self.encode_bucket_props(props, req)
mc = riak.pb.messages.MSG_CODE_SET_BUCKET_TYPE_REQ
rc = riak.pb.messages.MSG_CODE_SET_BUCKET_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_stream_mapred(self, content):
req = riak.pb.riak_kv_pb2.RpbMapRedReq()
req.request = str_to_bytes(content)
req.content_type = str_to_bytes("application/json")
mc = riak.pb.messages.MSG_CODE_MAP_RED_REQ
rc = riak.pb.messages.MSG_CODE_MAP_RED_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_delete(self, robj, rw=None, r=None,
w=None, dw=None, pr=None, pw=None,
timeout=None):
req = riak.pb.riak_kv_pb2.RpbDelReq()
if rw:
req.rw = self.encode_quorum(rw)
if r:
req.r = self.encode_quorum(r)
if w:
req.w = self.encode_quorum(w)
if dw:
req.dw = self.encode_quorum(dw)
if self._quorum_controls:
if pr:
req.pr = self.encode_quorum(pr)
if pw:
req.pw = self.encode_quorum(pw)
if self._client_timeouts and timeout:
req.timeout = timeout
use_vclocks = (self._tombstone_vclocks and
hasattr(robj, 'vclock') and robj.vclock)
if use_vclocks:
req.vclock = robj.vclock.encode('binary')
bucket = robj.bucket
req.bucket = str_to_bytes(bucket.name)
self._add_bucket_type(req, bucket.bucket_type)
req.key = str_to_bytes(robj.key)
mc = riak.pb.messages.MSG_CODE_DEL_REQ
rc = riak.pb.messages.MSG_CODE_DEL_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_timeseries_listkeysreq(self, table, timeout=None):
req = riak.pb.riak_ts_pb2.TsListKeysReq()
req.table = str_to_bytes(table.name)
if self._client_timeouts and timeout:
req.timeout = timeout
mc = riak.pb.messages.MSG_CODE_TS_LIST_KEYS_REQ
rc = riak.pb.messages.MSG_CODE_TS_LIST_KEYS_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_get_preflist(self, bucket, key):
req = riak.pb.riak_kv_pb2.RpbGetBucketKeyPreflistReq()
req.bucket = str_to_bytes(bucket.name)
req.key = str_to_bytes(key)
req.type = str_to_bytes(bucket.bucket_type.name)
mc = riak.pb.messages.MSG_CODE_GET_BUCKET_KEY_PREFLIST_REQ
rc = riak.pb.messages.MSG_CODE_GET_BUCKET_KEY_PREFLIST_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_index_req(self, bucket, index, startkey, endkey=None,
return_terms=None, max_results=None,
continuation=None, timeout=None, term_regex=None,
streaming=False):
"""
Encodes a secondary index request into the protobuf message.
:param bucket: the bucket whose index to query
:type bucket: string
:param index: the index to query
:type index: string
:param startkey: the value or beginning of the range
:type startkey: integer, string
:param endkey: the end of the range
:type endkey: integer, string
:param return_terms: whether to return the index term with the key
:type return_terms: bool
:param max_results: the maximum number of results to return (page size)
:type max_results: integer
:param continuation: the opaque continuation returned from a
previous paginated request
:type continuation: string
:param timeout: a timeout value in milliseconds, or 'infinity'
:type timeout: int
:param term_regex: a regular expression used to filter index terms
:type term_regex: string
:param streaming: encode as streaming request
:type streaming: bool
:rtype riak.pb.riak_kv_pb2.RpbIndexReq
"""
req = riak.pb.riak_kv_pb2.RpbIndexReq(
bucket=str_to_bytes(bucket.name),
index=str_to_bytes(index))
self._add_bucket_type(req, bucket.bucket_type)
if endkey is not None:
req.qtype = riak.pb.riak_kv_pb2.RpbIndexReq.range
req.range_min = str_to_bytes(str(startkey))
req.range_max = str_to_bytes(str(endkey))
else:
req.qtype = riak.pb.riak_kv_pb2.RpbIndexReq.eq
req.key = str_to_bytes(str(startkey))
if return_terms is not None:
req.return_terms = return_terms
if max_results:
req.max_results = max_results
if continuation:
req.continuation = str_to_bytes(continuation)
if timeout:
if timeout == 'infinity':
req.timeout = 0
else:
req.timeout = timeout
if term_regex:
req.term_regex = str_to_bytes(term_regex)
req.stream = streaming
mc = riak.pb.messages.MSG_CODE_INDEX_REQ
rc = riak.pb.messages.MSG_CODE_INDEX_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_timeseries_query(self, table, query, interpolations=None):
req = riak.pb.riak_ts_pb2.TsQueryReq()
q = query
if '{table}' in q:
q = q.format(table=table.name)
req.query.base = str_to_bytes(q)
mc = riak.pb.messages.MSG_CODE_TS_QUERY_REQ
rc = riak.pb.messages.MSG_CODE_TS_QUERY_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_stream_keys(self, bucket, timeout=None):
req = riak.pb.riak_kv_pb2.RpbListKeysReq()
req.bucket = str_to_bytes(bucket.name)
if self._client_timeouts and timeout:
req.timeout = timeout
self._add_bucket_type(req, bucket.bucket_type)
mc = riak.pb.messages.MSG_CODE_LIST_KEYS_REQ
rc = riak.pb.messages.MSG_CODE_LIST_KEYS_RESP
return Msg(mc, req.SerializeToString(), rc)
def 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, udef_a
mc = MSG_CODE_TS_TTB_MSG
rc = MSG_CODE_TS_TTB_MSG
return Msg(mc, encode(req), rc)
def encode_fetch_datatype(self, bucket, key, **kwargs):
req = riak.pb.riak_dt_pb2.DtFetchReq()
req.type = str_to_bytes(bucket.bucket_type.name)
req.bucket = str_to_bytes(bucket.name)
req.key = str_to_bytes(key)
self.encode_dt_options(req, **kwargs)
mc = riak.pb.messages.MSG_CODE_DT_FETCH_REQ
rc = riak.pb.messages.MSG_CODE_DT_FETCH_RESP
return Msg(mc, req.SerializeToString(), rc)
def get_server_info(self):
"""
Get information about the server
"""
# NB: can't do it this way due to recursion
# codec = self._get_codec(ttb_supported=False)
codec = PbufCodec()
msg = Msg(riak.pb.messages.MSG_CODE_GET_SERVER_INFO_REQ, None,
riak.pb.messages.MSG_CODE_GET_SERVER_INFO_RESP)
resp_code, resp = self._request(msg, codec)
return codec.decode_get_server_info(resp)
def encode_get_buckets(self, bucket_type, timeout=None, streaming=False):
# Bucket streaming landed in the same release as timeouts, so
# we don't need to check the capability.
req = riak.pb.riak_kv_pb2.RpbListBucketsReq()
req.stream = streaming
self._add_bucket_type(req, bucket_type)
if self._client_timeouts and timeout:
req.timeout = timeout
mc = riak.pb.messages.MSG_CODE_LIST_BUCKETS_REQ
rc = riak.pb.messages.MSG_CODE_LIST_BUCKETS_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_create_search_index(self, index, schema=None,
n_val=None, timeout=None):
index = str_to_bytes(index)
idx = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndex(name=index)
if schema:
idx.schema = str_to_bytes(schema)
if n_val:
idx.n_val = n_val
req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexPutReq(index=idx)
if timeout is not None:
req.timeout = timeout
mc = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_PUT_REQ
rc = riak.pb.messages.MSG_CODE_PUT_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_get_counter(self, bucket, key, **kwargs):
req = riak.pb.riak_kv_pb2.RpbCounterGetReq()
req.bucket = str_to_bytes(bucket.name)
req.key = str_to_bytes(key)
if kwargs.get('r') is not None:
req.r = self.encode_quorum(kwargs['r'])
if kwargs.get('pr') is not None:
req.pr = self.encode_quorum(kwargs['pr'])
if kwargs.get('basic_quorum') is not None:
req.basic_quorum = kwargs['basic_quorum']
if kwargs.get('notfound_ok') is not None:
req.notfound_ok = kwargs['notfound_ok']
mc = riak.pb.messages.MSG_CODE_COUNTER_GET_REQ
rc = riak.pb.messages.MSG_CODE_COUNTER_GET_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_update_counter(self, bucket, key, value, **kwargs):
req = riak.pb.riak_kv_pb2.RpbCounterUpdateReq()
req.bucket = str_to_bytes(bucket.name)
req.key = str_to_bytes(key)
req.amount = value
if kwargs.get('w') is not None:
req.w = self.encode_quorum(kwargs['w'])
if kwargs.get('dw') is not None:
req.dw = self.encode_quorum(kwargs['dw'])
if kwargs.get('pw') is not None:
req.pw = self.encode_quorum(kwargs['pw'])
if kwargs.get('returnvalue') is not None:
req.returnvalue = kwargs['returnvalue']
mc = riak.pb.messages.MSG_CODE_COUNTER_UPDATE_REQ
rc = riak.pb.messages.MSG_CODE_COUNTER_UPDATE_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_timeseries_keyreq(self, table, key, is_delete=False):
key_vals = None
if isinstance(key, list):
key_vals = key
else:
raise ValueError("key must be a list")
mc = MSG_CODE_TS_TTB_MSG
rc = MSG_CODE_TS_TTB_MSG
req_atom = tsgetreq_a
if is_delete:
req_atom = tsdelreq_a
# TODO FUTURE add timeout as last param
req = req_atom, table.name, \
[self.encode_to_ts_cell(k) for k in key_vals], udef_a
return Msg(mc, encode(req), rc)
def encode_update_datatype(self, datatype, **kwargs):
op = datatype.to_op()
type_name = datatype.type_name
if not op:
raise ValueError("No operation to send on datatype {!r}".
format(datatype))
req = riak.pb.riak_dt_pb2.DtUpdateReq()
req.bucket = str_to_bytes(datatype.bucket.name)
req.type = str_to_bytes(datatype.bucket.bucket_type.name)
if datatype.key:
req.key = str_to_bytes(datatype.key)
if datatype._context:
req.context = datatype._context
self.encode_dt_options(req, **kwargs)
self.encode_dt_op(type_name, req, op)
mc = riak.pb.messages.MSG_CODE_DT_UPDATE_REQ
rc = riak.pb.messages.MSG_CODE_DT_UPDATE_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_timeseries_keyreq(self, table, key, is_delete=False):
key_vals = None
if isinstance(key, list):
key_vals = key
else:
raise ValueError("key must be a list")
req = riak.pb.riak_ts_pb2.TsGetReq()
mc = riak.pb.messages.MSG_CODE_TS_GET_REQ
rc = riak.pb.messages.MSG_CODE_TS_GET_RESP
if is_delete:
req = riak.pb.riak_ts_pb2.TsDelReq()
mc = riak.pb.messages.MSG_CODE_TS_DEL_REQ
rc = riak.pb.messages.MSG_CODE_TS_DEL_RESP
req.table = str_to_bytes(table.name)
for cell in key_vals:
ts_cell = req.key.add()
self.encode_to_ts_cell(cell, ts_cell)
return Msg(mc, req.SerializeToString(), rc)
def encode_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 encode_ping(self):
return Msg(riak.pb.messages.MSG_CODE_PING_REQ, None,
riak.pb.messages.MSG_CODE_PING_RESP)
def encode_list_search_indexes(self):
req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexGetReq()
mc = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_REQ
rc = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_GET_RESP
return Msg(mc, req.SerializeToString(), rc)
def encode_delete_search_index(self, index):
req = riak.pb.riak_yokozuna_pb2.RpbYokozunaIndexDeleteReq(
name=str_to_bytes(index))
mc = riak.pb.messages.MSG_CODE_YOKOZUNA_INDEX_DELETE_REQ
rc = riak.pb.messages.MSG_CODE_DEL_RESP
return Msg(mc, req.SerializeToString(), rc)
