Skip to content

Commit 563a19f

Browse files
dhermesdhermes
committed
Merge pull request #444 from dhermes/implicit-dataset-from-environ
Implicit dataset from environ
2 parents 881f41c + 7cd77f2 commit 563a19f

File tree

12 files changed

+301
-31
lines changed

12 files changed

+301
-31
lines changed

‎gcloud/datastore/__init__.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,14 +44,41 @@
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_default_dataset(dataset_id=None):
62+
"""Determines auth settings from local enviroment.
63+
64+
Sets a default dataset either explicitly or via the local
65+
enviroment. Currently only supports enviroment variable but will
66+
implicitly support App Engine, Compute Engine and other environments
67+
in the future.
68+
69+
Local environment variable used is:
70+
- GCLOUD_DATASET_ID
71+
72+
:type dataset_id: :class:`str`.
73+
:param dataset_id: Optional. The dataset ID to use for the default
74+
dataset.
75+
"""
76+
if dataset_id is None:
77+
dataset_id = os.getenv(_DATASET_ENV_VAR_NAME)
78+
79+
if dataset_id is not None:
80+
_implicit_environ.DATASET = get_dataset(dataset_id)
81+
5582

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

‎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 directly 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: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,60 @@ def test_it(self):
3535
self.assertTrue(client._get_app_default_called)
3636

3737

38+
class Test_set_default_dataset(unittest2.TestCase):
39+
40+
def setUp(self):
41+
from gcloud.datastore import _implicit_environ
42+
self._replaced_dataset = _implicit_environ.DATASET
43+
_implicit_environ.DATASET = None
44+
45+
def tearDown(self):
46+
from gcloud.datastore import _implicit_environ
47+
_implicit_environ.DATASET = self._replaced_dataset
48+
49+
def _callFUT(self, dataset_id=None):
50+
from gcloud.datastore import set_default_dataset
51+
return set_default_dataset(dataset_id=dataset_id)
52+
53+
def _test_with_environ(self, environ, expected_result, dataset_id=None):
54+
import os
55+
from gcloud._testing import _Monkey
56+
from gcloud import datastore
57+
from gcloud.datastore import _implicit_environ
58+
59+
# Check the environment is unset.
60+
self.assertEqual(_implicit_environ.DATASET, None)
61+
62+
def custom_getenv(key):
63+
return environ.get(key)
64+
65+
def custom_get_dataset(local_dataset_id):
66+
return local_dataset_id
67+
68+
with _Monkey(os, getenv=custom_getenv):
69+
with _Monkey(datastore, get_dataset=custom_get_dataset):
70+
self._callFUT(dataset_id=dataset_id)
71+
72+
self.assertEqual(_implicit_environ.DATASET, expected_result)
73+
74+
def test_set_from_env_var(self):
75+
from gcloud.datastore import _DATASET_ENV_VAR_NAME
76+
77+
# Make a custom getenv function to Monkey.
78+
DATASET = 'dataset'
79+
VALUES = {
80+
_DATASET_ENV_VAR_NAME: DATASET,
81+
}
82+
self._test_with_environ(VALUES, DATASET)
83+
84+
def test_no_env_var_set(self):
85+
self._test_with_environ({}, None)
86+
87+
def test_set_explicit(self):
88+
DATASET_ID = 'DATASET'
89+
self._test_with_environ({}, DATASET_ID, dataset_id=DATASET_ID)
90+
91+
3892
class Test_get_dataset(unittest2.TestCase):
3993

