def test_DatastoreDeleteFn(self):
with patch.object(helper, 'get_client',
return_value=self._mock_client):
keys = [entity.key for entity in helper.create_entities(10)]
expected_keys = [key.to_client_key() for key in keys]
# Infer project from delete fn project arg.
key = Key(['k1', 1234], project=self._PROJECT)
expected_key = key.to_client_key()
key.project = None
keys.append(key)
expected_keys.append(expected_key)
all_batch_keys = []
self._mock_client.batch.side_effect = (
lambda: FakeBatch(all_batch_items=all_batch_keys))
datastore_delete_fn = DeleteFromDatastore._DatastoreDeleteFn(
self._PROJECT)
datastore_delete_fn.start_bundle()
for key in keys:
datastore_delete_fn.process(key)
datastore_delete_fn.finish_bundle()
self.assertListEqual(all_batch_keys, expected_keys)
def run():
import sys
args = sys.argv[1:]
options = DeleteOptions(args)
if not options.src.project:
options.src.project = options.project
p = beam.Pipeline(options=options)
sources = create_multi_datasource_reader(
p, options.src.project, options.src.namespace, options.src.kinds, keys_only=True)
sources | beam.Flatten() \
| 'EntityToKey' >> beam.ParDo(EntityToKey()) \
| 'DeleteFromDatastore' >> DeleteFromDatastore(options.src.project)
p.run().wait_until_finish()
def run(argv=None):
"""Main entry point."""
parser = argparse.ArgumentParser()
parser.add_argument('--kind',
dest='kind',
default='writereadtest',
help='Datastore Kind')
parser.add_argument('--num_entities',
dest='num_entities',
type=int,
required=True,
help='Number of entities to write')
parser.add_argument('--limit',
dest='limit',
type=int,
help='Limit of number of entities to write')
known_args, pipeline_args = parser.parse_known_args(argv)
pipeline_options = PipelineOptions(pipeline_args)
gcloud_options = pipeline_options.view_as(GoogleCloudOptions)
job_name = gcloud_options.job_name
kind = known_args.kind
num_entities = known_args.num_entities
project = gcloud_options.project
# Pipeline 1: Create and write the specified number of Entities to the
# Cloud Datastore.
ancestor_key = Key([kind, str(uuid.uuid4())], project=project)
_LOGGER.info('Writing %s entities to %s', num_entities, project)
p = new_pipeline_with_job_name(pipeline_options, job_name, '-write')
_ = (
p
| 'Input' >> beam.Create(list(range(num_entities)))
| 'To String' >> beam.Map(str)
|
'To Entity' >> beam.Map(EntityWrapper(kind, ancestor_key).make_entity)
| 'Write to Datastore' >> WriteToDatastore(project))
p.run()
query = Query(kind=kind, project=project, ancestor=ancestor_key)
# Optional Pipeline 2: If a read limit was provided, read it and confirm
# that the expected entities were read.
if known_args.limit is not None:
_LOGGER.info(
'Querying a limited set of %s entities and verifying count.',
known_args.limit)
p = new_pipeline_with_job_name(pipeline_options, job_name,
'-verify-limit')
query.limit = known_args.limit
entities = p | 'read from datastore' >> ReadFromDatastore(query)
assert_that(entities | beam.combiners.Count.Globally(),
equal_to([known_args.limit]))
p.run()
query.limit = None
# Pipeline 3: Query the written Entities and verify result.
_LOGGER.info('Querying entities, asserting they match.')
p = new_pipeline_with_job_name(pipeline_options, job_name, '-verify')
entities = p | 'read from datastore' >> ReadFromDatastore(query)
assert_that(entities | beam.combiners.Count.Globally(),
equal_to([num_entities]))
p.run()
# Pipeline 4: Delete Entities.
_LOGGER.info('Deleting entities.')
p = new_pipeline_with_job_name(pipeline_options, job_name, '-delete')
entities = p | 'read from datastore' >> ReadFromDatastore(query)
_ = (entities
| 'To Keys' >> beam.Map(lambda entity: entity.key)
| 'delete entities' >> DeleteFromDatastore(project))
p.run()
# Pipeline 5: Query the written Entities, verify no results.
_LOGGER.info(
'Querying for the entities to make sure there are none present.')
p = new_pipeline_with_job_name(pipeline_options, job_name,
'-verify-deleted')
entities = p | 'read from datastore' >> ReadFromDatastore(query)
assert_that(entities | beam.combiners.Count.Globally(), equal_to([0]))
p.run()
