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

Example 51 with TableRow

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

the class BigQueryServicesImplTest method testInsertDoesNotRetry.

/**
   * Tests that {@link DatasetServiceImpl#insertAll} does not retry non-rate-limited attempts.
   */
@Test
public void testInsertDoesNotRetry() throws Throwable {
    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 not-rate-limited, second response has valid payload but should not
    // be invoked.
    when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
    when(response.getStatusCode()).thenReturn(403).thenReturn(200);
    when(response.getContent()).thenReturn(toStream(errorWithReasonAndStatus("actually forbidden", 403))).thenReturn(toStream(new TableDataInsertAllResponse()));
    thrown.expect(GoogleJsonResponseException.class);
    thrown.expectMessage("actually forbidden");
    DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
    try {
        dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), new MockSleeper(), InsertRetryPolicy.alwaysRetry(), null);
        fail();
    } catch (RuntimeException e) {
        verify(response, times(1)).getStatusCode();
        verify(response, times(1)).getContent();
        verify(response, times(1)).getContentType();
        throw e.getCause();
    }
}

Example 52 with TableRow

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

the class BigQueryServicesImplTest method testIsTableEmptySucceeds.

@Test
public void testIsTableEmptySucceeds() throws Exception {
    TableReference tableRef = new TableReference().setProjectId("projectId").setDatasetId("datasetId").setTableId("tableId");
    TableDataList testDataList = new TableDataList().setRows(ImmutableList.of(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(testDataList));
    BigQueryServicesImpl.DatasetServiceImpl datasetService = new BigQueryServicesImpl.DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
    assertFalse(datasetService.isTableEmpty(tableRef, BackOff.ZERO_BACKOFF, Sleeper.DEFAULT));
    verify(response, times(2)).getStatusCode();
    verify(response, times(2)).getContent();
    verify(response, times(2)).getContentType();
}

Example 53 with TableRow

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

the class BigQueryIOTest method testBuildWriteWithSchema.

@Test
public void testBuildWriteWithSchema() {
    TableSchema schema = new TableSchema();
    BigQueryIO.Write<TableRow> write = BigQueryIO.<TableRow>write().to("foo.com:project:somedataset.sometable").withSchema(schema);
    checkWriteObject(write, "foo.com:project", "somedataset", "sometable", schema, CreateDisposition.CREATE_IF_NEEDED, WriteDisposition.WRITE_EMPTY, null);
}

Example 54 with TableRow

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

the class BigQueryIOTest method testWrite.

@Test
public void testWrite() throws Exception {
    BigQueryOptions bqOptions = TestPipeline.testingPipelineOptions().as(BigQueryOptions.class);
    bqOptions.setProject("defaultproject");
    bqOptions.setTempLocation(testFolder.newFolder("BigQueryIOTest").getAbsolutePath());
    FakeDatasetService datasetService = new FakeDatasetService();
    FakeBigQueryServices fakeBqServices = new FakeBigQueryServices().withJobService(new FakeJobService()).withDatasetService(datasetService);
    datasetService.createDataset("defaultproject", "dataset-id", "", "");
    Pipeline p = TestPipeline.create(bqOptions);
    p.apply(Create.of(new TableRow().set("name", "a").set("number", 1), new TableRow().set("name", "b").set("number", 2), new TableRow().set("name", "c").set("number", 3)).withCoder(TableRowJsonCoder.of())).apply(BigQueryIO.writeTableRows().to("dataset-id.table-id").withTableDescription(null).withCreateDisposition(CreateDisposition.CREATE_IF_NEEDED).withSchema(new TableSchema().setFields(ImmutableList.of(new TableFieldSchema().setName("name").setType("STRING"), new TableFieldSchema().setName("number").setType("INTEGER")))).withTestServices(fakeBqServices).withoutValidation());
    p.run();
    File tempDir = new File(bqOptions.getTempLocation());
    testNumFiles(tempDir, 0);
}

Example 55 with TableRow

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

the class BigQueryIOTest method testWriteValidatesDataset.

private void testWriteValidatesDataset(boolean unbounded) throws Exception {
    String projectId = "someproject";
    String datasetId = "somedataset";
    BigQueryOptions options = TestPipeline.testingPipelineOptions().as(BigQueryOptions.class);
    options.setProject(projectId);
    FakeBigQueryServices fakeBqServices = new FakeBigQueryServices().withJobService(new FakeJobService()).withDatasetService(new FakeDatasetService());
    Pipeline p = TestPipeline.create(options);
    TableReference tableRef = new TableReference();
    tableRef.setDatasetId(datasetId);
    tableRef.setTableId("sometable");
    PCollection<TableRow> tableRows;
    if (unbounded) {
        tableRows = p.apply(GenerateSequence.from(0)).apply(MapElements.via(new SimpleFunction<Long, TableRow>() {

            @Override
            public TableRow apply(Long input) {
                return null;
            }
        })).setCoder(TableRowJsonCoder.of());
    } else {
        tableRows = p.apply(Create.empty(TableRowJsonCoder.of()));
    }
    thrown.expect(RuntimeException.class);
    // Message will be one of following depending on the execution environment.
    thrown.expectMessage(Matchers.either(Matchers.containsString("Unable to confirm BigQuery dataset presence")).or(Matchers.containsString("BigQuery dataset not found for table")));
    tableRows.apply(BigQueryIO.writeTableRows().to(tableRef).withCreateDisposition(CreateDisposition.CREATE_IF_NEEDED).withSchema(new TableSchema()).withTestServices(fakeBqServices));
    p.run();
}