Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 17 additions & 12 deletions gcloud/datastore/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,8 +153,7 @@ def _require_connection(connection=None):
return connection


def get_entities(keys, missing=None, deferred=None,
connection=None, dataset_id=None):
def get_entities(keys, missing=None, deferred=None, connection=None):
"""Retrieves entities, along with their attributes.

:type keys: list of :class:`gcloud.datastore.key.Key`
Expand All @@ -173,14 +172,23 @@ def get_entities(keys, missing=None, deferred=None,
:type connection: :class:`gcloud.datastore.connection.Connection`
:param connection: Optional. The connection used to connect to datastore.

:type dataset_id: string
:param dataset_id: Optional. The ID of the dataset.

:rtype: list of :class:`gcloud.datastore.entity.Entity`
:returns: The requested entities.
:raises: :class:`ValueError` if the key dataset IDs don't agree.
"""
if not keys:
return []

connection = _require_connection(connection)
dataset_id = _require_dataset_id(dataset_id)
dataset_id = keys[0].dataset_id
# Rather than creating a list or set of all dataset IDs, we iterate
# and check. We could allow the backend to check this for us if IDs
# with no prefix worked (GoogleCloudPlatform/google-cloud-datastore#59)
# or if we made sure that a prefix s~ or e~ was on each key.
for key in keys[1:]:
if key.dataset_id != dataset_id:
raise ValueError('All keys in get_entities must be from the '
'same dataset.')

