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

@Test
public void wrapBigQueryInsertErrorReturnsValidJSON() {
  TableRow testRow =
      new TableRow()
          .set(NAME_KEY, testPerson.name)
          .set(AGE_KEY, testPerson.age);
  InsertErrors insertErrors = new TableDataInsertAllResponse.InsertErrors();
  ErrorProto errorProto = new ErrorProto().setMessage(ERROR_MESSAGE);
  insertErrors.setErrors(ImmutableList.of(errorProto));
  TableReference tableReference = new TableReference();
  BigQueryInsertError bigQueryInsertError =
      new BigQueryInsertError(testRow.clone(), insertErrors, tableReference);
  String expected = GSON.toJson(testPerson);

  FailsafeElement<String, String> wrappedValue =
      TextToBigQueryStreaming.wrapBigQueryInsertError(bigQueryInsertError);
  String actualOriginalPayload = wrappedValue.getOriginalPayload();
  String actualPayload = wrappedValue.getPayload();
  String actualErrorMessage = wrappedValue.getErrorMessage();

  assertThat(actualOriginalPayload).isEqualTo(expected);
  assertThat(actualPayload).isEqualTo(expected);
  assertThat(actualErrorMessage).isEqualTo(GSON.toJson(insertErrors));
}
 

@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();
}
 

@Test
public void testDecodeEncodeEqual() throws Exception {
  BigQueryInsertError value =
      new BigQueryInsertError(
          new TableRow().setF(Collections.singletonList(new TableCell().setV("Value"))),
          new TableDataInsertAllResponse.InsertErrors()
              .setIndex(0L)
              .setErrors(
                  Collections.singletonList(
                      new ErrorProto()
                          .setReason("a Reason")
                          .setLocation("A location")
                          .setMessage("A message")
                          .setDebugInfo("The debug info"))),
          new TableReference()
              .setProjectId("dummy-project-id")
              .setDatasetId("dummy-dataset-id")
              .setTableId("dummy-table-id"));

  CoderProperties.coderDecodeEncodeEqual(TEST_CODER, value);
}
 

Map<TableRow, List<TableDataInsertAllResponse.InsertErrors>> getInsertErrors() {
  Map<TableRow, List<TableDataInsertAllResponse.InsertErrors>> parsedInsertErrors =
      Maps.newHashMap();
  synchronized (tables) {
    for (Map.Entry<String, List<String>> entry : this.insertErrors.entrySet()) {
      TableRow tableRow = BigQueryHelpers.fromJsonString(entry.getKey(), TableRow.class);
      List<TableDataInsertAllResponse.InsertErrors> allErrors = Lists.newArrayList();
      for (String errorsString : entry.getValue()) {
        allErrors.add(
            BigQueryHelpers.fromJsonString(
                errorsString, TableDataInsertAllResponse.InsertErrors.class));
      }
      parsedInsertErrors.put(tableRow, allErrors);
    }
  }
  return parsedInsertErrors;
}
 

@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);
}
 

/**
 * Generates a {@link InsertErrors} used by {@link BigQueryInsertError}.
 *
 * @param error string to be added to {@link BigQueryInsertError}
 */
private static InsertErrors getInsertErrors(String error) {
  InsertErrors insertErrors = new TableDataInsertAllResponse.InsertErrors();
  ErrorProto errorProto = new ErrorProto().setMessage(error);
  insertErrors.setErrors(Lists.newArrayList(errorProto));
  return insertErrors;
}
 

