Examples with TableDataInsertAllResponse com.google.api.services.bigquery.model.TableDataInsertAllResponse used on opensource projects

Example 16 with TableDataInsertAllResponse

use of com.google.api.services.bigquery.model.TableDataInsertAllResponse in project beam by apache.

the class BigQueryServicesImplTest method testSimpleErrorRetrieval.

/**
 * Tests that {@link DatasetServiceImpl#insertAll} uses the supplied {@link ErrorContainer}.
 */
@Test
public void testSimpleErrorRetrieval() throws InterruptedException, IOException {
    TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows = ImmutableList.of(wrapValue(new TableRow().set("a", 1)), wrapValue(new TableRow().set("b", 2)));
    final List<ValueInSingleWindow<TableRow>> expected = ImmutableList.of(wrapErrorValue(new TableRow().set("a", 1)), wrapErrorValue(new TableRow().set("b", 2)));
    final TableDataInsertAllResponse failures = new TableDataInsertAllResponse().setInsertErrors(ImmutableList.of(new InsertErrors().setIndex(0L).setErrors(ImmutableList.of(new ErrorProto().setReason("timeout"))), new InsertErrors().setIndex(1L).setErrors(ImmutableList.of(new ErrorProto().setReason("invalid")))));
    setupMockResponses(response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(failures));
    });
    DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, null, PipelineOptionsFactory.create());
    List<ValueInSingleWindow<TableRow>> failedInserts = Lists.newArrayList();
    dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.neverRetry(), failedInserts, ErrorContainer.TABLE_ROW_ERROR_CONTAINER, false, false, false, null);
    assertThat(failedInserts, is(expected));
}

Example 17 with TableDataInsertAllResponse

use of com.google.api.services.bigquery.model.TableDataInsertAllResponse in project beam by apache.

the class BigQueryServicesImplTest method testInsertWithinRequestByteSizeLimitsErrorsOut.

/**
 * Tests that {@link DatasetServiceImpl#insertAll} does not go over limit of rows per request.
 */
@Test
public void testInsertWithinRequestByteSizeLimitsErrorsOut() throws Exception {
    TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows = ImmutableList.of(wrapValue(new TableRow().set("row", Strings.repeat("abcdefghi", 1024 * 1025))), wrapValue(new TableRow().set("row", "a")), wrapValue(new TableRow().set("row", "b")));
    List<String> insertIds = ImmutableList.of("a", "b", "c");
    final TableDataInsertAllResponse allRowsSucceeded = new TableDataInsertAllResponse();
    setupMockResponses(response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(allRowsSucceeded));
    }, response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(allRowsSucceeded));
    });
    DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, null, PipelineOptionsFactory.fromArgs("--maxStreamingBatchSize=15").create());
    List<ValueInSingleWindow<TableRow>> failedInserts = Lists.newArrayList();
    List<ValueInSingleWindow<TableRow>> successfulRows = Lists.newArrayList();
    RuntimeException e = assertThrows(RuntimeException.class, () -> dataService.<TableRow>insertAll(ref, rows, insertIds, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.alwaysRetry(), failedInserts, ErrorContainer.TABLE_ROW_ERROR_CONTAINER, false, false, false, successfulRows));
    assertThat(e.getMessage(), containsString("this row is too large."));
}

Example 18 with TableDataInsertAllResponse

use of com.google.api.services.bigquery.model.TableDataInsertAllResponse in project beam by apache.

the class BigQueryServicesImplTest method testFailInsertOtherRetry.

/**
 * Tests that {@link DatasetServiceImpl#insertAll} will not retry other non-rate-limited,
 * non-quota-exceeded attempts.
 */