entity_pbs = connection.lookup(
dataset_id=dataset_id,
Expand All @@ -205,7 +213,7 @@ def get_entities(keys, missing=None, deferred=None,
return entities


def allocate_ids(incomplete_key, num_ids, connection=None, dataset_id=None):
def allocate_ids(incomplete_key, num_ids, connection=None):
"""Allocates a list of IDs from a partial key.

:type incomplete_key: A :class:`gcloud.datastore.key.Key`
Expand All @@ -217,23 +225,20 @@ def allocate_ids(incomplete_key, num_ids, connection=None, dataset_id=None):
:type connection: :class:`gcloud.datastore.connection.Connection`
:param connection: Optional. The connection used to connect to datastore.

:type dataset_id: string
:param dataset_id: Optional. The ID of the dataset.

:rtype: list of :class:`gcloud.datastore.key.Key`
:returns: The (complete) keys allocated with ``incomplete_key`` as root.
:raises: :class:`ValueError` if ``incomplete_key`` is not a partial key.
"""
connection = _require_connection(connection)
dataset_id = _require_dataset_id(dataset_id)

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 = connection.allocate_ids(dataset_id, incomplete_key_pbs)
allocated_key_pbs = connection.allocate_ids(incomplete_key.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.completed_key(allocated_id)
Expand Down
3 changes: 1 addition & 2 deletions gcloud/datastore/key.py
Original file line number Diff line number Diff line change
Expand Up @@ -232,8 +232,7 @@ def get(self, connection=None):

# We allow partial keys to attempt a get, the backend will fail.
connection = connection or _implicit_environ.CONNECTION
entities = datastore.get_entities(
[self], connection=connection, dataset_id=self.dataset_id)
entities = datastore.get_entities([self], connection=connection)

if entities:
result = entities[0]
Expand Down
80 changes: 60 additions & 20 deletions gcloud/datastore/test___init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -213,11 +213,14 @@ def test__require_connection_implicit_set_passed_explicitly(self):

class Test_get_entities_function(unittest2.TestCase):

def _callFUT(self, keys, missing=None, deferred=None,
connection=None, dataset_id=None):
def _callFUT(self, keys, missing=None, deferred=None, connection=None):
from gcloud.datastore import get_entities
return get_entities(keys, missing=missing, deferred=deferred,
connection=connection, dataset_id=dataset_id)
connection=connection)

def test_get_entities_no_keys(self):
results = self._callFUT([])
self.assertEqual(results, [])

def test_get_entities_miss(self):
from gcloud.datastore.key import Key
Expand All @@ -226,8 +229,7 @@ def test_get_entities_miss(self):
DATASET_ID = 'DATASET'
connection = _Connection()
key = Key('Kind', 1234, dataset_id=DATASET_ID)
results = self._callFUT([key], connection=connection,
dataset_id=DATASET_ID)
results = self._callFUT([key], connection=connection)
self.assertEqual(results, [])

def test_get_entities_miss_w_missing(self):
Expand All @@ -252,8 +254,7 @@ def test_get_entities_miss_w_missing(self):

key = Key(KIND, ID, dataset_id=DATASET_ID)
missing = []
entities = self._callFUT([key], connection=connection,
dataset_id=DATASET_ID, missing=missing)
entities = self._callFUT([key], connection=connection, missing=missing)
self.assertEqual(entities, [])
self.assertEqual([missed.key.to_protobuf() for missed in missing],
[key.to_protobuf()])
Expand All @@ -271,22 +272,24 @@ def test_get_entities_miss_w_deferred(self):

deferred = []
entities = self._callFUT([key], connection=connection,
dataset_id=DATASET_ID, deferred=deferred)
deferred=deferred)
self.assertEqual(entities, [])
self.assertEqual([def_key.to_protobuf() for def_key in deferred],
[key.to_protobuf()])

def _make_entity_pb(self, dataset_id, kind, integer_id, name, str_val):
def _make_entity_pb(self, dataset_id, kind, integer_id,
name=None, str_val=None):
from gcloud.datastore import datastore_v1_pb2 as datastore_pb

entity_pb = datastore_pb.Entity()
entity_pb.key.partition_id.dataset_id = dataset_id
path_element = entity_pb.key.path_element.add()
path_element.kind = kind
path_element.id = integer_id
prop = entity_pb.property.add()
prop.name = name
prop.value.string_value = str_val
if name is not None and str_val is not None:
prop = entity_pb.property.add()
prop.name = name
prop.value.string_value = str_val

return entity_pb

Expand All @@ -307,8 +310,7 @@ def test_get_entities_hit(self):
connection = _Connection(entity_pb)

key = Key(KIND, ID, dataset_id=DATASET_ID)
result, = self._callFUT([key], connection=connection,
dataset_id=DATASET_ID)
result, = self._callFUT([key], connection=connection)
new_key = result.key

# Check the returned value is as expected.
Expand All @@ -318,6 +320,47 @@ def test_get_entities_hit(self):
self.assertEqual(list(result), ['foo'])
self.assertEqual(result['foo'], 'Foo')

def test_get_entities_hit_multiple_keys_same_dataset(self):
from gcloud.datastore.key import Key
from gcloud.datastore.test_connection import _Connection

DATASET_ID = 'DATASET'
KIND = 'Kind'
ID1 = 1234
ID2 = 2345

# Make a found entity pb to be returned from mock backend.
entity_pb1 = self._make_entity_pb(DATASET_ID, KIND, ID1)
entity_pb2 = self._make_entity_pb(DATASET_ID, KIND, ID2)

# Make a connection to return the entity pbs.
connection = _Connection(entity_pb1, entity_pb2)

key1 = Key(KIND, ID1, dataset_id=DATASET_ID)
key2 = Key(KIND, ID2, dataset_id=DATASET_ID)
retrieved1, retrieved2 = self._callFUT(
[key1, key2], connection=connection)

# Check values match.
self.assertEqual(retrieved1.key.path, key1.path)
self.assertEqual(dict(retrieved1), {})
self.assertEqual(retrieved2.key.path, key2.path)
self.assertEqual(dict(retrieved2), {})

def test_get_entities_hit_multiple_keys_different_dataset(self):
from gcloud.datastore.key import Key

DATASET_ID1 = 'DATASET'
DATASET_ID2 = 'DATASET-ALT'

# Make sure our IDs are actually different.
self.assertNotEqual(DATASET_ID1, DATASET_ID2)

key1 = Key('KIND', 1234, dataset_id=DATASET_ID1)
key2 = Key('KIND', 1234, dataset_id=DATASET_ID2)
with self.assertRaises(ValueError):
self._callFUT([key1, key2], connection=object())

def test_get_entities_implicit(self):
from gcloud.datastore import _implicit_environ
from gcloud.datastore.key import Key
Expand Down Expand Up @@ -358,11 +401,9 @@ def test_get_entities_implicit(self):

class Test_allocate_ids_function(unittest2.TestCase):

def _callFUT(self, incomplete_key, num_ids,
connection=None, dataset_id=None):
def _callFUT(self, incomplete_key, num_ids, connection=None):
from gcloud.datastore import allocate_ids
return allocate_ids(incomplete_key, num_ids, connection=connection,
dataset_id=dataset_id)
return allocate_ids(incomplete_key, num_ids, connection=connection)

def test_allocate_ids(self):
from gcloud.datastore.key import Key
Expand All @@ -372,8 +413,7 @@ def test_allocate_ids(self):
INCOMPLETE_KEY = Key('KIND', dataset_id=DATASET_ID)
CONNECTION = _Connection()
NUM_IDS = 2
result = self._callFUT(INCOMPLETE_KEY, NUM_IDS,
connection=CONNECTION, dataset_id=DATASET_ID)
result = self._callFUT(INCOMPLETE_KEY, NUM_IDS, connection=CONNECTION)

# Check the IDs returned match.
self.assertEqual([key.id for key in result], range(NUM_IDS))
Expand Down