How to setTimeout for JsonStreamWrite api

227 Views Asked by At

I am having a Java Spring boot application which uses JsonStreamWrite of Google BigQuery Storage Write API to write data to BigQuery. I wanted to timeout the stream write if it takes more than a minute to insert into BigQuery.

Here is my sample code

public void createWriteStream(String table, JsonArr arr) throws IOException, Descriptors.DescriptorValidationException, InterruptedException {
      BigQueryWriteClient bqClient = BigQueryWriteClient.create();
      WriteStream stream = WriteStream.newBuilder().setType(WriteStream.Type.COMMITTED).build();
      TableName tableName = TableName.of("ProjectId", "DataSet", table);
      CreateWriteStreamRequest createWriteStreamRequest =
              CreateWriteStreamRequest.newBuilder()
                      .setParent(tableName.toString())
                      .setWriteStream(stream)
                      .build();
      WriteStream writeStream = bqClient.createWriteStream(createWriteStreamRequest);

      JsonStreamWriter jsonStreamWriter = JsonStreamWriter
              .newBuilder(writeStream.getName(), writeStream.getTableSchema())
              .build();
     jsonStreamWriter.append(jsonArr);
  }

Does bigquery provide any such configuration to timeout the insert ?

1

There are 1 best solutions below

0
On

Try this code:

public void createWriteStream(String table, JsonArr arr) throws IOException, Descriptors.DescriptorValidationException, InterruptedException {

    BigQueryWriteSettings.Builder bigQueryWriteSettingsBuilder = BigQueryWriteSettings.newBuilder();
    bigQueryWriteSettingsBuilder
       .createWriteStreamSettings()
       .setRetrySettings(
            bigQueryWriteSettingsBuilder
               .createWriteStreamSettings()
               .getRetrySettings()
               .toBuilder()
               .setTotalTimeout(Duration.ofMinute(1))
               .build());
    BigQueryWriteSettings bigQueryWriteSettings = bigQueryWriteSettingsBuilder.build();
  
    BigQueryWriteClient bqClient = BigQueryWriteClient.create(bigQueryWriteSettingsBuilder);
       WriteStream stream = WriteStream.newBuilder().setType(WriteStream.Type.COMMITTED).build();
       TableName tableName = TableName.of("ProjectId", "DataSet", table);
       CreateWriteStreamRequest createWriteStreamRequest =
            CreateWriteStreamRequest.newBuilder()
                .setParent(tableName.toString())
                .setWriteStream(stream)
                .build();
       WriteStream writeStream = bqClient.createWriteStream(createWriteStreamRequest);
 
    JsonStreamWriter jsonStreamWriter = JsonStreamWriter
           .newBuilder(writeStream.getName(), writeStream.getTableSchema())
           .build();
    jsonStreamWriter.append(jsonArr);

   }

I referred to this documentation on how I implemented the BigQueryWriteSettings.