def test_execute(self, mock_hook):
source_project_dataset_table = '{}.{}'.format(TEST_DATASET,
TEST_TABLE_ID)
destination_cloud_storage_uris = ['gs://some-bucket/some-file.txt']
compression = 'NONE'
export_format = 'CSV'
field_delimiter = ','
print_header = True
labels = {'k1': 'v1'}
operator = BigQueryToGCSOperator(
task_id=TASK_ID,
source_project_dataset_table=source_project_dataset_table,
destination_cloud_storage_uris=destination_cloud_storage_uris,
compression=compression,
export_format=export_format,
field_delimiter=field_delimiter,
print_header=print_header,
labels=labels)
operator.execute(None)
mock_hook.return_value \
.get_conn.return_value \
.cursor.return_value \
.run_extract \
.assert_called_once_with(
source_project_dataset_table=source_project_dataset_table,
destination_cloud_storage_uris=destination_cloud_storage_uris,
compression=compression,
export_format=export_format,
field_delimiter=field_delimiter,
print_header=print_header,
labels=labels
)
DATASET_NAME),
use_legacy_sql=False,
)
copy_from_selected_data = BigQueryToBigQueryOperator(
task_id="copy_from_selected_data",
source_project_dataset_tables="{}.selected_data_from_external_table".
format(DATASET_NAME),
destination_project_dataset_table=
"{}.copy_of_selected_data_from_external_table".format(DATASET_NAME),
)
bigquery_to_gcs = BigQueryToGCSOperator(
task_id="bigquery_to_gcs",
source_project_dataset_table="{}.selected_data_from_external_table".
format(DATASET_NAME),
destination_cloud_storage_uris=[
"gs://{}/export-bigquery.csv".format(DATA_EXPORT_BUCKET_NAME)
],
)
create_dataset = BigQueryCreateEmptyDatasetOperator(
task_id="create-dataset", dataset_id=DATASET_NAME)
create_dataset_with_location = BigQueryCreateEmptyDatasetOperator(
task_id="create_dataset_with_location",
dataset_id=LOCATION_DATASET_NAME,
location=BQ_LOCATION)
create_table = BigQueryCreateEmptyTableOperator(
task_id="create_table",
dataset_id=DATASET_NAME,
with models.DAG(
"example_bigquery_transfer",
default_args=default_args,
schedule_interval=None, # Override to match your needs
tags=["example"],
) as dag:
copy_selected_data = BigQueryToBigQueryOperator(
task_id="copy_selected_data",
source_project_dataset_tables=f"{DATASET_NAME}.{ORIGIN}",
destination_project_dataset_table=f"{DATASET_NAME}.{TARGET}",
)
bigquery_to_gcs = BigQueryToGCSOperator(
task_id="bigquery_to_gcs",
source_project_dataset_table=f"{DATASET_NAME}.{ORIGIN}",
destination_cloud_storage_uris=[
f"gs://{DATA_EXPORT_BUCKET_NAME}/export-bigquery.csv"
],
)
create_dataset = BigQueryCreateEmptyDatasetOperator(
task_id="create_dataset", dataset_id=DATASET_NAME
)
for table in [ORIGIN, TARGET]:
create_table = BigQueryCreateEmptyTableOperator(
task_id=f"create_{table}_table",
dataset_id=DATASET_NAME,
table_id=table,
schema_fields=[
{"name": "emp_name", "type": "STRING", "mode": "REQUIRED"},