@Test
public void testFailuresNoRetryPolicy() throws Exception {
  TableRow row1 = new TableRow().set("name", "a").set("number", "1");
  TableRow row2 = new TableRow().set("name", "b").set("number", "2");
  TableRow row3 = new TableRow().set("name", "c").set("number", "3");

  TableDataInsertAllResponse.InsertErrors ephemeralError =
      new TableDataInsertAllResponse.InsertErrors()
          .setErrors(ImmutableList.of(new ErrorProto().setReason("timeout")));

  fakeDatasetService.failOnInsert(
      ImmutableMap.of(
          row1, ImmutableList.of(ephemeralError, ephemeralError),
          row2, ImmutableList.of(ephemeralError, ephemeralError)));

  p.apply(Create.of(row1, row2, row3))
      .apply(
          BigQueryIO.writeTableRows()
              .to("project-id:dataset-id.table-id")
              .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED)
              .withMethod(BigQueryIO.Write.Method.STREAMING_INSERTS)
              .withSchema(
                  new TableSchema()
                      .setFields(
                          ImmutableList.of(
                              new TableFieldSchema().setName("name").setType("STRING"),
                              new TableFieldSchema().setName("number").setType("INTEGER"))))
              .withTestServices(fakeBqServices)
              .withoutValidation());
  p.run();

  assertThat(
      fakeDatasetService.getAllRows("project-id", "dataset-id", "table-id"),
      containsInAnyOrder(row1, row2, row3));
}
 

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));
}
 

private TableDataInsertAllResponse.InsertErrors generateErrorAmongMany(
    int numErrors, String baseReason, String exceptionalReason) {
  // The retry policies are expected to search through the entire list of ErrorProtos to determine
  // whether to retry. Stick the exceptionalReason in a random position to exercise this.
  List<ErrorProto> errorProtos = Lists.newArrayListWithExpectedSize(numErrors);
  int exceptionalPosition = ThreadLocalRandom.current().nextInt(numErrors);
  for (int i = 0; i < numErrors; ++i) {
    ErrorProto error = new ErrorProto();
    error.setReason((i == exceptionalPosition) ? exceptionalReason : baseReason);
    errorProtos.add(error);
  }
  TableDataInsertAllResponse.InsertErrors errors = new TableDataInsertAllResponse.InsertErrors();
  errors.setErrors(errorProtos);
  return errors;
}
 

/**
 * Tests that {@link DatasetServiceImpl#insertAll} retries other non-rate-limited,
 * non-quota-exceeded attempts.
 */
@Test
public void testInsertOtherRetry() throws Throwable {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows = new ArrayList<>();
  rows.add(wrapValue(new TableRow()));

  // First response is 403 non-{rate-limited, quota-exceeded}, 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()));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.alwaysRetry(),
      null,
      null,
      false,
      false,
      false);
  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
  expectedLogs.verifyInfo("BigQuery insertAll error, retrying:");
}
 

/** Tests that {@link DatasetServiceImpl#insertAll} retries quota exceeded attempts. */
@Test
public void testInsertQuotaExceededRetry() throws Exception {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows = new ArrayList<>();
  rows.add(wrapValue(new TableRow()));

  // First response is 403 quota exceeded, 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("quotaExceeded", 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,
      null,
      false,
      false,
      false);
  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
  expectedLogs.verifyInfo("BigQuery insertAll error, retrying:");
}
 

