def test_poke(self, mock_msg_to_dict, mock_hook):
op = BigQueryDataTransferServiceTransferRunSensor(
transfer_config_id=TRANSFER_CONFIG_ID,
run_id=RUN_ID,
task_id="id",
project_id=PROJECT_ID,
expected_statuses={"success"},
)
op.poke(None)
mock_hook.return_value.get_transfer_run.assert_called_once_with(
transfer_config_id=TRANSFER_CONFIG_ID,
run_id=RUN_ID,
project_id=PROJECT_ID,
metadata=None,
retry=None,
timeout=None,
)
def test_poke_returns_true(self, mock_hook):
op = BigQueryDataTransferServiceTransferRunSensor(
transfer_config_id=TRANSFER_CONFIG_ID,
run_id=RUN_ID,
task_id="id",
project_id=PROJECT_ID,
expected_statuses={"SUCCEEDED"},
)
result = op.poke({})
self.assertEqual(result, True)
mock_hook.return_value.get_transfer_run.assert_called_once_with(
transfer_config_id=TRANSFER_CONFIG_ID,
run_id=RUN_ID,
project_id=PROJECT_ID,
metadata=None,
retry=None,
timeout=None,
)
)
# [END howto_bigquery_create_data_transfer]
# [START howto_bigquery_start_transfer]
gcp_bigquery_start_transfer = BigQueryDataTransferServiceStartTransferRunsOperator(
task_id="gcp_bigquery_start_transfer",
transfer_config_id=transfer_config_id,
requested_run_time={"seconds": int(time.time() + 60)},
)
run_id = "{{ task_instance.xcom_pull('gcp_bigquery_start_transfer', key='run_id') }}"
# [END howto_bigquery_start_transfer]
# [START howto_bigquery_dts_sensor]
gcp_run_sensor = BigQueryDataTransferServiceTransferRunSensor(
task_id="gcp_run_sensor",
transfer_config_id=transfer_config_id,
run_id=run_id,
expected_statuses={"SUCCEEDED"},
)
# [END howto_bigquery_dts_sensor]
# [START howto_bigquery_delete_data_transfer]
gcp_bigquery_delete_transfer = BigQueryDeleteDataTransferConfigOperator(
transfer_config_id=transfer_config_id, task_id="gcp_bigquery_delete_transfer"
)
# [END howto_bigquery_delete_data_transfer]
(
gcp_bigquery_create_transfer # noqa
>> gcp_bigquery_start_transfer # noqa
>> gcp_run_sensor # noqa
>> gcp_bigquery_delete_transfer # noqa
transfer_config_id = gcp_bigquery_create_transfer.output[
"transfer_config_id"]
# [END howto_bigquery_create_data_transfer]
# [START howto_bigquery_start_transfer]
gcp_bigquery_start_transfer = BigQueryDataTransferServiceStartTransferRunsOperator(
task_id="gcp_bigquery_start_transfer",
transfer_config_id=transfer_config_id,
requested_run_time={"seconds": int(time.time() + 60)},
)
# [END howto_bigquery_start_transfer]
# [START howto_bigquery_dts_sensor]
gcp_run_sensor = BigQueryDataTransferServiceTransferRunSensor(
task_id="gcp_run_sensor",
transfer_config_id=transfer_config_id,
run_id=gcp_bigquery_start_transfer.output["run_id"],
expected_statuses={"SUCCEEDED"},
)
# [END howto_bigquery_dts_sensor]
# [START howto_bigquery_delete_data_transfer]
gcp_bigquery_delete_transfer = BigQueryDeleteDataTransferConfigOperator(
transfer_config_id=transfer_config_id,
task_id="gcp_bigquery_delete_transfer")
# [END howto_bigquery_delete_data_transfer]
gcp_run_sensor >> gcp_bigquery_delete_transfer
# Task dependencies created via `XComArgs`:
# gcp_bigquery_create_transfer >> gcp_bigquery_start_transfer
# gcp_bigquery_create_transfer >> gcp_run_sensor