def test_export_table(self, mock_wait: mock.MagicMock) -> None:
"""Test that client().instances().export() is called and
wait_until_operation_finished is called.
"""
cloud_sql_to_gcs_export.export_table(self.mock_table_id,
self.mock_bq_refresh_config)
mock_export = self.mock_client.instances.return_value.export
mock_export.assert_called_with(project=self.mock_project_id,
instance=self.mock_instance_id,
body=mock.ANY)
mock_wait.assert_called()
def test_export_table_create_export_context(
self, mock_wait: mock.MagicMock) -> None:
"""Test that create_export_context is called with the schema_type, export_query, and export_uri."""
with mock.patch(
f"{CLOUD_SQL_TO_GCS_EXPORT_PACKAGE_NAME}.create_export_context"
) as mock_export_context:
cloud_sql_to_gcs_export.export_table(self.mock_table_id,
self.mock_bq_refresh_config)
mock_export_context.assert_called_with(self.schema_type,
self.mock_export_uri,
self.mock_table_query)
mock_wait.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_export_table_api_fail(self, mock_wait: mock.MagicMock) -> None:
"""Test that export_table fails if the export API request fails."""
mock_export = self.mock_client.instances.return_value.export
mock_export_request = mock_export.return_value
mock_export_request.execute.side_effect = \
googleapiclient.errors.HttpError('', content=b'')
with self.assertLogs(level='ERROR'):
success = cloud_sql_to_gcs_export.export_table(
self.mock_table_id, self.mock_bq_refresh_config)
self.assertFalse(success)
mock_wait.assert_not_called()