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

Example 31 with TableReference

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

the class WriteRename method startWriteRename.

private PendingJobData startWriteRename(TableDestination finalTableDestination, Iterable<WriteTables.Result> tempTableNames, ProcessContext c, BoundedWindow window) throws Exception {
    // The pane may have advanced either here due to triggering or due to an upstream trigger. We
    // check the upstream
    // trigger to handle the case where an earlier pane triggered the single-partition path. If this
    // happened, then the
    // table will already exist so we want to append to the table.
    boolean isFirstPane = Iterables.getFirst(tempTableNames, null).isFirstPane() && c.pane().isFirst();
    WriteDisposition writeDisposition = isFirstPane ? firstPaneWriteDisposition : WriteDisposition.WRITE_APPEND;
    CreateDisposition createDisposition = isFirstPane ? firstPaneCreateDisposition : CreateDisposition.CREATE_NEVER;
    List<TableReference> tempTables = StreamSupport.stream(tempTableNames.spliterator(), false).map(result -> BigQueryHelpers.fromJsonString(result.getTableName(), TableReference.class)).collect(Collectors.toList());
    // Make sure each destination table gets a unique job id.
    String jobIdPrefix = BigQueryResourceNaming.createJobIdWithDestination(c.sideInput(jobIdToken), finalTableDestination, -1, c.pane().getIndex());
    BigQueryHelpers.PendingJob retryJob = startCopy(bqServices.getJobService(c.getPipelineOptions().as(BigQueryOptions.class)), getDatasetService(c.getPipelineOptions().as(BigQueryOptions.class)), jobIdPrefix, finalTableDestination.getTableReference(), tempTables, writeDisposition, createDisposition, kmsKey, loadJobProjectId);
    return new PendingJobData(retryJob, finalTableDestination, tempTables, window);
}

Example 32 with TableReference

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

the class BigQueryServicesImplTest method testInsertRetrySelectRows.

/**
 * Tests that {@link DatasetServiceImpl#insertAll} retries selected rows on failure.
 */
@Test
public void testInsertRetrySelectRows() throws Exception {
    TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows = ImmutableList.of(wrapValue(new TableRow().set("row", "a")), wrapValue(new TableRow().set("row", "b")));
    List<String> insertIds = ImmutableList.of("a", "b");
    final TableDataInsertAllResponse bFailed = new TableDataInsertAllResponse().setInsertErrors(ImmutableList.of(new InsertErrors().setIndex(1L).setErrors(ImmutableList.of(new ErrorProto()))));
    final TableDataInsertAllResponse allRowsSucceeded = new TableDataInsertAllResponse();
    setupMockResponses(response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(bFailed));
    }, 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());
    dataService.insertAll(ref, rows, insertIds, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.alwaysRetry(), null, null, false, false, false, null);
    verifyAllResponsesAreRead();
    verifyWriteMetricWasSet("project", "dataset", "table", "unknown", 1);
    verifyWriteMetricWasSet("project", "dataset", "table", "ok", 1);
}

Example 33 with TableReference

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

the class BigQueryServicesImplTest method testSkipInvalidRowsIgnoreUnknownIgnoreInsertIdsValuesStreaming.

/**
 * Tests that {@link DatasetServiceImpl#insertAll} respects the skipInvalidRows,
 * ignoreUnknownValues and ignoreInsertIds parameters.
 */
@Test
public void testSkipInvalidRowsIgnoreUnknownIgnoreInsertIdsValuesStreaming() 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()));
    final TableDataInsertAllResponse allRowsSucceeded = new TableDataInsertAllResponse();
    // Return a 200 response each time
    MockSetupFunction allRowsSucceededResponseFunction = response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(allRowsSucceeded));
    };
    setupMockResponses(allRowsSucceededResponseFunction, allRowsSucceededResponseFunction);
    DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, null, PipelineOptionsFactory.create());
    // First, test with all flags disabled
    dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.neverRetry(), Lists.newArrayList(), ErrorContainer.TABLE_ROW_ERROR_CONTAINER, false, false, false, null);
    TableDataInsertAllRequest parsedRequest = fromString(request.getContentAsString(), TableDataInsertAllRequest.class);
    assertFalse(parsedRequest.getSkipInvalidRows());
    assertFalse(parsedRequest.getIgnoreUnknownValues());
    // Then with all enabled
    dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.neverRetry(), Lists.newArrayList(), ErrorContainer.TABLE_ROW_ERROR_CONTAINER, true, true, true, null);
    parsedRequest = fromString(request.getContentAsString(), TableDataInsertAllRequest.class);
    assertTrue(parsedRequest.getSkipInvalidRows());
    assertTrue(parsedRequest.getIgnoreUnknownValues());
    assertNull(parsedRequest.getRows().get(0).getInsertId());
    assertNull(parsedRequest.getRows().get(1).getInsertId());
    verifyWriteMetricWasSet("project", "dataset", "table", "ok", 2);
}

Example 34 with TableReference

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

the class BigQueryServicesImplTest method testCreateTableDoesNotRetry.

/**
 * Tests that {@link BigQueryServicesImpl} does not retry non-rate-limited attempts.
 */
@Test
public void testCreateTableDoesNotRetry() throws IOException {
    TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    Table testTable = new Table().setTableReference(ref);
    // First response is 403 not-rate-limited, 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(testTable));
    });
    thrown.expect(GoogleJsonResponseException.class);
    thrown.expectMessage("actually forbidden");
    BigQueryServicesImpl.DatasetServiceImpl services = new BigQueryServicesImpl.DatasetServiceImpl(bigquery, null, PipelineOptionsFactory.create());
    try {
        services.tryCreateTable(testTable, new RetryBoundedBackOff(BackOff.ZERO_BACKOFF, 3), Sleeper.DEFAULT);
        fail();
    } catch (IOException e) {
        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();
        throw e;
    }
}

Example 35 with TableReference

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

the class BigQueryServicesImplTest method testExtendedErrorRetrieval.

/**
 * Tests that {@link DatasetServiceImpl#insertAll} uses the supplied {@link ErrorContainer}.
 */
@Test
public void testExtendedErrorRetrieval() 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 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")))));
    final List<ValueInSingleWindow<BigQueryInsertError>> expected = ImmutableList.of(wrapErrorValue(new BigQueryInsertError(rows.get(0).getValue(), failures.getInsertErrors().get(0), ref)), wrapErrorValue(new BigQueryInsertError(rows.get(1).getValue(), failures.getInsertErrors().get(1), ref)));
    setupMockResponses(response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getContent()).thenReturn(toStream(failures));
    });
    DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, null, PipelineOptionsFactory.create());
    List<ValueInSingleWindow<BigQueryInsertError>> failedInserts = Lists.newArrayList();
    dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), TEST_BACKOFF, new MockSleeper(), InsertRetryPolicy.neverRetry(), failedInserts, ErrorContainer.BIG_QUERY_INSERT_ERROR_ERROR_CONTAINER, false, false, false, null);
    assertThat(failedInserts, is(expected));
}