Search in sources :

Example 1 with TableDataInsertAllResponse

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

the class BigQueryServicesImplTest method testInsertRetry.

/**
   * Tests that {@link DatasetServiceImpl#insertAll} retries quota rate limited attempts.
   */
@Test
public void testInsertRetry() throws Exception {
    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 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(new TableDataInsertAllResponse()));
    DatasetServiceImpl dataService = new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
    dataService.insertAll(ref, rows, null, BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()), new MockSleeper(), InsertRetryPolicy.alwaysRetry(), null);
    verify(response, times(2)).getStatusCode();
    verify(response, times(2)).getContent();
    verify(response, times(2)).getContentType();
    expectedLogs.verifyInfo("BigQuery insertAll exceeded rate limit, retrying");
}
Also used : TableReference(com.google.api.services.bigquery.model.TableReference) DatasetServiceImpl(org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.DatasetServiceImpl) TableRow(com.google.api.services.bigquery.model.TableRow) TableDataInsertAllResponse(com.google.api.services.bigquery.model.TableDataInsertAllResponse) ArrayList(java.util.ArrayList) ValueInSingleWindow(org.apache.beam.sdk.values.ValueInSingleWindow) MockSleeper(com.google.api.client.testing.util.MockSleeper) Test(org.junit.Test)

Example 2 with TableDataInsertAllResponse

use of com.google.api.services.bigquery.model.TableDataInsertAllResponse in project google-cloud-java by GoogleCloudPlatform.

the class InsertAllResponse method fromPb.

static InsertAllResponse fromPb(TableDataInsertAllResponse responsePb) {
    Map<Long, List<BigQueryError>> insertErrors = null;
    if (responsePb.getInsertErrors() != null) {
        List<InsertErrors> errorsPb = responsePb.getInsertErrors();
        insertErrors = Maps.newHashMapWithExpectedSize(errorsPb.size());
        for (InsertErrors errorPb : errorsPb) {
            insertErrors.put(errorPb.getIndex(), Lists.transform(errorPb.getErrors() != null ? errorPb.getErrors() : ImmutableList.<ErrorProto>of(), BigQueryError.FROM_PB_FUNCTION));
        }
    }
    return new InsertAllResponse(insertErrors);
}
Also used : InsertErrors(com.google.api.services.bigquery.model.TableDataInsertAllResponse.InsertErrors) List(java.util.List) ImmutableList(com.google.common.collect.ImmutableList) TableDataInsertAllResponse(com.google.api.services.bigquery.model.TableDataInsertAllResponse)

Example 3 with TableDataInsertAllResponse

use of com.google.api.services.bigquery.model.TableDataInsertAllResponse in project google-cloud-java by GoogleCloudPlatform.

the class BigQueryImplTest method testInsertAll.

