def test_model_pb_to_entity_pb_exclude_from_index_fields(self):
# type: () -> None
example_pb = example_pb2.ExampleDBModel()
example_pb.int32_key = 100
example_pb.string_key = 'string bar'
example_pb.bytes_key = b'foobarbytes'
example_pb.enum_key = 1 # type: ignore
# No exclude from index provided
entity_pb = model_pb_to_entity_pb(model_pb=example_pb)
entity_translated = datastore.helpers.entity_from_protobuf(entity_pb)
self.assertEqual(entity_translated.exclude_from_indexes, set([]))
entity_translated.key = self.client.key('ExampleModel',
'exclude_from_indexes_1')
self.client.put(entity_translated)
entity_translated_retrieved = self.client.get(entity_translated.key)
self.assertEqual(entity_translated, entity_translated_retrieved)
# Exclude from index provided for some fields
entity_pb = model_pb_to_entity_pb(
model_pb=example_pb, exclude_from_index=['int32_key', 'bytes_key'])
entity_translated = datastore.helpers.entity_from_protobuf(entity_pb)
self.assertEqual(entity_translated.exclude_from_indexes,
set(['int32_key', 'bytes_key']))
entity_translated.key = self.client.key('ExampleModel',
'exclude_from_indexes_2')
self.client.put(entity_translated)
entity_translated_retrieved = self.client.get(entity_translated.key)
self.assertEqual(entity_translated, entity_translated_retrieved)
def insert_db_model(fixture_path, model_name, primary_key):
# type: (str, str, str) -> bool
LOG.debug('Loading fixture from "%s"', fixture_path)
# 1. Load in JSON fixture
with open(fixture_path, 'r') as fp:
model_pb_json = fp.read()
# 2. Parse it in our custom Protobuf DB model type
module, model_class = get_module_and_class_for_model_name(model_name)
LOG.debug('Parsing JSON fixture as Protobuf message')
model_pb = json_format.Parse(model_pb_json, model_class())
# 3. Translate it into Entity PB
LOG.debug('Translating Protobuf PB to Entity PB')
entity_pb = model_pb_to_entity_pb(model_pb)
# 4. Store it in Datastore
if os.environ.get('DATASTORE_EMULATOR_HOST'):
client = datastore.Client(credentials=EmulatorCreds(),
_http=requests.Session())
else:
client = datastore.Client()
model_name = model_pb.DESCRIPTOR.name
key = client.key(model_name, primary_key)
key_pb = key.to_protobuf()
entity_pb.key.CopyFrom(key_pb) # pylint: disable=no-member
LOG.debug('Storing it in datastore under primary key "%s"', key)
entity_pb = entity = datastore.helpers.entity_from_protobuf(entity_pb)
client.put(entity)
return True
def test_model_pb_to_entity_pb_exclude_from_index_custom_extension_model_with_package(
self):
# type: () -> None
from tests.generated.models import example_with_options_pb2
model_pb = example_with_options_pb2.ExampleDBModelWithOptions1()
model_pb.string_key_one = 'one'
model_pb.string_key_two = 'two'
model_pb.string_key_three = 'three'
model_pb.string_key_four = 'four'
model_pb.int32_field_one = 111
model_pb.int32_field_two = 222
entity_pb = model_pb_to_entity_pb(model_pb=model_pb)
entity_translated = datastore.helpers.entity_from_protobuf(entity_pb)
self.assertEqual(
entity_translated.exclude_from_indexes,
set(['string_key_one', 'int32_field_two', 'string_key_three']))
entity_translated.key = self.client.key('ExampleModelWithOptions',
'exclude_from_index_1')
self.client.put(entity_translated)
entity_translated_retrieved = self.client.get(entity_translated.key)
self.assertEqual(entity_translated, entity_translated_retrieved)
def put_db_object(key):
# type: (str) -> Tuple[str, int, Dict[str, str]]
"""
Store arbitrary Protobuf object in Google Datastore.
NOTE: Request body needs to contain Protobuf model serialized as JSON.
"""
body = request.get_json()
# Fully qualified model name, e.g. "tests.generated.example_pb2.ExampleDBModel"
# NOTE: This module needs to be available in PYTHONPATH
model_name = body['model_name']
model_data = body['json_string']
_, model_class = get_module_and_class_for_model_name(model_name=model_name)
model_pb = json_format.Parse(model_data, model_class())
# 2. Convert it into entity object
entity_pb = model_pb_to_entity_pb(model_pb)
client = datastore.Client()
# Set PK on the object
key_pb = client.key(model_pb.DESCRIPTOR.name, key).to_protobuf()
entity_pb.key.CopyFrom(key_pb) # pylint: disable=no-member
# 3. Store it inside datastore
entity = datastore.helpers.entity_from_protobuf(entity_pb)
client.put(entity)
return '', 200, {}
def test_store_and_retrieve_populated_translated_object_from_datastore(
self):
# type: () -> None
"""
Test case which stores raw entity object in the datastore and verifies it matched the
same object which is stored using translated Protobuf definition.
"""
key_native = self.client.key('ExampleModel', 'native_entity_populated')
entity_native = datastore.Entity(key=key_native)
entity_native.update(EXAMPLE_DICT_POPULATED)
self.client.put(entity_native)
entity_native_retrieved = self.client.get(key_native)
self.assertTrue(entity_native_retrieved)
# Verify retrieved data matches the original input
self.assertEqual(entity_native_retrieved, EXAMPLE_DICT_POPULATED)
# Store custom Protobuf object in a datastore by translating it to Entity object
key_translated = self.client.key('ExampleModel',
'translated_entity_populated')
example_pb = EXAMPLE_PB_POPULATED
entity_pb_translated = model_pb_to_entity_pb(model_pb=example_pb)
# pylint: disable=no-member
entity_pb_translated.key.CopyFrom(key_translated.to_protobuf())
entity_translated = datastore.helpers.entity_from_protobuf(
entity_pb_translated)
self.client.put(entity_translated)
# Verify that the translated entity results in the same end result as using native
# entity object
entity_translated_retrieved = self.client.get(key_translated)
self.assertTrue(
entity_translated_retrieved.key != entity_native_retrieved.key)
# NOTE: key won't be the same so we clear it
entity_translated_retrieved.key = None
entity_native_retrieved.key = None
self.assertEqual(entity_translated_retrieved, entity_native_retrieved)
# If we translate retrieved entity back to the original Protobuf object definition, it
# should be the same as the original model (minus the key since the original model doesn't
# contain a key)
entity_pb_retrieved = datastore.helpers.entity_to_protobuf(
entity_translated_retrieved)
entity_pb_translated.ClearField('key')
self.assertEqual(entity_pb_translated, entity_pb_retrieved)
example_pb_retrieved = entity_pb_to_model_pb(
example_pb2.ExampleDBModel, entity_pb_retrieved)
self.assertEqual(example_pb_retrieved, example_pb)
def test_store_and_retrieve_default_values_and_translated_object_from_datastore(
self):
# type: () -> None
key_native = self.client.key('ExampleModel',
'native_entity_default_values')
entity_native = datastore.Entity(key=key_native)
entity_native.update(EXAMPLE_DICT_DEFAULT_VALUES)
self.client.put(entity_native)
entity_native_retrieved = self.client.get(key_native)
self.assertTrue(entity_native_retrieved)
# Verify retrieved data matches the original input
self.assertEqual(entity_native_retrieved, EXAMPLE_DICT_DEFAULT_VALUES)
# Store custom Protobuf object in a datastore by translating it to Entity object
key_translated = self.client.key('ExampleModel',
'translated_entity_default_values')
example_pb = EXAMPLE_PB_DEFAULT_VALUES
entity_pb_translated = model_pb_to_entity_pb(model_pb=example_pb)
# pylint: disable=no-member
entity_pb_translated.key.CopyFrom(key_translated.to_protobuf())
entity_translated = datastore.helpers.entity_from_protobuf(
entity_pb_translated)
self.client.put(entity_translated)
# Verify that the translated entity results in the same end result as using native
# entity object
entity_translated_retrieved = self.client.get(key_translated)
self.assertTrue(
entity_translated_retrieved.key != entity_native_retrieved.key)
# NOTE: key won't be the same so we clear it
entity_translated_retrieved.key = None
entity_native_retrieved.key = None
self.assertEqual(entity_translated_retrieved, entity_native_retrieved)
# If we translate retrieved entity back to the original Protobuf object definition, it
# should be the same as the original model (minus the key since the original model doesn't
# contain a key)
entity_pb_retrieved = datastore.helpers.entity_to_protobuf(
entity_translated_retrieved)
entity_pb_translated.ClearField('key')
self.assertEqual(entity_pb_translated, entity_pb_retrieved)
example_pb_retrieved = entity_pb_to_model_pb(
example_pb2.ExampleDBModel, entity_pb_retrieved)
self.assertEqual(example_pb_retrieved, example_pb)
# Storing and retrieving empty object should have the same end result
key_native_empty = self.client.key('ExampleModel',
'native_entity_empty')
entity_native_empty = datastore.Entity(key=key_native_empty)
entity_native_empty.update({})
self.client.put(entity_native_empty)
entity_native_empty_retrieved = self.client.get(key_native_empty)
self.assertTrue(entity_native_empty_retrieved is not None)
# Verify retrieved data matches the original input
self.assertEqual(entity_native_empty_retrieved, {})
# Store custom Protobuf object in a datastore by translating it to Entity object
key_translated_empty = self.client.key('ExampleModel',
'translated_entity_empty')
example_pb = example_pb2.ExampleDBModel()
entity_pb_translated_empty = model_pb_to_entity_pb(model_pb=example_pb)
# pylint: disable=no-member
entity_pb_translated_empty.key.CopyFrom(
key_translated_empty.to_protobuf())
entity_translated_empty = datastore.helpers.entity_from_protobuf(
entity_pb_translated_empty)
self.client.put(entity_translated_empty)
# Verify that the translated entity results in the same end result as using native
# entity object
entity_translated_empty_retrieved = self.client.get(
key_translated_empty)
self.assertTrue(entity_translated_empty_retrieved.key !=
entity_native_empty_retrieved.key)
# NOTE: key won't be the same so we clear it
entity_translated_empty_retrieved.key = None
entity_native_empty_retrieved.key = None
# self.assertEqual(entity_translated_empty_retrieved, entity_native_empty_retrieved)
# return
# If we translate retrieved entity back to the original Protobuf object definition, it
# should be the same as the original model (minus the key since the original model doesn't
# contain a key)
entity_pb_empty_retrieved = \
datastore.helpers.entity_to_protobuf(entity_translated_empty_retrieved)
entity_pb_translated_empty.ClearField('key')
entity_pb_empty_retrieved.ClearField('key')
self.assertEqual(entity_pb_translated_empty, entity_pb_empty_retrieved)
example_pb_empty_retrieved = entity_pb_to_model_pb(
example_pb2.ExampleDBModel, entity_pb_empty_retrieved)
self.assertEqual(example_pb_empty_retrieved, example_pb)
def measure_model_pb_to_entity_pb_with_exclude_field_from_index_simple_model():
return model_pb_to_entity_pb(model_pb=EXAMPLE_PB_WITH_OPTIONS_1)
def measure_model_pb_to_entity_pb_simple_model():
return model_pb_to_entity_pb(model_pb=simple_example_pb)
def measure_model_pb_to_entity_pb_complex_model():
return model_pb_to_entity_pb(model_pb=complex_example_pb)
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(BASE_DIR, '../'))
sys.path.append(os.path.join(BASE_DIR, '../tests/generated/'))
from protobuf_cloud_datastore_translator import model_pb_to_entity_pb
from protobuf_cloud_datastore_translator import entity_pb_to_model_pb
from tests.generated import example_pb2
from tests.mocks import EXAMPLE_PB_POPULATED
from tests.mocks import EXAMPLE_PB_DEFAULT_VALUES
from tests.mocks import EXAMPLE_PB_WITH_OPTIONS_1
complex_example_pb = EXAMPLE_PB_POPULATED
complex_entity_pb = model_pb_to_entity_pb(model_pb=complex_example_pb)
simple_example_pb = EXAMPLE_PB_DEFAULT_VALUES
simple_entity_pb = model_pb_to_entity_pb(model_pb=simple_example_pb)
def measure_model_pb_to_entity_pb_complex_model():
return model_pb_to_entity_pb(model_pb=complex_example_pb)
def measure_entity_pb_to_model_pb_complex_model():
return entity_pb_to_model_pb(example_pb2.ExampleDBModel, complex_entity_pb)
def measure_model_pb_to_entity_pb_simple_model():
return model_pb_to_entity_pb(model_pb=simple_example_pb)