com.google.cloud.bigquery.BigQueryError Java Examples

/**
 * Parses the {@link InsertAllResponse} object and constructs a mapping of each record in @{@link Records} that were
 * tried to sink in BQ and the error type {@link BQInsertionRecordsErrorType}.
 *
 * @param records - records that were tried with BQ insertion
 * @param bqResponse - the status of insertion for all records as returned by BQ
 * @return - map of each record and the associated list of error types.
 */
public Map<Record, List<BQInsertionRecordsErrorType>> parseBQResponse(final Records records, final InsertAllResponse bqResponse) {
    if (!bqResponse.hasErrors()) {
        return Collections.emptyMap();
    }
    Map<Record, List<BQInsertionRecordsErrorType>> parsedRecords = new HashMap<>();
    Map<Long, List<BigQueryError>> insertErrorsMap = bqResponse.getInsertErrors();
    for (long recordIndex : insertErrorsMap.keySet()) {
        final Record message = records.getRecords().get((int) recordIndex);
        final List<BQInsertionRecordsErrorType> errorTypeList = new ArrayList<>();
        parsedRecords.put(message, errorTypeList);
        for (final BigQueryError err : insertErrorsMap.get(recordIndex)) {
            errorTypeList.add(ErrorTypeFactory.getErrorType(err.getReason(), err.getMessage()));
        } //end of for each row
    }
    return parsedRecords;
}
 

@Override
public String toString() {
    String errString = "InsertStatus:";
    if (cause.getErrors().isEmpty()) {
        return  errString + " NO_ERRORS";
    }
    for (List<BigQueryError> valueList : this.cause.getErrors().values()) {
        for (BigQueryError err : valueList) {
            if (err.getMessage() != "") {
                errString += err.getMessage() + ",";
            }
        }
    }
    return errString;
}
 

@Before
public void setUp() {
    tableId = TableId.of("test-dataset", "test-table");
    builder = InsertAllRequest.newBuilder(tableId);
    bqRow = new BQRowWithInsertId();
    BQErrorHandler errorHandlerInstance = new OOBErrorHandler(new DefaultLogWriter());
    sink = new BqSink(bigquery, tableId, new BQResponseParser(), errorHandlerInstance, bqRow);
    when(successfulResponse.hasErrors()).thenReturn(false);
    when(bigquery.insertAll(any())).thenReturn(successfulResponse);
    when(failureResponse.hasErrors()).thenReturn(true);
    insertErrors = new HashMap<>();
    List<BigQueryError> columnError = Arrays.asList(new BigQueryError("failed since type mismatched", "column location", "message"));
    insertErrors.put(0L, columnError);
    when(failureResponse.getInsertErrors()).thenReturn(insertErrors);
}
 

@Test(expected = BigQuery.BigQueryErrors.class)
public void failsOnInsertErrors() throws Throwable {
  when(response.getErrorsFor(0)).thenReturn(ImmutableList.of(new BigQueryError("", "", "")));

  try {
    output.apply(EMPTY_MESSAGE).join();
  } catch (CompletionException e) {
    throw e.getCause();
  }
}
 

@Test
public void canDetectError() {
  // mock error result
  BigQueryError error = new BigQueryError("reason", "location", "message");
  when(jobStatus.getError()).thenReturn(error);
  when(jobStatus.getExecutionErrors()).thenReturn(null);
  assertEquals(Optional.of(ImmutableList.of(error)), load().map(e -> e.errors));
}
 

@Test
public void canDetectExecutionErrors() {
  // mock partial error result
  BigQueryError error = new BigQueryError("reason", "location", "message");
  when(jobStatus.getError()).thenReturn(null);
  when(jobStatus.getExecutionErrors()).thenReturn(ImmutableList.of(error));
  assertEquals(Optional.of(ImmutableList.of(error)), load().map(e -> e.errors));
}
 

@Test(expected = RuntimeException.class)
public void shouldFailWhenJobTerminatesWithError() throws InterruptedException {
  when(bigQuery.getDataset(DATASET_ID)).thenReturn(mock(Dataset.class));

  when(bigQuery.create(any(JobInfo.class))).thenReturn(job);
  when(job.waitFor(any(RetryOption.class))).thenReturn(job);
  when(job.getStatus()).thenReturn(mock(JobStatus.class));
  when(job.getStatus().getError()).thenReturn(new BigQueryError("", "", "job error"));

  BigQueryOutput.create(() -> floBigQueryClient, TABLE_ID).provide(null).publish();
}
 

