Skip to content

Commit

Permalink
Adding helper methods for implicit datastore environ.
Browse files Browse the repository at this point in the history
By setting two environment variables

- GCLOUD_DATASET_ID
- GOOGLE_APPLICATION_CREDENTIALS

the user can call convenience methods in `gcloud.datastore` directly
without worrying about auth or the name of the dataset. The
goal is that in places like App Engine and Compute Engine, the
dataset ID can be implied without **any** user intervention.

Partially addresses #337.

NOTE: This still needs to be documented, but it's unclear
      where is appropriate. We also need to have documentation
      for auth (outside of CONTRIBUTING.md).
  • Loading branch information
dhermes committed Dec 22, 2014
1 parent b343acf commit d85d200
Show file tree
Hide file tree
Showing 12 changed files with 306 additions and 30 deletions.
75 changes: 75 additions & 0 deletions gcloud/datastore/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,14 +44,34 @@
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_dataset_from_environ():
"""Determines auth settings from 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
"""
local_dataset_id = os.getenv(_DATASET_ENV_VAR_NAME)
if local_dataset_id is not None:
_implicit_environ.DATASET = get_dataset(local_dataset_id)


def get_connection():
"""Shortcut method to establish a connection to the Cloud Datastore.
Expand Down Expand Up @@ -97,3 +117,58 @@ 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 implied.')
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)


# Set DATASET if it can be implied from the environment.
_set_dataset_from_environ()
24 changes: 24 additions & 0 deletions gcloud/datastore/_implicit_environ.py
Original file line number Diff line number Diff line change
@@ -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
5 changes: 4 additions & 1 deletion gcloud/datastore/entity.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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 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:
Expand Down
5 changes: 3 additions & 2 deletions gcloud/datastore/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
142 changes: 142 additions & 0 deletions gcloud/datastore/test___init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,47 @@ def test_it(self):
self.assertTrue(client._get_app_default_called)


class Test__set_dataset_from_environ(unittest2.TestCase):

def _callFUT(self):
from gcloud.datastore import _set_dataset_from_environ
return _set_dataset_from_environ()

def _test_with_environ(self, environ, expected_result):
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(dataset_id):
return dataset_id

with _Monkey(os, getenv=custom_getenv):
with _Monkey(datastore, get_dataset=custom_get_dataset):
self._callFUT()

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)


class Test_get_dataset(unittest2.TestCase):

def _callFUT(self, dataset_id):
Expand All @@ -56,3 +97,104 @@ 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))

def test_set_DATASET(self):
import os
from gcloud._testing import _Monkey
from gcloud.test_credentials import _Client
from gcloud import credentials
from gcloud.datastore import _implicit_environ

# Make custom client for doing auth. Have to fake auth since we
# can't monkey patch `datastore.get_dataset` while reloading the
# `datastore.__init__` module.
client = _Client()

# Fake auth variables.
DATASET = 'dataset'

# Make a custom getenv function to Monkey.
VALUES = {
'GCLOUD_DATASET_ID': DATASET,
}

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 = _implicit_environ.DATASET
self.assertEqual(implicit_dataset.id(), DATASET)
# Check that the credentials on the implicit DATASET was set on the
# fake client.
cnxn_credentials = implicit_dataset.connection().credentials
self.assertTrue(cnxn_credentials is client._signed)
2 changes: 1 addition & 1 deletion gcloud/datastore/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
15 changes: 15 additions & 0 deletions gcloud/datastore/test_entity.py
Original file line number Diff line number Diff line change
Expand Up @@ -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=()):
Expand Down Expand Up @@ -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

Expand All @@ -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
Expand Down
3 changes: 3 additions & 0 deletions gcloud/datastore/test_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,11 @@ class Test_entity_from_protobuf(unittest2.TestCase):
_MARKER = object()

def _callFUT(self, val, dataset=_MARKER):
from gcloud.datastore import _implicit_environ
from gcloud.datastore.helpers import entity_from_protobuf

_implicit_environ.DATASET = None

if dataset is self._MARKER:
return entity_from_protobuf(val)

Expand Down
2 changes: 2 additions & 0 deletions gcloud/datastore/test_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,10 @@
class TestQuery(unittest2.TestCase):

def _getTargetClass(self):
from gcloud.datastore import _implicit_environ
from gcloud.datastore.query import Query

_implicit_environ.DATASET = None
return Query

def _makeOne(self, kind=None, dataset=None, namespace=None):
Expand Down
Loading

0 comments on commit d85d200

Please sign in to comment.