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')