def test_delete_non_existing_table(self): client = mock.Mock() client.tables.Delete.side_effect = HttpError( response={'status': '404'}, url='', content='') wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) wrapper._delete_table('', '', '') self.assertTrue(client.tables.Delete.called)
def test_delete_table_retries_for_timeouts(self, patched_time_sleep): client = mock.Mock() client.tables.Delete.side_effect = [ HttpError(response={'status': '408'}, url='', content=''), bigquery.BigqueryTablesDeleteResponse() ] wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) wrapper._delete_table('', '', '') self.assertTrue(client.tables.Delete.called)
def test_get_or_create_dataset_created(self): client = mock.Mock() client.datasets.Get.side_effect = HttpError( response={'status': '404'}, url='', content='') client.datasets.Insert.return_value = bigquery.Dataset( datasetReference=bigquery.DatasetReference( projectId='project_id', datasetId='dataset_id')) wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) new_dataset = wrapper.get_or_create_dataset('project_id', 'dataset_id') self.assertEqual(new_dataset.datasetReference.datasetId, 'dataset_id')
def test_no_table_and_create_never(self, patched_time_sleep): client = mock.Mock() client.tables.Get.side_effect = HttpError( response={'status': '404'}, url='', content='') create_disposition = beam.io.BigQueryDisposition.CREATE_NEVER with self.assertRaisesRegexp( RuntimeError, r'Table project:dataset\.table not found but create ' r'disposition is CREATE_NEVER'): with beam.io.BigQuerySink( 'project:dataset.table', create_disposition=create_disposition).writer(client): pass
def test_get_or_create_table_race_condition(self): client = mock.Mock() client.tables.Insert.side_effect = HttpError( response={'status': '409'}, url='', content='') client.tables.Get.side_effect = [None, 'table_id'] wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) new_table = wrapper.get_or_create_table( 'project_id', 'dataset_id', 'table_id', bigquery.TableSchema(fields=[ bigquery.TableFieldSchema(name='b', type='BOOLEAN', mode='REQUIRED')]), False, False) self.assertEqual(new_table, 'table_id')
def test_no_table_and_create_if_needed_and_no_schema( self, patched_time_sleep): client = mock.Mock() client.tables.Get.side_effect = HttpError( response={'status': '404'}, url='', content='') create_disposition = beam.io.BigQueryDisposition.CREATE_IF_NEEDED with self.assertRaisesRegexp( RuntimeError, r'Table project:dataset\.table requires a schema\. None ' r'can be inferred because the table does not exist'): with beam.io.BigQuerySink( 'project:dataset.table', create_disposition=create_disposition).writer(client): pass
def test_no_table_and_create_if_needed(self): client = mock.Mock() table = bigquery.Table(tableReference=bigquery.TableReference( projectId='project', datasetId='dataset', tableId='table'), schema=bigquery.TableSchema()) client.tables.Get.side_effect = HttpError(response={'status': '404'}, url='', content='') client.tables.Insert.return_value = table create_disposition = beam.io.BigQueryDisposition.CREATE_IF_NEEDED with beam.io.BigQuerySink( 'project:dataset.table', schema='somefield:INTEGER', create_disposition=create_disposition).writer(client): pass self.assertTrue(client.tables.Get.called) self.assertTrue(client.tables.Insert.called)