Merge pull request #444 from dhermes/implicit-dataset-from-environ

Implicit dataset from environ

gcloud/datastore/__init__.py

@@ -44,14 +44,41 @@
   which represents a lookup or search over the rows in the datastore.
 """
 
+import os
+
 from gcloud import credentials
+from gcloud.datastore import _implicit_environ
 from gcloud.datastore.connection import Connection
 
 
 SCOPE = ('https://www.googleapis.com/auth/datastore ',
          'https://www.googleapis.com/auth/userinfo.email')
 """The scope required for authenticating as a Cloud Datastore consumer."""
 
+_DATASET_ENV_VAR_NAME = 'GCLOUD_DATASET_ID'
+
+
+def set_default_dataset(dataset_id=None):
+    """Determines auth settings from local enviroment.
+
+    Sets a default dataset either explicitly or via the local
+    enviroment. Currently only supports enviroment variable but will
+    implicitly support App Engine, Compute Engine and other environments
+    in the future.
+
+    Local environment variable used is:
+    - GCLOUD_DATASET_ID
+
+    :type dataset_id: :class:`str`.
+    :param dataset_id: Optional. The dataset ID to use for the default
+                       dataset.
+    """
+    if dataset_id is None:
+        dataset_id = os.getenv(_DATASET_ENV_VAR_NAME)
+
+    if dataset_id is not None:
+        _implicit_environ.DATASET = get_dataset(dataset_id)
+
 
 def get_connection():
     """Shortcut method to establish a connection to the Cloud Datastore.
@@ -97,3 +124,54 @@ def get_dataset(dataset_id):
     """
     connection = get_connection()
     return connection.dataset(dataset_id)
+
+
+def _require_dataset():
+    """Convenience method to ensure DATASET is set.
+
+    :rtype: :class:`gcloud.datastore.dataset.Dataset`
+    :returns: A dataset based on the current environment.
+    :raises: :class:`EnvironmentError` if DATASET is not set.
+    """
+    if _implicit_environ.DATASET is None:
+        raise EnvironmentError('Dataset could not be inferred.')
+    return _implicit_environ.DATASET
+
+
+def get_entity(key):
+    """Retrieves entity from implicit dataset, along with its attributes.
+
+    :type key: :class:`gcloud.datastore.key.Key`
+    :param key: The name of the item to retrieve.
+
+    :rtype: :class:`gcloud.datastore.entity.Entity` or ``None``
+    :return: The requested entity, or ``None`` if there was no match found.
+    """
+    return _require_dataset().get_entity(key)
+
+
+def get_entities(keys):
+    """Retrieves entities from implied dataset, along with their attributes.
+
+    :type keys: list of :class:`gcloud.datastore.key.Key`
+    :param keys: The name of the item to retrieve.
+
+    :rtype: list of :class:`gcloud.datastore.entity.Entity`
+    :return: The requested entities.
+    """
+    return _require_dataset().get_entities(keys)
+
+
+def allocate_ids(incomplete_key, num_ids):
+    """Allocates a list of IDs from a partial key.
+
+    :type incomplete_key: A :class:`gcloud.datastore.key.Key`
+    :param incomplete_key: The partial key to use as base for allocated IDs.
+
+    :type num_ids: A :class:`int`.
+    :param num_ids: The number of IDs to allocate.
+
+    :rtype: list of :class:`gcloud.datastore.key.Key`
+    :return: The (complete) keys allocated with `incomplete_key` as root.
+    """
+    return _require_dataset().allocate_ids(incomplete_key, num_ids)

gcloud/datastore/_implicit_environ.py

@@ -0,0 +1,24 @@
+"""Module to provide implicit behavior based on enviroment.
+
+Acts as a mutable namespace to allow the datastore package to
+imply the current dataset from the enviroment.
+
+Also provides a base class for classes in the `datastore` package
+which could utilize the implicit enviroment.
+"""
+
+
+DATASET = None
+"""Module global to allow persistent implied dataset from enviroment."""
+
+
+class _DatastoreBase(object):
+    """Base for all classes in the datastore package.
+
+    Uses the implicit DATASET object as a default dataset attached
+    to the instances being created. Stores the dataset passed in
+    on the protected (i.e. non-public) attribute `_dataset`.
+    """
+
+    def __init__(self, dataset=None):
+        self._dataset = dataset or DATASET

gcloud/datastore/entity.py

@@ -14,6 +14,7 @@
 
 """Class for representing a single entity in the Cloud Datastore."""
 
+from gcloud.datastore import _implicit_environ
 from gcloud.datastore import datastore_v1_pb2 as datastore_pb
 from gcloud.datastore.key import Key
 
@@ -95,7 +96,9 @@ class Entity(dict):
 
     def __init__(self, dataset=None, kind=None, exclude_from_indexes=()):
         super(Entity, self).__init__()
-        self._dataset = dataset
+        # Does not inherit directly from object, so we don't use
+        # _implicit_environ._DatastoreBase to avoid split MRO.
+        self._dataset = dataset or _implicit_environ.DATASET
         if kind:
             self._key = Key().kind(kind)
         else:

gcloud/datastore/query.py

@@ -16,12 +16,13 @@
 
 import base64
 
+from gcloud.datastore import _implicit_environ
 from gcloud.datastore import datastore_v1_pb2 as datastore_pb
 from gcloud.datastore import helpers
 from gcloud.datastore.key import Key
 
 
-class Query(object):
+class Query(_implicit_environ._DatastoreBase):
     """A Query against the Cloud Datastore.
 
     This class serves as an abstraction for creating a query over data
