def testFlexTemplatesCreateJob(self): name = 'job name' gcs_path = 'gs://my_gcs_path' template_args = apis.TemplateArguments( job_name=name, gcs_location=gcs_path, worker_region=base.DEFAULT_REGION, project_id=self.Project()) launch_params = apis.Templates.FLEX_TEMPLATE_PARAMETER( jobName=template_args.job_name, containerSpecGcsPath=template_args.gcs_location, environment=apis.Templates.FLEX_TEMPLATE_ENVIRONMENT( workerRegion=base.DEFAULT_REGION, additionalExperiments=[]), parameters=None) request = apis.Templates.LAUNCH_FLEX_TEMPLATE_REQUEST( launchParameter=launch_params) wrapper_req = (apis.GetMessagesModule(). DataflowProjectsLocationsFlexTemplatesLaunchRequest) wrapped_request = wrapper_req(projectId=self.Project(), location=base.DEFAULT_REGION, launchFlexTemplateRequest=request) self.mocked_client.projects_locations_flexTemplates.Launch.Expect( wrapped_request, apis.GetMessagesModule().LaunchFlexTemplateResponse()) apis.Templates.CreateJobFromFlexTemplate(template_args)
def testTemplatesCreate(self): name = 'job name' gcs_path = 'gs://my_gcs_path' staging_path = 'gs://my_staging_path' request = apis.Templates.CREATE_REQUEST( jobName=name, location=base.DEFAULT_REGION, gcsPath=gcs_path, environment=apis.GetMessagesModule().RuntimeEnvironment( tempLocation=staging_path)) wrapper_req = (apis.GetMessagesModule(). DataflowProjectsLocationsTemplatesCreateRequest) wrapped_request = wrapper_req(createJobFromTemplateRequest=request, projectId=self.Project(), location=base.DEFAULT_REGION) self.mocked_client.projects_locations_templates.Create.Expect( wrapped_request, apis.GetMessagesModule().Job()) template_args = apis.TemplateArguments(job_name=name, gcs_location=gcs_path, project_id=self.Project(), staging_location=staging_path) apis.Templates.Create(template_args)
def testPrivateIPTemplatesCreate(self): name = 'job name' gcs_path = 'gs://my_gcs_path' staging_path = 'gs://my_staging_path' ip_configuration_enum = apis.GetMessagesModule( ).RuntimeEnvironment.IpConfigurationValueValuesEnum ip_private = ip_configuration_enum.WORKER_IP_PRIVATE request = apis.Templates.CREATE_REQUEST( jobName=name, location=base.DEFAULT_REGION, gcsPath=gcs_path, environment=apis.GetMessagesModule().RuntimeEnvironment( tempLocation=staging_path, ipConfiguration=ip_private)) wrapper_req = (apis.GetMessagesModule(). DataflowProjectsLocationsTemplatesCreateRequest) wrapped_request = wrapper_req(createJobFromTemplateRequest=request, projectId=self.Project(), location=base.DEFAULT_REGION) self.mocked_client.projects_locations_templates.Create.Expect( wrapped_request, apis.GetMessagesModule().Job()) template_args = apis.TemplateArguments(job_name=name, gcs_location=gcs_path, project_id=self.Project(), staging_location=staging_path, disable_public_ips=True) apis.Templates.Create(template_args)
def testLaunchDynamicTemplate(self): name = 'myjob' gcs_path = 'gs://dynamic_template_path' staging_path = 'gs://my_staging_path' parameters = apis.Templates.LAUNCH_TEMPLATE_PARAMETERS( jobName=name, environment=apis.GetMessagesModule().RuntimeEnvironment( tempLocation=staging_path), update=False) wrapper_req = (apis.GetMessagesModule(). DataflowProjectsLocationsTemplatesLaunchRequest) wrapped_request = wrapper_req( dynamicTemplate_gcsPath=gcs_path, dynamicTemplate_stagingLocation=staging_path, launchTemplateParameters=parameters, projectId=self.Project(), location=base.DEFAULT_REGION, validateOnly=False) self.mocked_client.projects_locations_templates.Launch.Expect( wrapped_request, apis.GetMessagesModule().LaunchTemplateResponse()) template_args = apis.TemplateArguments(job_name=name, gcs_location=gcs_path, project_id=self.Project(), staging_location=staging_path) apis.Templates.LaunchDynamicTemplate(template_args)
def _CommonRun(args): """Runs the command. Args: args: The arguments that were provided to this command invocation. Returns: A Job message. """ arguments = apis.TemplateArguments( project_id=properties.VALUES.core.project.Get(required=True), region_id=dataflow_util.GetRegion(args), job_name=args.job_name, gcs_location=args.gcs_location, zone=args.zone, max_workers=args.max_workers, num_workers=args.num_workers, network=args.network, subnetwork=args.subnetwork, worker_machine_type=args.worker_machine_type, staging_location=args.staging_location, kms_key_name=args.dataflow_kms_key, disable_public_ips=properties.VALUES.dataflow.disable_public_ips. GetBool(), parameters=args.parameters, service_account_email=args.service_account_email, worker_region=args.worker_region, worker_zone=args.worker_zone) return apis.Templates.Create(arguments)
def Run(self, args): region = dataflow_util.GetRegion(args) if args.sql_launcher_template: gcs_location = args.sql_launcher_template else: gcs_location = 'gs://dataflow-sql-templates-{}/latest/sql_launcher_template'.format( region) if args.parameters_file: query_parameters = sql_query_parameters.ParseParametersFile( args.parameters_file) elif args.parameter: query_parameters = sql_query_parameters.ParseParametersList( args.parameter) else: query_parameters = '[]' template_parameters = { 'dryRun': 'true' if args.dry_run else 'false', 'outputs': sql_util.ExtractOutputs(args), 'queryParameters': query_parameters, 'queryString': args.query, } arguments = apis.TemplateArguments( project_id=properties.VALUES.core.project.GetOrFail(), region_id=region, job_name=args.job_name, gcs_location=gcs_location, zone=args.worker_zone, max_workers=args.max_workers, disable_public_ips=properties.VALUES.dataflow.disable_public_ips. GetBool(), parameters=template_parameters, service_account_email=args.service_account_email) return apis.Templates.LaunchDynamicTemplate(arguments)
def Run(self, args): """Runs the command. Args: args: The arguments that were provided to this command invocation. Returns: A Job message. """ arguments = apis.TemplateArguments( project_id=properties.VALUES.core.project.Get(required=True), region_id=dataflow_util.GetRegion(args), job_name=args.job_name, gcs_location=args.template_file_gcs_location, max_workers=args.max_workers, num_workers=args.num_workers, network=args.network, subnetwork=args.subnetwork, worker_machine_type=args.worker_machine_type, kms_key_name=args.dataflow_kms_key, staging_location=args.staging_location, disable_public_ips=properties.VALUES.dataflow.disable_public_ips. GetBool(), service_account_email=args.service_account_email, worker_region=args.worker_region, worker_zone=args.worker_zone, enable_streaming_engine=properties.VALUES.dataflow. enable_streaming_engine.GetBool(), additional_experiments=args.additional_experiments, additional_user_labels=args.additional_user_labels, parameters=args.parameters) return apis.Templates.CreateJobFromFlexTemplate(arguments)
def _CommonRun(args): """Runs the command. Args: args: The arguments that were provided to this command invocation. Returns: A Job message. """ arguments = apis.TemplateArguments( project_id=properties.VALUES.core.project.Get(required=True), region_id=dataflow_util.GetRegion(args), job_name=args.job_name, gcs_location=args.template_file_gcs_location, parameters=args.parameters) return apis.Templates.CreateJobFromFlexTemplate(arguments)
def Run(self, args): use_dynamic_engine = (args.sql_launcher_template_engine == 'dynamic') region = dataflow_util.GetRegion(args) if args.sql_launcher_template: gcs_location = args.sql_launcher_template else: if use_dynamic_engine: suffix = 'sql_launcher_template' else: suffix = 'sql_launcher_flex_template' gcs_location = 'gs://dataflow-sql-templates-{}/latest/{}'.format( region, suffix) if args.parameters_file: query_parameters = sql_query_parameters.ParseParametersFile( args.parameters_file) elif args.parameter: query_parameters = sql_query_parameters.ParseParametersList( args.parameter) else: query_parameters = '[]' template_parameters = collections.OrderedDict([ ('dryRun', 'true' if args.dry_run else 'false'), ('outputs', sql_util.ExtractOutputs(args)), ('queryParameters', query_parameters), ('queryString', args.query), ]) arguments = apis.TemplateArguments( project_id=properties.VALUES.core.project.GetOrFail(), region_id=region, job_name=args.job_name, gcs_location=gcs_location, zone=args.zone, max_workers=args.max_workers, disable_public_ips=properties.VALUES.dataflow.disable_public_ips. GetBool(), parameters=template_parameters, service_account_email=args.service_account_email, kms_key_name=args.dataflow_kms_key, num_workers=args.num_workers, network=args.network, subnetwork=args.subnetwork, worker_machine_type=args.worker_machine_type, worker_region=args.worker_region, worker_zone=args.worker_zone) if use_dynamic_engine: return apis.Templates.LaunchDynamicTemplate(arguments) return apis.Templates.CreateJobFromFlexTemplate(arguments)
def _CommonRun(args, support_beta_features=False): """Runs the command. Args: args: The arguments that were provided to this command invocation. support_beta_features: True, if using the beta command. Returns: A Job message. """ num_workers = None worker_machine_type = None network = None subnetwork = None dataflow_kms_key = None if support_beta_features: num_workers = args.num_workers worker_machine_type = args.worker_machine_type network = args.network subnetwork = args.subnetwork dataflow_kms_key = args.dataflow_kms_key arguments = apis.TemplateArguments( project_id=properties.VALUES.core.project.Get(required=True), region_id=dataflow_util.GetRegion(args), job_name=args.job_name, gcs_location=args.gcs_location, zone=args.zone, max_workers=args.max_workers, num_workers=num_workers, network=network, subnetwork=subnetwork, worker_machine_type=worker_machine_type, staging_location=args.staging_location, kms_key_name=dataflow_kms_key, disable_public_ips=properties.VALUES.dataflow.disable_public_ips. GetBool(), parameters=args.parameters, service_account_email=args.service_account_email) flex_template = properties.VALUES.dataflow.flex_template.GetBool() if flex_template: return apis.Templates.CreateJobFromFlexTemplate(arguments) else: return apis.Templates.Create(arguments)
def _CommonRun(args): """Runs the command. Args: args: The arguments that were provided to this command invocation. Returns: A Job message. """ template_args = apis.TemplateArguments( max_workers=args.max_workers, num_workers=args.num_workers, network=args.network, subnetwork=args.subnetwork, worker_machine_type=args.worker_machine_type, kms_key_name=args.dataflow_kms_key, staging_location=args.staging_location, temp_location=args.temp_location, disable_public_ips=properties.VALUES.dataflow.disable_public_ips. GetBool(), service_account_email=args.service_account_email, worker_region=args.worker_region, worker_zone=args.worker_zone, enable_streaming_engine=properties.VALUES.dataflow. enable_streaming_engine.GetBool(), additional_experiments=args.additional_experiments, additional_user_labels=args.additional_user_labels) image_path = args.image if not args.image: image_path = args.image_gcr_path apis.Templates.BuildAndStoreFlexTemplateImage( args.image_gcr_path, args.flex_template_base_image, args.jar, args.py_path, args.env, args.sdk_language, args.gcs_log_dir) return apis.Templates.BuildAndStoreFlexTemplateFile( args.template_file_gcs_path, image_path, args.metadata_file, args.sdk_language, args.print_only, template_args, args.image_repository_username_secret_id, args.image_repository_password_secret_id, args.image_repository_cert_path)
def testLaunchDynamicTemplate_httpError(self): name = 'myjob' gcs_path = 'gs://dynamic_template_path' staging_path = 'gs://my_staging_path' parameters = apis.Templates.LAUNCH_TEMPLATE_PARAMETERS( jobName=name, environment=apis.GetMessagesModule().RuntimeEnvironment( tempLocation=staging_path), update=False) wrapper_req = (apis.GetMessagesModule(). DataflowProjectsLocationsTemplatesLaunchRequest) wrapped_request = wrapper_req( dynamicTemplate_gcsPath=gcs_path, dynamicTemplate_stagingLocation=staging_path, launchTemplateParameters=parameters, projectId=self.Project(), location=base.DEFAULT_REGION, validateOnly=False) self.mocked_client.projects_locations_templates.Launch.Expect( wrapped_request, exception=apitools_exceptions.HttpError.FromResponse( http_wrapper.Response( info={'status': 400}, content='{"error": ' '{"message": "The workflow could not be created."}}', request_url= 'https://dataflow.googleapis.com/v1b3/request-url'))) template_args = apis.TemplateArguments(job_name=name, gcs_location=gcs_path, project_id=self.Project(), staging_location=staging_path) with self.AssertRaisesExceptionMatches( exceptions.HttpException, 'The workflow could not be created.'): apis.Templates.LaunchDynamicTemplate(template_args)