#847: prepare for using connection as context mgr

gcloud/datastore/__init__.py

@@ -48,7 +48,6 @@
   when race conditions may occur.
 """
 
-from gcloud.datastore._implicit_environ import get_connection
 from gcloud.datastore._implicit_environ import get_default_connection
 from gcloud.datastore._implicit_environ import get_default_dataset_id
 from gcloud.datastore._implicit_environ import set_default_connection
@@ -87,3 +86,23 @@ def set_defaults(dataset_id=None, connection=None):
     """
     set_default_dataset_id(dataset_id=dataset_id)
     set_default_connection(connection=connection)
+
+
+def get_connection():
+    """Shortcut method to establish a connection to the Cloud Datastore.
+
+    Use this if you are going to access several datasets
+    with the same set of credentials (unlikely):
+
+    >>> from gcloud import datastore
+
+    >>> connection = datastore.get_connection()
+    >>> key1 = datastore.Key('Kind', 1234, dataset_id='dataset1')
+    >>> key2 = datastore.Key('Kind', 1234, dataset_id='dataset2')
+    >>> entity1 = datastore.get(key1, connection=connection)
+    >>> entity2 = datastore.get(key2, connection=connection)
+
+    :rtype: :class:`gcloud.datastore.connection.Connection`
+    :returns: A connection defined with the proper credentials.
+    """
+    return Connection.from_environment()

gcloud/datastore/_implicit_environ.py

@@ -24,6 +24,7 @@
 from gcloud._helpers import _compute_engine_id
 from gcloud._helpers import _lazy_property_deco
 from gcloud.datastore.connection import Connection
+from gcloud.datastore.connection import _CONNECTIONS
 
 
 _DATASET_ENV_VAR_NAME = 'GCLOUD_DATASET_ID'
@@ -104,24 +105,26 @@ def get_default_dataset_id():
     return _DEFAULTS.dataset_id
 
 
-def get_connection():
-    """Shortcut method to establish a connection to the Cloud Datastore.
+def _require_connection(connection=None):
+    """Infer a connection from the environment, if not passed explicitly.
 
-    Use this if you are going to access several datasets
-    with the same set of credentials (unlikely):
-
-    >>> from gcloud import datastore
-
-    >>> connection = datastore.get_connection()
-    >>> key1 = datastore.Key('Kind', 1234, dataset_id='dataset1')
-    >>> key2 = datastore.Key('Kind', 1234, dataset_id='dataset2')
-    >>> entity1 = datastore.get(key1, connection=connection)
-    >>> entity2 = datastore.get(key2, connection=connection)
+    :type connection: :class:`gcloud.datastore.connection.Connection`
+    :param connection: Optional.
 
     :rtype: :class:`gcloud.datastore.connection.Connection`
-    :returns: A connection defined with the proper credentials.
+    :returns: A connection based on the current environment.
+    :raises: :class:`EnvironmentError` if ``connection`` is ``None``, and
+             cannot be inferred from the environment.
     """
-    return Connection.from_environment()
+    if connection is None:
+        top = _CONNECTIONS.top
+        if top is not None:
+            connection = top.connection
+        else:
+            connection = get_default_connection()
+            if connection is None:
+                raise EnvironmentError('Connection could not be inferred.')
+    return connection
 
 
 def set_default_connection(connection=None):
@@ -130,8 +133,7 @@ def set_default_connection(connection=None):
     :type connection: :class:`gcloud.datastore.connection.Connection`
     :param connection: A connection provided to be the default.
     """
-    connection = connection or get_connection()
-    _DEFAULTS.connection = connection
+    _DEFAULTS.connection = connection or Connection.from_environment()
 
 
 def get_default_connection():
@@ -163,7 +165,7 @@ def dataset_id():
     @staticmethod
     def connection():
         """Return the implicit default connection.."""
-        return get_connection()
+        return Connection.from_environment()
 
     def __init__(self, connection=None, dataset_id=None, implicit=False):
         if connection is not None or not implicit:

