use of com.google.api.services.bigquery.model.TableDataInsertAllResponse in project beam by apache.
the class BigQueryServicesImplTest method testInsertRetry.
/**
* Tests that {@link DatasetServiceImpl#insertAll} retries quota rate limited attempts.
*/
@Test
public void testInsertRetry() throws Exception {
TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
List<ValueInSingleWindow<TableRow>> rows = new ArrayList<>();
rows.add(wrapTableRow(new TableRow()));
// First response is 403 rate limited, second response has valid payload.
when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
when(response.getStatusCode()).thenReturn(403).thenReturn(200);
when(response.getContent()).thenReturn(toStream(errorWithReasonAndStatus("rateLimitExceeded", 403))).thenReturn(toStream(new TableDataInsertAllResponse()));
DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), new MockSleeper(), InsertRetryPolicy.alwaysRetry(), null);
verify(response, times(2)).getStatusCode();
verify(response, times(2)).getContent();
verify(response, times(2)).getContentType();
expectedLogs.verifyInfo("BigQuery insertAll exceeded rate limit, retrying");
}
use of com.google.api.services.bigquery.model.TableDataInsertAllResponse in project google-cloud-java by GoogleCloudPlatform.
the class InsertAllResponse method fromPb.
static InsertAllResponse fromPb(TableDataInsertAllResponse responsePb) {
Map<Long, List<BigQueryError>> insertErrors = null;
if (responsePb.getInsertErrors() != null) {
List<InsertErrors> errorsPb = responsePb.getInsertErrors();
insertErrors = Maps.newHashMapWithExpectedSize(errorsPb.size());
for (InsertErrors errorPb : errorsPb) {
insertErrors.put(errorPb.getIndex(), Lists.transform(errorPb.getErrors() != null ? errorPb.getErrors() : ImmutableList.<ErrorProto>of(), BigQueryError.FROM_PB_FUNCTION));
}
}
return new InsertAllResponse(insertErrors);
}
use of com.google.api.services.bigquery.model.TableDataInsertAllResponse in project google-cloud-java by GoogleCloudPlatform.
the class BigQueryImplTest method testInsertAll.
@Test
public void testInsertAll() {
Map<String, Object> row1 = ImmutableMap.<String, Object>of("field", "value1");
Map<String, Object> row2 = ImmutableMap.<String, Object>of("field", "value2");
List<RowToInsert> rows = ImmutableList.of(new RowToInsert("row1", row1), new RowToInsert("row2", row2));
InsertAllRequest request = InsertAllRequest.newBuilder(TABLE_ID).setRows(rows).setSkipInvalidRows(false).setIgnoreUnknownValues(true).setTemplateSuffix("suffix").build();
TableDataInsertAllRequest requestPb = new TableDataInsertAllRequest().setRows(Lists.transform(rows, new Function<RowToInsert, TableDataInsertAllRequest.Rows>() {
@Override
public TableDataInsertAllRequest.Rows apply(RowToInsert rowToInsert) {
return new TableDataInsertAllRequest.Rows().setInsertId(rowToInsert.getId()).setJson(rowToInsert.getContent());
}
})).setSkipInvalidRows(false).setIgnoreUnknownValues(true).setTemplateSuffix("suffix");
TableDataInsertAllResponse responsePb = new TableDataInsertAllResponse().setInsertErrors(ImmutableList.of(new TableDataInsertAllResponse.InsertErrors().setIndex(0L).setErrors(ImmutableList.of(new ErrorProto().setMessage("ErrorMessage")))));
EasyMock.expect(bigqueryRpcMock.insertAll(PROJECT, DATASET, TABLE, requestPb)).andReturn(responsePb);
EasyMock.replay(bigqueryRpcMock);
bigquery = options.getService();
InsertAllResponse response = bigquery.insertAll(request);
assertNotNull(response.getErrorsFor(0L));
assertNull(response.getErrorsFor(1L));
assertEquals(1, response.getErrorsFor(0L).size());
assertEquals("ErrorMessage", response.getErrorsFor(0L).get(0).getMessage());
}
use of com.google.api.services.bigquery.model.TableDataInsertAllResponse in project wildfly-camel by wildfly-extras.
the class GoogleBigQueryIntegrationTest method setUp.
@Before
public void setUp() throws Exception {
configuration = new GoogleBigQueryConfiguration();
bigquery = Mockito.mock(Bigquery.class);
endpoint = Mockito.mock(GoogleBigQueryEndpoint.class);
tabledata = Mockito.mock(Bigquery.Tabledata.class);
mockInsertall = Mockito.mock(Bigquery.Tabledata.InsertAll.class);
Mockito.when(bigquery.tabledata()).thenReturn(tabledata);
Mockito.when(tabledata.insertAll(Mockito.anyString(), Mockito.anyString(), Mockito.anyString(), Mockito.any())).thenReturn(mockInsertall);
TableDataInsertAllResponse mockResponse = new TableDataInsertAllResponse();
Mockito.when(mockInsertall.execute()).thenReturn(mockResponse);
}
use of com.google.api.services.bigquery.model.TableDataInsertAllResponse in project beam by apache.
the class TestBigQuery method insertRows.
@Experimental(Kind.SCHEMAS)
public TableDataInsertAllResponse insertRows(Schema rowSchema, Row... rows) throws IOException {
List<Rows> bqRows = Arrays.stream(rows).map(row -> new Rows().setJson(BigQueryUtils.toTableRow(row))).collect(ImmutableList.toImmutableList());
Bigquery bq = newBigQueryClient(pipelineOptions);
return bq.tabledata().insertAll(pipelineOptions.getBigQueryProject() == null ? pipelineOptions.getProject() : pipelineOptions.getBigQueryProject(), pipelineOptions.getTargetDataset(), table.getTableReference().getTableId(), new TableDataInsertAllRequest().setRows(bqRows)).setPrettyPrint(false).execute();
}
Aggregations