Skip to content

Commit ced11fb

Browse files
dhermesdhermes
committed
Merge pull request #482 from dhermes/fix-477-part5
Address fifth part of 477: Moves Dataset.allocate_ids to connection module.
2 parents d171d44 + 24098aa commit ced11fb

6 files changed

Lines changed: 121 additions & 76 deletions

File tree

‎gcloud/datastore/__init__.py‎

Lines changed: 35 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,18 @@ def _require_dataset():
146146
return _implicit_environ.DATASET
147147

148148

149+
def _require_connection():
150+
"""Convenience method to ensure CONNECTION is set.
151+
152+
:rtype: :class:`gcloud.datastore.connection.Connection`
153+
:returns: A connection based on the current environment.
154+
:raises: :class:`EnvironmentError` if CONNECTION is not set.
155+
"""
156+
if _implicit_environ.CONNECTION is None:
157+
raise EnvironmentError('Connection could not be inferred.')
158+
return _implicit_environ.CONNECTION
159+
160+
149161
def get_entities(keys):
150162
"""Retrieves entities from implied dataset, along with their attributes.
151163
@@ -158,16 +170,36 @@ def get_entities(keys):
158170
return _require_dataset().get_entities(keys)
159171

160172

161-
def allocate_ids(incomplete_key, num_ids):
173+
def allocate_ids(incomplete_key, num_ids, connection=None, dataset_id=None):
162174
"""Allocates a list of IDs from a partial key.
163175
164176
:type incomplete_key: A :class:`gcloud.datastore.key.Key`
165-
:param incomplete_key: The partial key to use as base for allocated IDs.
177+
:param incomplete_key: Partial key to use as base for allocated IDs.
166178
167179
:type num_ids: A :class:`int`.
168180
:param num_ids: The number of IDs to allocate.
169181
182+
:type connection: :class:`gcloud.datastore.connection.Connection`
183+
:param connection: Optional. The connection used to allocate IDs.
184+
185+
:type dataset_id: :class:`str`.
186+
:param dataset_id: Optional. The ID of the dataset used to allocate.
187+
170188
:rtype: list of :class:`gcloud.datastore.key.Key`
171189
:returns: The (complete) keys allocated with `incomplete_key` as root.
190+
:raises: `ValueError` if `incomplete_key` is not a partial key.
172191
"""
173-
return _require_dataset().allocate_ids(incomplete_key, num_ids)
192+
connection = connection or _require_connection()
193+
dataset_id = dataset_id or _require_dataset().id()
194+
195+
if not incomplete_key.is_partial:
196+
raise ValueError(('Key is not partial.', incomplete_key))
197+
198+
incomplete_key_pb = incomplete_key.to_protobuf()
199+
incomplete_key_pbs = [incomplete_key_pb] * num_ids
200+
201+
allocated_key_pbs = connection.allocate_ids(dataset_id, incomplete_key_pbs)
202+
allocated_ids = [allocated_key_pb.path_element[-1].id
203+
for allocated_key_pb in allocated_key_pbs]
204+
return [incomplete_key.completed_key(allocated_id)
205+
for allocated_id in allocated_ids]

‎gcloud/datastore/connection.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,13 @@
1414

1515
"""Connections to gcloud datastore API servers."""
1616

17-
from gcloud import connection
17+
from gcloud import connection as base_connection
1818
from gcloud.datastore import datastore_v1_pb2 as datastore_pb
1919
from gcloud.datastore import helpers
2020
from gcloud.datastore.dataset import Dataset
2121

2222

23-
class Connection(connection.Connection):
23+
class Connection(base_connection.Connection):
2424
"""A connection to the Google Cloud Datastore via the Protobuf API.
2525
2626
This class should understand only the basic types (and protobufs)
@@ -125,7 +125,7 @@ def build_api_url(cls, dataset_id, method, base_url=None,
125125
api_version=(api_version or cls.API_VERSION),
126126
dataset_id=dataset_id, method=method)
127127