/** Tests that {@link DatasetServiceImpl#insertAll} retries rate limited attempts. */
@Test
public void testInsertRateLimitRetry() throws Exception {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows = new ArrayList<>();
  rows.add(wrapValue(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,
      null,
      false,
      false,
      false);
  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
  expectedLogs.verifyInfo("BigQuery insertAll error, retrying:");
}
 

@Override
public BigQueryInsertError decode(InputStream inStream) throws IOException {
  TableDataInsertAllResponse.InsertErrors err =
      MAPPER.readValue(
          StringUtf8Coder.of().decode(inStream), TableDataInsertAllResponse.InsertErrors.class);
  TableRow row = TableRowJsonCoder.of().decode(inStream);
  TableReference ref = BigQueryHelpers.parseTableSpec(StringUtf8Coder.of().decode(inStream));
  return new BigQueryInsertError(row, err, ref);
}
 

/**
 * Cause a given {@link TableRow} object to fail when it's inserted. The errors link the list will
 * be returned on subsequent retries, and the insert will succeed when the errors run out.
 */
public void failOnInsert(
    Map<TableRow, List<TableDataInsertAllResponse.InsertErrors>> insertErrors) {
  synchronized (tables) {
    for (Map.Entry<TableRow, List<TableDataInsertAllResponse.InsertErrors>> entry :
        insertErrors.entrySet()) {
      List<String> errorStrings = Lists.newArrayList();
      for (TableDataInsertAllResponse.InsertErrors errors : entry.getValue()) {
        errorStrings.add(BigQueryHelpers.toJsonString(errors));
      }
      this.insertErrors.put(BigQueryHelpers.toJsonString(entry.getKey()), errorStrings);
    }
  }
}
 

@Test
public void testNeverRetry() {
  assertFalse(
      InsertRetryPolicy.neverRetry()
          .shouldRetry(new Context(new TableDataInsertAllResponse.InsertErrors())));
}
 

void add(
List<ValueInSingleWindow<T>> failedInserts,
TableDataInsertAllResponse.InsertErrors error,
TableReference ref,
ValueInSingleWindow<TableRow> tableRow);
 

@Test
public void testExtendedErrorRetrieval() throws Exception {
  TableRow row1 = new TableRow().set("name", "a").set("number", "1");
  TableRow row2 = new TableRow().set("name", "b").set("number", "2");
  TableRow row3 = new TableRow().set("name", "c").set("number", "3");
  String tableSpec = "project-id:dataset-id.table-id";

  TableDataInsertAllResponse.InsertErrors ephemeralError =
      new TableDataInsertAllResponse.InsertErrors()
          .setErrors(ImmutableList.of(new ErrorProto().setReason("timeout")));
  TableDataInsertAllResponse.InsertErrors persistentError =
      new TableDataInsertAllResponse.InsertErrors()
          .setErrors(Lists.newArrayList(new ErrorProto().setReason("invalidQuery")));

  fakeDatasetService.failOnInsert(
      ImmutableMap.of(
          row1, ImmutableList.of(ephemeralError, ephemeralError),
          row2, ImmutableList.of(ephemeralError, ephemeralError, persistentError)));

  PCollection<BigQueryInsertError> failedRows =
      p.apply(Create.of(row1, row2, row3))
          .apply(
              BigQueryIO.writeTableRows()
                  .to(tableSpec)
                  .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED)
                  .withMethod(BigQueryIO.Write.Method.STREAMING_INSERTS)
                  .withSchema(
                      new TableSchema()
                          .setFields(
                              ImmutableList.of(
                                  new TableFieldSchema().setName("name").setType("STRING"),
                                  new TableFieldSchema().setName("number").setType("INTEGER"))))
                  .withFailedInsertRetryPolicy(InsertRetryPolicy.retryTransientErrors())
                  .withTestServices(fakeBqServices)
                  .withoutValidation()
                  .withExtendedErrorInfo())
          .getFailedInsertsWithErr();

  // row2 finally fails with a non-retryable error, so we expect to see it in the collection of
  // failed rows.
  PAssert.that(failedRows)
      .containsInAnyOrder(
          new BigQueryInsertError(
              row2, persistentError, BigQueryHelpers.parseTableSpec(tableSpec)));
  p.run();

  // Only row1 and row3 were successfully inserted.
  assertThat(
      fakeDatasetService.getAllRows("project-id", "dataset-id", "table-id"),
      containsInAnyOrder(row1, row3));
}
 

@Test
public void testRetryPolicy() throws Exception {
  TableRow row1 = new TableRow().set("name", "a").set("number", "1");
  TableRow row2 = new TableRow().set("name", "b").set("number", "2");
  TableRow row3 = new TableRow().set("name", "c").set("number", "3");

  TableDataInsertAllResponse.InsertErrors ephemeralError =
      new TableDataInsertAllResponse.InsertErrors()
          .setErrors(ImmutableList.of(new ErrorProto().setReason("timeout")));
  TableDataInsertAllResponse.InsertErrors persistentError =
      new TableDataInsertAllResponse.InsertErrors()
          .setErrors(ImmutableList.of(new ErrorProto().setReason("invalidQuery")));

  fakeDatasetService.failOnInsert(
      ImmutableMap.of(
          row1, ImmutableList.of(ephemeralError, ephemeralError),
          row2, ImmutableList.of(ephemeralError, ephemeralError, persistentError)));

  PCollection<TableRow> failedRows =
      p.apply(Create.of(row1, row2, row3))
          .apply(
              BigQueryIO.writeTableRows()
                  .to("project-id:dataset-id.table-id")
                  .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED)
                  .withMethod(BigQueryIO.Write.Method.STREAMING_INSERTS)
                  .withSchema(
                      new TableSchema()
                          .setFields(
                              ImmutableList.of(
                                  new TableFieldSchema().setName("name").setType("STRING"),
                                  new TableFieldSchema().setName("number").setType("INTEGER"))))
                  .withFailedInsertRetryPolicy(InsertRetryPolicy.retryTransientErrors())
                  .withTestServices(fakeBqServices)
                  .withoutValidation())
          .getFailedInserts();
  // row2 finally fails with a non-retryable error, so we expect to see it in the collection of
  // failed rows.
  PAssert.that(failedRows).containsInAnyOrder(row2);
  p.run();

  // Only row1 and row3 were successfully inserted.
  assertThat(
      fakeDatasetService.getAllRows("project-id", "dataset-id", "table-id"),
      containsInAnyOrder(row1, row3));
}
 

