Adding Bigtable row mutations for setting cell value.

gcloud/bigtable/row.py

@@ -15,11 +15,18 @@
 """User friendly container for Google Cloud Bigtable Row."""
 
 
+import struct
+
+import six
+
 from gcloud._helpers import _microseconds_from_datetime
 from gcloud._helpers import _to_bytes
 from gcloud.bigtable._generated import bigtable_data_pb2 as data_pb2
 
 
+_PACK_I64 = struct.Struct('>q').pack
+
+
 class Row(object):
     """Representation of a Google Cloud Bigtable Row.
 
@@ -92,6 +99,65 @@ def _get_mutations(self, state=None):
             else:
                 return self._false_pb_mutations
 
+    def set_cell(self, column_family_id, column, value, timestamp=None,
+                 state=None):
+        """Sets a value in this row.
+
+        The cell is determined by the ``row_key`` of the :class:`Row` and the
+        ``column``. The ``column`` must be in an existing
+        :class:`.column_family.ColumnFamily` (as determined by
+        ``column_family_id``).
+
+        .. note::
+
+            This method adds a mutation to the accumulated mutations on this
+            :class:`Row`, but does not make an API request. To actually
+            send an API request (with the mutations) to the Google Cloud
+            Bigtable API, call :meth:`commit`.
+
+        :type column_family_id: str
+        :param column_family_id: The column family that contains the column.
+                                 Must be of the form
+                                 ``[_a-zA-Z0-9][-_.a-zA-Z0-9]*``.
+
+        :type column: bytes
+        :param column: The column within the column family where the cell
+                       is located.
+
+        :type value: bytes or :class:`int`
+        :param value: The value to set in the cell. If an integer is used,
+                      will be interpreted as a 64-bit big-endian signed
+                      integer (8 bytes).
+
+        :type timestamp: :class:`datetime.datetime`
+        :param timestamp: (Optional) The timestamp of the operation.
+
+        :type state: bool
+        :param state: (Optional) The state that the mutation should be
+                      applied in. Unset if the mutation is not conditional,
+                      otherwise :data:`True` or :data:`False`.
+        """
+        column = _to_bytes(column)
+        if isinstance(value, six.integer_types):
+            value = _PACK_I64(value)
+        value = _to_bytes(value)
+        if timestamp is None:
+            # Use -1 for current Bigtable server time.
+            timestamp_micros = -1
+        else:
+            timestamp_micros = _microseconds_from_datetime(timestamp)
+            # Truncate to millisecond granularity.
+            timestamp_micros -= (timestamp_micros % 1000)
+
+        mutation_val = data_pb2.Mutation.SetCell(
+            family_name=column_family_id,
+            column_qualifier=column,
+            timestamp_micros=timestamp_micros,
+            value=value,
+        )
+        mutation_pb = data_pb2.Mutation(set_cell=mutation_val)
+        self._get_mutations(state).append(mutation_pb)
+
     def append_cell_value(self, column_family_id, column, value):
         """Appends a value to an existing cell.
 

gcloud/bigtable/test_row.py

@@ -89,6 +89,69 @@ def test__get_mutations_with_filter_bad_state(self):
         with self.assertRaises(ValueError):
             self._get_mutations_helper(filter_=filter_, state=state)
 
+    def _set_cell_helper(self, column=None, column_bytes=None,
+                         value=b'foobar', timestamp=None,
+                         timestamp_micros=-1):
+        import six
+        import struct
+        from gcloud.bigtable._generated import bigtable_data_pb2 as data_pb2
+
+        row_key = b'row_key'
+        column_family_id = u'column_family_id'
+        if column is None:
+            column = b'column'
+        table = object()
+        row = self._makeOne(row_key, table)
+        self.assertEqual(row._pb_mutations, [])
+        row.set_cell(column_family_id, column,
+                     value, timestamp=timestamp)
+
+        if isinstance(value, six.integer_types):
+            value = struct.pack('>q', value)
+        expected_pb = data_pb2.Mutation(
+            set_cell=data_pb2.Mutation.SetCell(
+                family_name=column_family_id,
+                column_qualifier=column_bytes or column,
+                timestamp_micros=timestamp_micros,
+                value=value,
+            ),
+        )
+        self.assertEqual(row._pb_mutations, [expected_pb])
+
+    def test_set_cell(self):
+        self._set_cell_helper()
+
+    def test_set_cell_with_string_column(self):
+        column_bytes = b'column'
+        column_non_bytes = u'column'
+        self._set_cell_helper(column=column_non_bytes,
+                              column_bytes=column_bytes)
+
+    def test_set_cell_with_integer_value(self):
+        value = 1337
+        self._set_cell_helper(value=value)
+
+    def test_set_cell_with_non_bytes_value(self):
+        row_key = b'row_key'
+        column = b'column'
+        column_family_id = u'column_family_id'
+        table = object()
+
+        row = self._makeOne(row_key, table)
+        value = object()  # Not bytes
+        with self.assertRaises(TypeError):
+            row.set_cell(column_family_id, column, value)
+
+    def test_set_cell_with_non_null_timestamp(self):
+        import datetime
+        from gcloud._helpers import _EPOCH
+
+        microseconds = 898294371
+        millis_granularity = microseconds - (microseconds % 1000)
+        timestamp = _EPOCH + datetime.timedelta(microseconds=microseconds)
+        self._set_cell_helper(timestamp=timestamp,
+                              timestamp_micros=millis_granularity)
+
     def test_append_cell_value(self):
         from gcloud.bigtable._generated import bigtable_data_pb2 as data_pb2