Ideas commit for better (implicit) auth story in datastore.

Also
- Integrates this idea into the datastore regression tests.
- Creates a (temporarily ignored) cyclic import
- Has an uncovered statement that confuses nosetests
- Implements a friendlier (more gcloud-node like) version of
  allocated_ids
- DOES NOT update CONTRIBUTING to explain change to regression
  tests or ditch unneeded helper / environ code for regressions
- Creates a datastore cleanup script for when tests get
  broken

Author: dhermes

CONTRIBUTING.rst

@@ -168,6 +168,9 @@ Running Regression Tests
     `docs <https://cloud.google.com/storage/docs/authentication#generating-a-private-key>`__
     for explanation on how to get a private key.
 
+  DJH: THIS PART NEEDS TO BE UPDATED AFTER DISCUSSION OF IMPLICIT ENVIRON
+       USE IN PRODUCTION CODE.
+
 - Examples of these can be found in ``regression/local_test_setup.sample``. We
   recommend copying this to ``regression/local_test_setup``, editing the values
   and sourcing them into your environment::

gcloud/datastore/__init__.py

@@ -32,11 +32,38 @@
   which represents a lookup or search over the rows in the datastore.
 """
 
+import os
+
 __version__ = '0.1.2'
 
 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 = None
+"""Module global which allows users to only optionally use a dataset."""
+
+
+def get_local_dataset_settings():
+    """Determines auth settings from local enviroment.
+
+    Currently only supports enviroment variables but will implicitly
+    support App Engine, Compute Engine and other environments in
+    the future.
+
+    Local environment variables used are:
+    - GCLOUD_DATASET_ID
+    - GCLOUD_CLIENT_EMAIL
+    - GCLOUD_KEY_FILE
+    """
+    local_dataset_settings = (
+        os.getenv('GCLOUD_DATASET_ID'),
+        os.getenv('GCLOUD_CLIENT_EMAIL'),
+        os.getenv('GCLOUD_KEY_FILE'),
+    )
+    if None in local_dataset_settings:
+        return None
+    else:
+        return local_dataset_settings
 
 
 def get_connection(client_email, private_key_path):
@@ -102,3 +129,73 @@ def get_dataset(dataset_id, client_email, private_key_path):
     """
     connection = get_connection(client_email, private_key_path)
     return connection.dataset(dataset_id)
+
+
+def _require_dataset():
+    """Convenience method to ensure DATASET is set.
+
+    :raises: :class:`EnvironmentError` if DATASET is not set.
+    """
+    if DATASET is None:
+        raise EnvironmentError('Dataset could not be implied.')
+
+
+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.
+    """
+    _require_dataset()
+    return 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.
+    """
+    _require_dataset()
+    return 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.
+    """
+    _require_dataset()
+
+    if not incomplete_key.is_partial():
+        raise ValueError(('Key is not partial.', incomplete_key))
+
+    incomplete_key_pb = incomplete_key.to_protobuf()
+    incomplete_key_pbs = [incomplete_key_pb] * num_ids
+
+    allocated_key_pbs = DATASET.connection().allocate_ids(
+        DATASET.id(), incomplete_key_pbs)
+    allocated_ids = [allocated_key_pb.path_element[-1].id
+                     for allocated_key_pb in allocated_key_pbs]
+    return [incomplete_key.id(allocated_id)
+            for allocated_id in allocated_ids]
+
+
+# Set DATASET if it can be implied from the environment.
+LOCAL_DATASET_SETTINGS = get_local_dataset_settings()
+if LOCAL_DATASET_SETTINGS is not None:
+    DATASET = get_dataset(*LOCAL_DATASET_SETTINGS)
+del LOCAL_DATASET_SETTINGS

gcloud/datastore/dataset.py

@@ -106,7 +106,7 @@ def get_entity(self, key):
         """Retrieves entity from the dataset, along with its attributes.
 
         :type key: :class:`gcloud.datastore.key.Key`
-        :param item_name: The name of the item to retrieve.
+        :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.
@@ -118,8 +118,8 @@ def get_entity(self, key):
     def get_entities(self, keys):
         """Retrieves entities from the dataset, along with their attributes.
 
-        :type key: list of :class:`gcloud.datastore.key.Key`
-        :param item_name: The name of the item to retrieve.
+        :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.

gcloud/datastore/entity.py

@@ -15,6 +15,7 @@
 delete or persist the data stored on the entity.
 """
 
