kstone_client/
query.rs

1/// Remote query builder and response types
2use crate::convert::*;
3use crate::error::Result;
4use bytes::Bytes;
5use kstone_core::Item;
6use kstone_proto::{self as proto, keystone_db_client::KeystoneDbClient};
7use tonic::transport::Channel;
8
9/// Remote query builder
10pub struct RemoteQuery {
11    partition_key: Vec<u8>,
12    sort_key_condition: Option<proto::SortKeyCondition>,
13    limit: Option<u32>,
14    exclusive_start_key: Option<proto::LastKey>,
15    scan_forward: Option<bool>,
16    index_name: Option<String>,
17}
18
19impl RemoteQuery {
20    /// Create a new query for a partition key
21    pub fn new(pk: &[u8]) -> Self {
22        Self {
23            partition_key: pk.to_vec(),
24            sort_key_condition: None,
25            limit: None,
26            exclusive_start_key: None,
27            scan_forward: None,
28            index_name: None,
29        }
30    }
31
32    /// Add a sort key equals condition
33    pub fn sk_eq(mut self, sk: &[u8]) -> Self {
34        self.sort_key_condition = Some(proto::SortKeyCondition {
35            condition: Some(proto::sort_key_condition::Condition::EqualTo(
36                value_to_proto_bytes(sk),
37            )),
38        });
39        self
40    }
41
42    /// Add a sort key less than condition
43    pub fn sk_lt(mut self, sk: &[u8]) -> Self {
44        self.sort_key_condition = Some(proto::SortKeyCondition {
45            condition: Some(proto::sort_key_condition::Condition::LessThan(
46                value_to_proto_bytes(sk),
47            )),
48        });
49        self
50    }
51
52    /// Add a sort key less than or equal condition
53    pub fn sk_lte(mut self, sk: &[u8]) -> Self {
54        self.sort_key_condition = Some(proto::SortKeyCondition {
55            condition: Some(proto::sort_key_condition::Condition::LessThanOrEqual(
56                value_to_proto_bytes(sk),
57            )),
58        });
59        self
60    }
61
62    /// Add a sort key greater than condition
63    pub fn sk_gt(mut self, sk: &[u8]) -> Self {
64        self.sort_key_condition = Some(proto::SortKeyCondition {
65            condition: Some(proto::sort_key_condition::Condition::GreaterThan(
66                value_to_proto_bytes(sk),
67            )),
68        });
69        self
70    }
71
72    /// Add a sort key greater than or equal condition
73    pub fn sk_gte(mut self, sk: &[u8]) -> Self {
74        self.sort_key_condition = Some(proto::SortKeyCondition {
75            condition: Some(proto::sort_key_condition::Condition::GreaterThanOrEqual(
76                value_to_proto_bytes(sk),
77            )),
78        });
79        self
80    }
81
82    /// Add a sort key between condition
83    pub fn sk_between(mut self, sk1: &[u8], sk2: &[u8]) -> Self {
84        self.sort_key_condition = Some(proto::SortKeyCondition {
85            condition: Some(proto::sort_key_condition::Condition::Between(
86                proto::BetweenCondition {
87                    lower: Some(value_to_proto_bytes(sk1)),
88                    upper: Some(value_to_proto_bytes(sk2)),
89                },
90            )),
91        });
92        self
93    }
94
95    /// Add a sort key begins_with condition
96    pub fn sk_begins_with(mut self, prefix: &[u8]) -> Self {
97        self.sort_key_condition = Some(proto::SortKeyCondition {
98            condition: Some(proto::sort_key_condition::Condition::BeginsWith(
99                value_to_proto_bytes(prefix),
100            )),
101        });
102        self
103    }
104
105    /// Set the scan direction (default: forward)
106    pub fn forward(mut self, forward: bool) -> Self {
107        self.scan_forward = Some(forward);
108        self
109    }
110
111    /// Set the maximum number of items to return
112    pub fn limit(mut self, limit: usize) -> Self {
113        self.limit = Some(limit as u32);
114        self
115    }
116
117    /// Set the exclusive start key for pagination
118    pub fn start_after(mut self, pk: &[u8], sk: Option<&[u8]>) -> Self {
119        self.exclusive_start_key = Some(proto::LastKey {
120            partition_key: pk.to_vec(),
121            sort_key: sk.map(|s| s.to_vec()),
122        });
123        self
124    }
125
126    /// Query a Local Secondary Index instead of the base table
127    pub fn index(mut self, index_name: impl Into<String>) -> Self {
128        self.index_name = Some(index_name.into());
129        self
130    }
131
132    /// Execute the query
133    pub async fn execute(
134        self,
135        client: &mut KeystoneDbClient<Channel>,
136    ) -> Result<RemoteQueryResponse> {
137        let request = proto::QueryRequest {
138            partition_key: self.partition_key,
139            sort_key_condition: self.sort_key_condition,
140            filter_expression: None,
141            expression_values: std::collections::HashMap::new(),
142            index_name: self.index_name,
143            limit: self.limit,
144            exclusive_start_key: self.exclusive_start_key,
145            scan_forward: self.scan_forward,
146        };
147
148        let response = client
149            .query(request)
150            .await?
151            .into_inner();
152
153        // Convert protobuf response to Rust types
154        let items: Vec<Item> = response
155            .items
156            .into_iter()
157            .map(|proto_item| proto_item_to_ks(proto_item).expect("Invalid item from server"))
158            .collect();
159
160        let last_key = response.last_evaluated_key.map(|key| {
161            let (pk, sk) = proto_last_key_to_ks(key);
162            (pk, sk)
163        });
164
165        Ok(RemoteQueryResponse {
166            items,
167            count: response.count as usize,
168            scanned_count: response.scanned_count as usize,
169            last_key,
170        })
171    }
172}
173
174/// Query response
175pub struct RemoteQueryResponse {
176    /// Items found
177    pub items: Vec<Item>,
178    /// Number of items returned
179    pub count: usize,
180    /// Last evaluated key (for pagination)
181    pub last_key: Option<(Bytes, Option<Bytes>)>,
182    /// Number of items examined
183    pub scanned_count: usize,
184}
185
186/// Helper function to convert bytes to protobuf Value (for sort key conditions)
187fn value_to_proto_bytes(bytes: &[u8]) -> proto::Value {
188    proto::Value {
189        value: Some(proto::value::Value::BinaryValue(bytes.to_vec())),
190    }
191}
kstone_client/query.rs

kstone_client/
query.rs
