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