Skip to content

Commit d85d200

Browse files
dhermesdhermes
committed
Adding helper methods for implicit datastore environ.
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).
1 parent b343acf commit d85d200

File tree

12 files changed

+306
-30
lines changed

12 files changed

+306
-30
lines changed

‎gcloud/datastore/__init__.py‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,14 +44,34 @@
4444
which represents a lookup or search over the rows in the datastore.
4545
"""
4646

47+
import os
48+
4749
from gcloud import credentials
50+
from gcloud.datastore import _implicit_environ
4851
from gcloud.datastore.connection import Connection
4952

5053

5154
SCOPE = ('https://www.googleapis.com/auth/datastore ',
5255
'https://www.googleapis.com/auth/userinfo.email')
5356
"""The scope required for authenticating as a Cloud Datastore consumer."""
5457

58+
_DATASET_ENV_VAR_NAME = 'GCLOUD_DATASET_ID'
59+
60+
61+
def _set_dataset_from_environ():
62+
"""Determines auth settings from local enviroment.
63+
64+
Currently only supports enviroment variable but will implicitly
65+
support App Engine, Compute Engine and other environments in
66+
the future.
67+
68+
Local environment variable used is:
69+
- GCLOUD_DATASET_ID
70+
"""
71+
local_dataset_id = os.getenv(_DATASET_ENV_VAR_NAME)
72+
if local_dataset_id is not None:
73+
_implicit_environ.DATASET = get_dataset(local_dataset_id)
74+
5575

5676
def get_connection():
5777
"""Shortcut method to establish a connection to the Cloud Datastore.
@@ -97,3 +117,58 @@ def get_dataset(dataset_id):
97117
"""
98118
connection = get_connection()
99119
return connection.dataset(dataset_id)
120+
121+
122+
def _require_dataset():
123+
"""Convenience method to ensure DATASET is set.
124+
125+
:rtype: :class:`gcloud.datastore.dataset.Dataset`
126+
:returns: A dataset based on the current environment.
127+
:raises: :class:`EnvironmentError` if DATASET is not set.
128+
"""
129+
if _implicit_environ.DATASET is None:
130+
raise EnvironmentError('Dataset could not be implied.')
131+
return _implicit_environ.DATASET
132+
133+
134+
def get_entity(key):
135+
"""Retrieves entity from implicit dataset, along with its attributes.
136+
137+
:type key: :class:`gcloud.datastore.key.Key`
138+
:param key: The name of the item to retrieve.
139+
140+
:rtype: :class:`gcloud.datastore.entity.Entity` or ``None``
141+
:return: The requested entity, or ``None`` if there was no match found.
142+
"""
143+
return _require_dataset().get_entity(key)
144+
145+
146+
def get_entities(keys):
147+
"""Retrieves entities from implied dataset, along with their attributes.
148+
149+
:type keys: list of :class:`gcloud.datastore.key.Key`
150+
:param keys: The name of the item to retrieve.
151+
152+
:rtype: list of :class:`gcloud.datastore.entity.Entity`
153+
:return: The requested entities.
154+
"""
155+
return _require_dataset().get_entities(keys)
156+
157+
158+
def allocate_ids(incomplete_key, num_ids):
159+
"""Allocates a list of IDs from a partial key.
160+
161+
:type incomplete_key: A :class:`gcloud.datastore.key.Key`
162+
:param incomplete_key: The partial key to use as base for allocated IDs.
163+
164+
:type num_ids: A :class:`int`.
165+
:param num_ids: The number of IDs to allocate.
166+
167+
:rtype: list of :class:`gcloud.datastore.key.Key`
168+
:return: The (complete) keys allocated with `incomplete_key` as root.
169+
"""
170+
return _require_dataset().allocate_ids(incomplete_key, num_ids)
171+
172+
173+
# Set DATASET if it can be implied from the environment.
174+
_set_dataset_from_environ()

‎gcloud/datastore/_implicit_environ.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
"""Module to provide implicit behavior based on enviroment.
2+
3+
Acts as a mutable namespace to allow the datastore package to
4+
imply the current dataset from the enviroment.
5+
6+
Also provides a base class for classes in the `datastore` package
7+
which could utilize the implicit enviroment.
8+
"""
9+
10+
11+
DATASET = None
12+
"""Module global to allow persistent implied dataset from enviroment."""
13+
14+
15+
class _DatastoreBase(object):
16+
"""Base for all classes in the datastore package.
17+
18+
Uses the implicit DATASET object as a default dataset attached
19+
to the instances being created. Stores the dataset passed in
20+
on the protected (i.e. non-public) attribute `_dataset`.
21+
"""
22+
23+
def __init__(self, dataset=None):
24+
self._dataset = dataset or DATASET

‎gcloud/datastore/entity.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
"""Class for representing a single entity in the Cloud Datastore."""
1616

17+
from gcloud.datastore import _implicit_environ
1718
from gcloud.datastore import datastore_v1_pb2 as datastore_pb
1819
from gcloud.datastore.key import Key
1920

@@ -95,7 +96,9 @@ class Entity(dict):
9596

9697
def __init__(self, dataset=None, kind=None, exclude_from_indexes=()):
9798
super(Entity, self).__init__()
98-
self._dataset = dataset
99+
# Does not inherit from object, so we don't use
100+
# _implicit_environ._DatastoreBase to avoid split MRO.
101+
self._dataset = dataset or _implicit_environ.DATASET
99102
if kind:
100103
self._key = Key().kind(kind)
101104
else:

‎gcloud/datastore/query.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,13 @@
1616

1717
import base64
1818

19+
from gcloud.datastore import _implicit_environ
1920
from gcloud.datastore import datastore_v1_pb2 as datastore_pb
2021
from gcloud.datastore import helpers
2122
from gcloud.datastore.key import Key
2223

2324

24-
class Query(object):
25+
class Query(_implicit_environ._DatastoreBase):
2526
"""A Query against the Cloud Datastore.
2627
2728
This class serves as an abstraction for creating a query over data
@@ -71,7 +72,7 @@ class Query(object):
7172
"""Mapping of operator strings and their protobuf equivalents."""
7273