@@ -71,7 +72,7 @@ class Query(object):
     """Mapping of operator strings and their protobuf equivalents."""
 
     def __init__(self, kind=None, dataset=None, namespace=None):
-        self._dataset = dataset
+        super(Query, self).__init__(dataset=dataset)
         self._namespace = namespace
         self._pb = datastore_pb.Query()
         self._offset = 0

gcloud/datastore/test___init__.py

@@ -35,6 +35,60 @@ def test_it(self):
         self.assertTrue(client._get_app_default_called)
 
 
+class Test_set_default_dataset(unittest2.TestCase):
+
+    def setUp(self):
+        from gcloud.datastore import _implicit_environ
+        self._replaced_dataset = _implicit_environ.DATASET
+        _implicit_environ.DATASET = None
+
+    def tearDown(self):
+        from gcloud.datastore import _implicit_environ
+        _implicit_environ.DATASET = self._replaced_dataset
+
+    def _callFUT(self, dataset_id=None):
+        from gcloud.datastore import set_default_dataset
+        return set_default_dataset(dataset_id=dataset_id)
+
+    def _test_with_environ(self, environ, expected_result, dataset_id=None):
+        import os
+        from gcloud._testing import _Monkey
+        from gcloud import datastore
+        from gcloud.datastore import _implicit_environ
+
+        # Check the environment is unset.
+        self.assertEqual(_implicit_environ.DATASET, None)
+
+        def custom_getenv(key):
+            return environ.get(key)
+
+        def custom_get_dataset(local_dataset_id):
+            return local_dataset_id
+
+        with _Monkey(os, getenv=custom_getenv):
+            with _Monkey(datastore, get_dataset=custom_get_dataset):
+                self._callFUT(dataset_id=dataset_id)
+
+        self.assertEqual(_implicit_environ.DATASET, expected_result)
+
+    def test_set_from_env_var(self):
+        from gcloud.datastore import _DATASET_ENV_VAR_NAME
+
+        # Make a custom getenv function to Monkey.
+        DATASET = 'dataset'
+        VALUES = {
+            _DATASET_ENV_VAR_NAME: DATASET,
+        }
+        self._test_with_environ(VALUES, DATASET)
+
+    def test_no_env_var_set(self):
+        self._test_with_environ({}, None)
+
+    def test_set_explicit(self):
+        DATASET_ID = 'DATASET'
+        self._test_with_environ({}, DATASET_ID, dataset_id=DATASET_ID)
+
+
 class Test_get_dataset(unittest2.TestCase):
 
     def _callFUT(self, dataset_id):
@@ -56,3 +110,67 @@ def test_it(self):
         self.assertTrue(isinstance(found.connection(), Connection))
         self.assertEqual(found.id(), DATASET_ID)
         self.assertTrue(client._get_app_default_called)
+
+
+class Test_implicit_behavior(unittest2.TestCase):
+
+    def test__require_dataset(self):
+        import gcloud.datastore
+        from gcloud.datastore import _implicit_environ
+        original_dataset = _implicit_environ.DATASET
+
+        try:
+            _implicit_environ.DATASET = None
+            self.assertRaises(EnvironmentError,
+                              gcloud.datastore._require_dataset)
+            NEW_DATASET = object()
+            _implicit_environ.DATASET = NEW_DATASET
+            self.assertEqual(gcloud.datastore._require_dataset(), NEW_DATASET)
+        finally:
+            _implicit_environ.DATASET = original_dataset
+
+    def test_get_entity(self):
+        import gcloud.datastore
+        from gcloud.datastore import _implicit_environ
+        from gcloud.datastore.test_entity import _Dataset
+        from gcloud._testing import _Monkey
+
+        CUSTOM_DATASET = _Dataset()
+        DUMMY_KEY = object()
+        DUMMY_VAL = object()
+        CUSTOM_DATASET[DUMMY_KEY] = DUMMY_VAL
+        with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
+            result = gcloud.datastore.get_entity(DUMMY_KEY)
+        self.assertTrue(result is DUMMY_VAL)
+
+    def test_get_entities(self):
+        import gcloud.datastore
+        from gcloud.datastore import _implicit_environ
+        from gcloud.datastore.test_entity import _Dataset
+        from gcloud._testing import _Monkey
+
+        CUSTOM_DATASET = _Dataset()
+        DUMMY_KEYS = [object(), object()]
+        DUMMY_VALS = [object(), object()]
+        for key, val in zip(DUMMY_KEYS, DUMMY_VALS):
+            CUSTOM_DATASET[key] = val
+
+        with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
+            result = gcloud.datastore.get_entities(DUMMY_KEYS)
+        self.assertTrue(result == DUMMY_VALS)
+
+    def test_allocate_ids(self):
+        import gcloud.datastore
+        from gcloud.datastore import _implicit_environ
+        from gcloud.datastore.key import Key
+        from gcloud.datastore.test_entity import _Dataset
+        from gcloud._testing import _Monkey
+
+        CUSTOM_DATASET = _Dataset()
+        INCOMPLETE_KEY = Key()
+        NUM_IDS = 2
+        with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
+            result = gcloud.datastore.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
+
+        # Check the IDs returned.
+        self.assertEqual([key.id() for key in result], range(1, NUM_IDS + 1))

