async def put(self, robj, return_body=True):
bucket = robj.bucket
req = riak_pb.RpbPutReq()
if return_body:
req.return_body = 1
req.bucket = str_to_bytes(bucket.name)
self._add_bucket_type(req, bucket.bucket_type)
if robj.key:
req.key = str_to_bytes(robj.key)
if robj.vclock:
req.vclock = robj.vclock.encode('binary')
self._encode_content(robj, req.content)
msg_code, resp = await self._request(messages.MSG_CODE_PUT_REQ, req,
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
async def put(self,
robj,
w=None,
dw=None,
pw=None,
return_body=True,
if_none_match=False,
timeout=None):
bucket = robj.bucket
req = riak_kv_pb2.RpbPutReq()
if w:
req.w = self._encode_quorum(w)
if dw:
req.dw = self._encode_quorum(dw)
if pw:
req.pw = self._encode_quorum(pw)
if return_body:
req.return_body = 1
if if_none_match:
req.if_none_match = 1
if timeout:
req.timeout = timeout
req.bucket = str_to_bytes(bucket.name)
self._add_bucket_type(req, bucket.bucket_type)
if robj.key:
req.key = str_to_bytes(robj.key)
if robj.vclock:
req.vclock = robj.vclock.encode('binary')
self._encode_content(robj, req.content)
msg_code, resp = await self._request(messages.MSG_CODE_PUT_REQ, req,
messages.MSG_CODE_PUT_RESP)
if resp is not None:
if resp.HasField('key'):
robj.key = bytes_to_str(resp.key)
if resp.HasField('vclock'):
robj.vclock = VClock(resp.vclock, 'binary')
if resp.content:
self._decode_contents(resp.content, robj)
elif not robj.key:
raise RiakError("missing response object")
return robj
def _parse_msg(self):
''' Parse protobuf message'''
self._msg = self._data[self.HEADER_LENGTH:self.HEADER_LENGTH +
self._msglen]
self.msg_code, = struct.unpack("B", self._msg[:1])
if self.msg_code is messages.MSG_CODE_ERROR_RESP:
error = self._get_pb_msg(self.msg_code, self._msg[1:])
logger.error('Riak error message recieved: %s',
bytes_to_str(error.errmsg))
raise RiakError(bytes_to_str(error.errmsg))
elif self.msg_code in messages.MESSAGE_CLASSES:
logger.debug('Normal message with code %d received', self.msg_code)
self.msg = self._get_pb_msg(self.msg_code, self._msg[1:])
else:
logger.error('Unknown message received [%d]', self.msg_code)
async def get_pbo(self):
if self._parse_header():
if self._check_eof():
self._parse_msg()
self._grow_tail()
return self.msg_code, self.msg
while not self.at_eof():
chunk = await self._reader.read(MAX_CHUNK_SIZE)
self._data.extend(chunk)
self._parse_header()
if self._check_eof():
self._parse_msg()
self._grow_tail()
return self.msg_code, self.msg
raise RiakError('Incomplete message')
def _encode_content(self, robj, rpb_content):
'''
Fills an RpbContent message with the appropriate data and
metadata from a RiakObject.
:param robj: a RiakObject
:type robj: RiakObject
:param rpb_content: the protobuf message to fill
:type rpb_content: riak_pb2.RpbContent
'''
if robj.content_type:
rpb_content.content_type = robj.content_type.encode()
if robj.charset:
rpb_content.charset = robj.charset.encode()
if robj.content_encoding:
rpb_content.content_encoding = robj.content_encoding
for uk in robj.usermeta:
pair = rpb_content.usermeta.add()
pair.key = uk.encode()
pair.value = robj.usermeta[uk].encode()
for link in robj.links:
pb_link = rpb_content.links.add()
try:
bucket, key, tag = link
except ValueError:
raise RiakError("Invalid link tuple %s" % link)
pb_link.bucket = _encode_link_field(bucket)
pb_link.key = _encode_link_field(key)
if tag:
pb_link.tag = _encode_link_field(tag)
else:
pb_link.tag = b''
for field, value in robj.indexes:
if isinstance(value, int):
value = str(value)
pair = rpb_content.indexes.add()
pair.key = str_to_bytes(field)
pair.value = value.encode()
rpb_content.value = robj.encoded_data
def add_index(self, field, value):
'''
add_index(field, value)
Tag this object with the specified field/value pair for
indexing.
:param field: The index field.
:type field: string
:param value: The index value.
:type value: string or integer
:rtype: :class:`RiakObject <riak.riak_object.RiakObject>`
'''
if field[-4:] not in ("_bin", "_int"):
raise RiakError("Riak 2i fields must end with either '_bin'"
" or '_int'.")
self.indexes.add((field, value))
return self._robject
def remove_index(self, field=None, value=None):
'''
remove_index(field=None, value=None)
Remove the specified field/value pair as an index on this
object.
:param field: The index field.
:type field: string
:param value: The index value.
:type value: string or integer
:rtype: :class:`RiakObject <riak.riak_object.RiakObject>`
'''
if not field and not value:
self.indexes.clear()
elif field and not value:
for index in [x for x in self.indexes if x[0] == field]:
self.indexes.remove(index)
elif field and value:
self.indexes.remove((field, value))
else:
raise RiakError("Cannot pass value without a field"
" name while removing index")
return self._robject