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)