Visitar URL original
Merge pull request #1494 from dhermes/bigtable-partial-row-update-from · googleapis/google-cloud-python@0891906 · GitHub
Skip to content

Commit 0891906

Browse files
committed
Merge pull request #1494 from dhermes/bigtable-partial-row-update-from
Implementing Bigtable PartialRowData.update_from_read_rows().
2 parents 37b1a38 + 8a49056 commit 0891906

2 files changed

Lines changed: 121 additions & 0 deletions

File tree

‎gcloud/bigtable/row_data.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -208,6 +208,44 @@ def _handle_row_contents(self, chunk):
208208
column_cells = column_family_dict.setdefault(column_name, [])
209209
column_cells.extend(cells)
210210

211+
def update_from_read_rows(self, read_rows_response_pb):
212+
"""Updates the current row from a ``ReadRows`` response.
213+
214+
:type read_rows_response_pb:
215+
:class:`._generated.bigtable_service_messages_pb2.ReadRowsResponse`
216+
:param read_rows_response_pb: A response streamed back as part of a
217+
``ReadRows`` request.
218+
219+
:raises: :class:`ValueError <exceptions.ValueError>` if the current
220+
partial row has already been committed, if the row key on the
221+
response doesn't match the current one or if there is a chunk
222+
encountered with an unexpected ``ONEOF`` protobuf property.
223+
"""
224+
if self._committed:
225+
raise ValueError('The row has been committed')
226+
227+
if read_rows_response_pb.row_key != self.row_key:
228+
raise ValueError('Response row key (%r) does not match current '
229+
'one (%r).' % (read_rows_response_pb.row_key,
230+
self.row_key))
231+
232+
last_chunk_index = len(read_rows_response_pb.chunks) - 1
233+
for index, chunk in enumerate(read_rows_response_pb.chunks):
234+
chunk_property = chunk.WhichOneof('chunk')
235+
if chunk_property == 'row_contents':
236+
self._handle_row_contents(chunk)
237+
elif chunk_property == 'reset_row':
238+
self._handle_reset_row(chunk)
239+
elif chunk_property == 'commit_row':
240+
self._handle_commit_row(chunk, index, last_chunk_index)
241+
else:
242+
# NOTE: This includes chunk_property == None since we always
243+
# want a value to be set
244+
raise ValueError('Unexpected chunk property: %s' % (
245+
chunk_property,))
246+
247+
self._chunks_encountered = True
248+
211249

212250
class PartialRowsData(object):
213251
"""Convenience wrapper for consuming a ``ReadRows`` streaming response.

‎gcloud/bigtable/test_row_data.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,89 @@ def test__handle_row_contents(self):
296296
}
297297
self.assertEqual(partial_row_data.cells, expected_cells)
298298

299+
def test_update_from_read_rows(self):
300+
from gcloud.bigtable._generated import bigtable_data_pb2 as data_pb2
301+
from gcloud.bigtable._generated import (
302+
bigtable_service_messages_pb2 as messages_pb2)
303+
304+
row_key = b'row-key'
305+
partial_row_data = self._makeOne(row_key)
306+
307+
# Set-up chunk1, some data that will be reset by chunk2.
308+
ignored_family_name = u'ignore-name'
309+
row_contents = data_pb2.Family(name=ignored_family_name)
310+
chunk1 = messages_pb2.ReadRowsResponse.Chunk(row_contents=row_contents)
311+
312+
# Set-up chunk2, a reset row.
313+
chunk2 = messages_pb2.ReadRowsResponse.Chunk(reset_row=True)
314+
315+
# Set-up chunk3, a column family with no columns.
316+
family_name = u'name'
317+
row_contents = data_pb2.Family(name=family_name)
318+
chunk3 = messages_pb2.ReadRowsResponse.Chunk(row_contents=row_contents)
319+
320+
# Set-up chunk4, a commit row.
321+
chunk4 = messages_pb2.ReadRowsResponse.Chunk(commit_row=True)
322+
323+
# Prepare request and make sure PartialRowData is empty before.
324+
read_rows_response_pb = messages_pb2.ReadRowsResponse(
325+
row_key=row_key, chunks=[chunk1, chunk2, chunk3, chunk4])
326+
self.assertEqual(partial_row_data.cells, {})
327+
self.assertFalse(partial_row_data.committed)
328+
self.assertFalse(partial_row_data._chunks_encountered)
329+
330+
# Parse the response and make sure the cells took place.
331+
partial_row_data.update_from_read_rows(read_rows_response_pb)
332+
self.assertEqual(partial_row_data.cells, {family_name: {}})
333+
self.assertFalse(ignored_family_name in partial_row_data.cells)
334+
self.assertTrue(partial_row_data.committed)
335+
self.assertTrue(partial_row_data._chunks_encountered)
336+
337+
def test_update_from_read_rows_while_committed(self):
338+
partial_row_data = self._makeOne(None)
339+
partial_row_data._committed = True
340+
self.assertFalse(partial_row_data._chunks_encountered)
341+
342+
with self.assertRaises(ValueError):
343+
partial_row_data.update_from_read_rows(None)
344+
345+
self.assertFalse(partial_row_data._chunks_encountered)
346+
347+
def test_update_from_read_rows_row_key_disagree(self):
348+
from gcloud.bigtable._generated import (
349+
bigtable_service_messages_pb2 as messages_pb2)
350+
351+
row_key1 = b'row-key1'
352+
row_key2 = b'row-key2'
353+
partial_row_data = self._makeOne(row_key1)
354+
self.assertFalse(partial_row_data._chunks_encountered)
355+
356+
self.assertNotEqual(row_key1, row_key2)
357+
read_rows_response_pb = messages_pb2.ReadRowsResponse(row_key=row_key2)
358+
with self.assertRaises(ValueError):
359+
partial_row_data.update_from_read_rows(read_rows_response_pb)
360+
361+
self.assertFalse(partial_row_data._chunks_encountered)
362+
363+
def test_update_from_read_rows_empty_chunk(self):
364+
from gcloud.bigtable._generated import (
365+
bigtable_service_messages_pb2 as messages_pb2)
366+
367+
row_key = b'row-key'
368+
partial_row_data = self._makeOne(row_key)
369+
self.assertFalse(partial_row_data._chunks_encountered)
370+
371+
chunk = messages_pb2.ReadRowsResponse.Chunk()
372+
read_rows_response_pb = messages_pb2.ReadRowsResponse(
373+
row_key=row_key, chunks=[chunk])
374+
375+
# This makes it an "empty" chunk.
376+
self.assertEqual(chunk.WhichOneof('chunk'), None)
377+
with self.assertRaises(ValueError):
378+
partial_row_data.update_from_read_rows(read_rows_response_pb)
379+
380+
self.assertFalse(partial_row_data._chunks_encountered)
381+
299382

300383
class TestPartialRowsData(unittest2.TestCase):
301384

0 commit comments

Comments
 (0)