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

Example 91 with TableReference

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

the class BigQueryUtilTest method testInsertAll.

@Test
public void testInsertAll() throws Exception {
    // Build up a list of indices to fail on each invocation. This should result in
    // 5 calls to insertAll.
    List<List<Long>> errorsIndices = new ArrayList<>();
    errorsIndices.add(Arrays.asList(0L, 5L, 10L, 15L, 20L));
    errorsIndices.add(Arrays.asList(0L, 2L, 4L));
    errorsIndices.add(Arrays.asList(0L, 2L));
    errorsIndices.add(new ArrayList<>());
    onInsertAll(errorsIndices);
    TableReference ref = BigQueryHelpers.parseTableSpec("project:dataset.table");
    DatasetServiceImpl datasetService = new DatasetServiceImpl(mockClient, null, options, 5);
    List<FailsafeValueInSingleWindow<TableRow, TableRow>> rows = new ArrayList<>();
    List<String> ids = new ArrayList<>();
    for (int i = 0; i < 25; ++i) {
        rows.add(FailsafeValueInSingleWindow.of(rawRow("foo", 1234), GlobalWindow.TIMESTAMP_MAX_VALUE, GlobalWindow.INSTANCE, PaneInfo.ON_TIME_AND_ONLY_FIRING, rawRow("foo", 1234)));
        ids.add("");
    }
    long totalBytes = datasetService.insertAll(ref, rows, ids, InsertRetryPolicy.alwaysRetry(), null, null, false, false, false, null);
    verifyInsertAll(5);
    // Each of the 25 rows has 1 byte for length and 30 bytes: '{"f":[{"v":"foo"},{"v":1234}]}'
    assertEquals("Incorrect byte count", 25L * 31L, totalBytes);
}

Example 92 with TableReference

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

the class BigQueryUtilTest method testTableGet.

@Test
public void testTableGet() throws InterruptedException, IOException {
    onTableGet(basicTableSchema());
    TableDataList dataList = new TableDataList().setTotalRows(0L);
    onTableList(dataList);
    BigQueryServicesImpl.DatasetServiceImpl services = new BigQueryServicesImpl.DatasetServiceImpl(mockClient, null, options);
    services.getTable(new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table"));
    verifyTableGet();
}

Example 93 with TableReference

use of com.google.api.services.bigquery.model.TableReference 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);
}

Example 94 with TableReference

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

the class BigQueryServicesImplTest method testCreateTableSucceeds.

@Test
public void testCreateTableSucceeds() throws IOException {
    TableReference ref = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    Table testTable = new Table().setTableReference(ref);
    setupMockResponses(response -> {
        when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
        when(response.getStatusCode()).thenReturn(200);
        when(response.getContent()).thenReturn(toStream(testTable));
    });
    BigQueryServicesImpl.DatasetServiceImpl services = new BigQueryServicesImpl.DatasetServiceImpl(bigquery, null, PipelineOptionsFactory.create());
    Table ret = services.tryCreateTable(testTable, new RetryBoundedBackOff(BackOff.ZERO_BACKOFF, 0), Sleeper.DEFAULT);
    assertEquals(testTable, ret);
    verifyAllResponsesAreRead();
}

Example 95 with TableReference

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

the class BigQueryHllSketchCompatibilityIT method prepareDatasetAndDataTables.

@BeforeClass
public static void prepareDatasetAndDataTables() throws Exception {
    BIGQUERY_CLIENT.createNewDataset(PROJECT_ID, DATASET_ID);
    TableSchema dataTableSchema = new TableSchema().setFields(Collections.singletonList(new TableFieldSchema().setName(DATA_FIELD_NAME).setType(DATA_FIELD_TYPE)));
    Table dataTableNonEmpty = new Table().setSchema(dataTableSchema).setTableReference(new TableReference().setProjectId(PROJECT_ID).setDatasetId(DATASET_ID).setTableId(DATA_TABLE_ID_NON_EMPTY));
    BIGQUERY_CLIENT.createNewTable(PROJECT_ID, DATASET_ID, dataTableNonEmpty);
    // Prepopulates dataTableNonEmpty with TEST_DATA
    List<Map<String, Object>> rows = TEST_DATA.stream().map(v -> Collections.singletonMap(DATA_FIELD_NAME, (Object) v)).collect(Collectors.toList());
    BIGQUERY_CLIENT.insertDataToTable(PROJECT_ID, DATASET_ID, DATA_TABLE_ID_NON_EMPTY, rows);
    Table dataTableEmpty = new Table().setSchema(dataTableSchema).setTableReference(new TableReference().setProjectId(PROJECT_ID).setDatasetId(DATASET_ID).setTableId(DATA_TABLE_ID_EMPTY));
    BIGQUERY_CLIENT.createNewTable(PROJECT_ID, DATASET_ID, dataTableEmpty);
}