Examples with BigQueryConnectorException com.google.cloud.bigquery.connector.common.BigQueryConnectorException used on opensource projects

Example 1 with BigQueryConnectorException

use of com.google.cloud.bigquery.connector.common.BigQueryConnectorException in project spark-bigquery-connector by GoogleCloudDataproc.

the class BigQueryDirectDataSourceWriterContext method commit.

/**
 * This function will determine, based on the WritingMode: if in IGNORE_INPUTS mode, no work is to
 * be done; otherwise all streams will be batch committed using the BigQuery Storage Write API,
 * and then: if in OVERWRITE mode, the overwriteDestinationWithTemporary function from
 * BigQueryClient will be called to replace the destination table with all the data from the
 * temporary table; if in ALL_ELSE mode no more work needs to be done.
 *
 * @see WritingMode
 * @see BigQueryClient#overwriteDestinationWithTemporary(TableId temporaryTableId, TableId
 *     destinationTableId)
 * @param messages the BigQueryWriterCommitMessage array returned by the BigQueryDataWriter's.
 */
@Override
public void commit(WriterCommitMessageContext[] messages) {
    if (writingMode.equals(WritingMode.IGNORE_INPUTS))
        return;
    logger.info("BigQuery DataSource writer {} committed with messages:\n{}", writeUUID, Arrays.toString(messages));
    BatchCommitWriteStreamsRequest.Builder batchCommitWriteStreamsRequest = BatchCommitWriteStreamsRequest.newBuilder().setParent(tablePathForBigQueryStorage);
    for (WriterCommitMessageContext message : messages) {
        batchCommitWriteStreamsRequest.addWriteStreams(((BigQueryDirectWriterCommitMessageContext) message).getWriteStreamName());
    }
    BatchCommitWriteStreamsResponse batchCommitWriteStreamsResponse = writeClient.batchCommitWriteStreams(batchCommitWriteStreamsRequest.build());
    if (!batchCommitWriteStreamsResponse.hasCommitTime()) {
        throw new BigQueryConnectorException("DataSource writer failed to batch commit its BigQuery write-streams");
    }
    logger.info("BigQuery DataSource writer has committed at time: {}", batchCommitWriteStreamsResponse.getCommitTime());
    if (writingMode.equals(WritingMode.OVERWRITE)) {
        Job overwriteJob = bigQueryClient.overwriteDestinationWithTemporary(tableToWrite.getTableId(), destinationTableId);
        BigQueryClient.waitForJob(overwriteJob);
        Preconditions.checkState(bigQueryClient.deleteTable(tableToWrite.getTableId()), new BigQueryConnectorException(String.format("Could not delete temporary table %s from BigQuery", tableToWrite)));
    }
}