gcloud/datastore/_testing.py

@@ -31,3 +31,36 @@ def _setup_defaults(test_case, *args, **kwargs):
 
 def _tear_down_defaults(test_case):
     _implicit_environ._DEFAULTS = test_case._replaced_defaults
+
+
+class _NoCommitBatch(object):
+
+    def __init__(self, dataset_id, connection):
+        from gcloud.datastore.batch import Batch
+        self._batch = Batch(dataset_id, connection)
+
+    def __enter__(self):
+        from gcloud.datastore.connection import _CONNECTIONS
+        _CONNECTIONS.push(self._batch)
+        return self._batch
+
+    def __exit__(self, *args):
+        from gcloud.datastore.connection import _CONNECTIONS
+        _CONNECTIONS.pop()
+
+
+class _NoCommitTransaction(object):
+
+    def __init__(self, dataset_id, connection, transaction_id='TRANSACTION'):
+        from gcloud.datastore.transaction import Transaction
+        xact = self._transaction = Transaction(dataset_id, connection)
+        xact._id = transaction_id
+
+    def __enter__(self):
+        from gcloud.datastore.connection import _CONNECTIONS
+        _CONNECTIONS.push(self._transaction)
+        return self._transaction
+
+    def __exit__(self, *args):
+        from gcloud.datastore.connection import _CONNECTIONS
+        _CONNECTIONS.pop()

gcloud/datastore/api.py

@@ -20,6 +20,7 @@
 
 from gcloud.datastore import _implicit_environ
 from gcloud.datastore.batch import Batch
+from gcloud.datastore.connection import _CONNECTIONS
 from gcloud.datastore.entity import Entity
 from gcloud.datastore.transaction import Transaction
 from gcloud.datastore import helpers
@@ -53,7 +54,7 @@ def _require_dataset_id(dataset_id=None, first_key=None):
     """
     if dataset_id is not None:
         return dataset_id
-    top = Batch.current()
+    top = _CONNECTIONS.top
     if top is not None:
         return top.dataset_id
     if first_key is not None:
@@ -65,28 +66,6 @@ def _require_dataset_id(dataset_id=None, first_key=None):
     return dataset_id
 
 
-def _require_connection(connection=None):
-    """Infer a connection from the environment, if not passed explicitly.
-
-    :type connection: :class:`gcloud.datastore.connection.Connection`
-    :param connection: Optional.
-
-    :rtype: :class:`gcloud.datastore.connection.Connection`
-    :returns: A connection based on the current environment.
-    :raises: :class:`EnvironmentError` if ``connection`` is ``None``, and
-             cannot be inferred from the environment.
-    """
-    if connection is None:
-        top = Batch.current()
-        if top is not None:
-            connection = top.connection
-        else:
-            connection = _implicit_environ.get_default_connection()
-            if connection is None:
-                raise EnvironmentError('Connection could not be inferred.')
-    return connection
-
-
 def _extended_lookup(connection, dataset_id, key_pbs,
                      missing=None, deferred=None,
                      eventual=False, transaction_id=None):
@@ -200,7 +179,7 @@ def get(keys, missing=None, deferred=None, connection=None, dataset_id=None):
     if not keys:
         return []
 
-    connection = _require_connection(connection)
+    connection = _implicit_environ._require_connection(connection)
     dataset_id = _require_dataset_id(dataset_id, keys[0])
 
     if list(set([key.dataset_id for key in keys])) != [dataset_id]:
@@ -259,10 +238,10 @@ def put(entities, connection=None, dataset_id=None):
     if not entities:
         return
 
-    connection = _require_connection(connection)
+    connection = _implicit_environ._require_connection(connection)
     dataset_id = _require_dataset_id(dataset_id, entities[0].key)
 
-    current = Batch.current()
+    current = _CONNECTIONS.top
     in_batch = current is not None
     if not in_batch:
         current = Batch(dataset_id=dataset_id, connection=connection)
@@ -294,11 +273,11 @@ def delete(keys, connection=None, dataset_id=None):
     if not keys:
         return
 
-    connection = _require_connection(connection)
+    connection = _implicit_environ._require_connection(connection)
     dataset_id = _require_dataset_id(dataset_id, keys[0])
 
     # We allow partial keys to attempt a delete, the backend will fail.
-    current = Batch.current()
+    current = _CONNECTIONS.top
     in_batch = current is not None
     if not in_batch:
         current = Batch(dataset_id=dataset_id, connection=connection)
@@ -324,7 +303,7 @@ def allocate_ids(incomplete_key, num_ids, connection=None):
     :returns: The (complete) keys allocated with ``incomplete_key`` as root.
     :raises: :class:`ValueError` if ``incomplete_key`` is not a partial key.
     """