@Test
public void testInsertAll() {
    Map<String, Object> row1 = ImmutableMap.<String, Object>of("field", "value1");
    Map<String, Object> row2 = ImmutableMap.<String, Object>of("field", "value2");
    List<RowToInsert> rows = ImmutableList.of(new RowToInsert("row1", row1), new RowToInsert("row2", row2));
    InsertAllRequest request = InsertAllRequest.newBuilder(TABLE_ID).setRows(rows).setSkipInvalidRows(false).setIgnoreUnknownValues(true).setTemplateSuffix("suffix").build();
    TableDataInsertAllRequest requestPb = new TableDataInsertAllRequest().setRows(Lists.transform(rows, new Function<RowToInsert, TableDataInsertAllRequest.Rows>() {

        @Override
        public TableDataInsertAllRequest.Rows apply(RowToInsert rowToInsert) {
            return new TableDataInsertAllRequest.Rows().setInsertId(rowToInsert.getId()).setJson(rowToInsert.getContent());
        }
    })).setSkipInvalidRows(false).setIgnoreUnknownValues(true).setTemplateSuffix("suffix");
    TableDataInsertAllResponse responsePb = new TableDataInsertAllResponse().setInsertErrors(ImmutableList.of(new TableDataInsertAllResponse.InsertErrors().setIndex(0L).setErrors(ImmutableList.of(new ErrorProto().setMessage("ErrorMessage")))));
    EasyMock.expect(bigqueryRpcMock.insertAll(PROJECT, DATASET, TABLE, requestPb)).andReturn(responsePb);
    EasyMock.replay(bigqueryRpcMock);
    bigquery = options.getService();
    InsertAllResponse response = bigquery.insertAll(request);
    assertNotNull(response.getErrorsFor(0L));
    assertNull(response.getErrorsFor(1L));
    assertEquals(1, response.getErrorsFor(0L).size());
    assertEquals("ErrorMessage", response.getErrorsFor(0L).get(0).getMessage());
}
Also used : ErrorProto(com.google.api.services.bigquery.model.ErrorProto) TableDataInsertAllResponse(com.google.api.services.bigquery.model.TableDataInsertAllResponse) TableDataInsertAllResponse(com.google.api.services.bigquery.model.TableDataInsertAllResponse) TableDataInsertAllRequest(com.google.api.services.bigquery.model.TableDataInsertAllRequest) RowToInsert(com.google.cloud.bigquery.InsertAllRequest.RowToInsert) TableDataInsertAllRequest(com.google.api.services.bigquery.model.TableDataInsertAllRequest) Test(org.junit.Test)

Example 4 with TableDataInsertAllResponse

use of com.google.api.services.bigquery.model.TableDataInsertAllResponse in project wildfly-camel by wildfly-extras.

the class GoogleBigQueryIntegrationTest method setUp.

@Before
public void setUp() throws Exception {
    configuration = new GoogleBigQueryConfiguration();
    bigquery = Mockito.mock(Bigquery.class);
    endpoint = Mockito.mock(GoogleBigQueryEndpoint.class);
    tabledata = Mockito.mock(Bigquery.Tabledata.class);
    mockInsertall = Mockito.mock(Bigquery.Tabledata.InsertAll.class);
    Mockito.when(bigquery.tabledata()).thenReturn(tabledata);
    Mockito.when(tabledata.insertAll(Mockito.anyString(), Mockito.anyString(), Mockito.anyString(), Mockito.any())).thenReturn(mockInsertall);
    TableDataInsertAllResponse mockResponse = new TableDataInsertAllResponse();
    Mockito.when(mockInsertall.execute()).thenReturn(mockResponse);
}
Also used : GoogleBigQueryEndpoint(org.apache.camel.component.google.bigquery.GoogleBigQueryEndpoint) Bigquery(com.google.api.services.bigquery.Bigquery) TableDataInsertAllResponse(com.google.api.services.bigquery.model.TableDataInsertAllResponse) GoogleBigQueryConfiguration(org.apache.camel.component.google.bigquery.GoogleBigQueryConfiguration) Before(org.junit.Before)

Example 5 with TableDataInsertAllResponse

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

the class TestBigQuery method insertRows.