gcloud/datastore/test_dataset.py

@@ -227,7 +227,7 @@ def test_allocate_ids(self):
         DATASET = self._makeOne(DATASET_ID, connection=CONNECTION)
         result = DATASET.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
 
-        # Check the IDs returned match _PathElementProto.
+        # Check the IDs returned match.
         self.assertEqual([key._id for key in result], range(NUM_IDS))
 
         # Check connection is called correctly.

gcloud/datastore/test_entity.py

@@ -23,8 +23,10 @@
 class TestEntity(unittest2.TestCase):
 
     def _getTargetClass(self):
+        from gcloud.datastore import _implicit_environ
         from gcloud.datastore.entity import Entity
 
+        _implicit_environ.DATASET = None
         return Entity
 
     def _makeOne(self, dataset=_MARKER, kind=_KIND, exclude_from_indexes=()):
@@ -265,6 +267,13 @@ def __init__(self, connection=None):
         super(_Dataset, self).__init__()
         self._connection = connection
 
+    def __bool__(self):
+        # Make sure the objects are Truth-y since an empty
+        # dict with _connection set will still be False-y.
+        return True
+
+    __nonzero__ = __bool__
+
     def id(self):
         return _DATASET_ID
 
@@ -274,6 +283,12 @@ def connection(self):
     def get_entity(self, key):
         return self.get(key)
 
+    def get_entities(self, keys):
+        return [self.get(key) for key in keys]
+
+    def allocate_ids(self, incomplete_key, num_ids):
+        return [incomplete_key.id(i + 1) for i in range(num_ids)]
+
 
 class _Connection(object):
     _transaction = _saved = _deleted = None

gcloud/datastore/test_helpers.py

@@ -19,6 +19,15 @@ class Test_entity_from_protobuf(unittest2.TestCase):
 
     _MARKER = object()
 
+    def setUp(self):
+        from gcloud.datastore import _implicit_environ
+        self._replaced_dataset = _implicit_environ.DATASET
+        _implicit_environ.DATASET = None
+
+    def tearDown(self):
+        from gcloud.datastore import _implicit_environ
+        _implicit_environ.DATASET = self._replaced_dataset
+
     def _callFUT(self, val, dataset=_MARKER):
         from gcloud.datastore.helpers import entity_from_protobuf
 

gcloud/datastore/test_query.py

@@ -17,9 +17,17 @@
 
 class TestQuery(unittest2.TestCase):
 
+    def setUp(self):
+        from gcloud.datastore import _implicit_environ
+        self._replaced_dataset = _implicit_environ.DATASET
+        _implicit_environ.DATASET = None
+
+    def tearDown(self):
+        from gcloud.datastore import _implicit_environ
+        _implicit_environ.DATASET = self._replaced_dataset
+
     def _getTargetClass(self):
         from gcloud.datastore.query import Query
-
         return Query
 
     def _makeOne(self, kind=None, dataset=None, namespace=None):

gcloud/datastore/test_transaction.py

@@ -38,6 +38,15 @@ def test_ctor(self):
         self.assertEqual(len(xact._auto_id_entities), 0)
         self.assertTrue(xact.connection() is connection)
 
+    def test_ctor_with_env(self):
+        SENTINEL_VAL = object()
+
+        from gcloud.datastore import _implicit_environ
+        _implicit_environ.DATASET = SENTINEL_VAL
+
+        transaction = self._makeOne(dataset=None)
+        self.assertEqual(transaction.dataset(), SENTINEL_VAL)
+
     def test_add_auto_id_entity(self):
         entity = _Entity()
         _DATASET = 'DATASET'