@Test
public void testFailInsertOtherRetry() throws Exception {
    TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows = new ArrayList<>();
    rows.add(wrapValue(new TableRow()));
    // First response is 403 non-{rate-limited, quota-exceeded}, second response has valid payload
    // but should not be invoked.
    setupMockResponses(response -> {
        when(response.getStatusCode()).thenReturn(403);
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getContent()).thenReturn(toStream(errorWithReasonAndStatus("actually forbidden", 403)));
    }, response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(new TableDataInsertAllResponse()));
    });
    DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, null, PipelineOptionsFactory.create());
    thrown.expect(RuntimeException.class);
    thrown.expectMessage("actually forbidden");
    try {
        dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.alwaysRetry(), null, null, false, false, false, null);
    } finally {
        verify(responses[0], atLeastOnce()).getStatusCode();
        verify(responses[0]).getContent();
        verify(responses[0]).getContentType();
        // It should not invoke 2nd response
        verify(responses[1], never()).getStatusCode();
        verify(responses[1], never()).getContent();
        verify(responses[1], never()).getContentType();
    }
    verifyWriteMetricWasSet("project", "dataset", "table", "actually forbidden", 1);
}

Example 19 with TableDataInsertAllResponse

use of com.google.api.services.bigquery.model.TableDataInsertAllResponse in project beam by apache.

the class BigQueryUtilTest method onInsertAll.

private void onInsertAll(List<List<Long>> errorIndicesSequence) throws Exception {
    when(mockClient.tabledata()).thenReturn(mockTabledata);
    final List<TableDataInsertAllResponse> responses = new ArrayList<>();
    for (List<Long> errorIndices : errorIndicesSequence) {
        List<TableDataInsertAllResponse.InsertErrors> errors = new ArrayList<>();
        for (long i : errorIndices) {
            TableDataInsertAllResponse.InsertErrors error = new TableDataInsertAllResponse.InsertErrors();
            error.setIndex(i);
        }
        TableDataInsertAllResponse response = new TableDataInsertAllResponse();
        response.setInsertErrors(errors);
        responses.add(response);
    }
    Bigquery.Tabledata.InsertAll mockInsertAll = mock(Bigquery.Tabledata.InsertAll.class);
    when(mockTabledata.insertAll(anyString(), anyString(), anyString(), any(TableDataInsertAllRequest.class))).thenReturn(mockInsertAll);
    when(mockInsertAll.setPrettyPrint(any())).thenReturn(mockInsertAll);
    when(mockInsertAll.execute()).thenReturn(responses.get(0), responses.subList(1, responses.size()).toArray(new TableDataInsertAllResponse[responses.size() - 1]));
}

Example 20 with TableDataInsertAllResponse

use of com.google.api.services.bigquery.model.TableDataInsertAllResponse in project beam by apache.

the class BigQueryServicesImplTest method testInsertRetryPolicy.

/**
 * Tests that {@link DatasetServiceImpl#insertAll} uses the supplied {@link InsertRetryPolicy},
 * and returns the list of rows not retried.
 */
@Test
public void testInsertRetryPolicy() throws InterruptedException, IOException {
    TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows = ImmutableList.of(wrapValue(new TableRow()), wrapValue(new TableRow()));
    // First time row0 fails with a retryable error, and row1 fails with a persistent error.
    final TableDataInsertAllResponse firstFailure = new TableDataInsertAllResponse().setInsertErrors(ImmutableList.of(new InsertErrors().setIndex(0L).setErrors(ImmutableList.of(new ErrorProto().setReason("timeout"))), new InsertErrors().setIndex(1L).setErrors(ImmutableList.of(new ErrorProto().setReason("invalid")))));
    // Second time there is only one row, which fails with a retryable error.
    final TableDataInsertAllResponse secondFialure = new TableDataInsertAllResponse().setInsertErrors(ImmutableList.of(new InsertErrors().setIndex(0L).setErrors(ImmutableList.of(new ErrorProto().setReason("timeout")))));
    // On the final attempt, no failures are returned.
    final TableDataInsertAllResponse allRowsSucceeded = new TableDataInsertAllResponse();
    setupMockResponses(response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        // Always return 200.
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(firstFailure));
    }, response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(secondFialure));
    }, response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(allRowsSucceeded));
    });
    DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, null, PipelineOptionsFactory.create());
    List<ValueInSingleWindow<TableRow>> failedInserts = Lists.newArrayList();
    dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.retryTransientErrors(), failedInserts, ErrorContainer.TABLE_ROW_ERROR_CONTAINER, false, false, false, null);
    assertEquals(1, failedInserts.size());
    expectedLogs.verifyInfo("Retrying 1 failed inserts to BigQuery");
    verifyWriteMetricWasSet("project", "dataset", "table", "timeout", 2);
}