def _disable_auto_scheduling(
config: Union[dict, TransferConfig]) -> TransferConfig:
"""
In the case of Airflow, the customer needs to create a transfer config
with the automatic scheduling disabled (UI, CLI or an Airflow operator) and
then trigger a transfer run using a specialized Airflow operator that will
call start_manual_transfer_runs.
:param config: Data transfer configuration to create.
:type config: Union[dict, google.cloud.bigquery_datatransfer_v1.types.TransferConfig]
"""
config = TransferConfig.to_dict(config) if isinstance(
config, TransferConfig) else config
new_config = copy(config)
schedule_options = new_config.get("schedule_options")
if schedule_options:
disable_auto_scheduling = schedule_options.get(
"disable_auto_scheduling", None)
if disable_auto_scheduling is None:
schedule_options["disable_auto_scheduling"] = True
else:
new_config["schedule_options"] = {"disable_auto_scheduling": True}
# HACK: TransferConfig.to_dict returns invalid representation
# See: https://github.com/googleapis/python-bigquery-datatransfer/issues/90
if isinstance(new_config.get('user_id'), str):
new_config['user_id'] = int(new_config['user_id'])
return TransferConfig(**new_config)
def _disable_auto_scheduling(
config: Union[dict, TransferConfig]) -> TransferConfig:
"""
In the case of Airflow, the customer needs to create a transfer config
with the automatic scheduling disabled (UI, CLI or an Airflow operator) and
then trigger a transfer run using a specialized Airflow operator that will
call start_manual_transfer_runs.
:param config: Data transfer configuration to create.
"""
config = TransferConfig.to_dict(config) if isinstance(
config, TransferConfig) else config
new_config = copy(config)
schedule_options = new_config.get("schedule_options")
if schedule_options:
disable_auto_scheduling = schedule_options.get(
"disable_auto_scheduling", None)
if disable_auto_scheduling is None:
schedule_options["disable_auto_scheduling"] = True
else:
new_config["schedule_options"] = {"disable_auto_scheduling": True}
return TransferConfig(**new_config)
"field_delimiter": ",",
"max_bad_records": "0",
"skip_leading_rows": "1",
"data_path_template": "bucket",
"destination_table_name_template": "name",
"file_format": "CSV",
}
TRANSFER_CONFIG = ParseDict(
{
"destination_dataset_id": "dataset",
"display_name": "GCS Test Config",
"data_source_id": "google_cloud_storage",
"params": PARAMS,
},
TransferConfig(),
)
TRANSFER_CONFIG_ID = "id1234"
class BigQueryDataTransferHookTestCase(unittest.TestCase):
def setUp(self) -> None:
with mock.patch(
"airflow.providers.google.cloud.hooks.bigquery_dts.CloudBaseHook.__init__",
new=mock_base_gcp_hook_no_default_project_id,
):
self.hook = BiqQueryDataTransferServiceHook()
self.hook._get_credentials = mock.MagicMock( # type: ignore
return_value=CREDENTIALS
)
CREDENTIALS = "test-creds"
PROJECT_ID = "id"
PARAMS = {
"field_delimiter": ",",
"max_bad_records": 0,
"skip_leading_rows": 1,
"data_path_template": "bucket",
"destination_table_name_template": "name",
"file_format": "CSV",
}
TRANSFER_CONFIG = TransferConfig(
destination_dataset_id="dataset",
display_name="GCS Test Config",
data_source_id="google_cloud_storage",
params=PARAMS,
)
TRANSFER_CONFIG_ID = "id1234"
class BigQueryDataTransferHookTestCase(unittest.TestCase):
def setUp(self) -> None:
with mock.patch(
"airflow.providers.google.cloud.hooks.bigquery_dts.GoogleBaseHook.__init__",
new=mock_base_gcp_hook_no_default_project_id,
):
self.hook = BiqQueryDataTransferServiceHook()
self.hook._get_credentials = mock.MagicMock(
return_value=CREDENTIALS) # type: ignore