7374
def __init__(self, kind=None, dataset=None, namespace=None):
74-
self._dataset = dataset
75+
super(Query, self).__init__(dataset=dataset)
7576
self._namespace = namespace
7677
self._pb = datastore_pb.Query()
7778
self._offset = 0

‎gcloud/datastore/test___init__.py‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,47 @@ def test_it(self):
3535
self.assertTrue(client._get_app_default_called)
3636

3737

38+
class Test__set_dataset_from_environ(unittest2.TestCase):
39+
40+
def _callFUT(self):
41+
from gcloud.datastore import _set_dataset_from_environ
42+
return _set_dataset_from_environ()
43+
44+
def _test_with_environ(self, environ, expected_result):
45+
import os
46+
from gcloud._testing import _Monkey
47+
from gcloud import datastore
48+
from gcloud.datastore import _implicit_environ
49+
50+
# Check the environment is unset.
51+
self.assertEqual(_implicit_environ.DATASET, None)
52+
53+
def custom_getenv(key):
54+
return environ.get(key)
55+
56+
def custom_get_dataset(dataset_id):
57+
return dataset_id
58+
59+
with _Monkey(os, getenv=custom_getenv):
60+
with _Monkey(datastore, get_dataset=custom_get_dataset):
61+
self._callFUT()
62+
63+
self.assertEqual(_implicit_environ.DATASET, expected_result)
64+
65+
def test_set_from_env_var(self):
66+
from gcloud.datastore import _DATASET_ENV_VAR_NAME
67+
68+
# Make a custom getenv function to Monkey.
69+
DATASET = 'dataset'
70+
VALUES = {
71+
_DATASET_ENV_VAR_NAME: DATASET,
72+
}
73+
self._test_with_environ(VALUES, DATASET)
74+
75+
def test_no_env_var_set(self):
76+
self._test_with_environ({}, None)
77+
78+
3879
class Test_get_dataset(unittest2.TestCase):
3980