128-
def transaction(self, transaction=connection.Connection._EMPTY):
128+
def transaction(self, transaction=base_connection.Connection._EMPTY):
129129
"""Getter/setter for the connection's transaction object.
130130
131131
:type transaction: :class:`gcloud.datastore.transaction.Transaction`,

‎gcloud/datastore/dataset.py‎

Lines changed: 0 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -112,29 +112,3 @@ def get_entities(self, keys, missing=None, deferred=None):
112112
entities.append(helpers.entity_from_protobuf(
113113
entity_pb, dataset=self))
114114
return entities
115-
116-
def allocate_ids(self, incomplete_key, num_ids):
117-
"""Allocates a list of IDs from a partial key.
118-
119-
:type incomplete_key: A :class:`gcloud.datastore.key.Key`
120-
:param incomplete_key: Partial key to use as base for allocated IDs.
121-
122-
:type num_ids: A :class:`int`.
123-
:param num_ids: The number of IDs to allocate.
124-
125-
:rtype: list of :class:`gcloud.datastore.key.Key`
126-
:returns: The (complete) keys allocated with `incomplete_key` as root.
127-
:raises: `ValueError` if `incomplete_key` is not a partial key.
128-
"""
129-
if not incomplete_key.is_partial:
130-
raise ValueError(('Key is not partial.', incomplete_key))
131-
132-
incomplete_key_pb = incomplete_key.to_protobuf()
133-
incomplete_key_pbs = [incomplete_key_pb] * num_ids
134-
135-
allocated_key_pbs = self.connection().allocate_ids(
136-
self.id(), incomplete_key_pbs)
137-
allocated_ids = [allocated_key_pb.path_element[-1].id
138-
for allocated_key_pb in allocated_key_pbs]
139-
return [incomplete_key.completed_key(allocated_id)
140-
for allocated_id in allocated_ids]

‎gcloud/datastore/test___init__.py‎

Lines changed: 83 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -151,20 +151,43 @@ def test_it(self):
151151

152152
class Test_implicit_behavior(unittest2.TestCase):
153153

154-
def test__require_dataset(self):
154+
def test__require_dataset_value_unset(self):
155155
import gcloud.datastore
156156
from gcloud.datastore import _implicit_environ
157-
original_dataset = _implicit_environ.DATASET
158-
159-
try:
160-
_implicit_environ.DATASET = None
161-
self.assertRaises(EnvironmentError,
162-
gcloud.datastore._require_dataset)
163-
NEW_DATASET = object()
164-
_implicit_environ.DATASET = NEW_DATASET
165-
self.assertEqual(gcloud.datastore._require_dataset(), NEW_DATASET)
166-
finally:
167-
_implicit_environ.DATASET = original_dataset
157+
from gcloud._testing import _Monkey
158+
159+
with _Monkey(_implicit_environ, DATASET=None):
160+
with self.assertRaises(EnvironmentError):
161+
gcloud.datastore._require_dataset()
162+
163+
def test__require_dataset_value_set(self):
164+
import gcloud.datastore
165+
from gcloud.datastore import _implicit_environ
166+
from gcloud._testing import _Monkey
167+
168+
FAKE_DATASET = object()
169+
with _Monkey(_implicit_environ, DATASET=FAKE_DATASET):
170+
stored_dataset = gcloud.datastore._require_dataset()
171+
self.assertTrue(stored_dataset is FAKE_DATASET)
172+
173+
def test__require_connection_value_unset(self):
174+
import gcloud.datastore
175+
from gcloud.datastore import _implicit_environ
176+
from gcloud._testing import _Monkey
177+
178+
with _Monkey(_implicit_environ, CONNECTION=None):
179+
with self.assertRaises(EnvironmentError):
180+
gcloud.datastore._require_connection()
181+
182+
def test__require_connection_value_set(self):
183+
import gcloud.datastore
184+
from gcloud.datastore import _implicit_environ
185+
from gcloud._testing import _Monkey
186+
187+
FAKE_CONNECTION = object()
188+
with _Monkey(_implicit_environ, CONNECTION=FAKE_CONNECTION):
189+
stored_connection = gcloud.datastore._require_connection()
190+
self.assertTrue(stored_connection is FAKE_CONNECTION)
168191

169192
def test_get_entities(self):
170193
import gcloud.datastore
@@ -182,18 +205,62 @@ def test_get_entities(self):
182205
result = gcloud.datastore.get_entities(DUMMY_KEYS)
183206
self.assertTrue(result == DUMMY_VALS)
184207

208+
209+
class Test_allocate_ids_function(unittest2.TestCase):
210+
211+
def _callFUT(self, incomplete_key, num_ids,
212+
connection=None, dataset_id=None):
213+
from gcloud.datastore import allocate_ids
214+
return allocate_ids(incomplete_key, num_ids, connection=connection,
215+
dataset_id=dataset_id)
216+
185217
def test_allocate_ids(self):
186-
import gcloud.datastore
218+
from gcloud.datastore.key import Key
219+
from gcloud.datastore.test_dataset import _Connection
220+
221+
DATASET_ID = 'DATASET'
222+
INCOMPLETE_KEY = Key('KIND', dataset_id=DATASET_ID)
223+
CONNECTION = _Connection()
224+
NUM_IDS = 2
225+
result = self._callFUT(INCOMPLETE_KEY, NUM_IDS,
226+
connection=CONNECTION, dataset_id=DATASET_ID)
227+
228+
# Check the IDs returned match.
229+
self.assertEqual([key.id for key in result], range(NUM_IDS))
230+
231+
# Check connection is called correctly.
232+
self.assertEqual(CONNECTION._called_dataset_id, DATASET_ID)
233+
self.assertEqual(len(CONNECTION._called_key_pbs), NUM_IDS)
234+
235+
def test_allocate_ids_implicit(self):
187236
from gcloud.datastore import _implicit_environ
188237
from gcloud.datastore.key import Key
238+
from gcloud.datastore.test_dataset import _Connection
189239
from gcloud.datastore.test_entity import _Dataset
190240
from gcloud._testing import _Monkey
191241

192242
CUSTOM_DATASET = _Dataset()
243+
CUSTOM_CONNECTION = _Connection()
193244
NUM_IDS = 2
194-
with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET):
245+
with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET,
246+
CONNECTION=CUSTOM_CONNECTION):
195247
INCOMPLETE_KEY = Key('KIND')
196-
result = gcloud.datastore.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
248+
result = self._callFUT(INCOMPLETE_KEY, NUM_IDS)
197249

198250
# Check the IDs returned.
199-
self.assertEqual([key.id for key in result], range(1, NUM_IDS + 1))
251+
self.assertEqual([key.id for key in result], range(NUM_IDS))
252+
253+
def test_allocate_ids_with_complete(self):
254+
from gcloud.datastore import _implicit_environ
255+
from gcloud.datastore.key import Key
256+
from gcloud.datastore.test_dataset import _Connection
257+
from gcloud.datastore.test_entity import _Dataset
258+
from gcloud._testing import _Monkey
259+
260+
CUSTOM_DATASET = _Dataset()
261+
CUSTOM_CONNECTION = _Connection()
262+
with _Monkey(_implicit_environ, DATASET=CUSTOM_DATASET,
263+
CONNECTION=CUSTOM_CONNECTION):
264+
COMPLETE_KEY = Key('KIND', 1234)
265+
self.assertRaises(ValueError, self._callFUT,
266+
COMPLETE_KEY, 2)

‎gcloud/datastore/test_dataset.py‎

Lines changed: 0 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -108,31 +108,6 @@ def test_get_entities_hit(self):
108108
self.assertEqual(list(result), ['foo'])
109109
self.assertEqual(result['foo'], 'Foo')
110110

111-
def test_allocate_ids(self):
112-
from gcloud.datastore.key import Key
113-
114-
DATASET_ID = 'DATASET'
115-
INCOMPLETE_KEY = Key('KIND', dataset_id=DATASET_ID)
116-
CONNECTION = _Connection()
117-
NUM_IDS = 2
118-
DATASET = self._makeOne(DATASET_ID, connection=CONNECTION)
119-
result = DATASET.allocate_ids(INCOMPLETE_KEY, NUM_IDS)
120-
121-
# Check the IDs returned match.
122-
self.assertEqual([key.id for key in result], range(NUM_IDS))
123-
124-
# Check connection is called correctly.
125-
self.assertEqual(CONNECTION._called_dataset_id, DATASET_ID)
126-
self.assertEqual(len(CONNECTION._called_key_pbs), NUM_IDS)
127-
128-
def test_allocate_ids_with_complete(self):
129-
from gcloud.datastore.test_entity import _Key
130-
131-
COMPLETE_KEY = _Key()
132-
DATASET = self._makeOne(None)
133-
self.assertRaises(ValueError, DATASET.allocate_ids,
134-
COMPLETE_KEY, 2)
135-
136111

137112
class _Connection(object):
138113
_called_with = None

‎gcloud/datastore/test_entity.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -274,9 +274,6 @@ def connection(self):
274274
def get_entities(self, keys):
275275
return [self.get(key) for key in keys]
276276

277-
def allocate_ids(self, incomplete_key, num_ids):
278-
return [incomplete_key.completed_key(i + 1) for i in range(num_ids)]
279-
280277

281278
class _Connection(object):
282279
_transaction = _saved = _deleted = None

0 commit comments

Comments
 (0)