+import gcloud.datastore
 from gcloud.datastore import datastore_v1_pb2 as datastore_pb
 from gcloud.datastore.key import Key
 
@@ -73,6 +74,8 @@ class Entity(dict):
     def __init__(self, dataset=None, kind=None):
         super(Entity, self).__init__()
         self._dataset = dataset
+        if self._dataset is None:
+            self._dataset = gcloud.datastore.DATASET
         if kind:
             self._key = Key().kind(kind)
         else:

gcloud/datastore/query.py

@@ -2,6 +2,7 @@
 
 import base64
 
+import gcloud.datastore
 from gcloud.datastore import datastore_v1_pb2 as datastore_pb
 from gcloud.datastore import helpers
 from gcloud.datastore.key import Key
@@ -56,6 +57,8 @@ class Query(object):
 
     def __init__(self, kind=None, dataset=None, namespace=None):
         self._dataset = dataset
+        if self._dataset is None:
+            self._dataset = gcloud.datastore.DATASET
         self._namespace = namespace
         self._pb = datastore_pb.Query()
         self._cursor = None

gcloud/datastore/test___init__.py

@@ -33,6 +33,50 @@ def test_it(self):
         self.assertEqual(client._called_with, expected_called_with)
 
 
+class Test_get_local_dataset_settings(unittest2.TestCase):
+
+    def _callFUT(self):
+        from gcloud.datastore import get_local_dataset_settings
+        return get_local_dataset_settings()
+
+    def _test_with_environ(self, environ, expected_result):
+        import os
+        from gcloud._testing import _Monkey
+
+        def custom_getenv(key):
+            return environ.get(key)
+
+        with _Monkey(os, getenv=custom_getenv):
+            result = self._callFUT()
+
+        self.assertEqual(result, expected_result)
+
+    def test_all_set(self):
+        # Fake auth variables.
+        DATASET = 'dataset'
+        CLIENT_EMAIL = '[email protected]'
+        TEMP_PATH = 'fakepath'
+
+        # Make a custom getenv function to Monkey.
+        VALUES = {
+            'GCLOUD_DATASET_ID': DATASET,
+            'GCLOUD_CLIENT_EMAIL': CLIENT_EMAIL,
+            'GCLOUD_KEY_FILE': TEMP_PATH,
+        }
+        expected_result = (DATASET, CLIENT_EMAIL, TEMP_PATH)
+        self._test_with_environ(VALUES, expected_result)
+
+    def test_partial_set(self):
+        # Fake auth variables.
+        DATASET = 'dataset'
+
+        # Make a custom getenv function to Monkey.
+        VALUES = {
+            'GCLOUD_DATASET_ID': DATASET,
+        }
+        self._test_with_environ(VALUES, None)
+
+
 class Test_get_dataset(unittest2.TestCase):
 
     def _callFUT(self, dataset_id, client_email, private_key_path):
@@ -66,3 +110,160 @@ def test_it(self):
             'scope': SCOPE,
         }
         self.assertEqual(client._called_with, expected_called_with)