4081
def _callFUT(self, dataset_id):
@@ -56,3 +97,104 @@ def test_it(self):
5697
self.assertTrue(isinstance(found.connection(), Connection))
5798
self.assertEqual(found.id(), DATASET_ID)
5899
self.assertTrue(client._get_app_default_called)
100+
101+
102+
class Test_implicit_behavior(unittest2.TestCase):
103+
104+
def test__require_dataset(self):
105+
import gcloud.datastore
106+
from gcloud.datastore import _implicit_environ
107+
original_dataset = _implicit_environ.DATASET
108+
109+
try:
110+
_implicit_environ.DATASET = None
111+
self.assertRaises(EnvironmentError,
112+
gcloud.datastore._require_dataset)
113+
NEW_DATASET = object()
114+
_implicit_environ.DATASET = NEW_DATASET
115+
self.assertEqual(gcloud.datastore._require_dataset(), NEW_DATASET)
116+
finally:
117+
_implicit_environ.DATASET = original_dataset
118+
119+
def test_get_entity(self):
120+
import gcloud.datastore
121+
from gcloud.datastore import _implicit_environ
122+
from gcloud.datastore.test_entity import _Dataset
123+
from gcloud._testing import _Monkey
124+
125+
CUSTOM_DATASET = _Dataset()
126+
DUMMY_KEY = object()
127+
DUMMY_VAL = object()
128+
CUSTOM_DATASET[DUMMY_KEY] = DUMMY_VAL
129+
with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
130+
result = gcloud.datastore.get_entity(DUMMY_KEY)
131+
self.assertTrue(result is DUMMY_VAL)
132+
133+
def test_get_entities(self):
134+
import gcloud.datastore
135+
from gcloud.datastore import _implicit_environ
136+
from gcloud.datastore.test_entity import _Dataset
137+
from gcloud._testing import _Monkey
138+
139+
CUSTOM_DATASET = _Dataset()
140+
DUMMY_KEYS = [object(), object()]
141+
DUMMY_VALS = [object(), object()]
142+
for key, val in zip(DUMMY_KEYS, DUMMY_VALS):
143+
CUSTOM_DATASET[key] = val
144+
145+
with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
146+
result = gcloud.datastore.get_entities(DUMMY_KEYS)
147+
self.assertTrue(result == DUMMY_VALS)
148+
149+
def test_allocate_ids(self):
150+
import gcloud.datastore
151+
from gcloud.datastore import _implicit_environ
152+
from gcloud.datastore.key import Key
153+
from gcloud.datastore.test_entity import _Dataset
154+
from gcloud._testing import _Monkey
155+
156+
CUSTOM_DATASET = _Dataset()
157+
INCOMPLETE_KEY = Key()
158+
NUM_IDS = 2
159+
with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
160+
result = gcloud.datastore.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
161+
162+
# Check the IDs returned.
163+
self.assertEqual([key.id() for key in result], range(1, NUM_IDS + 1))
164+
165+
def test_set_DATASET(self):
166+
import os
167+
from gcloud._testing import _Monkey
168+
from gcloud.test_credentials import _Client
169+
from gcloud import credentials
170+
from gcloud.datastore import _implicit_environ
171+
172+
# Make custom client for doing auth. Have to fake auth since we
173+
# can't monkey patch `datastore.get_dataset` while reloading the
174+
# `datastore.__init__` module.
175+
client = _Client()
176+
177+
# Fake auth variables.
178+
DATASET = 'dataset'
179+
180+
# Make a custom getenv function to Monkey.
181+
VALUES = {
182+
'GCLOUD_DATASET_ID': DATASET,
183+
}
184+
185+
def custom_getenv(key):
186+
return VALUES.get(key)
187+
188+
# Perform the import again with our test patches.
189+
with _Monkey(credentials, client=client):
190+
with _Monkey(os, getenv=custom_getenv):
191+
import gcloud.datastore
192+
reload(gcloud.datastore)
193+
194+
# Check that the DATASET was correctly implied from the environ.
195+
implicit_dataset = _implicit_environ.DATASET
196+
self.assertEqual(implicit_dataset.id(), DATASET)
197+
# Check that the credentials on the implicit DATASET was set on the
198+
# fake client.
199+
cnxn_credentials = implicit_dataset.connection().credentials
200+
self.assertTrue(cnxn_credentials is client._signed)

‎gcloud/datastore/test_dataset.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -227,7 +227,7 @@ def test_allocate_ids(self):
227227
DATASET = self._makeOne(DATASET_ID, connection=CONNECTION)
228228
result = DATASET.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
229229

230-
# Check the IDs returned match _PathElementProto.
230+
# Check the IDs returned match.
231231
self.assertEqual([key._id for key in result], range(NUM_IDS))
232232

233233
# Check connection is called correctly.

‎gcloud/datastore/test_entity.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,10 @@
2323
class TestEntity(unittest2.TestCase):
2424

2525
def _getTargetClass(self):
26+
from gcloud.datastore import _implicit_environ
2627
from gcloud.datastore.entity import Entity
2728

29+
_implicit_environ.DATASET = None
2830
return Entity
2931

3032
def _makeOne(self, dataset=_MARKER, kind=_KIND, exclude_from_indexes=()):
@@ -265,6 +267,13 @@ def __init__(self, connection=None):
265267
super(_Dataset, self).__init__()
266268
self._connection = connection
267269

270+
def __bool__(self):
271+
# Make sure the objects are Truth-y since an empty
272+
# dict with _connection set will still be False-y.
273+
return True
274+
275+
__nonzero__ = __bool__
276+
268277
def id(self):
269278
return _DATASET_ID
270279

@@ -274,6 +283,12 @@ def connection(self):
274283
def get_entity(self, key):
275284
return self.get(key)
276285

286+
def get_entities(self, keys):
287+
return [self.get(key) for key in keys]
288+
289+
def allocate_ids(self, incomplete_key, num_ids):
290+
return [incomplete_key.id(i + 1) for i in range(num_ids)]
291+
277292

278293
class _Connection(object):
279294
_transaction = _saved = _deleted = None

‎gcloud/datastore/test_helpers.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,11 @@ class Test_entity_from_protobuf(unittest2.TestCase):
2020
_MARKER = object()
2121

2222
def _callFUT(self, val, dataset=_MARKER):
23+
from gcloud.datastore import _implicit_environ
2324
from gcloud.datastore.helpers import entity_from_protobuf
2425

26+
_implicit_environ.DATASET = None
27+
2528
if dataset is self._MARKER:
2629
return entity_from_protobuf(val)
2730

‎gcloud/datastore/test_query.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,10 @@
1818
class TestQuery(unittest2.TestCase):
1919

2020
def _getTargetClass(self):
21+
from gcloud.datastore import _implicit_environ
2122
from gcloud.datastore.query import Query
2223

24+
_implicit_environ.DATASET = None
2325
return Query
2426

2527
def _makeOne(self, kind=None, dataset=None, namespace=None):

0 commit comments

Comments
 (0)