1   /*
2    * Copyright 2010 The Apache Software Foundation
3    *
4    * Licensed to the Apache Software Foundation (ASF) under one
5    * or more contributor license agreements.  See the NOTICE file
6    * distributed with this work for additional information
7    * regarding copyright ownership.  The ASF licenses this file
8    * to you under the Apache License, Version 2.0 (the
9    * "License"); you may not use this file except in compliance
10   * with the License.  You may obtain a copy of the License at
11   *
12   *     http://www.apache.org/licenses/LICENSE-2.0
13   *
14   * Unless required by applicable law or agreed to in writing, software
15   * distributed under the License is distributed on an "AS IS" BASIS,
16   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17   * See the License for the specific language governing permissions and
18   * limitations under the License.
19   */
20  
21  package org.apache.hadoop.hbase.rest;
22  
23  import java.io.ByteArrayInputStream;
24  import java.io.IOException;
25  import java.io.StringWriter;
26  import java.util.Iterator;
27  import java.util.Random;
28  
29  import javax.xml.bind.JAXBContext;
30  import javax.xml.bind.JAXBException;
31  import javax.xml.bind.Marshaller;
32  import javax.xml.bind.Unmarshaller;
33  
34  import org.apache.commons.httpclient.Header;
35  import org.apache.hadoop.hbase.HColumnDescriptor;
36  import org.apache.hadoop.hbase.HTableDescriptor;
37  import org.apache.hadoop.hbase.KeyValue;
38  import org.apache.hadoop.hbase.client.HBaseAdmin;
39  import org.apache.hadoop.hbase.client.HTable;
40  import org.apache.hadoop.hbase.client.Put;
41  import org.apache.hadoop.hbase.rest.client.Client;
42  import org.apache.hadoop.hbase.rest.client.Cluster;
43  import org.apache.hadoop.hbase.rest.client.Response;
44  import org.apache.hadoop.hbase.rest.model.CellModel;
45  import org.apache.hadoop.hbase.rest.model.CellSetModel;
46  import org.apache.hadoop.hbase.rest.model.RowModel;
47  import org.apache.hadoop.hbase.rest.model.ScannerModel;
48  import org.apache.hadoop.hbase.util.Bytes;
49  
50  public class TestScannerResource extends HBaseRESTClusterTestBase {
51    static final String TABLE = "TestScannerResource";
52    static final String CFA = "a";
53    static final String CFB = "b";
54    static final String COLUMN_1 = CFA + ":1";
55    static final String COLUMN_2 = CFB + ":2";
56  
57    static int expectedRows1;
58    static int expectedRows2;
59  
60    Client client;
61    JAXBContext context;
62    Marshaller marshaller;
63    Unmarshaller unmarshaller;
64    HBaseAdmin admin;
65  
66    int insertData(String tableName, String column, double prob)
67        throws IOException {
68      Random rng = new Random();
69      int count = 0;
70      HTable table = new HTable(conf, tableName);
71      byte[] k = new byte[3];
72      byte [][] famAndQf = KeyValue.parseColumn(Bytes.toBytes(column));
73      for (byte b1 = 'a'; b1 < 'z'; b1++) {
74        for (byte b2 = 'a'; b2 < 'z'; b2++) {
75          for (byte b3 = 'a'; b3 < 'z'; b3++) {
76            if (rng.nextDouble() < prob) {
77              k[0] = b1;
78              k[1] = b2;
79              k[2] = b3;
80              Put put = new Put(k);
81              put.add(famAndQf[0], famAndQf[1], k);
82              table.put(put);
83              count++;
84            }
85          }
86        }
87      }
88      table.flushCommits();
89      return count;
90    }
91  
92    @Override
93    protected void setUp() throws Exception {
94      super.setUp();
95      context = JAXBContext.newInstance(
96          CellModel.class,
97          CellSetModel.class,
98          RowModel.class,
99          ScannerModel.class);
100     marshaller = context.createMarshaller();
101     unmarshaller = context.createUnmarshaller();
102     client = new Client(new Cluster().add("localhost", testServletPort));
103     admin = new HBaseAdmin(conf);
104     if (admin.tableExists(TABLE)) {
105       return;
106     }
107     HTableDescriptor htd = new HTableDescriptor(TABLE);
108     htd.addFamily(new HColumnDescriptor(CFA));
109     htd.addFamily(new HColumnDescriptor(CFB));
110     admin.createTable(htd);
111     expectedRows1 = insertData(TABLE, COLUMN_1, 1.0);
112     expectedRows2 = insertData(TABLE, COLUMN_2, 0.5);
113   }
114 
115   @Override
116   protected void tearDown() throws Exception {
117     client.shutdown();
118     super.tearDown();
119   }
120 
121   int countCellSet(CellSetModel model) {
122     int count = 0;
123     Iterator<RowModel> rows = model.getRows().iterator();
124     while (rows.hasNext()) {
125       RowModel row = rows.next();
126       Iterator<CellModel> cells = row.getCells().iterator();
127       while (cells.hasNext()) {
128         cells.next();
129         count++;
130       }
131     }
132     return count;
133   }
134 
135   void doTestSimpleScannerXML() throws IOException, JAXBException {
136     final int BATCH_SIZE = 5;
137     // new scanner
138     ScannerModel model = new ScannerModel();
139     model.setBatch(BATCH_SIZE);
140     model.addColumn(Bytes.toBytes(COLUMN_1));
141     StringWriter writer = new StringWriter();
142     marshaller.marshal(model, writer);
143     byte[] body = Bytes.toBytes(writer.toString());
144     Response response = client.put("/" + TABLE + "/scanner", MIMETYPE_XML,
145       body);
146     assertEquals(response.getCode(), 201);
147     String scannerURI = response.getLocation();
148     assertNotNull(scannerURI);
149 
150     // get a cell set
151     response = client.get(scannerURI, MIMETYPE_XML);
152     assertEquals(response.getCode(), 200);
153     CellSetModel cellSet = (CellSetModel)
154       unmarshaller.unmarshal(new ByteArrayInputStream(response.getBody()));
155     // confirm batch size conformance
156     assertEquals(countCellSet(cellSet), BATCH_SIZE);
157 
158     // delete the scanner
159     response = client.delete(scannerURI);
160     assertEquals(response.getCode(), 200);
161   }
162 
163   void doTestSimpleScannerPB() throws IOException {
164     final int BATCH_SIZE = 10;
165     // new scanner
166     ScannerModel model = new ScannerModel();
167     model.setBatch(BATCH_SIZE);
168     model.addColumn(Bytes.toBytes(COLUMN_1));
169     Response response = client.put("/" + TABLE + "/scanner",
170       MIMETYPE_PROTOBUF, model.createProtobufOutput());
171     assertEquals(response.getCode(), 201);
172     String scannerURI = response.getLocation();
173     assertNotNull(scannerURI);
174 
175     // get a cell set
176     response = client.get(scannerURI, MIMETYPE_PROTOBUF);
177     assertEquals(response.getCode(), 200);
178     CellSetModel cellSet = new CellSetModel();
179     cellSet.getObjectFromMessage(response.getBody());
180     // confirm batch size conformance
181     assertEquals(countCellSet(cellSet), BATCH_SIZE);
182 
183     // delete the scanner
184     response = client.delete(scannerURI);
185     assertEquals(response.getCode(), 200);
186   }
187 
188   void doTestSimpleScannerBinary() throws IOException {
189     // new scanner
190     ScannerModel model = new ScannerModel();
191     model.setBatch(1);
192     model.addColumn(Bytes.toBytes(COLUMN_1));
193     Response response = client.put("/" + TABLE + "/scanner",
194       MIMETYPE_PROTOBUF, model.createProtobufOutput());
195     assertEquals(response.getCode(), 201);
196     String scannerURI = response.getLocation();
197     assertNotNull(scannerURI);
198 
199     // get a cell
200     response = client.get(scannerURI, MIMETYPE_BINARY);
201     assertEquals(response.getCode(), 200);
202     // verify that data was returned
203     assertTrue(response.getBody().length > 0);
204     // verify that the expected X-headers are present
205     boolean foundRowHeader = false, foundColumnHeader = false,
206       foundTimestampHeader = false;
207     for (Header header: response.getHeaders()) {
208       if (header.getName().equals("X-Row")) {
209         foundRowHeader = true;
210       } else if (header.getName().equals("X-Column")) {
211         foundColumnHeader = true;
212       } else if (header.getName().equals("X-Timestamp")) {
213         foundTimestampHeader = true;
214       }
215     }
216     assertTrue(foundRowHeader);
217     assertTrue(foundColumnHeader);
218     assertTrue(foundTimestampHeader);
219 
220     // delete the scanner
221     response = client.delete(scannerURI);
222     assertEquals(response.getCode(), 200);
223   }
224 
225   int fullTableScan(ScannerModel model) throws IOException {
226     model.setBatch(100);
227     Response response = client.put("/" + TABLE + "/scanner",
228         MIMETYPE_PROTOBUF, model.createProtobufOutput());
229     assertEquals(response.getCode(), 201);
230     String scannerURI = response.getLocation();
231     assertNotNull(scannerURI);
232     int count = 0;
233     while (true) {
234       response = client.get(scannerURI, MIMETYPE_PROTOBUF);
235       assertTrue(response.getCode() == 200 || response.getCode() == 204);
236       if (response.getCode() == 200) {
237         CellSetModel cellSet = new CellSetModel();
238         cellSet.getObjectFromMessage(response.getBody());
239         Iterator<RowModel> rows = cellSet.getRows().iterator();
240         while (rows.hasNext()) {
241           RowModel row = rows.next();
242           Iterator<CellModel> cells = row.getCells().iterator();
243           while (cells.hasNext()) {
244             cells.next();
245             count++;
246           }
247         }
248       } else {
249         break;
250       }
251     }
252     // delete the scanner
253     response = client.delete(scannerURI);
254     assertEquals(response.getCode(), 200);
255     return count;
256   }
257 
258   void doTestFullTableScan() throws IOException {
259     ScannerModel model = new ScannerModel();
260     model.addColumn(Bytes.toBytes(COLUMN_1));
261     assertEquals(fullTableScan(model), expectedRows1);
262 
263     model = new ScannerModel();
264     model.addColumn(Bytes.toBytes(COLUMN_2));
265     assertEquals(fullTableScan(model), expectedRows2);
266   }
267 
268   public void testScannerResource() throws Exception {
269     doTestSimpleScannerXML();
270     doTestSimpleScannerPB();
271     doTestSimpleScannerBinary();
272     doTestFullTableScan();
273   }
274 }