def execute(self, context: 'Context'): hook = BiqQueryDataTransferServiceHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, location=self.location ) self.log.info("Creating DTS transfer config") response = hook.create_transfer_config( project_id=self.project_id, transfer_config=self.transfer_config, authorization_code=self.authorization_code, retry=self.retry, timeout=self.timeout, metadata=self.metadata, ) transfer_config = _get_transfer_config_details(response.name) BigQueryDataTransferConfigLink.persist( context=context, task_instance=self, region=transfer_config["region"], config_id=transfer_config["config_id"], project_id=transfer_config["project_id"], ) result = TransferConfig.to_dict(response) self.log.info("Created DTS transfer config %s", get_object_id(result)) self.xcom_push(context, key="transfer_config_id", value=get_object_id(result)) return result
def execute(self, context): hook = BiqQueryDataTransferServiceHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain ) self.log.info("Creating DTS transfer config") response = hook.create_transfer_config( project_id=self.project_id, transfer_config=self.transfer_config, authorization_code=self.authorization_code, retry=self.retry, timeout=self.timeout, metadata=self.metadata, ) result = TransferConfig.to_dict(response) self.log.info("Created DTS transfer config %s", get_object_id(result)) self.xcom_push(context, key="transfer_config_id", value=get_object_id(result)) return result
class BigQueryCreateDataTransferOperatorTestCase(unittest.TestCase): @mock.patch( "airflow.providers.google.cloud.operators.bigquery_dts.BiqQueryDataTransferServiceHook", **{'return_value.create_transfer_config.return_value': TransferConfig(name=TRANSFER_CONFIG_NAME)}, ) def test_execute(self, mock_hook): op = BigQueryCreateDataTransferOperator( transfer_config=TRANSFER_CONFIG, project_id=PROJECT_ID, task_id="id" ) ti = mock.MagicMock() op.execute({'ti': ti}) mock_hook.return_value.create_transfer_config.assert_called_once_with( authorization_code=None, metadata=None, transfer_config=TRANSFER_CONFIG, project_id=PROJECT_ID, retry=None, timeout=None, ) ti.xcom_push.assert_called_once_with(execution_date=None, key='transfer_config_id', value='1a2b3c')