@Test(expected = StageException.class)
public void runQueryHasErrors() throws Exception {
  QueryJobConfiguration queryRequest = QueryJobConfiguration.newBuilder("SELECT * FROM [sample:table] LIMIT 1000")
      .setUseQueryCache(true)
      .setUseLegacySql(useLegacySql)
      .build();

  TableResult mockQueryResponse = mock(TableResult.class);
  Job mockJob = mock(Job.class);
  JobStatus mockJobStatus = mock(JobStatus.class);

  // First pretend we haven't finished running the query, second time around its completed.
  when(mockJob.isDone()).thenReturn(true);
  when(mockJob.getJobId()).thenReturn(jobId);

  when(mockJob.getQueryResults()).thenReturn(mockQueryResponse);
  when(mockJobStatus.getError()).thenReturn(new BigQueryError(
      "Some Error",
      "Some Location",
      "Some Error Message"
  ));
  when(mockJob.getStatus()).thenReturn(mockJobStatus);

  when(mockBigquery.create((JobInfo)any())).thenReturn(mockJob);
  when(mockBigquery.cancel(jobId)).thenReturn(true);

  BigQueryDelegate delegate = new BigQueryDelegate(mockBigquery, useLegacySql);

  ErrorCode code = null;
  try {
    delegate.runQuery(queryRequest, 1000, 1000);
  } catch (StageException e) {
    code = e.getErrorCode();
    throw e;
  } finally {
    assertEquals(Errors.BIGQUERY_02, code);
  }
}
 

/** Example of inserting rows into a table without running a load job. */
// [TARGET insertAll(InsertAllRequest)]
// [VARIABLE "my_dataset_name"]
// [VARIABLE "my_table_name"]
public InsertAllResponse insertAll(String datasetName, String tableName) {
  // [START bigquery_table_insert_rows]
  TableId tableId = TableId.of(datasetName, tableName);
  // Values of the row to insert
  Map<String, Object> rowContent = new HashMap<>();
  rowContent.put("booleanField", true);
  // Bytes are passed in base64
  rowContent.put("bytesField", "Cg0NDg0="); // 0xA, 0xD, 0xD, 0xE, 0xD in base64
  // Records are passed as a map
  Map<String, Object> recordsContent = new HashMap<>();
  recordsContent.put("stringField", "Hello, World!");
  rowContent.put("recordField", recordsContent);
  InsertAllResponse response =
      bigquery.insertAll(
          InsertAllRequest.newBuilder(tableId)
              .addRow(rowContent)
              // More rows can be added in the same RPC by invoking .addRow() on the builder.
              // You can also supply optional unique row keys to support de-duplication scenarios.
              .build());
  if (response.hasErrors()) {
    // If any of the insertions failed, this lets you inspect the errors
    for (Entry<Long, List<BigQueryError>> entry : response.getInsertErrors().entrySet()) {
      // inspect row error
    }
  }
  // [END bigquery_table_insert_rows]
  return response;
}
 

static BigQueryException convertToBigQueryException(BigQueryError error) {
    return new BigQueryException(UNKNOWN_CODE, error.getMessage(), error);
}
 

static BigQueryException convertToBigQueryException(BigQueryError error)
{
    return new BigQueryException(UNKNOWN_CODE, error.getMessage(), error);
}
 

public InsertStatus(boolean success, Map<Long, List<BigQueryError>> insertErrors) {
    this.success = success;
    this.cause = new BqInsertErrors(insertErrors);
}
 

public Map<Long, List<BigQueryError>> getErrors() {
    return errors;
}
 

private BigQueryErrors(List<BigQueryError> errors) {
  super(errors.toString());
  this.errors = errors;
}
 

