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

Example 21 with TableDataInsertAllResponse

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

the class BigqueryClient method insertDataToTable.

/**
 * Inserts rows to a table using a BigQuery streaming write.
 */
public void insertDataToTable(String projectId, String datasetId, String tableName, List<Map<String, Object>> rows) throws IOException, InterruptedException {
    Sleeper sleeper = Sleeper.DEFAULT;
    BackOff backoff = BackOffAdapter.toGcpBackOff(BACKOFF_FACTORY.backoff());
    IOException lastException = null;
    do {
        if (lastException != null) {
            LOG.warn("Retrying insert table ({}) after exception", tableName, lastException);
        }
        try {
            List<Rows> dataRows = rows.stream().map(row -> new Rows().setJson(row)).collect(Collectors.toList());
            TableDataInsertAllResponse response = this.bqClient.tabledata().insertAll(projectId, datasetId, tableName, new TableDataInsertAllRequest().setRows(dataRows)).execute();
            if (response != null && (response.getInsertErrors() == null || response.getInsertErrors().isEmpty())) {
                LOG.info("Successfully inserted data into table : " + tableName);
                return;
            } else {
                if (response == null || response.getInsertErrors() == null) {
                    lastException = new IOException("Expected valid response from insert data job, but received null.");
                } else {
                    lastException = new IOException(String.format("Got insertion error (%s)", response.getInsertErrors().toString()));
                }
            }
        } catch (IOException e) {
            // ignore and retry
            lastException = e;
        }
    } while (BackOffUtils.next(sleeper, backoff));
    throw new RuntimeException(String.format("Unable to get BigQuery response after retrying %d times for table (%s)", MAX_QUERY_RETRIES, tableName), lastException);
}

Example 22 with TableDataInsertAllResponse

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

the class BigQueryServicesImplTest method testInsertStoppedRetry.

/**
 * Tests that {@link DatasetServiceImpl#insertAll} can stop quotaExceeded retry attempts.
 */
@Test
public void testInsertStoppedRetry() throws Exception {
    TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows = new ArrayList<>();
    rows.add(wrapValue(new TableRow()));
    MockSetupFunction quotaExceededResponse = response -> {
        when(response.getStatusCode()).thenReturn(403);
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getContent()).thenReturn(toStream(errorWithReasonAndStatus("quotaExceeded", 403)));
    };
    // Respond 403 four times, then valid payload.
    setupMockResponses(quotaExceededResponse, quotaExceededResponse, quotaExceededResponse, quotaExceededResponse, response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(new TableDataInsertAllResponse()));
    });
    thrown.expect(RuntimeException.class);
    // Google-http-client 1.39.1 and higher does not read the content of the response with error
    // status code. How can we ensure appropriate exception is thrown?
    thrown.expectMessage("quotaExceeded");
    DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, null, PipelineOptionsFactory.create());
    dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.alwaysRetry(), null, null, false, false, false, null);
    verifyAllResponsesAreRead();
    verifyWriteMetricWasSet("project", "dataset", "table", "quotaexceeded", 1);
}

Example 23 with TableDataInsertAllResponse

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

the class BigQueryServicesImplTest method testInsertQuotaExceededRetry.

/**
 * Tests that {@link DatasetServiceImpl#insertAll} retries quota exceeded attempts.
 */
@Test
public void testInsertQuotaExceededRetry() 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 quota exceeded, second response has valid payload.
    setupMockResponses(response -> {
        when(response.getStatusCode()).thenReturn(403);
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getContent()).thenReturn(toStream(errorWithReasonAndStatus("quotaExceeded", 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());
    dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.alwaysRetry(), null, null, false, false, false, null);
    verifyAllResponsesAreRead();
    expectedLogs.verifyInfo("BigQuery insertAll error, retrying:");
    verifyWriteMetricWasSet("project", "dataset", "table", "quotaexceeded", 1);
}