def test_passing_arguments_to_hook(self, mock_hook):
task = BigQueryTableExistenceSensor(task_id='task-id',
project_id=TEST_PROJECT_ID,
dataset_id=TEST_DATASET_ID,
table_id=TEST_TABLE_ID,
bigquery_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO)
mock_hook.return_value.table_exists.return_value = True
results = task.poke(mock.MagicMock())
self.assertEqual(True, results)
mock_hook.assert_called_once_with(bigquery_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO)
mock_hook.return_value.table_exists.assert_called_once_with(
TEST_PROJECT_ID, TEST_DATASET_ID, TEST_TABLE_ID)
def test_passing_arguments_to_hook(self, mock_hook):
task = BigQueryTableExistenceSensor(
task_id='task-id',
project_id=TEST_PROJECT_ID,
dataset_id=TEST_DATASET_ID,
table_id=TEST_TABLE_ID,
bigquery_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_hook.return_value.table_exists.return_value = True
results = task.poke(mock.MagicMock())
assert results is True
mock_hook.assert_called_once_with(
bigquery_conn_id=TEST_GCP_CONN_ID,
delegate_to=TEST_DELEGATE_TO,
impersonation_chain=TEST_IMPERSONATION_CHAIN,
)
mock_hook.return_value.table_exists.assert_called_once_with(
project_id=TEST_PROJECT_ID, dataset_id=TEST_DATASET_ID, table_id=TEST_TABLE_ID
)
project_id=PROJECT_ID)
create_table = BigQueryCreateEmptyTableOperator(
task_id="create_table",
dataset_id=DATASET_NAME,
table_id=TABLE_NAME,
schema_fields=SCHEMA,
time_partitioning={
"type": "DAY",
"field": "ds",
},
)
# [START howto_sensor_bigquery_table]
check_table_exists = BigQueryTableExistenceSensor(
task_id="check_table_exists",
project_id=PROJECT_ID,
dataset_id=DATASET_NAME,
table_id=TABLE_NAME)
# [END howto_sensor_bigquery_table]
execute_insert_query = BigQueryExecuteQueryOperator(
task_id="execute_insert_query",
sql=INSERT_ROWS_QUERY,
use_legacy_sql=False)
# [START howto_sensor_bigquery_table_partition]
check_table_partition_exists = BigQueryTablePartitionExistenceSensor(
task_id="check_table_partition_exists",
project_id=PROJECT_ID,
dataset_id=DATASET_NAME,
table_id=TABLE_NAME,
fivetran_conn_id='fivetran_default',
connector_id='{{ var.value.connector_id }}')
fivetran_sync_wait = FivetranSensor(
task_id='fivetran-sensor',
fivetran_conn_id='fivetran_default',
connector_id='{{ var.value.connector_id }}',
poke_interval=5)
"""
#### BigQuery row validation task
Ensure that data was copied to BigQuery correctly, i.e. the table and dataset
exist.
"""
validate_bigquery = BigQueryTableExistenceSensor(
task_id='validate_bigquery',
project_id='{{ var.value.gcp_project_id }}',
dataset_id=DATASET,
table_id='forestfires',
)
"""
#### Row-level data quality check
Run a data quality check on a few rows, ensuring that the data in BigQuery
matches the ground truth in the correspoding JSON file.
"""
check_bq_row_count = BigQueryValueCheckOperator(
task_id="check_row_count",
sql=f"SELECT COUNT(*) FROM {DATASET}.{TABLE}",
pass_value=516,
use_legacy_sql=False,
)
done = DummyOperator(task_id='done')