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)