View Javadoc
1   /*
2    * Copyright 2021 TiKV Project Authors.
3    *
4    * Licensed under the Apache License, Version 2.0 (the "License");
5    * you may not use this file except in compliance with the License.
6    * You may obtain a copy of the License at
7    *
8    * http://www.apache.org/licenses/LICENSE-2.0
9    *
10   * Unless required by applicable law or agreed to in writing, software
11   * distributed under the License is distributed on an "AS IS" BASIS,
12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13   * See the License for the specific language governing permissions and
14   * limitations under the License.
15   *
16   */
17  
18  package org.tikv.common;
19  
20  import static org.tikv.common.operation.iterator.CoprocessorIterator.getHandleIterator;
21  import static org.tikv.common.operation.iterator.CoprocessorIterator.getRowIterator;
22  import static org.tikv.common.operation.iterator.CoprocessorIterator.getTiChunkIterator;
23  
24  import com.google.protobuf.ByteString;
25  import java.util.ArrayList;
26  import java.util.Iterator;
27  import java.util.List;
28  import java.util.stream.Collectors;
29  import javax.annotation.Nonnull;
30  import org.tikv.common.columnar.TiChunk;
31  import org.tikv.common.key.Key;
32  import org.tikv.common.meta.TiDAGRequest;
33  import org.tikv.common.meta.TiTimestamp;
34  import org.tikv.common.operation.iterator.ConcreteScanIterator;
35  import org.tikv.common.operation.iterator.IndexScanIterator;
36  import org.tikv.common.row.Row;
37  import org.tikv.common.util.ConcreteBackOffer;
38  import org.tikv.common.util.RangeSplitter;
39  import org.tikv.common.util.RangeSplitter.RegionTask;
40  import org.tikv.kvproto.Kvrpcpb.KvPair;
41  
42  public class Snapshot {
43    private final TiTimestamp timestamp;
44    private final TiSession session;
45  
46    public Snapshot(@Nonnull TiTimestamp timestamp, TiSession session) {
47      this.timestamp = timestamp;
48      this.session = session;
49    }
50  
51    public TiSession getSession() {
52      return session;
53    }
54  
55    public long getVersion() {
56      return timestamp.getVersion();
57    }
58  
59    public TiTimestamp getTimestamp() {
60      return timestamp;
61    }
62  
63    public byte[] get(byte[] key) {
64      ByteString keyString = ByteString.copyFrom(key);
65      ByteString value = get(keyString);
66      return value.toByteArray();
67    }
68  
69    public ByteString get(ByteString key) {
70      try (KVClient client = new KVClient(session, session.getRegionStoreClientBuilder())) {
71        return client.get(key, timestamp.getVersion());
72      }
73    }
74  
75    public List<org.tikv.common.BytePairWrapper> batchGet(int backOffer, List<byte[]> keys) {
76      List<ByteString> list = new ArrayList<>();
77      for (byte[] key : keys) {
78        list.add(ByteString.copyFrom(key));
79      }
80      try (KVClient client = new KVClient(session, session.getRegionStoreClientBuilder())) {
81        List<KvPair> kvPairList =
82            client.batchGet(
83                ConcreteBackOffer.newCustomBackOff(backOffer, session.getPDClient().getClusterId()),
84                list,
85                timestamp.getVersion());
86        return kvPairList
87            .stream()
88            .map(
89                kvPair ->
90                    new org.tikv.common.BytePairWrapper(
91                        kvPair.getKey().toByteArray(), kvPair.getValue().toByteArray()))
92            .collect(Collectors.toList());
93      }
94    }
95  
96    public Iterator<TiChunk> tableReadChunk(
97        TiDAGRequest dagRequest, List<RegionTask> tasks, int numOfRows) {
98      if (dagRequest.isDoubleRead()) {
99        throw new UnsupportedOperationException(
100           "double read case should first read handle in row-wise fashion");
101     } else {
102       return getTiChunkIterator(dagRequest, tasks, getSession(), numOfRows);
103     }
104   }
105   /**
106    * Issue a table read request
107    *
108    * @param dagRequest DAG request for coprocessor
109    * @return a Iterator that contains all result from this select request.
110    */
111   public Iterator<Row> tableReadRow(TiDAGRequest dagRequest, long physicalId) {
112     return tableReadRow(
113         dagRequest,
114         RangeSplitter.newSplitter(session.getRegionManager())
115             .splitRangeByRegion(
116                 dagRequest.getRangesByPhysicalId(physicalId), dagRequest.getStoreType()));
117   }
118 
119   /**
120    * Below is lower level API for env like Spark which already did key range split Perform table
121    * scan
122    *
123    * @param dagRequest DAGRequest for coprocessor
124    * @param tasks RegionTasks of the coprocessor request to send
125    * @return Row iterator to iterate over resulting rows
126    */
127   private Iterator<Row> tableReadRow(TiDAGRequest dagRequest, List<RegionTask> tasks) {
128     if (dagRequest.isDoubleRead()) {
129       Iterator<Long> iter = getHandleIterator(dagRequest, tasks, getSession());
130       return new IndexScanIterator(this, dagRequest, iter);
131     } else {
132       return getRowIterator(dagRequest, tasks, getSession());
133     }
134   }
135 
136   /**
137    * Below is lower level API for env like Spark which already did key range split Perform handle
138    * scan
139    *
140    * @param dagRequest DAGRequest for coprocessor
141    * @param tasks RegionTask of the coprocessor request to send
142    * @return Row iterator to iterate over resulting rows
143    */
144   public Iterator<Long> indexHandleRead(TiDAGRequest dagRequest, List<RegionTask> tasks) {
145     return getHandleIterator(dagRequest, tasks, session);
146   }
147 
148   /**
149    * scan all keys after startKey, inclusive
150    *
151    * @param startKey start of keys
152    * @return iterator of kvPair
153    */
154   public Iterator<KvPair> scan(ByteString startKey) {
155     return new ConcreteScanIterator(
156         session.getConf(),
157         session.getRegionStoreClientBuilder(),
158         startKey,
159         timestamp.getVersion(),
160         Integer.MAX_VALUE);
161   }
162 
163   /**
164    * scan all keys with prefix
165    *
166    * @param prefix prefix of keys
167    * @return iterator of kvPair
168    */
169   public Iterator<KvPair> scanPrefix(ByteString prefix) {
170     ByteString nextPrefix = Key.toRawKey(prefix).nextPrefix().toByteString();
171     return new ConcreteScanIterator(
172         session.getConf(),
173         session.getRegionStoreClientBuilder(),
174         prefix,
175         nextPrefix,
176         timestamp.getVersion());
177   }
178 
179   public TiConfiguration getConf() {
180     return this.session.getConf();
181   }
182 }