def test_refresh_bq_table_from_gcs_export_synchronous(
self, mock_load_gcs: mock.MagicMock,
mock_load_stale_data: mock.MagicMock,
mock_delete_temp_table: mock.MagicMock) -> None:
"""Test that refresh_bq_table_from_gcs_export_synchronous calls the following load jobs in order:
1. load_table_from_gcs_and_wait
2. load_rows_excluded_from_refresh_into_temp_table_and_wait
3. big_query_client.load_table_from_table_async
4. delete_temp_table_if_exists
"""
temp_table_name = f'{self.mock_table_id}_temp'
mock_parent = mock.Mock()
mock_parent.attach_mock(mock_load_gcs, 'load_from_gcs')
mock_parent.attach_mock(mock_load_stale_data, 'load_stale_data')
mock_parent.attach_mock(mock_delete_temp_table, 'delete_temp_table')
expected_calls = [
mock.call.load_from_gcs(self.mock_bq_client,
self.mock_table_id,
self.mock_bq_refresh_config,
destination_table_id=temp_table_name),
mock.call.load_stale_data(self.mock_bq_client,
self.mock_table_id,
self.mock_bq_refresh_config,
destination_table_id=temp_table_name),
mock.call.delete_temp_table(self.mock_bq_client, temp_table_name,
self.mock_bq_refresh_config)
]
bq_refresh.refresh_bq_table_from_gcs_export_synchronous(
self.mock_bq_client, self.mock_table_id,
self.mock_bq_refresh_config)
mock_parent.assert_has_calls(expected_calls)
self.mock_bq_client.load_table_from_table_async.assert_called()
def export_table_then_load_table(
big_query_client: BigQueryClient, table: str,
cloud_sql_to_bq_config: CloudSqlToBQConfig) -> None:
"""Exports a Cloud SQL table to CSV, then loads it into BigQuery.
If a table excludes some region codes, it first loads all the GCS and the excluded region's data to a temp table.
See for details: load_table_with_excluded_regions
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 metadata_base class for its corresponding SchemaType.
cloud_sql_to_bq_config: The config class for the given SchemaType.
Returns:
True if load succeeds, else False.
"""
export_success = cloud_sql_to_gcs_export.export_table(
table, cloud_sql_to_bq_config)
if not export_success:
raise ValueError(
f"Failure to export CloudSQL table to GCS, skipping BigQuery load of table [{table}]."
)
bq_refresh.refresh_bq_table_from_gcs_export_synchronous(
big_query_client, table, cloud_sql_to_bq_config)
def test_refresh_bq_table_from_gcs_export_synchronous_load_table_async_fails(
self, mock_wait: mock.MagicMock,
mock_delete: mock.MagicMock) -> None:
"""Test that refresh_bq_table_from_gcs_export_synchronous raises a ValueError if wait_for_table_load fails
and does not all delete_temp_table_if_exists
"""
mock_wait.return_value = False
with self.assertRaises(ValueError):
bq_refresh.refresh_bq_table_from_gcs_export_synchronous(
self.mock_bq_client, self.mock_table_id,
self.mock_bq_refresh_config)
mock_delete.assert_not_called()
def test_refresh_bq_table_from_gcs_export_synchronous_create_table_fails(
self, mock_load_stale_data: mock.MagicMock) -> None:
"""Test if a create table fails, it raises a ValueError."""
with self.assertRaises(ValueError):
self.mock_bq_refresh_config.dataset_id = self.mock_dataset_id
self.mock_bq_refresh_config.get_bq_schema_for_table.return_value = self.mock_table_schema
self.mock_bq_client.table_exists.return_value = False
self.mock_bq_client.create_table_with_schema.return_value = None
bq_refresh.refresh_bq_table_from_gcs_export_synchronous(
self.mock_bq_client, 'new_table', self.mock_bq_refresh_config)
mock_load_stale_data.assert_not_called()
self.mock_bq_client.load_table_from_table_async.assert_not_called()
def test_refresh_bq_table_from_gcs_export_synchronous_load_stale_data_fails(
self, mock_load_gcs: mock.MagicMock,
mock_load_stale_data: mock.MagicMock,
mock_delete_temp_table: mock.MagicMock) -> None:
"""Test that refresh_bq_table_from_gcs_export_synchronous does not continue when
load_rows_excluded_from_refresh_into_temp_table_and_wait raises an error."""
mock_load_stale_data.side_effect = ValueError
with self.assertRaises(ValueError):
bq_refresh.refresh_bq_table_from_gcs_export_synchronous(
self.mock_bq_client, self.mock_table_id,
self.mock_bq_refresh_config)
mock_load_gcs.assert_called()
mock_load_stale_data.assert_called()
mock_delete_temp_table.assert_not_called()
self.mock_bq_client.create_table_with_schema.assert_not_called()
self.mock_bq_client.load_table_from_table_async.assert_not_called()
def test_refresh_bq_table_from_gcs_export_synchronous_load_from_gcs_fails(
self,
mock_load_gcs: mock.MagicMock,
mock_load_stale_data: mock.MagicMock,
mock_delete_temp_table: mock.MagicMock,
) -> None:
"""Test that refresh_bq_table_from_gcs_export_synchronous raises an error and exits
if load_from_gcs_and_wait fails
"""
mock_load_gcs.side_effect = ValueError
with self.assertRaises(ValueError):
bq_refresh.refresh_bq_table_from_gcs_export_synchronous(
self.mock_bq_client, self.mock_table_id,
self.mock_bq_refresh_config)
mock_load_gcs.assert_called()
mock_load_stale_data.assert_not_called()
mock_delete_temp_table.assert_not_called()
self.mock_bq_client.create_table_with_schema.assert_not_called()
self.mock_bq_client.load_table_from_table_async.assert_not_called()
def test_refresh_from_or_create_table_table_doesnt_exist(
self, mock_load_stale_data: mock.MagicMock) -> None:
"""Test if a destination table does not exist, it does not call
load_rows_excluded_from_refresh_into_temp_table_and_wait and it creates the table with the schema.
"""
with self.assertLogs(level='INFO'):
self.mock_bq_refresh_config.dataset_id = self.mock_dataset_id
self.mock_bq_refresh_config.get_bq_schema_for_table.return_value = self.mock_table_schema
self.mock_bq_client.table_exists.return_value = False
bq_refresh.refresh_bq_table_from_gcs_export_synchronous(
self.mock_bq_client, 'new_table', self.mock_bq_refresh_config)
self.mock_bq_client.create_table_with_schema.assert_called_with(
dataset_id=self.mock_dataset_id,
table_id='new_table',
schema_fields=self.mock_table_schema)
mock_load_stale_data.assert_not_called()
self.mock_bq_client.load_table_from_table_async.assert_called()