Adding helper to parse a Bigtable ReadModifyWriteRow response.

dhermes · dhermes · commit 97b85a45ec8b · 2016-02-15T14:23:38.000-08:00
diff --git a/gcloud/bigtable/row.py b/gcloud/bigtable/row.py
@@ -1195,6 +1195,46 @@ def to_pb(self):
         return data_pb2.RowFilter(condition=condition)
 
 
+def _parse_rmw_row_response(row_response):
+    """Parses the response to a ``ReadModifyWriteRow`` request.
+
+    :type row_response: :class:`.data_pb2.Row`
+    :param row_response: The response row (with only modified cells) from a
+                         ``ReadModifyWriteRow`` request.
+
+    :rtype: dict
+    :returns: The new contents of all modified cells. Returned as a
+              dictionary of column families, each of which holds a
+              dictionary of columns. Each column contains a list of cells
+              modified. Each cell is represented with a two-tuple with the
+              value (in bytes) and the timestamp for the cell. For example:
+
+              .. code:: python
+
+                  {
+                      u'col-fam-id': {
+                          b'col-name1': [
+                              (b'cell-val', datetime.datetime(...)),
+                              (b'cell-val-newer', datetime.datetime(...)),
+                          ],
+                          b'col-name2': [
+                              (b'altcol-cell-val', datetime.datetime(...)),
+                          ],
+                      },
+                      u'col-fam-id2': {
+                          b'col-name3-but-other-fam': [
+                              (b'foo', datetime.datetime(...)),
+                          ],
+                      },
+                  }
+    """
+    result = {}
+    for column_family in row_response.families:
+        column_family_id, curr_family = _parse_family_pb(column_family)
+        result[column_family_id] = curr_family
+    return result
+
+
 def _parse_family_pb(family_pb):
     """Parses a Family protobuf into a dictionary.
 
diff --git a/gcloud/bigtable/test_row.py b/gcloud/bigtable/test_row.py
@@ -1535,6 +1535,92 @@ def test_to_pb_false_only(self):
         self.assertEqual(filter_pb, expected_pb)
 
 
+class Test__parse_rmw_row_response(unittest2.TestCase):
+
+    def _callFUT(self, row_response):
+        from gcloud.bigtable.row import _parse_rmw_row_response
+        return _parse_rmw_row_response(row_response)
+
+    def test_it(self):
+        from gcloud._helpers import _datetime_from_microseconds
+        from gcloud.bigtable._generated import bigtable_data_pb2 as data_pb2
+
+        col_fam1 = u'col-fam-id'
+        col_fam2 = u'col-fam-id2'
+        col_name1 = b'col-name1'
+        col_name2 = b'col-name2'
+        col_name3 = b'col-name3-but-other-fam'
+        cell_val1 = b'cell-val'
+        cell_val2 = b'cell-val-newer'
+        cell_val3 = b'altcol-cell-val'
+        cell_val4 = b'foo'
+
+        microseconds = 1000871
+        timestamp = _datetime_from_microseconds(microseconds)
+        expected_output = {
+            col_fam1: {
+                col_name1: [
+                    (cell_val1, timestamp),
+                    (cell_val2, timestamp),
+                ],
+                col_name2: [
+                    (cell_val3, timestamp),
+                ],
+            },
+            col_fam2: {
+                col_name3: [
+                    (cell_val4, timestamp),
+                ],
+            },
+        }
+        sample_input = data_pb2.Row(
+            families=[
+                data_pb2.Family(
+                    name=col_fam1,
+                    columns=[
+                        data_pb2.Column(
+                            qualifier=col_name1,
+                            cells=[
+                                data_pb2.Cell(
+                                    value=cell_val1,
+                                    timestamp_micros=microseconds,
+                                ),
+                                data_pb2.Cell(
+                                    value=cell_val2,
+                                    timestamp_micros=microseconds,
+                                ),
+                            ],
+                        ),
+                        data_pb2.Column(
+                            qualifier=col_name2,
+                            cells=[
+                                data_pb2.Cell(
+                                    value=cell_val3,
+                                    timestamp_micros=microseconds,
+                                ),
+                            ],
+                        ),
+                    ],
+                ),
+                data_pb2.Family(
+                    name=col_fam2,
+                    columns=[
+                        data_pb2.Column(
+                            qualifier=col_name3,
+                            cells=[
+                                data_pb2.Cell(
+                                    value=cell_val4,
+                                    timestamp_micros=microseconds,
+                                ),
+                            ],
+                        ),
+                    ],
+                ),
+            ],
+        )
+        self.assertEqual(expected_output, self._callFUT(sample_input))
+
+
 class Test__parse_family_pb(unittest2.TestCase):
 
     def _callFUT(self, family_pb):
