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)