4094
def _callFUT(self, dataset_id):
@@ -56,3 +110,67 @@ def test_it(self):
56110
self.assertTrue(isinstance(found.connection(), Connection))
57111
self.assertEqual(found.id(), DATASET_ID)
58112
self.assertTrue(client._get_app_default_called)
113+
114+
115+
class Test_implicit_behavior(unittest2.TestCase):
116+
117+
def test__require_dataset(self):
118+
import gcloud.datastore
119+
from gcloud.datastore import _implicit_environ
120+
original_dataset = _implicit_environ.DATASET
121+
122+
try:
123+
_implicit_environ.DATASET = None
124+
self.assertRaises(EnvironmentError,
125+
gcloud.datastore._require_dataset)
126+
NEW_DATASET = object()
127+
_implicit_environ.DATASET = NEW_DATASET
128+
self.assertEqual(gcloud.datastore._require_dataset(), NEW_DATASET)
129+
finally:
130+
_implicit_environ.DATASET = original_dataset
131+
132+
def test_get_entity(self):
133+
import gcloud.datastore
134+
from gcloud.datastore import _implicit_environ
135+
from gcloud.datastore.test_entity import _Dataset
136+
from gcloud._testing import _Monkey
137+
138+
CUSTOM_DATASET = _Dataset()
139+
DUMMY_KEY = object()
140+
DUMMY_VAL = object()
141+
CUSTOM_DATASET[DUMMY_KEY] = DUMMY_VAL
142+
with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
143+
result = gcloud.datastore.get_entity(DUMMY_KEY)
144+
self.assertTrue(result is DUMMY_VAL)
145+
146+
def test_get_entities(self):
147+
import gcloud.datastore
148+
from gcloud.datastore import _implicit_environ
149+
from gcloud.datastore.test_entity import _Dataset
150+
from gcloud._testing import _Monkey
151+
152+
CUSTOM_DATASET = _Dataset()
153+
DUMMY_KEYS = [object(), object()]
154+
DUMMY_VALS = [object(), object()]
155+
for key, val in zip(DUMMY_KEYS, DUMMY_VALS):
156+
CUSTOM_DATASET[key] = val
157+
158+
with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
159+
result = gcloud.datastore.get_entities(DUMMY_KEYS)
160+
self.assertTrue(result == DUMMY_VALS)
161+
162+
def test_allocate_ids(self):
163+
import gcloud.datastore
164+
from gcloud.datastore import _implicit_environ
165+
from gcloud.datastore.key import Key
166+
from gcloud.datastore.test_entity import _Dataset
167+
from gcloud._testing import _Monkey
168+
169+
CUSTOM_DATASET = _Dataset()
170+
INCOMPLETE_KEY = Key()
171+
NUM_IDS = 2
172+
with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
173+
result = gcloud.datastore.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
174+
175+
# Check the IDs returned.
176+
self.assertEqual([key.id() for key in result], range(1, NUM_IDS + 1))

‎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: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,15 @@ class Test_entity_from_protobuf(unittest2.TestCase):
1919

2020
_MARKER = object()
2121

22+
def setUp(self):
23+
from gcloud.datastore import _implicit_environ
24+
self._replaced_dataset = _implicit_environ.DATASET
25+
_implicit_environ.DATASET = None
26+
27+
def tearDown(self):
28+
from gcloud.datastore import _implicit_environ
29+
_implicit_environ.DATASET = self._replaced_dataset
30+
2231
def _callFUT(self, val, dataset=_MARKER):
2332
from gcloud.datastore.helpers import entity_from_protobuf
2433

‎gcloud/datastore/test_query.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,17 @@
1717

1818
class TestQuery(unittest2.TestCase):
1919

20+
def setUp(self):
21+
from gcloud.datastore import _implicit_environ
22+
self._replaced_dataset = _implicit_environ.DATASET
23+
_implicit_environ.DATASET = None
24+
25+
def tearDown(self):
26+
from gcloud.datastore import _implicit_environ
27+
_implicit_environ.DATASET = self._replaced_dataset
28+
2029
def _getTargetClass(self):
2130
from gcloud.datastore.query import Query
22-
2331
return Query
2432

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

‎gcloud/datastore/test_transaction.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,15 @@ def test_ctor(self):
3838
self.assertEqual(len(xact._auto_id_entities), 0)
3939
self.assertTrue(xact.connection() is connection)
4040

41+
def test_ctor_with_env(self):
42+
SENTINEL_VAL = object()
43+
44+
from gcloud.datastore import _implicit_environ
45+
_implicit_environ.DATASET = SENTINEL_VAL
46+
47+
transaction = self._makeOne(dataset=None)
48+
self.assertEqual(transaction.dataset(), SENTINEL_VAL)
49+
4150
def test_add_auto_id_entity(self):
4251
entity = _Entity()
4352
_DATASET = 'DATASET'

0 commit comments

Comments
 (0)