Adding HappyBase batch send() and batch context manager.

gcloud/bigtable/happybase/batch.py

@@ -95,3 +95,41 @@ def __init__(self, table, timestamp=None, batch_size=None,
         # Internal state for tracking mutations.
         self._row_map = {}
         self._mutation_count = 0
+
+    def send(self):
+        """Send / commit the batch of mutations to the server."""
+        for row in self._row_map.values():
+            # commit() does nothing if row hasn't accumulated any mutations.
+            row.commit()
+
+        self._row_map.clear()
+        self._mutation_count = 0
+
+    def __enter__(self):
+        """Enter context manager, no set-up required."""
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Exit context manager, no set-up required.
+
+        :type exc_type: type
+        :param exc_type: The type of the exception if one occurred while the
+                         context manager was active. Otherwise, :data:`None`.
+
+        :type exc_value: :class:`Exception <exceptions.Exception>`
+        :param exc_value: An instance of ``exc_type`` if an exception occurred
+                          while the context was active.
+                          Otherwise, :data:`None`.
+
+        :type traceback: ``traceback`` type
+        :param traceback: The traceback where the exception occurred (if one
+                          did occur). Otherwise, :data:`None`.
+        """
+        # If the context manager encountered an exception and the batch is
+        # transactional, we don't commit the mutations.
+        if self._transaction and exc_type is not None:
+            return
+
+        # NOTE: For non-transactional batches, this will even commit mutations
+        #       if an error occurred during the context manager.
+        self.send()

gcloud/bigtable/happybase/test_batch.py

@@ -91,3 +91,102 @@ def test_constructor_with_batch_size_and_transactional(self):
         with self.assertRaises(TypeError):
             self._makeOne(table, batch_size=batch_size,
                           transaction=transaction)
+
+    def test_send(self):
+        table = object()
+        batch = self._makeOne(table)
+
+        batch._row_map = row_map = _MockRowMap()
+        row_map['row-key1'] = row1 = _MockRow()
+        row_map['row-key2'] = row2 = _MockRow()
+        batch._mutation_count = 1337
+
+        self.assertEqual(row_map.clear_count, 0)
+        self.assertEqual(row1.commits, 0)
+        self.assertEqual(row2.commits, 0)
+        self.assertNotEqual(batch._mutation_count, 0)
+        self.assertNotEqual(row_map, {})
+
+        batch.send()
+        self.assertEqual(row_map.clear_count, 1)
+        self.assertEqual(row1.commits, 1)
+        self.assertEqual(row2.commits, 1)
+        self.assertEqual(batch._mutation_count, 0)
+        self.assertEqual(row_map, {})
+
+    def test_context_manager(self):
+        klass = self._getTargetClass()
+
+        class BatchWithSend(_SendMixin, klass):
+            pass
+
+        table = object()
+        batch = BatchWithSend(table)
+        self.assertFalse(batch._send_called)
+
+        with batch:
+            pass
+
+        self.assertTrue(batch._send_called)
+
+    def test_context_manager_with_exception_non_transactional(self):
+        klass = self._getTargetClass()
+
+        class BatchWithSend(_SendMixin, klass):
+            pass
+
+        table = object()
+        batch = BatchWithSend(table)
+        self.assertFalse(batch._send_called)
+
+        with self.assertRaises(ValueError):
+            with batch:
+                raise ValueError('Something bad happened')
+
+        self.assertTrue(batch._send_called)
+
+    def test_context_manager_with_exception_transactional(self):
+        klass = self._getTargetClass()
+
+        class BatchWithSend(_SendMixin, klass):
+            pass
+
+        table = object()
+        batch = BatchWithSend(table, transaction=True)
+        self.assertFalse(batch._send_called)
+
+        with self.assertRaises(ValueError):
+            with batch:
+                raise ValueError('Something bad happened')
+
+        self.assertFalse(batch._send_called)
+
+        # Just to make sure send() actually works (and to make cover happy).
+        batch.send()
+        self.assertTrue(batch._send_called)
+
+
+class _SendMixin(object):
+
+    _send_called = False
+
+    def send(self):
+        self._send_called = True
+
+
+class _MockRowMap(dict):
+
+    clear_count = 0
+
+    def clear(self):
+        self.clear_count += 1
+        super(_MockRowMap, self).clear()
+
+
+class _MockRow(object):
+
+    def __init__(self):
+        self.commits = 0
+
+    def commit(self):
+        self.commits += 1