-    connection = _require_connection(connection)
+    connection = _implicit_environ._require_connection(connection)
 
     if not incomplete_key.is_partial:
         raise ValueError(('Key is not partial.', incomplete_key))

gcloud/datastore/batch.py

@@ -21,16 +21,13 @@
 https://cloud.google.com/datastore/docs/concepts/entities#Datastore_Batch_operations
 """
 
-from gcloud._helpers import _LocalStack
 from gcloud.datastore import _implicit_environ
 from gcloud.datastore import helpers
+from gcloud.datastore.connection import _CONNECTIONS
 from gcloud.datastore.key import _dataset_ids_equal
 from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
 
 
-_BATCHES = _LocalStack()
-
-
 class Batch(object):
     """An abstraction representing a collected group of updates / deletes.
 
@@ -90,7 +87,7 @@ def __init__(self, dataset_id=None, connection=None):
     @staticmethod
     def current():
         """Return the topmost batch / transaction, or None."""
-        return _BATCHES.top
+        return _CONNECTIONS.top
 
     @property
     def dataset_id(self):
@@ -229,7 +226,7 @@ def rollback(self):
         pass
 
     def __enter__(self):
-        _BATCHES.push(self)
+        _CONNECTIONS.push(self)
         self.begin()
         return self
 
@@ -240,7 +237,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             else:
                 self.rollback()
         finally:
-            _BATCHES.pop()
+            _CONNECTIONS.pop()
 
 
 def _assign_entity_to_mutation(mutation_pb, entity, auto_id_entities):

gcloud/datastore/connection.py

@@ -17,10 +17,14 @@
 import os
 
 from gcloud import connection
+from gcloud._helpers import _LocalStack
 from gcloud.exceptions import make_exception
 from gcloud.datastore import _datastore_v1_pb2 as datastore_pb
 
 
+_CONNECTIONS = _LocalStack()
+
+
 SCOPE = ('https://www.googleapis.com/auth/datastore',
          'https://www.googleapis.com/auth/userinfo.email')
 """The scopes required for authenticating as a Cloud Datastore consumer."""

gcloud/datastore/test___init__.py

@@ -44,3 +44,25 @@ def call_set_connection(connection=None):
 
         self.assertEqual(SET_DATASET_CALLED, [DATASET_ID])
         self.assertEqual(SET_CONNECTION_CALLED, [CONNECTION])
+
+
+class Test_get_connection(unittest2.TestCase):
+
+    def _callFUT(self):
+        from gcloud.datastore import get_connection
+        return get_connection()
+
+    def test_it(self):
+        from gcloud import credentials
+        from gcloud.datastore.connection import SCOPE
+        from gcloud.datastore.connection import Connection
+        from gcloud.test_credentials import _Client
+        from gcloud._testing import _Monkey
+
+        client = _Client()
+        with _Monkey(credentials, client=client):
+            found = self._callFUT()
+        self.assertTrue(isinstance(found, Connection))
+        self.assertTrue(found._credentials is client._signed)
+        self.assertEqual(found._credentials._scopes, SCOPE)
+        self.assertTrue(client._get_app_default_called)