Merge pull request #2000 from dhermes/counter_set-happybase-just-put

Make counter_set in HappyBase agree with HBase impl. semantics

Author: dhermes

gcloud/bigtable/happybase/__init__.py

@@ -21,26 +21,25 @@
 -------------------------
 
 Some concepts from HBase/Thrift do not map directly to the Cloud
-Bigtable API. As a result, the following instance methods and functions
-could not be implemented:
+Bigtable API. As a result
 
+* :meth:`Table.regions() <gcloud.bigtable.happybase.table.Table.regions>`
+  could not be implemented since tables in Cloud Bigtable do not expose
+  internal storage details
 * :meth:`Connection.enable_table() \
-      <gcloud.bigtable.happybase.connection.Connection.enable_table>` - no
-  concept of enabled/disabled
+      <gcloud.bigtable.happybase.connection.Connection.enable_table>`
+  does nothing since Cloud Bigtable has no concept of enabled/disabled
 * :meth:`Connection.disable_table() \
-      <gcloud.bigtable.happybase.connection.Connection.disable_table>` - no
-  concept of enabled/disabled
+      <gcloud.bigtable.happybase.connection.Connection.disable_table>`
+  does nothing since Cloud Bigtable has no concept of enabled/disabled
 * :meth:`Connection.is_table_enabled() \
       <gcloud.bigtable.happybase.connection.Connection.is_table_enabled>`
-  - no concept of enabled/disabled
+  always returns :data:`True` since Cloud Bigtable has no concept of
+  enabled/disabled
 * :meth:`Connection.compact_table() \
-      <gcloud.bigtable.happybase.connection.Connection.compact_table>` -
-  table storage is opaque to user
-* :meth:`Table.regions() <gcloud.bigtable.happybase.table.Table.regions>`
-  - tables in Cloud Bigtable do not expose internal storage details
-* :meth:`Table.counter_set() \
-      <gcloud.bigtable.happybase.table.Table.counter_set>` - method can't
-  be atomic, so we disable it
+      <gcloud.bigtable.happybase.connection.Connection.compact_table>`
+  does nothing since Cloud Bigtable handles table compactions automatically
+  and does not expose an API for it
 * The ``__version__`` value for the HappyBase package is :data:`None`.
   However, it's worth nothing this implementation was based off HappyBase
   0.9.

gcloud/bigtable/happybase/table.py

@@ -42,6 +42,7 @@
 
 
 _WARN = warnings.warn
+_PACK_I64 = struct.Struct('>q').pack
 _UNPACK_I64 = struct.Struct('>q').unpack
 _SIMPLE_GC_RULES = (MaxAgeGCRule, MaxVersionsGCRule)
 
@@ -531,9 +532,11 @@ def counter_get(self, row, column):
     def counter_set(self, row, column, value=0):
         """Set a counter column to a specific value.
 
-        This method is provided in HappyBase, but we do not provide it here
-        because it defeats the purpose of using atomic increment and decrement
-        of a counter.
+        .. note::
+
+            Be careful using this method. It can be useful for setting the
+            initial value of a counter, but it defeats the purpose of using
+            atomic increment and decrement.
 
         :type row: str
         :param row: Row key for the row we are setting a counter in.
@@ -544,13 +547,8 @@ def counter_set(self, row, column, value=0):
 
         :type value: int
         :param value: Value to set the counter to.
-
-        :raises: :class:`NotImplementedError <exceptions.NotImplementedError>`
-                 always
         """
-        raise NotImplementedError('Table.counter_set will not be implemented. '
-                                  'Instead use the increment/decrement '
-                                  'methods along with counter_get.')
+        self.put(row, {column: _PACK_I64(value)})
 
     def counter_inc(self, row, column, value=1):
         """Atomically increment a counter column.

gcloud/bigtable/happybase/test_table.py

@@ -889,15 +889,45 @@ def _counter_inc_helper(self, row, column, value, commit_result):
                          {tuple(column.split(':')): incremented_value})
 
     def test_counter_set(self):
+        import struct
+        from gcloud._testing import _Monkey
+        from gcloud.bigtable.happybase import table as MUT
+
         name = 'table-name'
         connection = None
         table = self._makeOne(name, connection)
+        batches_created = []
+
+        def make_batch(*args, **kwargs):
+            result = _MockBatch(*args, **kwargs)
+            batches_created.append(result)
+            return result
 
         row = 'row-key'
         column = 'fam:col1'
         value = 42
-        with self.assertRaises(NotImplementedError):
-            table.counter_set(row, column, value=value)
+        with _Monkey(MUT, Batch=make_batch):
+            result = table.counter_set(row, column, value=value)
+
+        # There is no return value.
+        self.assertEqual(result, None)
+
+        # Check how the batch was created and used.
+        batch, = batches_created
+        self.assertTrue(isinstance(batch, _MockBatch))
+        self.assertEqual(batch.args, (table,))
+        expected_kwargs = {
+            'timestamp': None,
+            'batch_size': None,
+            'transaction': False,
+            'wal': MUT._WAL_SENTINEL,
+        }
+        self.assertEqual(batch.kwargs, expected_kwargs)
+        # Make sure it was a successful context manager
+        self.assertEqual(batch.exit_vals, [(None, None, None)])
+        data = {column: struct.Struct('>q').pack(value)}
+        self.assertEqual(batch.put_args, [(row, data)])
+        self.assertEqual(batch.delete_args, [])
 
     def test_counter_inc(self):
         import struct