feat(!): add helper method get_row_range_from_prefix (#99)

liufuyang · Oct 23, 2024 · c88bc67 · c88bc67
1 parent a9ce0bf
commit c88bc67
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 24 deletions.
diff --git a/README.md b/README.md
@@ -65,7 +65,7 @@ You can use the library as follows:
 
 ```toml
 [dependencies]
-bigtable_rs = "0.2.13"
+bigtable_rs = "0.2.14"
 tokio = { version = "1.0", features = ["rt-multi-thread"] }
 env_logger = "0.11.1"
 ```

diff --git a/bigtable_rs/Cargo.toml b/bigtable_rs/Cargo.toml
@@ -1,7 +1,7 @@
 [package]
 name = "bigtable_rs"
 description = "A very simple Google Bigtable client lib in Rust"
-version = "0.2.13"
+version = "0.2.14"
 authors = ["Fuyang Liu <[email protected]>"]
 edition = "2021"
 license = "MIT"

diff --git a/bigtable_rs/src/bigtable.rs b/bigtable_rs/src/bigtable.rs
@@ -99,14 +99,12 @@ use tower::ServiceBuilder;
 
 use crate::auth_service::AuthSvc;
 use crate::bigtable::read_rows::decode_read_rows_response;
-use crate::google::bigtable::v2::row_range::{EndKey, StartKey};
 use crate::google::bigtable::v2::{
     bigtable_client::BigtableClient, MutateRowRequest, MutateRowResponse, MutateRowsRequest,
-    MutateRowsResponse, ReadRowsRequest, RowRange, RowSet, SampleRowKeysRequest,
-    SampleRowKeysResponse,
+    MutateRowsResponse, ReadRowsRequest, RowSet, SampleRowKeysRequest, SampleRowKeysResponse,
 };
 use crate::google::bigtable::v2::{CheckAndMutateRowRequest, CheckAndMutateRowResponse};
-use crate::{root_ca_certificate, util::get_end_key_for_prefix};
+use crate::{root_ca_certificate, util::get_row_range_from_prefix};
 
 pub mod read_rows;
 
@@ -486,14 +484,10 @@ impl BigTable {
         mut request: ReadRowsRequest,
         prefix: Vec<u8>,
     ) -> Result<Vec<(RowKey, Vec<RowCell>)>> {
-        let end_key =
-            get_end_key_for_prefix(prefix.as_ref()).map(|end_key| EndKey::EndKeyOpen(end_key));
+        let row_range = get_row_range_from_prefix(prefix);
         request.rows = Some(RowSet {
             row_keys: vec![], // use this field to put keys for reading specific rows
-            row_ranges: vec![RowRange {
-                start_key: Some(StartKey::StartKeyClosed(prefix)),
-                end_key,
-            }],
+            row_ranges: vec![row_range],
         });
         let response = self.client.read_rows(request).await?.into_inner();
         decode_read_rows_response(self.timeout.as_ref(), response).await

diff --git a/bigtable_rs/src/util.rs b/bigtable_rs/src/util.rs
@@ -1,4 +1,15 @@
-pub fn get_end_key_for_prefix(start_key: &[u8]) -> Option<Vec<u8>> {
+use crate::google::bigtable::v2::row_range::{EndKey, StartKey};
+use crate::google::bigtable::v2::RowRange;
+
+pub fn get_row_range_from_prefix(prefix: Vec<u8>) -> RowRange {
+    let end_key = get_end_key_for_prefix(prefix.as_ref()).map(EndKey::EndKeyOpen);
+    RowRange {
+        start_key: Some(StartKey::StartKeyClosed(prefix)),
+        end_key,
+    }
+}
+
+fn get_end_key_for_prefix(start_key: &[u8]) -> Option<Vec<u8>> {
     let size = start_key.len();
     if size < 1 {
         return None;

diff --git a/examples/src/custom_query.rs b/examples/src/custom_query.rs
@@ -1,9 +1,8 @@
 use bigtable_rs::bigtable;
 use bigtable_rs::bigtable::read_rows::decode_read_rows_response;
 use bigtable_rs::google::bigtable::v2::row_filter::{Chain, Filter};
-use bigtable_rs::google::bigtable::v2::row_range::{EndKey, StartKey};
-use bigtable_rs::google::bigtable::v2::{ReadRowsRequest, RowFilter, RowRange, RowSet};
-use bigtable_rs::util::get_end_key_for_prefix;
+use bigtable_rs::google::bigtable::v2::{ReadRowsRequest, RowFilter, RowSet};
+use bigtable_rs::util::get_row_range_from_prefix;
 use env_logger;
 use std::error::Error;
 use std::time::Duration;
@@ -40,14 +39,8 @@ async fn main() -> Result<(), Box<dyn Error>> {
         rows: Some(RowSet {
             row_keys: vec![],
             row_ranges: vec![
-                RowRange {
-                    start_key: Some(StartKey::StartKeyClosed(prefix_a.clone().into_bytes())),
-                    end_key: get_end_key_for_prefix(prefix_a.as_ref()).map(EndKey::EndKeyOpen),
-                },
-                RowRange {
-                    start_key: Some(StartKey::StartKeyClosed(prefix_b.clone().into_bytes())),
-                    end_key: get_end_key_for_prefix(prefix_b.as_ref()).map(EndKey::EndKeyOpen),
-                },
+                get_row_range_from_prefix(prefix_a.into_bytes()),
+                get_row_range_from_prefix(prefix_b.into_bytes()),
             ],
         }),
         filter: Some(RowFilter {