com.google.api.services.bigquery.model.TableDataInsertAllRequest Java Examples

@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.getProject(),
          pipelineOptions.getTargetDataset(),
          table.getTableReference().getTableId(),
          new TableDataInsertAllRequest().setRows(bqRows))
      .execute();
}
 

private void onInsertAll(List<List<Long>> errorIndicesSequence) throws Exception {
  when(mockClient.tabledata()).thenReturn(mockTabledata);

  final List<TableDataInsertAllResponse> responses = new ArrayList<>();
  for (List<Long> errorIndices : errorIndicesSequence) {
    List<TableDataInsertAllResponse.InsertErrors> errors = new ArrayList<>();
    for (long i : errorIndices) {
      TableDataInsertAllResponse.InsertErrors error =
          new TableDataInsertAllResponse.InsertErrors();
      error.setIndex(i);
    }
    TableDataInsertAllResponse response = new TableDataInsertAllResponse();
    response.setInsertErrors(errors);
    responses.add(response);
  }

  doAnswer(
          invocation -> {
            Bigquery.Tabledata.InsertAll mockInsertAll = mock(Bigquery.Tabledata.InsertAll.class);
            when(mockInsertAll.execute())
                .thenReturn(
                    responses.get(0),
                    responses
                        .subList(1, responses.size())
                        .toArray(new TableDataInsertAllResponse[responses.size() - 1]));
            return mockInsertAll;
          })
      .when(mockTabledata)
      .insertAll(anyString(), anyString(), anyString(), any(TableDataInsertAllRequest.class));
}
 

@Test
public void sendMessage() throws Exception {
    createProducer().process(createExchangeWithBody(new DefaultCamelContext(), new HashMap<>()));
    ArgumentCaptor<TableDataInsertAllRequest> dataCaptor = ArgumentCaptor.forClass(TableDataInsertAllRequest.class);
    Mockito.verify(tabledata).insertAll(Mockito.eq(TEST_PROJECT_ID), Mockito.eq(TEST_DATASET_ID), Mockito.eq(TEST_TABLE_ID), dataCaptor.capture());
    List<TableDataInsertAllRequest> requests = dataCaptor.getAllValues();
    Assert.assertEquals(1, requests.size());
    Assert.assertEquals(1, requests.get(0).getRows().size());
    Assert.assertNull(requests.get(0).getRows().get(0).getInsertId());
}
 

@Test
public void sendMessageWithTableId() throws Exception {
    Exchange exchange = createExchangeWithBody(new DefaultCamelContext(), new HashMap<>());
    exchange.getIn().setHeader(GoogleBigQueryConstants.TABLE_ID, "exchange_table_id");
    createProducer().process(exchange);
    ArgumentCaptor<TableDataInsertAllRequest> dataCaptor = ArgumentCaptor.forClass(TableDataInsertAllRequest.class);
    Mockito.verify(tabledata).insertAll(Mockito.eq(TEST_PROJECT_ID), Mockito.eq(TEST_DATASET_ID), Mockito.eq("exchange_table_id"), dataCaptor.capture());
    List<TableDataInsertAllRequest> requests = dataCaptor.getAllValues();
    Assert.assertEquals(1, requests.size());
    Assert.assertEquals(1, requests.get(0).getRows().size());
    Assert.assertNull(requests.get(0).getRows().get(0).getInsertId());
}
 

@Test
public void useAsInsertIdConfig() throws Exception {
    configuration.setUseAsInsertId("row1");
    Map<String, String> map = new HashMap<>();
    map.put("row1", "value1");
    createProducer().process(createExchangeWithBody(new DefaultCamelContext(), map));
    ArgumentCaptor<TableDataInsertAllRequest> dataCaptor = ArgumentCaptor.forClass(TableDataInsertAllRequest.class);
    Mockito.verify(tabledata).insertAll(Mockito.eq(TEST_PROJECT_ID), Mockito.eq(TEST_DATASET_ID), Mockito.eq(TEST_TABLE_ID), dataCaptor.capture());
    List<TableDataInsertAllRequest> requests = dataCaptor.getAllValues();
    Assert.assertEquals(1, requests.size());
    Assert.assertEquals(1, requests.get(0).getRows().size());
    Assert.assertEquals("value1", requests.get(0).getRows().get(0).getInsertId());
}
 