@Test
public void testErrorInIngestingMultipleTables() throws Exception {
  List<Record> records = new ArrayList<>();

  Record record1 = createRecord(ImmutableMap.of("a", 1));
  Record record2 = createRecord(ImmutableMap.of("a",  2));
  Record record3 = createRecord(ImmutableMap.of("a",  3));

  record1.getHeader().setAttribute("table", "table1");
  record2.getHeader().setAttribute("table", "table2");
  record3.getHeader().setAttribute("table", "table3");

  records.add(record1);
  records.add(record2);
  records.add(record3);

  Answer<InsertAllResponse> insertAllResponseAnswer = invocationOnMock -> {
    InsertAllRequest request = (InsertAllRequest) (invocationOnMock.getArguments()[0]);
    InsertAllResponse response = PowerMockito.mock(InsertAllResponse.class);
    if (request.getTable().getTable().equals("table2")) {
      BigQueryError bigQueryError = PowerMockito.mock(BigQueryError.class);
      Mockito.doReturn("Error in bigquery").when(bigQueryError).getMessage();
      Mockito.doReturn("Error in bigquery").when(bigQueryError).getReason();
      Mockito.doReturn(
          ImmutableMap.of(0L, Collections.singletonList(bigQueryError))
      ).when(response).getInsertErrors();
      Mockito.doReturn(true).when(response).hasErrors();
    } else {
      Mockito.doReturn(Collections.emptyMap()).when(response).getInsertErrors();
      Mockito.doReturn(false).when(response).hasErrors();
    }
    return response;
  };

  PowerMockito.doAnswer(insertAllResponseAnswer).when(bigQuery).insertAll(Mockito.any(InsertAllRequest.class));
  PowerMockito.doAnswer((Answer<Table>) invocationOnMock -> Mockito.mock(Table.class))
      .when(bigQuery).getTable(Mockito.any(TableId.class));

  BigQueryTargetConfigBuilder configBuilder = new BigQueryTargetConfigBuilder();
  configBuilder.datasetEL("sample");
  configBuilder.tableNameEL("${record:attribute('table')}");
  TargetRunner targetRunner = createAndRunner(configBuilder.build(), records);

  Assert.assertEquals(1, targetRunner.getErrorRecords().size());
  for (Record errorRecord : targetRunner.getErrorRecords()) {
    String errorCode = errorRecord.getHeader().getErrorCode();
    Assert.assertNotNull(errorCode);
    Assert.assertEquals("table2", errorRecord.getHeader().getAttribute("table"));
  }

}
 

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }

    final String projectId = context.getProperty(PROJECT_ID).evaluateAttributeExpressions().getValue();
    final String dataset = context.getProperty(DATASET).evaluateAttributeExpressions(flowFile).getValue();
    final String tableName = context.getProperty(TABLE_NAME).evaluateAttributeExpressions(flowFile).getValue();

    final TableId tableId;
    if (StringUtils.isEmpty(projectId)) {
        tableId = TableId.of(dataset, tableName);
    } else {
        tableId = TableId.of(projectId, dataset, tableName);
    }

    try {

        InsertAllRequest.Builder request = InsertAllRequest.newBuilder(tableId);
        int nbrecord = 0;

        try (final InputStream in = session.read(flowFile)) {
            final RecordReaderFactory readerFactory = context.getProperty(RECORD_READER).asControllerService(RecordReaderFactory.class);
            try (final RecordReader reader = readerFactory.createRecordReader(flowFile, in, getLogger());) {
                Record currentRecord;
                while ((currentRecord = reader.nextRecord()) != null) {
                    request.addRow(convertMapRecord(currentRecord.toMap()));
                    nbrecord++;
                }
            }
        }

        request.setIgnoreUnknownValues(context.getProperty(IGNORE_UNKNOWN).evaluateAttributeExpressions(flowFile).asBoolean());
        request.setSkipInvalidRows(context.getProperty(SKIP_INVALID_ROWS).evaluateAttributeExpressions(flowFile).asBoolean());

        InsertAllResponse response = getCloudService().insertAll(request.build());

        final Map<String, String> attributes = new HashMap<>();

        if (response.hasErrors()) {
            getLogger().log(LogLevel.WARN, "Failed to insert {} of {} records into BigQuery {} table.", new Object[] { response.getInsertErrors().size(), nbrecord, tableName });
            if (getLogger().isDebugEnabled()) {
                for (long index : response.getInsertErrors().keySet()) {
                    for (BigQueryError e : response.getInsertErrors().get(index)) {
                        getLogger().log(LogLevel.DEBUG, "Failed to insert record #{}: {}", new Object[] { index, e.getMessage() });
                    }
                }
            }

            attributes.put(BigQueryAttributes.JOB_NB_RECORDS_ATTR, Long.toString(nbrecord - response.getInsertErrors().size()));

            flowFile = session.penalize(flowFile);
            flowFile = session.putAllAttributes(flowFile, attributes);
            session.transfer(flowFile, REL_FAILURE);
        } else {
            attributes.put(BigQueryAttributes.JOB_NB_RECORDS_ATTR, Long.toString(nbrecord));
            flowFile = session.putAllAttributes(flowFile, attributes);
            session.transfer(flowFile, REL_SUCCESS);
        }

    } catch (Exception ex) {
        getLogger().log(LogLevel.ERROR, ex.getMessage(), ex);
        flowFile = session.penalize(flowFile);
        session.transfer(flowFile, REL_FAILURE);
    }
}