def test_query_with_not_equal_filter_raises_type_error(self) -> None:
query = datastore_services.Query(filters=BarModel.prop != 1)
with self.assertRaisesRegex(
TypeError,
'forbidden filter'): # type: ignore[no-untyped-call]
job_utils.get_beam_query_from_ndb_query(query)
def test_query_everything(self) -> None:
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) -> None:
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 expand(
self, pbegin: pvalue.PBegin
) -> beam.PCollection[datastore_services.Model]:
"""Returns a PCollection with models matching the corresponding query.
This overrides the expand() method from the parent class.
Args:
pbegin: 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=pbegin.pipeline.options.namespace)
return (pbegin.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) -> 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_multiple_orders(self) -> None:
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_query_with_filter(self) -> None:
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) -> None:
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) -> None:
query = base_models.BaseModel.query()
beam_query = job_utils.get_beam_query_from_ndb_query(query)
self.assertEqual(beam_query.kind, 'BaseModel')