+
+
+class Test_implicit_behavior(unittest2.TestCase):
+
+    def test__require_dataset(self):
+        import gcloud.datastore
+        original_dataset = gcloud.datastore.DATASET
+
+        try:
+            gcloud.datastore.DATASET = None
+            self.assertRaises(EnvironmentError,
+                              gcloud.datastore._require_dataset)
+            gcloud.datastore.DATASET = object()
+            self.assertEqual(gcloud.datastore._require_dataset(), None)
+        finally:
+            gcloud.datastore.DATASET = original_dataset
+
+    def test_get_entity(self):
+        import gcloud.datastore
+        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(gcloud.datastore, 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.test_entity import _Dataset
+        from gcloud._testing import _Monkey
+
+        class _ExtendedDataset(_Dataset):
+            def get_entities(self, keys):
+                return [self.get(key) for key in keys]
+
+        CUSTOM_DATASET = _ExtendedDataset()
+        DUMMY_KEYS = [object(), object()]
+        DUMMY_VALS = [object(), object()]
+        for key, val in zip(DUMMY_KEYS, DUMMY_VALS):
+            CUSTOM_DATASET[key] = val
+
+        with _Monkey(gcloud.datastore, 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.test_entity import _Connection
+        from gcloud.datastore.test_entity import _DATASET_ID
+        from gcloud.datastore.test_entity import _Dataset
+        from gcloud.datastore.test_entity import _Key
+        from gcloud._testing import _Monkey
+
+        class _PathElementProto(object):
+            COUNTER = 0
+
+            def __init__(self):
+                _PathElementProto.COUNTER += 1
+                self.id = _PathElementProto.COUNTER
+
+        class _KeyProto(object):
+
+            def __init__(self):
+                self.path_element = [_PathElementProto()]
+
+        class _ExtendedKey(_Key):
+            def id(self, id_to_set):
+                self._called_id = id_to_set
+                return id_to_set
+
+        INCOMPLETE_KEY = _ExtendedKey()
+        INCOMPLETE_KEY._key = _KeyProto()
+        INCOMPLETE_KEY._partial = True
+        NUM_IDS = 2
+
+        class _ExtendedConnection(_Connection):
+            def allocate_ids(self, dataset_id, key_pbs):
+                self._called_dataset_id = dataset_id
+                self._called_key_pbs = key_pbs
+                return key_pbs
+
+        CUSTOM_CONNECTION = _ExtendedConnection()
+        CUSTOM_DATASET = _Dataset(connection=CUSTOM_CONNECTION)
+        with _Monkey(gcloud.datastore, DATASET=CUSTOM_DATASET):
+            result = gcloud.datastore.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
+
+        self.assertEqual(_PathElementProto.COUNTER, 1)
+        self.assertEqual(result, [1, 1])
+        self.assertEqual(CUSTOM_CONNECTION._called_dataset_id, _DATASET_ID)
+        self.assertEqual(len(CUSTOM_CONNECTION._called_key_pbs), 2)
+        key_paths = [key_pb.path_element[-1].id
+                     for key_pb in CUSTOM_CONNECTION._called_key_pbs]
+        self.assertEqual(key_paths, [1, 1])
+
+    def test_allocate_ids_with_complete(self):
+        import gcloud.datastore
+        from gcloud.datastore.test_entity import _Connection
+        from gcloud.datastore.test_entity import _Dataset
+        from gcloud.datastore.test_entity import _Key
+        from gcloud._testing import _Monkey
+
+        COMPLETE_KEY = _Key()
+        NUM_IDS = 2
+        CUSTOM_CONNECTION = _Connection()
+        CUSTOM_DATASET = _Dataset(connection=CUSTOM_CONNECTION)
+        with _Monkey(gcloud.datastore, DATASET=CUSTOM_DATASET):
+            self.assertRaises(ValueError, gcloud.datastore.allocate_ids,
+                              COMPLETE_KEY, NUM_IDS)
+
+    def test_set_DATASET(self):
+        import os
+        import tempfile
+        from gcloud import credentials
+        from gcloud.test_credentials import _Client
+        from gcloud._testing import _Monkey
+
+        # Make custom client for doing auth.
+        client = _Client()
+
+        # Fake auth variables.
+        CLIENT_EMAIL = '[email protected]'
+        PRIVATE_KEY = 'SEEkR1t'
+        DATASET = 'dataset'
+
+        # Write the fake key to a temp file.
+        TEMP_PATH = tempfile.mktemp()
+        with open(TEMP_PATH, 'w') as file_obj:
+            file_obj.write(PRIVATE_KEY)
+            file_obj.flush()
+
+        # Make a custom getenv function to Monkey.
+        VALUES = {
+            'GCLOUD_DATASET_ID': DATASET,
+            'GCLOUD_CLIENT_EMAIL': CLIENT_EMAIL,
+            'GCLOUD_KEY_FILE': TEMP_PATH,
+        }
+
+        def custom_getenv(key):
+            return VALUES.get(key)
+
+        # Perform the import again with our test patches.
+        with _Monkey(credentials, client=client):
+            with _Monkey(os, getenv=custom_getenv):
+                import gcloud.datastore
+                reload(gcloud.datastore)
+
+        # Check that the DATASET was correctly implied from the environ.
+        implicit_dataset = gcloud.datastore.DATASET
+        self.assertEqual(implicit_dataset.id(), DATASET)
+        # Check that the credentials on the implicit DATASET was set on the
+        # fake client.
+        credentials = implicit_dataset.connection().credentials
+        self.assertTrue(credentials is client._signed)

gcloud/datastore/test_entity.py

@@ -9,6 +9,9 @@
 class TestEntity(unittest2.TestCase):
 
     def _getTargetClass(self):
+        import gcloud.datastore
+        gcloud.datastore.DATASET = None
+
         from gcloud.datastore.entity import Entity
 
         return Entity