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 create_entities(count, id_or_name=False): """Creates a list of entities with random keys.""" if id_or_name: ids_or_names = [uuid.uuid4().int & ((1 << 63) - 1) for _ in range(count)] else: ids_or_names = [str(uuid.uuid4()) for _ in range(count)] keys = [types.Key(['EntityKind', x], project='project') for x in ids_or_names] return [types.Entity(key) for key in keys]
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_get_beam_key_from_ndb_key(self) -> None: beam_key = beam_datastore_types.Key(('FooModel', 'abc'), project=feconf.OPPIA_PROJECT_ID, namespace=self.namespace) # We use private _from_ds_key here because it provides functionality # for obtaining an NDB key from a Beam key, and writing it ourselves # would be too complicated. ndb_key = datastore_services.Key._from_ds_key(beam_key.to_client_key()) # pylint: disable=protected-access self.assertEqual(job_utils.get_beam_key_from_ndb_key(ndb_key), beam_key)
def get_beam_key_from_ndb_key(ndb_key): """Returns an Apache Beam key equivalent to the given NDB key. Args: ndb_key: datastore_services.Key. The NDB key. Returns: beam_datastore_types.Key. The Apache Beam key. """ return beam_datastore_types.Key( ndb_key.flat(), project=ndb_key.app(), namespace=ndb_key.namespace())
def process(self, element, timestamp=None, *args, **kwargs): from apache_beam.io.gcp.datastore.v1new import types kind_id = timestamp.strftime('%Y%m%d-%H%M%S') + '-' + str(element.pop('indice')) key = types.Key([self.kind, kind_id], project=self.project) entity = types.Entity(key=key) entity.set_properties({ 'user_id': element['user_id'], 'item_id': element['item_id'], 'created': timestamp, 'rule': 'RULE_NAME', }) yield 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 CreateEntity(elem): """Make a SignalQualitySampleStats entity from summarised stats.""" key = ds_types.Key( [ 'SignalQuality', '{bot_group}/{bot}/{measurement}'.format(**elem[0].as_dict()), 'SignalQualitySampleStats', '0', # version ], project='chromeperf', ) entity = ds_types.Entity(key) properties = elem[1].as_dict() properties['updated_time'] = datetime.datetime.now() entity.set_properties(properties) return 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))