Remove confusion between 'connection' and '_batch_stack'.

gcloud/datastore/client.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Convenience wrapper for invoking APIs/factories w/ a dataset ID."""
 
+from gcloud._helpers import _LocalStack
 from gcloud.datastore import helpers
 from gcloud.datastore.batch import Batch
 from gcloud.datastore.entity import Entity
@@ -126,43 +127,42 @@ def __init__(self, dataset_id=None, namespace=None, connection=None):
         self.dataset_id = dataset_id
         if connection is None:
             connection = get_connection()
-        self._connection_stack = [connection]
+        self.connection = connection
+        self._batch_stack = _LocalStack()
         self.namespace = namespace
 
-    def _push_connection(self, connection):
-        """Push a connection/batch/transaction onto our stack.
+    def _push_batch(self, batch):
+        """Push a batch/transaction onto our stack.
 
         "Protected", intended for use by batch / transaction context mgrs.
 
-        :type connection: :class:`gcloud.datastore.connection.Connection`,
-                          or an object implementing its API.
-        :param connection: newly-active connection/batch/transaction to
-                           pass to proxied API methods
+        :type batch: :class:`gcloud.datastore.batch.Batch`, or an object
+                     implementing its API.
+        :param batch: newly-active batch/batch/transaction.
         """
-        self._connection_stack.append(connection)
+        self._batch_stack.push(batch)
 
-    def _pop_connection(self):
-        """Pop a connection/batch/transaction from our stack.
+    def _pop_batch(self):
+        """Pop a batch/transaction from our stack.
 
         "Protected", intended for use by batch / transaction context mgrs.
 
         :raises: IndexError if the stack is empty.
-        :rtype: :class:`gcloud.datastore.connection.Connection`, or
-                an object implementing its API.
-        :returns: the top-most connection/batch/transaction, after removing it.
+        :rtype: :class:`gcloud.datastore.batch.Batch`, or an object
+                 implementing its API.
+        :returns: the top-most batch/transaction, after removing it.
         """
-        return self._connection_stack.pop()
+        return self._batch_stack.pop()
 
     @property
-    def connection(self):
-        """Currently-active connection.
+    def current_batch(self):
+        """Currently-active batch.
 
-        :rtype: :class:`gcloud.datastore.connection.Connection`, or
-                an object implementing its API.
-        :returns: The connection/batch/transaction at the toop of the
-                  connection stack.
+        :rtype: :class:`gcloud.datastore.batch.Batch`, or an object
+                implementing its API, or ``NoneType`` (if no batch is active).
+        :returns: The batch/transaction at the toop of the batch stack.
         """
-        return self._connection_stack[-1]
+        return self._batch_stack.top
 
     def get(self, key, missing=None, deferred=None):
         """Retrieve an entity from a single key (if it exists).

gcloud/datastore/test_client.py

@@ -63,6 +63,7 @@ def test_ctor_w_implicit_inputs(self):
         self.assertEqual(client.dataset_id, OTHER)
         self.assertEqual(client.namespace, None)
         self.assertTrue(client.connection is conn)
+        self.assertTrue(client.current_batch is None)
 
     def test_ctor_w_explicit_inputs(self):
         OTHER = 'other'
@@ -74,18 +75,25 @@ def test_ctor_w_explicit_inputs(self):
         self.assertEqual(client.dataset_id, OTHER)
         self.assertEqual(client.namespace, NAMESPACE)
         self.assertTrue(client.connection is conn)
-        self.assertEqual(list(client._connection_stack), [conn])
+        self.assertTrue(client.current_batch is None)
+        self.assertEqual(list(client._batch_stack), [])
 
     def test__push_connection_and__pop_connection(self):
         conn = object()
-        new_conn = object()
+        batch1 = object()
+        batch2 = object()
         client = self._makeOne(connection=conn)
-        client._push_connection(new_conn)
-        self.assertTrue(client.connection is new_conn)
-        self.assertEqual(list(client._connection_stack), [conn, new_conn])
-        self.assertTrue(client._pop_connection() is new_conn)
-        self.assertTrue(client.connection is conn)
-        self.assertEqual(list(client._connection_stack), [conn])
+        client._push_batch(batch1)
+        self.assertEqual(list(client._batch_stack), [batch1])
+        self.assertTrue(client.current_batch is batch1)
+        client._push_batch(batch2)
+        self.assertTrue(client.current_batch is batch2)
+        # list(_LocalStack) returns in reverse order.
+        self.assertEqual(list(client._batch_stack), [batch2, batch1])
+        self.assertTrue(client._pop_batch() is batch2)
+        self.assertEqual(list(client._batch_stack), [batch1])
+        self.assertTrue(client._pop_batch() is batch1)
+        self.assertEqual(list(client._batch_stack), [])
 
     def test_get_miss(self):
         _called_with = []