public TableDataInsertAllResponse.InsertErrors getInsertErrors() {
  return errors;
}
 

public Context(TableDataInsertAllResponse.InsertErrors errors) {
  this.errors = errors;
}
 

public BigQueryInsertError(
    TableRow row, TableDataInsertAllResponse.InsertErrors error, TableReference table) {
  this.row = row;
  this.error = error;
  this.table = table;
}
 

@Test
public void testAlwaysRetry() {
  assertTrue(
      InsertRetryPolicy.alwaysRetry()
          .shouldRetry(new Context(new TableDataInsertAllResponse.InsertErrors())));
}
 

/** 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<ValueInSingleWindow<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(
          wrapValue(
              new BigQueryInsertError(
                  rows.get(0).getValue(), failures.getInsertErrors().get(0), ref)),
          wrapValue(
              new BigQueryInsertError(
                  rows.get(1).getValue(), failures.getInsertErrors().get(1), ref)));

  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, PipelineOptionsFactory.create());

  List<ValueInSingleWindow<BigQueryInsertError>> failedInserts = Lists.newArrayList();
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.neverRetry(),
      failedInserts,
      ErrorContainer.BIG_QUERY_INSERT_ERROR_ERROR_CONTAINER,
      false,
      false,
      false);

  assertThat(failedInserts, is(expected));
}
 

/** Tests that {@link DatasetServiceImpl#insertAll} uses the supplied {@link ErrorContainer}. */
@Test
public void testSimpleErrorRetrieval() throws InterruptedException, IOException {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<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")))));

  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, PipelineOptionsFactory.create());

  List<ValueInSingleWindow<TableRow>> failedInserts = Lists.newArrayList();
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.neverRetry(),
      failedInserts,
      ErrorContainer.TABLE_ROW_ERROR_CONTAINER,
      false,
      false,
      false);

  assertThat(failedInserts, is(rows));
}
 

