Java tutorial
/** * This file is part of riak-java-pb-client * * Copyright (c) 2010 by Trifork * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. **/ package com.basho.riak.pbc; import java.io.DataOutputStream; import java.io.IOException; import java.net.InetAddress; import java.security.NoSuchAlgorithmException; import java.security.SecureRandom; import java.util.Collections; import java.util.Date; import java.util.HashMap; import java.util.Map; import java.util.UUID; import java.util.prefs.BackingStoreException; import java.util.prefs.Preferences; import org.apache.commons.codec.binary.Base64; import org.json.JSONObject; import com.basho.riak.client.http.util.Constants; import com.basho.riak.client.util.CharsetUtils; import com.basho.riak.protobuf.RiakKvPB; import com.basho.riak.protobuf.RiakPB; import com.basho.riak.protobuf.RiakKvPB.RpbDelReq; import com.basho.riak.protobuf.RiakKvPB.RpbGetClientIdResp; import com.basho.riak.protobuf.RiakKvPB.RpbGetReq; import com.basho.riak.protobuf.RiakKvPB.RpbGetResp; import com.basho.riak.protobuf.RiakPB.RpbGetServerInfoResp; import com.basho.riak.protobuf.RiakKvPB.RpbListBucketsResp; import com.basho.riak.protobuf.RiakKvPB.RpbMapRedReq; import com.basho.riak.protobuf.RiakKvPB.RpbPutReq; import com.basho.riak.protobuf.RiakKvPB.RpbPutResp; import com.google.protobuf.ByteString; import java.util.ArrayList; import java.util.List; /** * Low level protocol buffers client. */ public class RiakClient implements RiakMessageCodes { private static final int BUFFER_SIZE_KB; static { BUFFER_SIZE_KB = Integer.parseInt(System.getProperty("com.basho.riak.client.pbc,buffer", "16")); } private static final RiakObject[] NO_RIAK_OBJECTS = new RiakObject[0]; private static final ByteString[] NO_BYTE_STRINGS = new ByteString[0]; //private static final String[] NO_STRINGS = new String[0]; //private static final MapReduceResponse[] NO_MAP_REDUCE_RESPONSES = new MapReduceResponse[0]; private String node; private String serverVersion; /** * if this has been set (or gotten) then it will be applied to new * connections */ private volatile byte[] clientId; private final RiakConnectionPool pool; public RiakClient(String host) throws IOException { this(host, RiakConnection.DEFAULT_RIAK_PB_PORT); } public RiakClient(String host, int port) throws IOException { this(InetAddress.getByName(host), port); } public RiakClient(RiakConnectionPool pool) { this.pool = pool; } public RiakClient(InetAddress addr, int port) throws IOException { this.pool = new RiakConnectionPool(0, RiakConnectionPool.LIMITLESS, addr, port, 1000, BUFFER_SIZE_KB, 1000, 0); this.pool.start(); } public RiakClient(String host, int port, int bufferSizeKb) throws IOException { this.pool = new RiakConnectionPool(0, RiakConnectionPool.LIMITLESS, InetAddress.getByName(host), port, 1000, bufferSizeKb, 1000, 0); this.pool.start(); } RiakConnection getConnection() throws IOException { RiakConnection c = pool.getConnection(clientId); return c; } void release(RiakConnection c) { pool.releaseConnection(c); } /** * helper method to use a reasonable default client id * beware, it caches the client id. If you call it multiple times on the same client * you get the *same* id (not good for reusing a client with different ids) * * @throws IOException */ public void prepareClientID() throws IOException { Preferences prefs = Preferences.userNodeForPackage(RiakClient.class); String clid = prefs.get("client_id", null); if (clid == null) { SecureRandom sr; try { sr = SecureRandom.getInstance("SHA1PRNG"); // Not totally secure, but doesn't need to be // and 100% less prone to 30 second hangs on linux jdk5 sr.setSeed(UUID.randomUUID().getLeastSignificantBits() + new Date().getTime()); } catch (NoSuchAlgorithmException e) { throw new RuntimeException(e); } byte[] data = new byte[6]; sr.nextBytes(data); clid = CharsetUtils.asString(Base64.encodeBase64Chunked(data), CharsetUtils.ISO_8859_1); prefs.put("client_id", clid); try { prefs.flush(); } catch (BackingStoreException e) { throw new IOException(e.toString()); } } setClientID(clid); } public void ping() throws IOException { RiakConnection c = getConnection(); try { c.send(MSG_PingReq); c.receive_code(MSG_PingResp); } finally { release(c); } } /** * Warning: the riak client id is 4 bytes. This method silently truncates anymore bytes than that. * Be aware that if you have two client Ids, "boris1" and "boris2" this method will leave you with 1 client id, "bori". * Use {@link RiakClient#prepareClientID()} to generate a reasonably unique Id. * @see RiakClient#prepareClientID() * @param id * @throws IOException */ public void setClientID(String id) throws IOException { if (id == null || id.length() < Constants.RIAK_CLIENT_ID_LENGTH) { throw new IllegalArgumentException( "Client ID must be at least " + Constants.RIAK_CLIENT_ID_LENGTH + " bytes long"); } setClientID(ByteString.copyFrom(CharsetUtils.utf8StringToBytes(id), 0, Constants.RIAK_CLIENT_ID_LENGTH)); } // ///////////////////// public void setClientID(ByteString id) throws IOException { if (id.size() > Constants.RIAK_CLIENT_ID_LENGTH) { id = ByteString.copyFrom(id.toByteArray(), 0, Constants.RIAK_CLIENT_ID_LENGTH); } this.clientId = id.toByteArray(); } public String getClientID() throws IOException { RiakConnection c = getConnection(); try { c.send(MSG_GetClientIdReq); byte[] data = c.receive(MSG_GetClientIdResp); if (data == null) return null; RpbGetClientIdResp res = RiakKvPB.RpbGetClientIdResp.parseFrom(data); clientId = res.getClientId().toByteArray(); return CharsetUtils.asUTF8String(clientId); } finally { release(c); } } public Map<String, String> getServerInfo() throws IOException { RiakConnection c = getConnection(); try { c.send(MSG_GetServerInfoReq); byte[] data = c.receive(MSG_GetServerInfoResp); if (data == null) return Collections.emptyMap(); RpbGetServerInfoResp res = RiakPB.RpbGetServerInfoResp.parseFrom(data); if (res.hasNode()) { this.node = res.getNode().toStringUtf8(); } if (res.hasServerVersion()) { this.serverVersion = res.getServerVersion().toStringUtf8(); } Map<String, String> result = new HashMap<String, String>(); result.put("node", node); result.put("server_version", serverVersion); return result; } finally { release(c); } } // ///////////////////// public RiakObject[] fetch(String bucket, String key, int readQuorum) throws IOException { return fetch(ByteString.copyFromUtf8(bucket), ByteString.copyFromUtf8(key), readQuorum); } public RiakObject[] fetch(ByteString bucket, ByteString key, int readQuorum) throws IOException { RpbGetReq req = RiakKvPB.RpbGetReq.newBuilder().setBucket(bucket).setKey(key).setR(readQuorum).build(); RiakConnection c = getConnection(); try { c.send(MSG_GetReq, req); return processFetchReply(c, bucket, key).getObjects(); } finally { release(c); } } // All the fetch parameters public FetchResponse fetch(String bucket, String key, FetchMeta fetchMeta) throws IOException { return fetch(ByteString.copyFromUtf8(bucket), ByteString.copyFromUtf8(key), fetchMeta); } // All the fetch parameters public FetchResponse fetch(ByteString bucket, ByteString key, FetchMeta fetchMeta) throws IOException { RpbGetReq.Builder b = RiakKvPB.RpbGetReq.newBuilder().setBucket(bucket).setKey(key); fetchMeta.write(b); RiakConnection c = getConnection(); try { c.send(MSG_GetReq, b.build()); return processFetchReply(c, bucket, key); } finally { release(c); } } public RiakObject[] fetch(String bucket, String key) throws IOException { return fetch(ByteString.copyFromUtf8(bucket), ByteString.copyFromUtf8(key)); } public RiakObject[] fetch(ByteString bucket, ByteString key) throws IOException { RpbGetReq req = RiakKvPB.RpbGetReq.newBuilder().setBucket(bucket).setKey(key).build(); RiakConnection c = getConnection(); try { c.send(MSG_GetReq, req); return processFetchReply(c, bucket, key).getObjects(); } finally { release(c); } } private FetchResponse processFetchReply(RiakConnection c, ByteString bucket, ByteString key) throws IOException { byte[] rep = c.receive(MSG_GetResp); if (rep == null) { return new FetchResponse(NO_RIAK_OBJECTS, false, null); } RpbGetResp resp = RiakKvPB.RpbGetResp.parseFrom(rep); int count = resp.getContentCount(); // To unify the behavior of having just a tombstone vs. siblings // that include a tombstone, we create an empty object and mark // it deleted if (count == 0) { RiakKvPB.RpbGetResp.Builder responseBuilder = resp.toBuilder(); RiakKvPB.RpbContent.Builder contentBuilder = RiakKvPB.RpbContent.getDefaultInstance().toBuilder(); contentBuilder.setDeleted(true).setValue(ByteString.EMPTY); resp = responseBuilder.addContent(contentBuilder.build()).build(); count = 1; } RiakObject[] out = new RiakObject[count]; ByteString vclock = resp.getVclock(); for (int i = 0; i < count; i++) { out[i] = new RiakObject(vclock, bucket, key, resp.getContent(i)); } boolean unchanged = resp.getUnchanged(); return new FetchResponse(out, unchanged, vclock.toByteArray()); } // ///////////////////// /** * * @param bucket * the bucket name * @param indexName * the name of the index (e.g. 'user_bin') * @param value * the index value * @return List<String> * list of keys for the index * @throws IOException */ public List<String> index(String bucket, String indexName, String value) throws IOException { RiakKvPB.RpbIndexReq req = RiakKvPB.RpbIndexReq.newBuilder().setBucket(ByteString.copyFromUtf8(bucket)) .setIndex(ByteString.copyFromUtf8(indexName)).setKey(ByteString.copyFromUtf8(value)) .setQtype(RiakKvPB.RpbIndexReq.IndexQueryType.eq).build(); RiakConnection c = getConnection(); try { c.send(MSG_IndexReq, req); return processIndexReply(c); } finally { release(c); } } /** * * @param bucket * the bucket name * @param indexName * the name of the index (e.g. 'user_bin') * @param start * the start value in a range (e.g 'a') * @param end * the end value in a range (e.g. 'z') * @return * a List<String> of keys for this index range * @throws IOException */ public List<String> index(String bucket, String indexName, String start, String end) throws IOException { RiakKvPB.RpbIndexReq req = RiakKvPB.RpbIndexReq.newBuilder().setBucket(ByteString.copyFromUtf8(bucket)) .setIndex(ByteString.copyFromUtf8(indexName)).setRangeMin(ByteString.copyFromUtf8(start)) .setRangeMax(ByteString.copyFromUtf8(end)).setQtype(RiakKvPB.RpbIndexReq.IndexQueryType.range) .build(); RiakConnection c = getConnection(); try { c.send(MSG_IndexReq, req); return processIndexReply(c); } finally { release(c); } } /** * * @param bucket * the bucket name * @param indexName * the index (e.g. 'age_int') * @param value * a long for the index value * @return * a List<String> of keys for this index value * @throws IOException */ public List<String> index(String bucket, String indexName, long value) throws IOException { RiakKvPB.RpbIndexReq req = RiakKvPB.RpbIndexReq.newBuilder().setBucket(ByteString.copyFromUtf8(bucket)) .setIndex(ByteString.copyFromUtf8(indexName)).setKey(ByteString.copyFromUtf8(String.valueOf(value))) .setQtype(RiakKvPB.RpbIndexReq.IndexQueryType.eq).build(); RiakConnection c = getConnection(); try { c.send(MSG_IndexReq, req); return processIndexReply(c); } finally { release(c); } } /** * * @param bucket * the bucket name * @param indexName * the index (e.g. 'age_int') * @param start * the start value in a range (e.g 16) * @param end * the end value in a range (e.g. 32) * @return * List<String> of the keys for this index range * @throws IOException */ public List<String> index(String bucket, String indexName, long start, long end) throws IOException { RiakKvPB.RpbIndexReq req = RiakKvPB.RpbIndexReq.newBuilder().setBucket(ByteString.copyFromUtf8(bucket)) .setIndex(ByteString.copyFromUtf8(indexName)) .setRangeMin(ByteString.copyFromUtf8(String.valueOf(start))) .setRangeMax(ByteString.copyFromUtf8(String.valueOf(end))) .setQtype(RiakKvPB.RpbIndexReq.IndexQueryType.range).build(); RiakConnection c = getConnection(); try { c.send(MSG_IndexReq, req); return processIndexReply(c); } finally { release(c); } } private List<String> processIndexReply(RiakConnection c) throws IOException { byte[] rep = c.receive(MSG_IndexResp); if (null == rep) { return Collections.EMPTY_LIST; } RiakKvPB.RpbIndexResp resp = RiakKvPB.RpbIndexResp.parseFrom(rep); List<String> keys = new ArrayList<String>(resp.getKeysCount()); for (ByteString bs : resp.getKeysList()) { keys.add(bs.toStringUtf8()); } return keys; } // ///////////////////// public ByteString[] store(RiakObject[] values, RequestMeta meta) throws IOException { RiakConnection c = getConnection(); try { BulkReader reader = new BulkReader(c, values.length); Thread worker = new Thread(reader); worker.start(); DataOutputStream dout = c.getOutputStream(); for (int i = 0; i < values.length; i++) { RiakObject value = values[i]; RiakKvPB.RpbPutReq.Builder builder = RiakKvPB.RpbPutReq.newBuilder().setBucket(value.getBucketBS()) .setKey(value.getKeyBS()).setContent(value.buildContent()); if (value.getVclock() != null) { builder.setVclock(value.getVclock()); } builder.setReturnBody(true); if (meta != null) { if (meta.writeQuorum != null) { builder.setW(meta.writeQuorum.intValue()); } if (meta.durableWriteQuorum != null) { builder.setDw(meta.durableWriteQuorum.intValue()); } } RpbPutReq req = builder.build(); int len = req.getSerializedSize(); dout.writeInt(len + 1); dout.write(MSG_PutReq); req.writeTo(dout); } dout.flush(); try { worker.join(); } catch (InterruptedException e) { e.printStackTrace(); } return reader.vclocks; } finally { release(c); } } class BulkReader implements Runnable { private ByteString[] vclocks; private final RiakConnection c; public BulkReader(RiakConnection c, int count) { this.c = c; this.vclocks = new ByteString[count]; } public void run() { try { for (int i = 0; i < vclocks.length; i++) { byte[] data = c.receive(MSG_PutResp); if (data != null) { RpbPutResp resp = RiakKvPB.RpbPutResp.parseFrom(data); if (resp.hasVclock()) { vclocks[i] = resp.getVclock(); } } } } catch (IOException e) { // TODO at least log it e.printStackTrace(); } } } public void store(RiakObject value) throws IOException { store(value, null); } public RiakObject[] store(RiakObject value, IRequestMeta meta) throws IOException { RiakKvPB.RpbPutReq.Builder builder = RiakKvPB.RpbPutReq.newBuilder().setBucket(value.getBucketBS()) .setContent(value.buildContent()); if (value.getKeyBS() != null) { builder.setKey(value.getKeyBS()); } if (value.getVclock() != null) { builder.setVclock(value.getVclock()); } if (meta != null) { meta.preparePut(builder); } RiakConnection c = getConnection(); try { c.send(MSG_PutReq, builder.build()); byte[] r = c.receive(MSG_PutResp); if (r == null) { return NO_RIAK_OBJECTS; } RpbPutResp resp = RiakKvPB.RpbPutResp.parseFrom(r); RiakObject[] res = new RiakObject[resp.getContentCount()]; ByteString vclock = resp.getVclock(); // The key parameter will be set only if the server generated a // key for the object so we check and set it accordingly for (int i = 0; i < res.length; i++) { res[i] = new RiakObject(vclock, value.getBucketBS(), (resp.hasKey()) ? resp.getKey() : value.getKeyBS(), resp.getContent(i)); } return res; } finally { release(c); } } // ///////////////////// public void delete(String bucket, String key, DeleteMeta deleteMeta) throws IOException { delete(ByteString.copyFromUtf8(bucket), ByteString.copyFromUtf8(key), deleteMeta); } public void delete(ByteString bucket, ByteString key, DeleteMeta deleteMeta) throws IOException { RpbDelReq.Builder builder = RiakKvPB.RpbDelReq.newBuilder().setBucket(bucket).setKey(key); deleteMeta.write(builder); RiakConnection c = getConnection(); try { c.send(MSG_DelReq, builder.build()); c.receive_code(MSG_DelResp); } finally { release(c); } } public void delete(String bucket, String key, int rw) throws IOException { delete(ByteString.copyFromUtf8(bucket), ByteString.copyFromUtf8(key), rw); } public void delete(ByteString bucket, ByteString key, int rw) throws IOException { RpbDelReq req = RiakKvPB.RpbDelReq.newBuilder().setBucket(bucket).setKey(key).setRw(rw).build(); RiakConnection c = getConnection(); try { c.send(MSG_DelReq, req); c.receive_code(MSG_DelResp); } finally { release(c); } } public void delete(String bucket, String key) throws IOException { delete(ByteString.copyFromUtf8(bucket), ByteString.copyFromUtf8(key)); } public void delete(ByteString bucket, ByteString key) throws IOException { RpbDelReq req = RiakKvPB.RpbDelReq.newBuilder().setBucket(bucket).setKey(key).build(); RiakConnection c = getConnection(); try { c.send(MSG_DelReq, req); c.receive_code(MSG_DelResp); } finally { release(c); } } public ByteString[] listBuckets() throws IOException { byte[] data; RiakConnection c = getConnection(); try { c.send(MSG_ListBucketsReq); data = c.receive(MSG_ListBucketsResp); if (data == null) { return NO_BYTE_STRINGS; } } finally { release(c); } RpbListBucketsResp resp = RiakKvPB.RpbListBucketsResp.parseFrom(data); ByteString[] out = new ByteString[resp.getBucketsCount()]; for (int i = 0; i < out.length; i++) { out[i] = resp.getBuckets(i); } return out; } public BucketProperties getBucketProperties(ByteString bucket) throws IOException { RiakConnection c = getConnection(); try { c.send(MSG_GetBucketReq, RiakKvPB.RpbGetBucketReq.newBuilder().setBucket(bucket).build()); byte[] data = c.receive(MSG_GetBucketResp); BucketProperties bp = new BucketProperties(); if (data == null) { return bp; } bp.init(RiakKvPB.RpbGetBucketResp.parseFrom(data)); return bp; } finally { release(c); } } public void setBucketProperties(ByteString bucket, BucketProperties props) throws IOException { RiakKvPB.RpbSetBucketReq req = RiakKvPB.RpbSetBucketReq.newBuilder().setBucket(bucket) .setProps(props.build()).build(); RiakConnection c = getConnection(); try { c.send(MSG_SetBucketReq, req); c.receive_code(MSG_SetBucketResp); } finally { release(c); } } // ///////////////////// public KeySource listKeys(ByteString bucket) throws IOException { RiakConnection c = getConnection(); c.send(MSG_ListKeysReq, RiakKvPB.RpbListKeysReq.newBuilder().setBucket(bucket).build()); return new KeySource(this, c); } // ///////////////////// public MapReduceResponseSource mapReduce(JSONObject obj) throws IOException { return mapReduce(ByteString.copyFromUtf8(obj.toString()), new RequestMeta().contentType("application/json")); } public MapReduceResponseSource mapReduce(String request, IRequestMeta meta) throws IOException { return mapReduce(ByteString.copyFromUtf8(request), meta); } public MapReduceResponseSource mapReduce(ByteString request, IRequestMeta meta) throws IOException { RiakConnection c = getConnection(); ByteString contentType = meta.getContentType(); if (contentType == null) { throw new IllegalArgumentException("no content type"); } RpbMapRedReq req = RiakKvPB.RpbMapRedReq.newBuilder().setRequest(request) .setContentType(meta.getContentType()).build(); c.send(MSG_MapRedReq, req); return new MapReduceResponseSource(this, c, contentType); } public void shutdown() { pool.shutdown(); } }