def execute(self, context):
hook = DataFlowHook(gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
poll_sleep=self.poll_sleep)
hook.start_template_dataflow(self.job_name, self.dataflow_default_options,
self.parameters, self.template)
def execute(self, context):
hook = DataFlowHook(gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
poll_sleep=self.poll_sleep)
hook.start_template_dataflow(self.task_id, self.dataflow_default_options,
self.parameters, self.template)
def execute(
self,
# Some context about the context: https://bcb.github.io/airflow/execute-context
context: Dict[str, Any] # pylint: disable=unused-argument
) -> None:
hook = DataFlowHook(gcp_conn_id=self.gcp_conn_id,
delegate_to=self.delegate_to,
poll_sleep=self.poll_sleep)
# In DataflowTemplateOperator, start_template_dataflow has the default append_job_name set to True
# so it adds a unique-id to the end of the job name. This overwrites that default argument.
hook.start_template_dataflow(self.task_id, self.dataflow_default_options,
self.parameters, self.template, append_job_name=False)
class DataFlowTemplateHookTest(unittest.TestCase):
def setUp(self):
with mock.patch(BASE_STRING.format('GoogleCloudBaseHook.__init__'),
new=mock_init):
self.dataflow_hook = DataFlowHook(gcp_conn_id='test')
@mock.patch(DATAFLOW_STRING.format('DataFlowHook._start_template_dataflow'))
def test_start_template_dataflow(self, internal_dataflow_mock):
self.dataflow_hook.start_template_dataflow(
task_id=TASK_ID, variables=DATAFLOW_OPTIONS_TEMPLATE, parameters=PARAMETERS,
dataflow_template=TEMPLATE)
internal_dataflow_mock.assert_called_once_with(
mock.ANY, DATAFLOW_OPTIONS_TEMPLATE, PARAMETERS, TEMPLATE)
class DataFlowTemplateHookTest(unittest.TestCase):
def setUp(self):
with mock.patch(BASE_STRING.format('GoogleCloudBaseHook.__init__'),
new=mock_init):
self.dataflow_hook = DataFlowHook(gcp_conn_id='test')
@mock.patch(DATAFLOW_STRING.format('DataFlowHook._start_template_dataflow')
)
def test_start_template_dataflow(self, internal_dataflow_mock):
self.dataflow_hook.start_template_dataflow(
job_name=JOB_NAME,
variables=DATAFLOW_OPTIONS_TEMPLATE,
parameters=PARAMETERS,
dataflow_template=TEMPLATE)
options_with_region = {'region': 'us-central1'}
options_with_region.update(DATAFLOW_OPTIONS_TEMPLATE)
internal_dataflow_mock.assert_called_once_with(mock.ANY,
options_with_region,
PARAMETERS, TEMPLATE)
@mock.patch(DATAFLOW_STRING.format('_DataflowJob'))
@mock.patch(DATAFLOW_STRING.format('DataFlowHook.get_conn'))
def test_start_template_dataflow_with_runtime_env(self, mock_conn,
mock_dataflowjob):
dataflow_options_template = copy.deepcopy(DATAFLOW_OPTIONS_TEMPLATE)
options_with_runtime_env = copy.deepcopy(RUNTIME_ENV)
options_with_runtime_env.update(dataflow_options_template)
dataflowjob_instance = mock_dataflowjob.return_value
dataflowjob_instance.wait_for_done.return_value = None
method = (mock_conn.return_value.projects.return_value.locations.
return_value.templates.return_value.launch)
self.dataflow_hook.start_template_dataflow(
job_name=JOB_NAME,
variables=options_with_runtime_env,
parameters=PARAMETERS,
dataflow_template=TEMPLATE)
body = {
"jobName": mock.ANY,
"parameters": PARAMETERS,
"environment": RUNTIME_ENV
}
method.assert_called_once_with(
projectId=options_with_runtime_env['project'],
location='us-central1',
gcsPath=TEMPLATE,
body=body,
)
class DataFlowTemplateHookTest(unittest.TestCase):
def setUp(self):
with mock.patch(BASE_STRING.format('GoogleCloudBaseHook.__init__'),
new=mock_init):
self.dataflow_hook = DataFlowHook(gcp_conn_id='test')
@mock.patch(DATAFLOW_STRING.format('DataFlowHook._start_template_dataflow'))
def test_start_template_dataflow(self, internal_dataflow_mock):
self.dataflow_hook.start_template_dataflow(
job_name=JOB_NAME, variables=DATAFLOW_OPTIONS_TEMPLATE, parameters=PARAMETERS,
dataflow_template=TEMPLATE)
options_with_region = {'region': 'us-central1'}
options_with_region.update(DATAFLOW_OPTIONS_TEMPLATE)
internal_dataflow_mock.assert_called_once_with(
mock.ANY, options_with_region, PARAMETERS, TEMPLATE)
@mock.patch(DATAFLOW_STRING.format('_DataflowJob'))
@mock.patch(DATAFLOW_STRING.format('DataFlowHook.get_conn'))
def test_start_template_dataflow_with_runtime_env(self, mock_conn, mock_dataflowjob):
dataflow_options_template = copy.deepcopy(DATAFLOW_OPTIONS_TEMPLATE)
options_with_runtime_env = copy.deepcopy(RUNTIME_ENV)
options_with_runtime_env.update(dataflow_options_template)
dataflowjob_instance = mock_dataflowjob.return_value
dataflowjob_instance.wait_for_done.return_value = None
method = (mock_conn.return_value
.projects.return_value
.locations.return_value
.templates.return_value
.launch)
self.dataflow_hook.start_template_dataflow(
job_name=JOB_NAME,
variables=options_with_runtime_env,
parameters=PARAMETERS,
dataflow_template=TEMPLATE
)
body = {"jobName": mock.ANY,
"parameters": PARAMETERS,
"environment": RUNTIME_ENV
}
method.assert_called_once_with(
projectId=options_with_runtime_env['project'],
location='us-central1',
gcsPath=TEMPLATE,
body=body,
)
class DataFlowTemplateHookTest(unittest.TestCase):
def setUp(self):
with mock.patch(BASE_STRING.format('GoogleCloudBaseHook.__init__'),
new=mock_init):
self.dataflow_hook = DataFlowHook(gcp_conn_id='test')
@mock.patch(DATAFLOW_STRING.format('DataFlowHook._start_template_dataflow')
)
def test_start_template_dataflow(self, internal_dataflow_mock):
self.dataflow_hook.start_template_dataflow(
task_id=TASK_ID,
variables=DATAFLOW_OPTIONS_TEMPLATE,
parameters=PARAMETERS,
dataflow_template=TEMPLATE)
internal_dataflow_mock.assert_called_once_with(
mock.ANY, DATAFLOW_OPTIONS_TEMPLATE, PARAMETERS, TEMPLATE)
class DataFlowTemplateHookTest(unittest.TestCase):
def setUp(self):
with mock.patch(BASE_STRING.format('GoogleCloudBaseHook.__init__'),
new=mock_init):
self.dataflow_hook = DataFlowHook(gcp_conn_id='test')
@mock.patch(DATAFLOW_STRING.format('DataFlowHook._start_template_dataflow'))
def test_start_template_dataflow(self, internal_dataflow_mock):
self.dataflow_hook.start_template_dataflow(
job_name=JOB_NAME, variables=DATAFLOW_OPTIONS_TEMPLATE, parameters=PARAMETERS,
dataflow_template=TEMPLATE)
options_with_region = {'region': 'us-central1'}
options_with_region.update(DATAFLOW_OPTIONS_TEMPLATE)
internal_dataflow_mock.assert_called_once_with(
mock.ANY, options_with_region, PARAMETERS, TEMPLATE)
