401 lines
13 KiB
Python
401 lines
13 KiB
Python
|
# Copyright 2014 Google Inc. All rights reserved.
|
||
|
|
#
|
||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
# you may not use this file except in compliance with the License.
|
||
|
|
# You may obtain a copy of the License at
|
||
|
|
#
|
||
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
#
|
||
|
|
# Unless required by applicable law or agreed to in writing, software
|
||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
# See the License for the specific language governing permissions and
|
||
|
|
# limitations under the License.
|
||
|
|
|
||
|
|
import unittest2
|
||
|
|
|
||
|
|
|
||
|
|
class TestBatch(unittest2.TestCase):
|
||
|
|
|
||
|
|
def _getTargetClass(self):
|
||
|
|
from gcloud.datastore.batch import Batch
|
||
|
|
|
||
|
|
return Batch
|
||
|
|
|
||
|
|
def _makeOne(self, client):
|
||
|
|
return self._getTargetClass()(client)
|
||
|
|
|
||
|
|
def test_ctor(self):
|
||
|
|
from gcloud.datastore._generated import datastore_pb2
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
_NAMESPACE = 'NAMESPACE'
|
||
|
|
connection = _Connection()
|
||
|
|
client = _Client(_PROJECT, connection, _NAMESPACE)
|
||
|
|
batch = self._makeOne(client)
|
||
|
|
|
||
|
|
self.assertEqual(batch.project, _PROJECT)
|
||
|
|
self.assertEqual(batch.connection, connection)
|
||
|
|
self.assertEqual(batch.namespace, _NAMESPACE)
|
||
|
|
self.assertTrue(batch._id is None)
|
||
|
|
self.assertEqual(batch._status, batch._INITIAL)
|
||
|
|
self.assertTrue(isinstance(batch._commit_request,
|
||
|
|
datastore_pb2.CommitRequest))
|
||
|
|
self.assertTrue(batch.mutations is batch._commit_request.mutations)
|
||
|
|
self.assertEqual(batch._partial_key_entities, [])
|
||
|
|
|
||
|
|
def test_current(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
connection = _Connection()
|
||
|
|
client = _Client(_PROJECT, connection)
|
||
|
|
batch1 = self._makeOne(client)
|
||
|
|
batch2 = self._makeOne(client)
|
||
|
|
self.assertTrue(batch1.current() is None)
|
||
|
|
self.assertTrue(batch2.current() is None)
|
||
|
|
with batch1:
|
||
|
|
self.assertTrue(batch1.current() is batch1)
|
||
|
|
self.assertTrue(batch2.current() is batch1)
|
||
|
|
with batch2:
|
||
|
|
self.assertTrue(batch1.current() is batch2)
|
||
|
|
self.assertTrue(batch2.current() is batch2)
|
||
|
|
self.assertTrue(batch1.current() is batch1)
|
||
|
|
self.assertTrue(batch2.current() is batch1)
|
||
|
|
self.assertTrue(batch1.current() is None)
|
||
|
|
self.assertTrue(batch2.current() is None)
|
||
|
|
|
||
|
|
def test_put_entity_wo_key(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
connection = _Connection()
|
||
|
|
client = _Client(_PROJECT, connection)
|
||
|
|
batch = self._makeOne(client)
|
||
|
|
|
||
|
|
self.assertRaises(ValueError, batch.put, _Entity())
|
||
|
|
|
||
|
|
def test_put_entity_w_key_wrong_project(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
connection = _Connection()
|
||
|
|
client = _Client(_PROJECT, connection)
|
||
|
|
batch = self._makeOne(client)
|
||
|
|
entity = _Entity()
|
||
|
|
entity.key = _Key('OTHER')
|
||
|
|
|
||
|
|
self.assertRaises(ValueError, batch.put, entity)
|
||
|
|
|
||
|
|
def test_put_entity_w_partial_key(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
_PROPERTIES = {'foo': 'bar'}
|
||
|
|
connection = _Connection()
|
||
|
|
client = _Client(_PROJECT, connection)
|
||
|
|
batch = self._makeOne(client)
|
||
|
|
entity = _Entity(_PROPERTIES)
|
||
|
|
key = entity.key = _Key(_PROJECT)
|
||
|
|
key._id = None
|
||
|
|
|
||
|
|
batch.put(entity)
|
||
|
|
|
||
|
|
mutated_entity = _mutated_pb(self, batch.mutations, 'insert')
|
||
|
|
self.assertEqual(mutated_entity.key, key._key)
|
||
|
|
self.assertEqual(batch._partial_key_entities, [entity])
|
||
|
|
|
||
|
|
def test_put_entity_w_completed_key(self):
|
||
|
|
from gcloud.datastore.helpers import _property_tuples
|
||
|
|
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
_PROPERTIES = {
|
||
|
|
'foo': 'bar',
|
||
|
|
'baz': 'qux',
|
||
|
|
'spam': [1, 2, 3],
|
||
|
|
'frotz': [], # will be ignored
|
||
|
|
}
|
||
|
|
connection = _Connection()
|
||
|
|
client = _Client(_PROJECT, connection)
|
||
|
|
batch = self._makeOne(client)
|
||
|
|
entity = _Entity(_PROPERTIES)
|
||
|
|
entity.exclude_from_indexes = ('baz', 'spam')
|
||
|
|
key = entity.key = _Key(_PROJECT)
|
||
|
|
|
||
|
|
batch.put(entity)
|
||
|
|
|
||
|
|
mutated_entity = _mutated_pb(self, batch.mutations, 'upsert')
|
||
|
|
self.assertEqual(mutated_entity.key, key._key)
|
||
|
|
|
||
|
|
prop_dict = dict(_property_tuples(mutated_entity))
|
||
|
|
self.assertEqual(len(prop_dict), 3)
|
||
|
|
self.assertFalse(prop_dict['foo'].exclude_from_indexes)
|
||
|
|
self.assertTrue(prop_dict['baz'].exclude_from_indexes)
|
||
|
|
self.assertFalse(prop_dict['spam'].exclude_from_indexes)
|
||
|
|
spam_values = prop_dict['spam'].array_value.values
|
||
|
|
self.assertTrue(spam_values[0].exclude_from_indexes)
|
||
|
|
self.assertTrue(spam_values[1].exclude_from_indexes)
|
||
|
|
self.assertTrue(spam_values[2].exclude_from_indexes)
|
||
|
|
self.assertFalse('frotz' in prop_dict)
|
||
|
|
|
||
|
|
def test_delete_w_partial_key(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
connection = _Connection()
|
||
|
|
client = _Client(_PROJECT, connection)
|
||
|
|
batch = self._makeOne(client)
|
||
|
|
key = _Key(_PROJECT)
|
||
|
|
key._id = None
|
||
|
|
|
||
|
|
self.assertRaises(ValueError, batch.delete, key)
|
||
|
|
|
||
|
|
def test_delete_w_key_wrong_project(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
connection = _Connection()
|
||
|
|
client = _Client(_PROJECT, connection)
|
||
|
|
batch = self._makeOne(client)
|
||
|
|
key = _Key('OTHER')
|
||
|
|
|
||
|
|
self.assertRaises(ValueError, batch.delete, key)
|
||
|
|
|
||
|
|
def test_delete_w_completed_key(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
connection = _Connection()
|
||
|
|
client = _Client(_PROJECT, connection)
|
||
|
|
batch = self._makeOne(client)
|
||
|
|
key = _Key(_PROJECT)
|
||
|
|
|
||
|
|
batch.delete(key)
|
||
|
|
|
||
|
|
mutated_key = _mutated_pb(self, batch.mutations, 'delete')
|
||
|
|
self.assertEqual(mutated_key, key._key)
|
||
|
|
|
||
|
|
def test_begin(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
client = _Client(_PROJECT, None)
|
||
|
|
batch = self._makeOne(client)
|
||
|
|
self.assertEqual(batch._status, batch._INITIAL)
|
||
|
|
batch.begin()
|
||
|
|
self.assertEqual(batch._status, batch._IN_PROGRESS)
|
||
|
|
|
||
|
|
def test_begin_fail(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
client = _Client(_PROJECT, None)
|
||
|
|
batch = self._makeOne(client)
|
||
|
|
batch._status = batch._IN_PROGRESS
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
batch.begin()
|
||
|
|
|
||
|
|
def test_rollback(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
client = _Client(_PROJECT, None)
|
||
|
|
batch = self._makeOne(client)
|
||
|
|
self.assertEqual(batch._status, batch._INITIAL)
|
||
|
|
batch.rollback()
|
||
|
|
self.assertEqual(batch._status, batch._ABORTED)
|
||
|
|
|
||
|
|
def test_commit(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
connection = _Connection()
|
||
|
|
client = _Client(_PROJECT, connection)
|
||
|
|
batch = self._makeOne(client)
|
||
|
|
|
||
|
|
self.assertEqual(batch._status, batch._INITIAL)
|
||
|
|
batch.commit()
|
||
|
|
self.assertEqual(batch._status, batch._FINISHED)
|
||
|
|
|
||
|
|
self.assertEqual(connection._committed,
|
||
|
|
[(_PROJECT, batch._commit_request, None)])
|
||
|
|
|
||
|
|
def test_commit_w_partial_key_entities(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
_NEW_ID = 1234
|
||
|
|
connection = _Connection(_NEW_ID)
|
||
|
|
client = _Client(_PROJECT, connection)
|
||
|
|
batch = self._makeOne(client)
|
||
|
|
entity = _Entity({})
|
||
|
|
key = entity.key = _Key(_PROJECT)
|
||
|
|
key._id = None
|
||
|
|
batch._partial_key_entities.append(entity)
|
||
|
|
|
||
|
|
self.assertEqual(batch._status, batch._INITIAL)
|
||
|
|
batch.commit()
|
||
|
|
self.assertEqual(batch._status, batch._FINISHED)
|
||
|
|
|
||
|
|
self.assertEqual(connection._committed,
|
||
|
|
[(_PROJECT, batch._commit_request, None)])
|
||
|
|
self.assertFalse(entity.key.is_partial)
|
||
|
|
self.assertEqual(entity.key._id, _NEW_ID)
|
||
|
|
|
||
|
|
def test_as_context_mgr_wo_error(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
_PROPERTIES = {'foo': 'bar'}
|
||
|
|
connection = _Connection()
|
||
|
|
entity = _Entity(_PROPERTIES)
|
||
|
|
key = entity.key = _Key(_PROJECT)
|
||
|
|
|
||
|
|
client = _Client(_PROJECT, connection)
|
||
|
|
self.assertEqual(list(client._batches), [])
|
||
|
|
|
||
|
|
with self._makeOne(client) as batch:
|
||
|
|
self.assertEqual(list(client._batches), [batch])
|
||
|
|
batch.put(entity)
|
||
|
|
|
||
|
|
self.assertEqual(list(client._batches), [])
|
||
|
|
|
||
|
|
mutated_entity = _mutated_pb(self, batch.mutations, 'upsert')
|
||
|
|
self.assertEqual(mutated_entity.key, key._key)
|
||
|
|
self.assertEqual(connection._committed,
|
||
|
|
[(_PROJECT, batch._commit_request, None)])
|
||
|
|
|
||
|
|
def test_as_context_mgr_nested(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
_PROPERTIES = {'foo': 'bar'}
|
||
|
|
connection = _Connection()
|
||
|
|
entity1 = _Entity(_PROPERTIES)
|
||
|
|
key1 = entity1.key = _Key(_PROJECT)
|
||
|
|
entity2 = _Entity(_PROPERTIES)
|
||
|
|
key2 = entity2.key = _Key(_PROJECT)
|
||
|
|
|
||
|
|
client = _Client(_PROJECT, connection)
|
||
|
|
self.assertEqual(list(client._batches), [])
|
||
|
|
|
||
|
|
with self._makeOne(client) as batch1:
|
||
|
|
self.assertEqual(list(client._batches), [batch1])
|
||
|
|
batch1.put(entity1)
|
||
|
|
with self._makeOne(client) as batch2:
|
||
|
|
self.assertEqual(list(client._batches), [batch2, batch1])
|
||
|
|
batch2.put(entity2)
|
||
|
|
|
||
|
|
self.assertEqual(list(client._batches), [batch1])
|
||
|
|
|
||
|
|
self.assertEqual(list(client._batches), [])
|
||
|
|
|
||
|
|
mutated_entity1 = _mutated_pb(self, batch1.mutations, 'upsert')
|
||
|
|
self.assertEqual(mutated_entity1.key, key1._key)
|
||
|
|
|
||
|
|
mutated_entity2 = _mutated_pb(self, batch2.mutations, 'upsert')
|
||
|
|
self.assertEqual(mutated_entity2.key, key2._key)
|
||
|
|
|
||
|
|
self.assertEqual(connection._committed,
|
||
|
|
[(_PROJECT, batch2._commit_request, None),
|
||
|
|
(_PROJECT, batch1._commit_request, None)])
|
||
|
|
|
||
|
|
def test_as_context_mgr_w_error(self):
|
||
|
|
_PROJECT = 'PROJECT'
|
||
|
|
_PROPERTIES = {'foo': 'bar'}
|
||
|
|
connection = _Connection()
|
||
|
|
entity = _Entity(_PROPERTIES)
|
||
|
|
key = entity.key = _Key(_PROJECT)
|
||
|
|
|
||
|
|
client = _Client(_PROJECT, connection)
|
||
|
|
self.assertEqual(list(client._batches), [])
|
||
|
|
|
||
|
|
try:
|
||
|
|
with self._makeOne(client) as batch:
|
||
|
|
self.assertEqual(list(client._batches), [batch])
|
||
|
|
batch.put(entity)
|
||
|
|
raise ValueError("testing")
|
||
|
|
except ValueError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
self.assertEqual(list(client._batches), [])
|
||
|
|
|
||
|
|
mutated_entity = _mutated_pb(self, batch.mutations, 'upsert')
|
||
|
|
self.assertEqual(mutated_entity.key, key._key)
|
||
|
|
self.assertEqual(connection._committed, [])
|
||
|
|
|
||
|
|
|
||
|
|
class _PathElementPB(object):
|
||
|
|
|
||
|
|
def __init__(self, id_):
|
||
|
|
self.id = id_
|
||
|
|
|
||
|
|
|
||
|
|
class _KeyPB(object):
|
||
|
|
|
||
|
|
def __init__(self, id_):
|
||
|
|
self.path = [_PathElementPB(id_)]
|
||
|
|
|
||
|
|
|
||
|
|
class _Connection(object):
|
||
|
|
_marker = object()
|
||
|
|
_save_result = (False, None)
|
||
|
|
|
||
|
|
def __init__(self, *new_keys):
|
||
|
|
self._completed_keys = [_KeyPB(key) for key in new_keys]
|
||
|
|
self._committed = []
|
||
|
|
self._index_updates = 0
|
||
|
|
|
||
|
|
def commit(self, project, commit_request, transaction_id):
|
||
|
|
self._committed.append((project, commit_request, transaction_id))
|
||
|
|
return self._index_updates, self._completed_keys
|
||
|
|
|
||
|
|
|
||
|
|
class _Entity(dict):
|
||
|
|
key = None
|
||
|
|
exclude_from_indexes = ()
|
||
|
|
_meanings = {}
|
||
|
|
|
||
|
|
|
||
|
|
class _Key(object):
|
||
|
|
_MARKER = object()
|
||
|
|
_kind = 'KIND'
|
||
|
|
_key = 'KEY'
|
||
|
|
_path = None
|
||
|
|
_id = 1234
|
||
|
|
_stored = None
|
||
|
|
|
||
|
|
def __init__(self, project):
|
||
|
|
self.project = project
|
||
|
|
|
||
|
|
@property
|
||
|
|
def is_partial(self):
|
||
|
|
return self._id is None
|
||
|
|
|
||
|
|
def to_protobuf(self):
|
||
|
|
from gcloud.datastore._generated import entity_pb2
|
||
|
|
key = self._key = entity_pb2.Key()
|
||
|
|
# Don't assign it, because it will just get ripped out
|
||
|
|
# key.partition_id.project_id = self.project
|
||
|
|
|
||
|
|
element = key.path.add()
|
||
|
|
element.kind = self._kind
|
||
|
|
if self._id is not None:
|
||
|
|
element.id = self._id
|
||
|
|
|
||
|
|
return key
|
||
|
|
|
||
|
|
def completed_key(self, new_id):
|
||
|
|
assert self.is_partial
|
||
|
|
new_key = self.__class__(self.project)
|
||
|
|
new_key._id = new_id
|
||
|
|
return new_key
|
||
|
|
|
||
|
|
|
||
|
|
class _Client(object):
|
||
|
|
|
||
|
|
def __init__(self, project, connection, namespace=None):
|
||
|
|
self.project = project
|
||
|
|
self.connection = connection
|
||
|
|
self.namespace = namespace
|
||
|
|
self._batches = []
|
||
|
|
|
||
|
|
def _push_batch(self, batch):
|
||
|
|
self._batches.insert(0, batch)
|
||
|
|
|
||
|
|
def _pop_batch(self):
|
||
|
|
return self._batches.pop(0)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def current_batch(self):
|
||
|
|
if self._batches:
|
||
|
|
return self._batches[0]
|
||
|
|
|
||
|
|
|
||
|
|
def _assert_num_mutations(test_case, mutation_pb_list, num_mutations):
|
||
|
|
test_case.assertEqual(len(mutation_pb_list), num_mutations)
|
||
|
|
|
||
|
|
|
||
|
|
def _mutated_pb(test_case, mutation_pb_list, mutation_type):
|
||
|
|
# Make sure there is only one mutation.
|
||
|
|
_assert_num_mutations(test_case, mutation_pb_list, 1)
|
||
|
|
|
||
|
|
# We grab the only mutation.
|
||
|
|
mutated_pb = mutation_pb_list[0]
|
||
|
|
# Then check if it is the correct type.
|
||
|
|
test_case.assertEqual(mutated_pb.WhichOneof('operation'),
|
||
|
|
mutation_type)
|
||
|
|
|
||
|
|
return getattr(mutated_pb, mutation_type)
|