def test_from_and_then_to_model(self):
model = FooModel(id='abc', project=feconf.OPPIA_PROJECT_ID, prop='123')
self.assertEqual(
model,
job_utils.get_ndb_model_from_beam_entity(
job_utils.get_beam_entity_from_ndb_model(model)))
def test_get_model_from_beam_entity(self):
beam_entity = beam_datastore_types.Entity(
beam_datastore_types.Key(('FooModel', 'abc'),
project=feconf.OPPIA_PROJECT_ID))
beam_entity.set_properties({'prop': '123'})
self.assertEqual(
FooModel(id='abc', app=feconf.OPPIA_PROJECT_ID, prop='123'),
job_utils.get_ndb_model_from_beam_entity(beam_entity))
def test_from_and_then_to_beam_entity(self):
beam_entity = beam_datastore_types.Entity(
beam_datastore_types.Key(('FooModel', 'abc'),
project=feconf.OPPIA_PROJECT_ID))
beam_entity.set_properties({'prop': '123'})
self.assertEqual(
beam_entity,
job_utils.get_beam_entity_from_ndb_model(
job_utils.get_ndb_model_from_beam_entity(beam_entity)))
def test_from_and_then_to_beam_entity(self):
beam_entity = beam_datastore_types.Entity(
beam_datastore_types.Key(
('CoreModel', 'abc'), project=feconf.OPPIA_PROJECT_ID))
beam_entity.set_properties({
'prop': 123,
'created_on': None,
'last_updated': None,
'deleted': False,
})
self.assertEqual(
beam_entity,
job_utils.get_beam_entity_from_ndb_model(
job_utils.get_ndb_model_from_beam_entity(beam_entity)))
def test_get_model_from_beam_entity_with_time(self):
utcnow = datetime.datetime.utcnow()
beam_entity = beam_datastore_types.Entity(
beam_datastore_types.Key(
('CoreModel', 'abc'), project=feconf.OPPIA_PROJECT_ID,
namespace=self.namespace))
beam_entity.set_properties({
'prop': 3.14,
'created_on': utcnow.replace(tzinfo=datetime.timezone.utc),
'last_updated': None,
'deleted': False,
})
self.assertEqual(
CoreModel(
id='abc', project=feconf.OPPIA_PROJECT_ID, prop=3.14,
created_on=utcnow),
job_utils.get_ndb_model_from_beam_entity(beam_entity))
def _write_to_datastore_handler(self, pickled_beam_entities):
"""XML-RPC handler for a WriteToDatastore request.
IMPORTANT: This operation must be idempotent!
Args:
pickled_beam_entities: str. The list of models to put into the
datastore, encoded as a pickled list of Apache Beam entities.
"""
ndb_models_to_put = [
job_utils.get_ndb_model_from_beam_entity(model)
for model in pickle.loads(pickled_beam_entities)
]
datastore_services.update_timestamps_multi(
# The caller is responsible for updating timestamps, not the stub.
ndb_models_to_put, update_last_updated_time=False)
with self._datastore_lock:
datastore_services.put_multi(ndb_models_to_put)