@Test
public void listOfMessages() throws Exception {
    List<Map<String, String>> messages = new ArrayList<>();
    messages.add(new HashMap<>());
    messages.add(new HashMap<>());
    createProducer().process(createExchangeWithBody(new DefaultCamelContext(), messages));
    ArgumentCaptor<TableDataInsertAllRequest> dataCaptor = ArgumentCaptor.forClass(TableDataInsertAllRequest.class);
    Mockito.verify(tabledata).insertAll(Mockito.eq(TEST_PROJECT_ID), Mockito.eq(TEST_DATASET_ID), Mockito.eq(TEST_TABLE_ID), dataCaptor.capture());
    List<TableDataInsertAllRequest> requests = dataCaptor.getAllValues();
    Assert.assertEquals(1, requests.size());
    Assert.assertEquals(2, requests.get(0).getRows().size());
}
 

/**
 * 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<ValueInSingleWindow<TableRow>> rows =
      ImmutableList.of(wrapValue(new TableRow()), wrapValue(new TableRow()));

  final TableDataInsertAllResponse allRowsSucceeded = new TableDataInsertAllResponse();

  // Return a 200 response each time
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContent())
      .thenReturn(toStream(allRowsSucceeded))
      .thenReturn(toStream(allRowsSucceeded));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());

  // First, test with all flags disabled
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.neverRetry(),
      Lists.newArrayList(),
      ErrorContainer.TABLE_ROW_ERROR_CONTAINER,
      false,
      false,
      false);

  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()),
      new MockSleeper(),
      InsertRetryPolicy.neverRetry(),
      Lists.newArrayList(),
      ErrorContainer.TABLE_ROW_ERROR_CONTAINER,
      true,
      true,
      true);

  parsedRequest = fromString(request.getContentAsString(), TableDataInsertAllRequest.class);

  assertTrue(parsedRequest.getSkipInvalidRows());
  assertTrue(parsedRequest.getIgnoreUnknownValues());
  assertNull(parsedRequest.getRows().get(0).getInsertId());
  assertNull(parsedRequest.getRows().get(1).getInsertId());
}
 

private void verifyInsertAll(int expectedRetries) throws IOException {
  verify(mockClient, times(expectedRetries)).tabledata();
  verify(mockTabledata, times(expectedRetries))
      .insertAll(anyString(), anyString(), anyString(), any(TableDataInsertAllRequest.class));
}
 

@Test
public void testExtract()
        throws Exception
{
    assumeThat(GCS_TEST_BUCKET, not(isEmptyOrNullString()));

    // Create source table
    String tableId = "data";
    String datasetId = BQ_TAG + "_extract_test";
    Dataset dataset = new Dataset().setDatasetReference(new DatasetReference()
            .setProjectId(gcpProjectId)
            .setDatasetId(datasetId));
    retryExecutor.run(() -> bq.datasets().insert(gcpProjectId, dataset)
            .execute());
    Table table = new Table().setTableReference(new TableReference()
            .setProjectId(gcpProjectId)
            .setTableId(tableId))
            .setSchema(new TableSchema()
                    .setFields(ImmutableList.of(
                            new TableFieldSchema().setName("foo").setType("STRING"),
                            new TableFieldSchema().setName("bar").setType("STRING")
                    )));
    retryExecutor.run(() -> bq.tables().insert(gcpProjectId, datasetId, table)
            .execute());

    // Populate source table
    TableDataInsertAllRequest content = new TableDataInsertAllRequest()
            .setRows(ImmutableList.of(
                    new TableDataInsertAllRequest.Rows().setJson(ImmutableMap.of(
                            "foo", "a",
                            "bar", "b")),
                    new TableDataInsertAllRequest.Rows().setJson(ImmutableMap.of(
                            "foo", "c",
                            "bar", "d"))));
    retryExecutor.run(() -> bq.tabledata().insertAll(gcpProjectId, datasetId, tableId, content)
            .execute());

    // Run extract
    String objectName = GCS_PREFIX + "test.csv";
    addWorkflow(projectDir, "acceptance/bigquery/extract.dig");
    Id attemptId = pushAndStart(server.endpoint(), projectDir, "extract", ImmutableMap.of(
            "src_dataset", datasetId,
            "src_table", tableId,
            "dst_bucket", GCS_TEST_BUCKET,
            "dst_object", objectName,
            "outfile", outfile.toString()));
    expect(Duration.ofMinutes(5), attemptSuccess(server.endpoint(), attemptId));
    assertThat(Files.exists(outfile), is(true));

    // Check that destination file was created
    StorageObject metadata = retryExecutor.run(() -> gcs.objects().get(GCS_TEST_BUCKET, objectName)
            .execute());
    assertThat(metadata.getName(), is(objectName));
    ByteArrayOutputStream data = new ByteArrayOutputStream();
    retryExecutor.run(() -> {
        try {
            gcs.objects().get(GCS_TEST_BUCKET, objectName)
                    .executeMediaAndDownloadTo(data);
        }
        catch (IOException e) {
            throw Throwables.propagate(e);
        }
    });
}