def test_start_table_load_and_wait_not_called(self, mock_start, mock_wait):
"""Test that start_table_load_and_wait doesn't call wait_for_table_load.
Should be the case if start_table_load fails."""
mock_start.return_value = None
success = bq_load.start_table_load_and_wait(self.mock_dataset,
self.mock_table_id,
self.schema_type)
mock_start.assert_called()
mock_wait.assert_not_called()
self.assertFalse(success)
def test_start_table_load_and_wait(self, mock_start, mock_wait):
"""Test that start_table_load and wait_for_table_load are called."""
mock_start.return_value = (self.mock_load_job, self.mock_table)
mock_wait.return_value = True
success = bq_load.start_table_load_and_wait(self.mock_dataset,
self.mock_table_id,
self.schema_type)
mock_start.assert_called()
mock_wait.assert_called()
self.assertTrue(success)
def export_table_then_load_table(
big_query_client: BigQueryClient, table: str,
dataset_ref: bigquery.dataset.DatasetReference,
schema_type: SchemaType) -> bool:
"""Exports a Cloud SQL table to CSV, then loads it into BigQuery.
Waits until the BigQuery load is completed.
Args:
big_query_client: A BigQueryClient.
table: Table to export then import. Table must be defined
in the TABLES_TO_EXPORT for its corresponding schema.
dataset_ref: The BigQuery dataset to load the table into.
Gets created if it does not already exist.
schema_type: The schema, SchemaType.COUNTY, SchemaType.STATE,
or SchemaType.OPERATIONS, where this table lives.
Returns:
True if load succeeds, else False.
"""
if schema_type == SchemaType.JAILS:
export_queries = export_config.COUNTY_TABLE_EXPORT_QUERIES
elif schema_type == SchemaType.STATE:
export_queries = export_config.STATE_TABLE_EXPORT_QUERIES
elif schema_type == SchemaType.OPERATIONS:
export_queries = export_config.OPERATIONS_TABLE_EXPORT_QUERIES
else:
logging.error("Unknown schema_type: %s", schema_type)
return False
try:
export_query = export_queries[table]
except KeyError:
logging.exception(
"Unknown table name [%s]. Is it listed in "
"the TABLES_TO_EXPORT for the %s schema_type?", table, schema_type)
return False
export_success = cloudsql_export.export_table(schema_type, table,
export_query)
if export_success: # pylint: disable=no-else-return
load_success = bq_load.start_table_load_and_wait(
big_query_client, dataset_ref, table, schema_type)
return load_success
else:
logging.error(
"Skipping BigQuery load of table [%s], "
"which failed to export.", table)
return False