Add request ctor functions taking high-level types

Signed-off-by: Nick Cameron <[email protected]>

Author: nrc

src/lib.rs

@@ -103,14 +103,14 @@ pub use crate::backoff::Backoff;
 #[doc(inline)]
 pub use crate::kv::{BoundRange, IntoOwnedRange, Key, KvPair, Value};
 #[doc(inline)]
-pub use crate::raw::{Client as RawClient, ColumnFamily};
+pub use crate::raw::{lowering::*, Client as RawClient, ColumnFamily};
 #[doc(inline)]
 pub use crate::request::RetryOptions;
 #[doc(inline)]
 pub use crate::timestamp::{Timestamp, TimestampExt};
 #[doc(inline)]
 pub use crate::transaction::{
-    CheckLevel, Client as TransactionClient, Snapshot, Transaction, TransactionOptions,
+    lowering::*, CheckLevel, Client as TransactionClient, Snapshot, Transaction, TransactionOptions,
 };
 #[doc(inline)]
 pub use config::Config;

src/raw/client.rs

@@ -2,10 +2,10 @@
 
 use tikv_client_common::Error;
 
-use super::requests;
 use crate::{
     config::Config,
     pd::PdRpcClient,
+    raw::lowering::*,
     request::{KvRequest, RetryOptions},
     BoundRange, ColumnFamily, Key, KvPair, Result, Value,
 };
