From 83a423fcd585666f217d60538f305a8a2300930d Mon Sep 17 00:00:00 2001 From: jaguarx Date: Sun, 28 Apr 2019 20:54:56 +0800 Subject: [PATCH] expose columnfamily as parameter --- .../operation/iterator/RawScanIterator.java | 5 +- .../tikv/common/region/RegionStoreClient.java | 54 +++++++++++++++++-- src/main/java/org/tikv/raw/RawKVClient.java | 27 +++++++--- 3 files changed, 73 insertions(+), 13 deletions(-) diff --git a/src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java b/src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java index 81d60694d04..4956de73586 100644 --- a/src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java +++ b/src/main/java/org/tikv/common/operation/iterator/RawScanIterator.java @@ -27,14 +27,17 @@ import org.tikv.common.util.ConcreteBackOffer; public class RawScanIterator extends ScanIterator { + private ByteString cf; public RawScanIterator( TiConfiguration conf, RegionStoreClientBuilder builder, + ByteString cf, ByteString startKey, ByteString endKey, int limit) { super(conf, builder, startKey, endKey, limit); + this.cf = cf; } TiRegion loadCurrentRegionToCache() throws Exception { @@ -47,7 +50,7 @@ TiRegion loadCurrentRegionToCache() throws Exception { } else { while (true) { try { - currentCache = client.rawScan(backOffer, startKey, limit); + currentCache = client.rawScan(backOffer, startKey, cf, limit); break; } catch (final TiKVException e) { backOffer.doBackOff(BackOffFunction.BackOffFuncType.BoRegionMiss, e); diff --git a/src/main/java/org/tikv/common/region/RegionStoreClient.java b/src/main/java/org/tikv/common/region/RegionStoreClient.java index 13d735c506a..1ad4c94e084 100644 --- a/src/main/java/org/tikv/common/region/RegionStoreClient.java +++ b/src/main/java/org/tikv/common/region/RegionStoreClient.java @@ -334,6 +334,26 @@ public void rawPut(BackOffer backOffer, ByteString key, ByteString value) { rawPutHelper(resp); } + public void rawPut(BackOffer backOffer, ByteString key, ByteString value, ByteString cf) { + Supplier factory = + () -> + RawPutRequest.newBuilder() + .setContext(region.getContext()) + .setCfBytes(cf) + .setKey(key) + .setValue(value) + .build(); + + KVErrorHandler handler = + new KVErrorHandler<>( + regionManager, + this, + region, + resp -> resp.hasRegionError() ? resp.getRegionError() : null); + RawPutResponse resp = callWithRetry(backOffer, TikvGrpc.METHOD_RAW_PUT, factory, handler); + rawPutHelper(resp); + } + private void rawPutHelper(RawPutResponse resp) { if (resp == null) { this.regionManager.onRequestFail(region); @@ -369,6 +389,28 @@ public void rawBatchPut(BackOffer backOffer, List kvPairs) { handleRawBatchPut(resp); } + public void rawBatchPut(BackOffer backOffer, List kvPairs, ByteString cf) { + if (kvPairs.isEmpty()) { + return; + } + Supplier factory = + () -> + RawBatchPutRequest.newBuilder() + .setContext(region.getContext()) + .setCfBytes(cf) + .addAllPairs(kvPairs) + .build(); + KVErrorHandler handler = + new KVErrorHandler<>( + regionManager, + this, + region, + resp -> resp.hasRegionError() ? resp.getRegionError() : null); + RawBatchPutResponse resp = + callWithRetry(backOffer, TikvGrpc.METHOD_RAW_BATCH_PUT, factory, handler); + handleRawBatchPut(resp); + } + private void handleRawBatchPut(RawBatchPutResponse resp) { if (resp == null) { this.regionManager.onRequestFail(region); @@ -388,11 +430,13 @@ private void handleRawBatchPut(RawBatchPutResponse resp) { * @param keyOnly true if value of KvPair is not needed * @return KvPair list */ - private List rawScan(BackOffer backOffer, ByteString key, int limit, boolean keyOnly) { + private List rawScan( + BackOffer backOffer, ByteString key, ByteString cf, int limit, boolean keyOnly) { Supplier factory = () -> RawScanRequest.newBuilder() .setContext(region.getContext()) + .setCfBytes(cf) .setStartKey(key) .setKeyOnly(keyOnly) .setLimit(limit) @@ -408,12 +452,12 @@ private List rawScan(BackOffer backOffer, ByteString key, int limit, boo return rawScanHelper(resp); } - public List rawScan(BackOffer backOffer, ByteString key) { - return rawScan(backOffer, key, getConf().getScanBatchSize()); + public List rawScan(BackOffer backOffer, ByteString key, ByteString cf) { + return rawScan(backOffer, key, cf, getConf().getScanBatchSize()); } - public List rawScan(BackOffer backOffer, ByteString key, int limit) { - return rawScan(backOffer, key, limit, false); + public List rawScan(BackOffer backOffer, ByteString key, ByteString cf, int limit) { + return rawScan(backOffer, key, cf, limit, false); } private List rawScanHelper(RawScanResponse resp) { diff --git a/src/main/java/org/tikv/raw/RawKVClient.java b/src/main/java/org/tikv/raw/RawKVClient.java index 72401b9042e..2eb4cb22f98 100644 --- a/src/main/java/org/tikv/raw/RawKVClient.java +++ b/src/main/java/org/tikv/raw/RawKVClient.java @@ -135,13 +135,17 @@ public ByteString get(ByteString key) { * @param endKey raw end key, exclusive * @return list of key-value pairs in range */ - public List scan(ByteString startKey, ByteString endKey) { - Iterator iterator = rawScanIterator(conf, clientBuilder, startKey, endKey); + public List scan(ByteString cf, ByteString startKey, ByteString endKey) { + Iterator iterator = rawScanIterator(conf, clientBuilder, cf, startKey, endKey); List result = new ArrayList<>(); iterator.forEachRemaining(result::add); return result; } + public List scan(ByteString startKey, ByteString endKey) { + return scan(ByteString.EMPTY, startKey, endKey); + } + /** * Scan raw key-value pairs from TiKV in range [startKey, endKey) * @@ -149,13 +153,17 @@ public List scan(ByteString startKey, ByteString endKey) { * @param limit limit of key-value pairs * @return list of key-value pairs in range */ - public List scan(ByteString startKey, int limit) { - Iterator iterator = rawScanIterator(conf, clientBuilder, startKey, limit); + public List scan(ByteString cf, ByteString startKey, int limit) { + Iterator iterator = rawScanIterator(conf, clientBuilder, cf, startKey, limit); List result = new ArrayList<>(); iterator.forEachRemaining(result::add); return result; } + public List scan(ByteString startKey, int limit) { + return scan(ByteString.EMPTY, startKey, limit); + } + /** * Delete a raw key-value pair from TiKV if key exists * @@ -294,14 +302,19 @@ private void sendBatchPut(BackOffer backOffer, List batches) { private Iterator rawScanIterator( TiConfiguration conf, RegionStoreClientBuilder builder, + ByteString cf, ByteString startKey, ByteString endKey) { - return new RawScanIterator(conf, builder, startKey, endKey, Integer.MAX_VALUE); + return new RawScanIterator(conf, builder, cf, startKey, endKey, Integer.MAX_VALUE); } private Iterator rawScanIterator( - TiConfiguration conf, RegionStoreClientBuilder builder, ByteString startKey, int limit) { - return new RawScanIterator(conf, builder, startKey, ByteString.EMPTY, limit); + TiConfiguration conf, + RegionStoreClientBuilder builder, + ByteString cf, + ByteString startKey, + int limit) { + return new RawScanIterator(conf, builder, cf, startKey, ByteString.EMPTY, limit); } private BackOffer defaultBackOff() {