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()