def main():
logging.basicConfig()
logging.getLogger().setLevel(logging.INFO)
# parse arguments
args, beam_args = parse_args()
# Perform the export (unless we are skipping it).
if not args.skip_export:
export.export_to_gcs_content_types(args.parent, args.gcs_destination,
args.content_types,
args.asset_types)
# Perform the import, via template or beam runner.
launch_location = args.template_job_launch_location
if launch_location:
final_state = pipeline_runner.run_pipeline_template(
args.template_job_project, args.template_job_region,
launch_location, args.input, args.group_by, args.write_disposition,
args.dataset, args.stage, args.load_time,
args.template_job_runtime_environment_json)
else:
final_state = pipeline_runner.run_pipeline_beam_runner(
None, None, args.input, args.group_by, args.write_disposition,
args.dataset, args.stage, args.load_time, beam_args)
if not pipeline_runner.is_successful_state(final_state):
sys.exit(1)
def test_export_to_gcs_content_types(self, mock_export_to_gcs):
export.export_to_gcs_content_types('parent', 'gcs_prefix',
['RESOURCE'], ['a', 'b'])
mock_export_to_gcs.assert_has_calls([
mock.call('parent', 'gcs_prefix/RESOURCE.json', 'RESOURCE',
['a', 'b']),
])
self.assertEqual(mock_export_to_gcs.call_count, 1)
def run_export_import():
"""Run export and import process."""
# Optionally require that we are invoked by a cron tasks and not any old
# pulic HTTP request.
if CONFIG.restrict_to_cron_tasks:
if 'X-Appengine-Cron' not in request.headers:
return '', 403
# perform export
export_arguments = get_export_arguments()
export_result = export.export_to_gcs_content_types(*export_arguments)
# perform import
final_state, pipeline_result = run_import()
# Return 200 if everything worked.
status_code = 200
if not pipeline_runner.is_successful_state(final_state):
status_code = 500
return pprint.pformat({
'export_result': export_result,
'pipeline_result': pipeline_result
}), status_code
def test_export_to_gcs_all_content_types(self, mock_export_to_gcs):
export.export_to_gcs_content_types('parent', 'gcs_prefix', None, None)
self.assertEqual(mock_export_to_gcs.call_count, 2)