def test_query_with_not_equal_filter_raises_type_error(self) -> None:
query = datastore_services.Query(filters=BarModel.prop != 1)
with self.assertRaisesRegexp(
TypeError,
'forbidden filter'): # type: ignore[no-untyped-call]
job_utils.get_beam_query_from_ndb_query(query)
def test_query_everything(self):
query = datastore_services.query_everything()
beam_query = job_utils.get_beam_query_from_ndb_query(query)
self.assertIsNone(beam_query.kind)
self.assertEqual(beam_query.order, ('__key__',))
def test_query_with_range_like_filter(self):
query = datastore_services.Query(filters=datastore_services.all_of(
BarModel.prop >= 3, BarModel.prop < 6))
beam_query = job_utils.get_beam_query_from_ndb_query(query)
self.assertEqual(
beam_query.filters, [('prop', '>=', 3), ('prop', '<', 6)])
def test_stub_is_initially_empty(self):
query_everything = job_utils.get_beam_query_from_ndb_query(
datastore_services.query_everything())
self.assertItemsEqual(self.get_everything(), [])
with self.stub.context():
self.assert_pcoll_empty(
self.pipeline | self.stub.ReadFromDatastore(query_everything))
def test_read_after_delete_raises_error(self):
query_everything = job_utils.get_beam_query_from_ndb_query(
datastore_services.query_everything())
with self.stub.context():
empty_pcoll = self.pipeline | beam.Create([])
self.assert_pcoll_empty(
empty_pcoll
| self.stub.DeleteFromDatastore(feconf.OPPIA_PROJECT_ID))
self.assertRaisesRegexp(
RuntimeError, 'Cannot read from datastore after a mutation',
lambda: (self.pipeline
| self.stub.ReadFromDatastore(query_everything)))
def expand(self, pbegin):
"""Returns a PCollection containing the queried models.
Args:
pbegin: PValue. The initial PValue of the pipeline, used to anchor
the models to its underlying pipeline.
Returns:
PCollection. The PCollection of models.
"""
return (pbegin
| 'Reading %r from the datastore' % self.query >>
(self.datastoreio.ReadFromDatastore(
job_utils.get_beam_query_from_ndb_query(self.query)))
| 'Transforming %r into NDB models' % self.query >>
(beam.Map(job_utils.get_ndb_model_from_beam_entity)))
def test_read_from_datastore(self):
query_everything = job_utils.get_beam_query_from_ndb_query(
datastore_services.query_everything())
model_list = [
self.create_model(base_models.BaseModel, id='a'),
self.create_model(base_models.BaseModel, id='b'),
self.create_model(base_models.BaseModel, id='c'),
]
self.put_multi(model_list)
with self.stub.context():
model_pcoll = (self.pipeline
| self.stub.ReadFromDatastore(query_everything))
self.assert_pcoll_equal(model_pcoll, [
job_utils.get_beam_entity_from_ndb_model(m) for m in model_list
])
def expand(self, initial_pipeline):
"""Returns a PCollection with models matching the corresponding query.
This overrides the expand() method from the parent class.
Args:
initial_pipeline: PValue. The initial pipeline. This pipeline
is used to anchor the models to itself.
Returns:
PCollection. The PCollection of models.
"""
query = job_utils.get_beam_query_from_ndb_query(
self.query, namespace=initial_pipeline.pipeline.options.namespace)
return (initial_pipeline.pipeline
| 'Reading %r from the datastore' % self.query >>
(datastoreio.ReadFromDatastore(query))
| 'Transforming %r into NDB models' % self.query >>
(beam.Map(job_utils.get_ndb_model_from_beam_entity)))
def test_query_with_descending_order(self):
query = BarModel.query().order(-BarModel.prop)
beam_query = job_utils.get_beam_query_from_ndb_query(query)
self.assertEqual(beam_query.order, ('-prop',))
def test_query_with_in_filter_raises_type_error(self):
query = datastore_services.Query(filters=BarModel.prop.IN([1, 2, 3]))
with self.assertRaisesRegexp(TypeError, 'forbidden filter'):
job_utils.get_beam_query_from_ndb_query(query)
def test_query_with_multiple_orders(self):
query = BarModel.query().order(BarModel.prop, BarModel.prop)
beam_query = job_utils.get_beam_query_from_ndb_query(query)
self.assertEqual(beam_query.order, ('prop', 'prop'))
def test_read_from_datastore_without_acquiring_context_raises_error(self):
query_everything = job_utils.get_beam_query_from_ndb_query(
datastore_services.query_everything())
with self.assertRaisesRegexp(RuntimeError, 'Must enter context'):
self.stub.ReadFromDatastore(query_everything)
def test_query_with_or_filter_raises_type_error(self):
query = datastore_services.Query(filters=datastore_services.any_of(
BarModel.prop == 1, BarModel.prop == 2))
with self.assertRaisesRegexp(TypeError, 'forbidden filter'):
job_utils.get_beam_query_from_ndb_query(query)
def test_query_with_filter(self):
query = datastore_services.Query(filters=BarModel.prop >= 3)
beam_query = job_utils.get_beam_query_from_ndb_query(query)
self.assertEqual(beam_query.filters, [('prop', '>=', 3)])
def test_query_with_namespace(self):
query = datastore_services.Query(namespace='abc')
beam_query = job_utils.get_beam_query_from_ndb_query(query)
self.assertEqual(beam_query.namespace, 'abc')
def test_query_with_kind(self):
query = base_models.BaseModel.query()
beam_query = job_utils.get_beam_query_from_ndb_query(query)
self.assertEqual(beam_query.kind, 'BaseModel')
def test_query_with_order(self) -> None:
query = BarModel.query().order(BarModel.prop)
beam_query = job_utils.get_beam_query_from_ndb_query(query)
self.assertEqual(beam_query.order, ('prop', ))
def test_query_with_project(self):
query = datastore_services.Query(app='foo-project')
beam_query = job_utils.get_beam_query_from_ndb_query(query)
self.assertEqual(beam_query.project, 'foo-project')
