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')