forked from tikv/client-java
-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathRawKVClient.java
More file actions
120 lines (106 loc) · 3.88 KB
/
Copy pathRawKVClient.java
File metadata and controls
120 lines (106 loc) · 3.88 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
package org.tikv;
import com.google.protobuf.ByteString;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import org.tikv.kvproto.Kvrpcpb;
import org.tikv.kvproto.Metapb;
import org.tikv.operation.iterator.RawScanIterator;
import org.tikv.region.RegionManager;
import org.tikv.region.RegionStoreClient;
import org.tikv.region.TiRegion;
import org.tikv.util.BackOffer;
import org.tikv.util.ConcreteBackOffer;
import org.tikv.util.Pair;
public class RawKVClient {
private static final String DEFAULT_PD_ADDRESS = "127.0.0.1:2379";
private final TiSession session;
private final RegionManager regionManager;
private RawKVClient(String addresses) {
session = TiSession.create(TiConfiguration.createDefault(addresses));
regionManager = session.getRegionManager();
}
private RawKVClient() {
this(DEFAULT_PD_ADDRESS);
}
public static RawKVClient create() {
return new RawKVClient();
}
public static RawKVClient create(String address) {
return new RawKVClient(address);
}
/**
* Put a raw key-value pair to TiKV
*
* @param key raw key
* @param value raw value
*/
public void put(ByteString key, ByteString value) {
Pair<TiRegion, Metapb.Store> pair = regionManager.getRegionStorePairByRawKey(key);
RegionStoreClient client = RegionStoreClient.create(pair.first, pair.second, session);
client.rawPut(defaultBackOff(), key, value);
}
/**
* Get a raw key-value pair from TiKV if key exists
*
* @param key raw key
* @return a ByteString value if key exists, ByteString.EMPTY if key does not exist
*/
public ByteString get(ByteString key) {
Pair<TiRegion, Metapb.Store> pair = regionManager.getRegionStorePairByRawKey(key);
RegionStoreClient client = RegionStoreClient.create(pair.first, pair.second, session);
return client.rawGet(defaultBackOff(), key);
}
/**
* Scan raw key-value pairs from TiKV in range [startKey, endKey)
*
* @param startKey raw start key, inclusive
* @param endKey raw end key, exclusive
* @return list of key-value pairs in range
*/
public List<Kvrpcpb.KvPair> scan(ByteString startKey, ByteString endKey) {
Iterator<Kvrpcpb.KvPair> iterator = rawScanIterator(startKey, endKey);
List<Kvrpcpb.KvPair> result = new ArrayList<>();
iterator.forEachRemaining(result::add);
return result;
}
/**
* Scan raw key-value pairs from TiKV in range [startKey, endKey)
*
* @param startKey raw start key, inclusive
* @param limit limit of key-value pairs
* @return list of key-value pairs in range
*/
public List<Kvrpcpb.KvPair> scan(ByteString startKey, int limit) {
Iterator<Kvrpcpb.KvPair> iterator = rawScanIterator(startKey, limit);
List<Kvrpcpb.KvPair> result = new ArrayList<>();
iterator.forEachRemaining(result::add);
return result;
}
/**
* Delete a raw key-value pair from TiKV if key exists
*
* @param key raw key to be deleted
*/
public void delete(ByteString key) {
TiRegion region = regionManager.getRegionByRawKey(key);
Kvrpcpb.Context context =
Kvrpcpb.Context.newBuilder()
.setRegionId(region.getId())
.setRegionEpoch(region.getRegionEpoch())
.setPeer(region.getLeader())
.build();
Pair<TiRegion, Metapb.Store> pair = regionManager.getRegionStorePairByRawKey(key);
RegionStoreClient client = RegionStoreClient.create(pair.first, pair.second, session);
client.rawDelete(defaultBackOff(), key, context);
}
private Iterator<Kvrpcpb.KvPair> rawScanIterator(ByteString startKey, ByteString endKey) {
return new RawScanIterator(startKey, endKey, Integer.MAX_VALUE, session);
}
private Iterator<Kvrpcpb.KvPair> rawScanIterator(ByteString startKey, int limit) {
return new RawScanIterator(startKey, ByteString.EMPTY, limit, session);
}
private BackOffer defaultBackOff() {
return ConcreteBackOffer.newCustomBackOff(1000);
}
}