/**
 * 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());
}
 

/**
 * 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<ValueInSingleWindow<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();

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  // Always return 200.
  when(response.getStatusCode()).thenReturn(200);
  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200).thenReturn(200);

  // First fail
  when(response.getContent())
      .thenReturn(toStream(firstFailure))
      .thenReturn(toStream(secondFialure))
      .thenReturn(toStream(allRowsSucceeded));

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

  List<ValueInSingleWindow<TableRow>> failedInserts = Lists.newArrayList();
  dataService.insertAll(
      ref,
      rows,
      null,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.retryTransientErrors(),
      failedInserts,
      ErrorContainer.TABLE_ROW_ERROR_CONTAINER,
      false,
      false,
      false);
  assertEquals(1, failedInserts.size());
  expectedLogs.verifyInfo("Retrying 1 failed inserts to BigQuery");
}
 

public TableDataInsertAllResponse.InsertErrors getError() {
  return error;
}
 

/** Tests that {@link DatasetServiceImpl#insertAll} fails gracefully when persistent issues. */
@Test
public void testInsertFailsGracefully() throws Exception {
  TableReference ref =
      new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
  List<ValueInSingleWindow<TableRow>> rows =
      ImmutableList.of(wrapValue(new TableRow()), wrapValue(new TableRow()));

  final TableDataInsertAllResponse row1Failed =
      new TableDataInsertAllResponse()
          .setInsertErrors(ImmutableList.of(new InsertErrors().setIndex(1L)));

  final TableDataInsertAllResponse row0Failed =
      new TableDataInsertAllResponse()
          .setInsertErrors(ImmutableList.of(new InsertErrors().setIndex(0L)));

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  // Always return 200.
  when(response.getStatusCode()).thenReturn(200);
  // Return row 1 failing, then we retry row 1 as row 0, and row 0 persistently fails.
  when(response.getContent())
      .thenReturn(toStream(row1Failed))
      .thenAnswer(invocation -> toStream(row0Failed));

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

  // Expect it to fail.
  try {
    dataService.insertAll(
        ref,
        rows,
        null,
        BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
        new MockSleeper(),
        InsertRetryPolicy.alwaysRetry(),
        null,
        null,
        false,
        false,
        false);
    fail();
  } catch (IOException e) {
    assertThat(e, instanceOf(IOException.class));
    assertThat(e.getMessage(), containsString("Insert failed:"));
    assertThat(e.getMessage(), containsString("[{\"index\":0}]"));
  }

  // Verify the exact number of retries as well as log messages.
  verify(response, times(4)).getStatusCode();
  verify(response, times(4)).getContent();
  verify(response, times(4)).getContentType();
  expectedLogs.verifyInfo("Retrying 1 failed inserts to BigQuery");
}
 

/** 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<ValueInSingleWindow<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();

  when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
  when(response.getStatusCode()).thenReturn(200).thenReturn(200);
  when(response.getContent())
      .thenReturn(toStream(bFailed))
      .thenReturn(toStream(allRowsSucceeded));

  DatasetServiceImpl dataService =
      new DatasetServiceImpl(bigquery, PipelineOptionsFactory.create());
  dataService.insertAll(
      ref,
      rows,
      insertIds,
      BackOffAdapter.toGcpBackOff(TEST_BACKOFF.backoff()),
      new MockSleeper(),
      InsertRetryPolicy.alwaysRetry(),
      null,
      null,
      false,
      false,
      false);
  verify(response, times(2)).getStatusCode();
  verify(response, times(2)).getContent();
  verify(response, times(2)).getContentType();
}
 

@Override
public <T> long insertAll(
    TableReference ref,
    List<ValueInSingleWindow<TableRow>> rowList,
    @Nullable List<String> insertIdList,
    InsertRetryPolicy retryPolicy,
    List<ValueInSingleWindow<T>> failedInserts,
    ErrorContainer<T> errorContainer,
    boolean skipInvalidRows,
    boolean ignoreUnknownValues,
    boolean ignoreInsertIds)
    throws IOException, InterruptedException {
  Map<TableRow, List<TableDataInsertAllResponse.InsertErrors>> insertErrors = getInsertErrors();
  synchronized (tables) {
    if (ignoreInsertIds) {
      insertIdList = null;
    }

    if (insertIdList != null) {
      assertEquals(rowList.size(), insertIdList.size());
    }

    long dataSize = 0;
    TableContainer tableContainer =
        getTableContainer(
            ref.getProjectId(),
            ref.getDatasetId(),
            BigQueryHelpers.stripPartitionDecorator(ref.getTableId()));
    for (int i = 0; i < rowList.size(); ++i) {
      TableRow row = rowList.get(i).getValue();
      List<TableDataInsertAllResponse.InsertErrors> allErrors = insertErrors.get(row);
      boolean shouldInsert = true;
      if (allErrors != null) {
        for (TableDataInsertAllResponse.InsertErrors errors : allErrors) {
          if (!retryPolicy.shouldRetry(new Context(errors))) {
            shouldInsert = false;
          }
        }
      }
      if (shouldInsert) {
        if (insertIdList == null) {
          dataSize += tableContainer.addRow(row, null);
        } else {
          dataSize += tableContainer.addRow(row, insertIdList.get(i));
        }
      } else {
        errorContainer.add(
            failedInserts, allErrors.get(allErrors.size() - 1), ref, rowList.get(i));
      }
    }
    return dataSize;
  }
}