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)