@@ -110,7 +110,7 @@ impl Client {
     /// # });
     /// ```
     pub async fn get(&self, key: impl Into<Key>) -> Result<Option<Value>> {
-        requests::new_raw_get_request(key, self.cf.clone())
+        new_raw_get_request(key.into(), self.cf.clone())
             .execute(self.rpc.clone(), RetryOptions::default_optimistic())
             .await
     }
@@ -137,7 +137,7 @@ impl Client {
         &self,
         keys: impl IntoIterator<Item = impl Into<Key>>,
     ) -> Result<Vec<KvPair>> {
-        requests::new_raw_batch_get_request(keys, self.cf.clone())
+        new_raw_batch_get_request(keys.into_iter().map(Into::into), self.cf.clone())
             .execute(self.rpc.clone(), RetryOptions::default_optimistic())
             .await
     }
@@ -159,7 +159,7 @@ impl Client {
     /// # });
     /// ```
     pub async fn put(&self, key: impl Into<Key>, value: impl Into<Value>) -> Result<()> {
-        requests::new_raw_put_request(key, value, self.cf.clone())
+        new_raw_put_request(key.into(), value.into(), self.cf.clone())
             .execute(self.rpc.clone(), RetryOptions::default_optimistic())
             .await
     }
@@ -185,7 +185,7 @@ impl Client {
         &self,
         pairs: impl IntoIterator<Item = impl Into<KvPair>>,
     ) -> Result<()> {
-        requests::new_raw_batch_put_request(pairs, self.cf.clone())
+        new_raw_batch_put_request(pairs.into_iter().map(Into::into), self.cf.clone())
             .execute(self.rpc.clone(), RetryOptions::default_optimistic())
             .await
     }
@@ -208,7 +208,7 @@ impl Client {
     /// # });
     /// ```
     pub async fn delete(&self, key: impl Into<Key>) -> Result<()> {
-        requests::new_raw_delete_request(key, self.cf.clone())
+        new_raw_delete_request(key.into(), self.cf.clone())
             .execute(self.rpc.clone(), RetryOptions::default_optimistic())
             .await
     }
@@ -231,7 +231,7 @@ impl Client {
     /// # });
     /// ```
     pub async fn batch_delete(&self, keys: impl IntoIterator<Item = impl Into<Key>>) -> Result<()> {
-        requests::new_raw_batch_delete_request(keys, self.cf.clone())
+        new_raw_batch_delete_request(keys.into_iter().map(Into::into), self.cf.clone())
             .execute(self.rpc.clone(), RetryOptions::default_optimistic())
             .await
     }
@@ -252,7 +252,7 @@ impl Client {
     /// # });
     /// ```
     pub async fn delete_range(&self, range: impl Into<BoundRange>) -> Result<()> {
-        requests::new_raw_delete_range_request(range, self.cf.clone())
+        new_raw_delete_range_request(range.into(), self.cf.clone())
             .execute(self.rpc.clone(), RetryOptions::default_optimistic())
             .await
     }
@@ -277,7 +277,7 @@ impl Client {
     /// # });
     /// ```
     pub async fn scan(&self, range: impl Into<BoundRange>, limit: u32) -> Result<Vec<KvPair>> {
-        self.scan_inner(range, limit, false).await
+        self.scan_inner(range.into(), limit, false).await
     }
 
     /// Create a new 'scan' request that only returns the keys.
@@ -388,7 +388,7 @@ impl Client {
             });
         }
 
-        let res = requests::new_raw_scan_request(range, limit, key_only, self.cf.clone())
+        let res = new_raw_scan_request(range.into(), limit, key_only, self.cf.clone())
             .execute(self.rpc.clone(), RetryOptions::default_optimistic())
             .await;
         res.map(|mut s| {
@@ -410,8 +410,13 @@ impl Client {
             });
         }
 
-        requests::new_raw_batch_scan_request(ranges, each_limit, key_only, self.cf.clone())
-            .execute(self.rpc.clone(), RetryOptions::default_optimistic())
-            .await
+        new_raw_batch_scan_request(
+            ranges.into_iter().map(Into::into),
+            each_limit,
+            key_only,
+            self.cf.clone(),
+        )
+        .execute(self.rpc.clone(), RetryOptions::default_optimistic())
+        .await
     }
 }

src/raw/lowering.rs

@@ -0,0 +1,78 @@
+// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
+
+/// This module provides constructor functions for requests which take arguments as high-level
+/// types (i.e., the types from the client crate) and converts these to the types used in the
+/// generated protobuf code, then calls the low-level ctor functions in the requests module.
+use crate::{raw::requests, BoundRange, ColumnFamily, Key, KvPair, Value};
+use std::iter::Iterator;
+use tikv_client_proto::kvrpcpb;
+
+pub fn new_raw_get_request(key: Key, cf: Option<ColumnFamily>) -> kvrpcpb::RawGetRequest {
+    requests::new_raw_get_request(key.into(), cf)
+}
+
+pub fn new_raw_batch_get_request(
+    keys: impl Iterator<Item = Key>,
+    cf: Option<ColumnFamily>,
+) -> kvrpcpb::RawBatchGetRequest {
+    requests::new_raw_batch_get_request(keys.map(Into::into).collect(), cf)
+}
+
+pub fn new_raw_put_request(
+    key: Key,
+    value: Value,
+    cf: Option<ColumnFamily>,
+) -> kvrpcpb::RawPutRequest {
+    requests::new_raw_put_request(key.into(), value, cf)
+}
+
+pub fn new_raw_batch_put_request(
+    pairs: impl Iterator<Item = KvPair>,
+    cf: Option<ColumnFamily>,
+) -> kvrpcpb::RawBatchPutRequest {
+    requests::new_raw_batch_put_request(pairs.map(Into::into).collect(), cf)
+}
+
+pub fn new_raw_delete_request(key: Key, cf: Option<ColumnFamily>) -> kvrpcpb::RawDeleteRequest {
+    requests::new_raw_delete_request(key.into(), cf)
+}
+
+pub fn new_raw_batch_delete_request(
+    keys: impl Iterator<Item = Key>,
+    cf: Option<ColumnFamily>,
+) -> kvrpcpb::RawBatchDeleteRequest {
+    requests::new_raw_batch_delete_request(keys.map(Into::into).collect(), cf)
+}
+
+pub fn new_raw_delete_range_request(
+    range: BoundRange,
+    cf: Option<ColumnFamily>,
+) -> kvrpcpb::RawDeleteRangeRequest {
+    let (start_key, end_key) = range.into_keys();
+    requests::new_raw_delete_range_request(start_key.into(), end_key.unwrap_or_default().into(), cf)
+}
+
+pub fn new_raw_scan_request(
+    range: BoundRange,
+    limit: u32,
+    key_only: bool,
+    cf: Option<ColumnFamily>,
+) -> kvrpcpb::RawScanRequest {
+    let (start_key, end_key) = range.into_keys();
+    requests::new_raw_scan_request(
+        start_key.into(),
+        end_key.unwrap_or_default().into(),
+        limit,
+        key_only,
+        cf,
+    )
+}
+
+pub fn new_raw_batch_scan_request(
+    ranges: impl Iterator<Item = BoundRange>,
+    each_limit: u32,
+    key_only: bool,
+    cf: Option<ColumnFamily>,
+) -> kvrpcpb::RawBatchScanRequest {
+    requests::new_raw_batch_scan_request(ranges.map(Into::into).collect(), each_limit, key_only, cf)
+}

src/raw/mod.rs

@@ -14,6 +14,7 @@ use crate::Error;
 use std::{convert::TryFrom, fmt};
 
 mod client;
+pub mod lowering;
 mod requests;
 
 /// A [`ColumnFamily`](ColumnFamily) is an optional parameter for [`raw::Client`](Client) requests.

src/raw/requests.rs

@@ -54,12 +54,9 @@ impl KvRequest for kvrpcpb::RawGetRequest {
     }
 }
 
-pub fn new_raw_get_request(
-    key: impl Into<Key>,
-    cf: Option<ColumnFamily>,
-) -> kvrpcpb::RawGetRequest {
+pub fn new_raw_get_request(key: Vec<u8>, cf: Option<ColumnFamily>) -> kvrpcpb::RawGetRequest {
     let mut req = kvrpcpb::RawGetRequest::default();
-    req.set_key(key.into().into());
+    req.set_key(key);
     req.maybe_set_cf(cf);
 
     req
@@ -98,11 +95,11 @@ impl KvRequest for kvrpcpb::RawBatchGetRequest {
 }
 
 pub fn new_raw_batch_get_request(
-    keys: impl IntoIterator<Item = impl Into<Key>>,
+    keys: Vec<Vec<u8>>,
     cf: Option<ColumnFamily>,
 ) -> kvrpcpb::RawBatchGetRequest {
     let mut req = kvrpcpb::RawBatchGetRequest::default();
-    req.set_keys(keys.into_iter().map(Into::into).map(Into::into).collect());
+    req.set_keys(keys);
     req.maybe_set_cf(cf);
 
     req
@@ -144,13 +141,13 @@ impl KvRequest for kvrpcpb::RawPutRequest {
 }
 
 pub fn new_raw_put_request(
-    key: impl Into<Key>,
-    value: impl Into<Value>,
+    key: Vec<u8>,
+    value: Vec<u8>,
     cf: Option<ColumnFamily>,
 ) -> kvrpcpb::RawPutRequest {
     let mut req = kvrpcpb::RawPutRequest::default();
-    req.set_key(key.into().into());
-    req.set_value(value.into());
+    req.set_key(key);
+    req.set_value(value);
     req.maybe_set_cf(cf);
 
     req
@@ -187,11 +184,11 @@ impl KvRequest for kvrpcpb::RawBatchPutRequest {
 }
 
 pub fn new_raw_batch_put_request(
-    pairs: impl IntoIterator<Item = impl Into<KvPair>>,
+    pairs: Vec<kvrpcpb::KvPair>,
     cf: Option<ColumnFamily>,
 ) -> kvrpcpb::RawBatchPutRequest {
     let mut req = kvrpcpb::RawBatchPutRequest::default();
-    req.set_pairs(pairs.into_iter().map(Into::into).map(Into::into).collect());
+    req.set_pairs(pairs);
     req.maybe_set_cf(cf);
 
     req
@@ -229,12 +226,9 @@ impl KvRequest for kvrpcpb::RawDeleteRequest {
     }
 }
 
-pub fn new_raw_delete_request(
-    key: impl Into<Key>,
-    cf: Option<ColumnFamily>,
-) -> kvrpcpb::RawDeleteRequest {
+pub fn new_raw_delete_request(key: Vec<u8>, cf: Option<ColumnFamily>) -> kvrpcpb::RawDeleteRequest {
     let mut req = kvrpcpb::RawDeleteRequest::default();
-    req.set_key(key.into().into());
+    req.set_key(key);
     req.maybe_set_cf(cf);
 
     req
@@ -271,11 +265,11 @@ impl KvRequest for kvrpcpb::RawBatchDeleteRequest {
 }
 
 pub fn new_raw_batch_delete_request(
-    keys: impl IntoIterator<Item = impl Into<Key>>,
+    keys: Vec<Vec<u8>>,
     cf: Option<ColumnFamily>,
 ) -> kvrpcpb::RawBatchDeleteRequest {
     let mut req = kvrpcpb::RawBatchDeleteRequest::default();
-    req.set_keys(keys.into_iter().map(Into::into).map(Into::into).collect());
+    req.set_keys(keys);
     req.maybe_set_cf(cf);
 
     req
@@ -316,13 +310,13 @@ impl KvRequest for kvrpcpb::RawDeleteRangeRequest {
 }
 
 pub fn new_raw_delete_range_request(
-    range: impl Into<BoundRange>,
+    start_key: Vec<u8>,
+    end_key: Vec<u8>,
     cf: Option<ColumnFamily>,
 ) -> kvrpcpb::RawDeleteRangeRequest {
-    let (start_key, end_key) = range.into().into_keys();
     let mut req = kvrpcpb::RawDeleteRangeRequest::default();
-    req.set_start_key(start_key.into());
-    req.set_end_key(end_key.unwrap_or_default().into());
+    req.set_start_key(start_key);
+    req.set_end_key(end_key);
     req.maybe_set_cf(cf);
 
     req
@@ -365,15 +359,15 @@ impl KvRequest for kvrpcpb::RawScanRequest {
 }
 
 pub fn new_raw_scan_request(
-    range: impl Into<BoundRange>,
+    start_key: Vec<u8>,
+    end_key: Vec<u8>,
     limit: u32,
     key_only: bool,
     cf: Option<ColumnFamily>,
 ) -> kvrpcpb::RawScanRequest {
-    let (start_key, end_key) = range.into().into_keys();
     let mut req = kvrpcpb::RawScanRequest::default();
-    req.set_start_key(start_key.into());
-    req.set_end_key(end_key.unwrap_or_default().into());
+    req.set_start_key(start_key);
+    req.set_end_key(end_key);
     req.set_limit(limit);
     req.set_key_only(key_only);
     req.maybe_set_cf(cf);
@@ -418,13 +412,13 @@ impl KvRequest for kvrpcpb::RawBatchScanRequest {
 }
 
 pub fn new_raw_batch_scan_request(
-    ranges: impl IntoIterator<Item = impl Into<BoundRange>>,
+    ranges: Vec<kvrpcpb::KeyRange>,
     each_limit: u32,
     key_only: bool,
     cf: Option<ColumnFamily>,
 ) -> kvrpcpb::RawBatchScanRequest {
     let mut req = kvrpcpb::RawBatchScanRequest::default();
-    req.set_ranges(ranges.into_iter().map(Into::into).map(Into::into).collect());
+    req.set_ranges(ranges);
     req.set_each_limit(each_limit);
     req.set_key_only(key_only);
     req.maybe_set_cf(cf);

src/transaction/buffer.rs

@@ -81,7 +81,7 @@ impl Buffer {
         f: F,
     ) -> Result<impl Iterator<Item = KvPair>>
     where
-        F: FnOnce(Vec<Key>) -> Fut,
+        F: FnOnce(Box<dyn Iterator<Item = Key>>) -> Fut,
         Fut: Future<Output = Result<Vec<KvPair>>>,
     {
         let (cached_results, undetermined_keys) = {
@@ -106,11 +106,11 @@ impl Buffer {
                 .into_iter()
                 .filter_map(|(k, v)| v.unwrap().map(|v| KvPair(k, v)));
 
-            let undetermined_keys = undetermined_keys.into_iter().map(|(k, _)| k).collect();
+            let undetermined_keys = undetermined_keys.into_iter().map(|(k, _)| k);
             (cached_results, undetermined_keys)
         };
 
-        let fetched_results = f(undetermined_keys).await?;
+        let fetched_results = f(Box::new(undetermined_keys)).await?;
         let mut mutations = self.mutations.lock().await;
         for kvpair in &fetched_results {
             let key = kvpair.0.clone();

src/transaction/client.rs

@@ -185,8 +185,8 @@ impl Client {
         let mut start_key = vec![];
         loop {
             let req = new_scan_lock_request(
-                mem::take(&mut start_key).into(),
-                safepoint.clone(),
+                mem::take(&mut start_key),
+                safepoint.version(),
                 SCAN_LOCK_BATCH_SIZE,
             );
             let res: Vec<kvrpcpb::LockInfo> = req