Examples with StreamAppendClient org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.StreamAppendClient used on opensource projects

Example 1 with StreamAppendClient

use of org.apache.beam.sdk.io.gcp.bigquery.BigQueryServices.StreamAppendClient in project beam by apache.

the class FakeDatasetService method getStreamAppendClient.

@Override
public StreamAppendClient getStreamAppendClient(String streamName, Descriptor descriptor) {
    return new StreamAppendClient() {

        private Descriptor protoDescriptor;

        {
            this.protoDescriptor = descriptor;
        }

        @Override
        public ApiFuture<AppendRowsResponse> appendRows(long offset, ProtoRows rows) throws Exception {
            synchronized (FakeDatasetService.class) {
                Stream stream = writeStreams.get(streamName);
                if (stream == null) {
                    throw new RuntimeException("No such stream: " + streamName);
                }
                List<TableRow> tableRows = Lists.newArrayListWithExpectedSize(rows.getSerializedRowsCount());
                for (ByteString bytes : rows.getSerializedRowsList()) {
                    DynamicMessage msg = DynamicMessage.parseFrom(protoDescriptor, bytes);
                    if (msg.getUnknownFields() != null && !msg.getUnknownFields().asMap().isEmpty()) {
                        throw new RuntimeException("Unknown fields set in append! " + msg.getUnknownFields());
                    }
                    tableRows.add(TableRowToStorageApiProto.tableRowFromMessage(DynamicMessage.parseFrom(protoDescriptor, bytes)));
                }
                stream.appendRows(offset, tableRows);
            }
            return ApiFutures.immediateFuture(AppendRowsResponse.newBuilder().build());
        }

        @Override
        public void close() throws Exception {
        }

        @Override
        public void pin() {
        }

        @Override
        public void unpin() throws Exception {
        }
    };
}