@Experimental(Kind.SCHEMAS)
public TableDataInsertAllResponse insertRows(Schema rowSchema, Row... rows) throws IOException {
    List<Rows> bqRows = Arrays.stream(rows).map(row -> new Rows().setJson(BigQueryUtils.toTableRow(row))).collect(ImmutableList.toImmutableList());
    Bigquery bq = newBigQueryClient(pipelineOptions);
    return bq.tabledata().insertAll(pipelineOptions.getBigQueryProject() == null ? pipelineOptions.getProject() : pipelineOptions.getBigQueryProject(), pipelineOptions.getTargetDataset(), table.getTableReference().getTableId(), new TableDataInsertAllRequest().setRows(bqRows)).setPrettyPrint(false).execute();
}
Also used : NullCredentialInitializer(org.apache.beam.sdk.extensions.gcp.auth.NullCredentialInitializer) Statement(org.junit.runners.model.Statement) Arrays(java.util.Arrays) TestPipelineOptions(org.apache.beam.sdk.testing.TestPipelineOptions) TestRule(org.junit.rules.TestRule) Experimental(org.apache.beam.sdk.annotations.Experimental) Duration(org.joda.time.Duration) Rows(com.google.api.services.bigquery.model.TableDataInsertAllRequest.Rows) HttpRequestInitializer(com.google.api.client.http.HttpRequestInitializer) Kind(org.apache.beam.sdk.annotations.Experimental.Kind) Transport(org.apache.beam.sdk.extensions.gcp.util.Transport) ThreadLocalRandom(java.util.concurrent.ThreadLocalRandom) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) TableRow(com.google.api.services.bigquery.model.TableRow) TableSchema(com.google.api.services.bigquery.model.TableSchema) MatcherAssert.assertThat(org.hamcrest.MatcherAssert.assertThat) Row(org.apache.beam.sdk.values.Row) DateTimeFormat(org.joda.time.format.DateTimeFormat) TableReference(com.google.api.services.bigquery.model.TableReference) RetryHttpRequestInitializer(org.apache.beam.sdk.extensions.gcp.util.RetryHttpRequestInitializer) DateTimeFormatter(org.joda.time.format.DateTimeFormatter) DateTime(org.joda.time.DateTime) Credentials(com.google.auth.Credentials) IOException(java.io.IOException) FieldType(org.apache.beam.sdk.schemas.Schema.FieldType) Description(org.junit.runner.Description) Seconds.secondsBetween(org.joda.time.Seconds.secondsBetween) Collectors(java.util.stream.Collectors) Schema(org.apache.beam.sdk.schemas.Schema) Table(com.google.api.services.bigquery.model.Table) List(java.util.List) Bigquery(com.google.api.services.bigquery.Bigquery) TableDataInsertAllResponse(com.google.api.services.bigquery.model.TableDataInsertAllResponse) TableDataInsertAllRequest(com.google.api.services.bigquery.model.TableDataInsertAllRequest) HttpCredentialsAdapter(com.google.auth.http.HttpCredentialsAdapter) Matcher(org.hamcrest.Matcher) Instant(org.joda.time.Instant) ChainingHttpRequestInitializer(com.google.cloud.hadoop.util.ChainingHttpRequestInitializer) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) Collections(java.util.Collections) TableDataInsertAllRequest(com.google.api.services.bigquery.model.TableDataInsertAllRequest) Bigquery(com.google.api.services.bigquery.Bigquery) Rows(com.google.api.services.bigquery.model.TableDataInsertAllRequest.Rows) Experimental(org.apache.beam.sdk.annotations.Experimental)

Aggregations

TableDataInsertAllResponse (com.google.api.services.bigquery.model.TableDataInsertAllResponse)23 TableReference (com.google.api.services.bigquery.model.TableReference)18 TableRow (com.google.api.services.bigquery.model.TableRow)18 Test (org.junit.Test)18 MockSleeper (com.google.api.client.testing.util.MockSleeper)16 DatasetServiceImpl (org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.DatasetServiceImpl)16 FailsafeValueInSingleWindow (org.apache.beam.sdk.values.FailsafeValueInSingleWindow)14 ValueInSingleWindow (org.apache.beam.sdk.values.ValueInSingleWindow)10 ErrorProto (com.google.api.services.bigquery.model.ErrorProto)9 ArrayList (java.util.ArrayList)9 TableDataInsertAllRequest (com.google.api.services.bigquery.model.TableDataInsertAllRequest)8 InsertErrors (com.google.api.services.bigquery.model.TableDataInsertAllResponse.InsertErrors)8 Matchers.containsString (org.hamcrest.Matchers.containsString)8 Bigquery (com.google.api.services.bigquery.Bigquery)6 List (java.util.List)6 Table (com.google.api.services.bigquery.model.Table)5 TableSchema (com.google.api.services.bigquery.model.TableSchema)5 IOException (java.io.IOException)5 Transport (org.apache.beam.sdk.extensions.gcp.util.Transport)5 ImmutableList (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList)5