001/*
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *     http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing, software
013 * distributed under the License is distributed on an "AS IS" BASIS,
014 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015 * See the License for the specific language governing permissions and
016 * limitations under the License.
017 */
018package org.apache.hadoop.hbase.rest;
019
020import java.io.IOException;
021import java.io.OutputStream;
022import java.util.List;
023import org.apache.hadoop.hbase.Cell;
024import org.apache.hadoop.hbase.CellUtil;
025import org.apache.hadoop.hbase.client.Result;
026import org.apache.hadoop.hbase.client.ResultScanner;
027import org.apache.hadoop.hbase.rest.model.CellModel;
028import org.apache.hadoop.hbase.rest.model.CellSetModel;
029import org.apache.hadoop.hbase.rest.model.RowModel;
030import org.apache.hadoop.hbase.util.Bytes;
031import org.apache.yetus.audience.InterfaceAudience;
032import org.slf4j.Logger;
033import org.slf4j.LoggerFactory;
034
035import org.apache.hbase.thirdparty.javax.ws.rs.WebApplicationException;
036import org.apache.hbase.thirdparty.javax.ws.rs.core.StreamingOutput;
037
038@InterfaceAudience.Private
039public class ProtobufStreamingOutput implements StreamingOutput {
040  private static final Logger LOG = LoggerFactory.getLogger(ProtobufStreamingOutput.class);
041
042  private String contentType;
043  private ResultScanner resultScanner;
044  private int limit;
045  private int fetchSize;
046
047  protected ProtobufStreamingOutput(ResultScanner scanner, String type, int limit, int fetchSize) {
048    this.resultScanner = scanner;
049    this.contentType = type;
050    this.limit = limit;
051    this.fetchSize = fetchSize;
052    if (LOG.isTraceEnabled()) {
053      LOG.trace("Created StreamingOutput with content type = " + this.contentType + " user limit : "
054        + this.limit + " scan fetch size : " + this.fetchSize);
055    }
056  }
057
058  @Override
059  public void write(OutputStream outStream) throws IOException, WebApplicationException {
060    Result[] rowsToSend;
061    if (limit < fetchSize) {
062      rowsToSend = this.resultScanner.next(limit);
063      writeToStream(createModelFromResults(rowsToSend), this.contentType, outStream);
064    } else {
065      int count = limit;
066      while (count > 0) {
067        if (count < fetchSize) {
068          rowsToSend = this.resultScanner.next(count);
069        } else {
070          rowsToSend = this.resultScanner.next(this.fetchSize);
071        }
072        if (rowsToSend.length == 0) {
073          break;
074        }
075        count = count - rowsToSend.length;
076        writeToStream(createModelFromResults(rowsToSend), this.contentType, outStream);
077      }
078    }
079  }
080
081  private void writeToStream(CellSetModel model, String contentType, OutputStream outStream)
082    throws IOException {
083    byte[] objectBytes = model.createProtobufOutput();
084    outStream.write(Bytes.toBytes((short) objectBytes.length));
085    outStream.write(objectBytes);
086    outStream.flush();
087    if (LOG.isTraceEnabled()) {
088      LOG.trace("Wrote " + model.getRows().size() + " rows to stream successfully.");
089    }
090  }
091
092  private CellSetModel createModelFromResults(Result[] results) {
093    CellSetModel cellSetModel = new CellSetModel();
094    for (Result rs : results) {
095      byte[] rowKey = rs.getRow();
096      RowModel rModel = new RowModel(rowKey);
097      List<Cell> kvs = rs.listCells();
098      for (Cell kv : kvs) {
099        rModel.addCell(new CellModel(CellUtil.cloneFamily(kv), CellUtil.cloneQualifier(kv),
100          kv.getTimestamp(), CellUtil.cloneValue(kv)));
101      }
102      cellSetModel.addRow(rModel);
103    }
104    return cellSetModel;
105  }
106}