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