kstone_client/
scan.rs

1/// Remote scan 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;
8use tonic::Streaming;
9
10/// Remote scan builder
11pub struct RemoteScan {
12    limit: Option<u32>,
13    exclusive_start_key: Option<proto::LastKey>,
14    index_name: Option<String>,
15    segment: Option<u32>,
16    total_segments: Option<u32>,
17}
18
19impl RemoteScan {
20    /// Create a new scan
21    pub fn new() -> Self {
22        Self {
23            limit: None,
24            exclusive_start_key: None,
25            index_name: None,
26            segment: None,
27            total_segments: None,
28        }
29    }
30
31    /// Set the maximum number of items to return
32    pub fn limit(mut self, limit: usize) -> Self {
33        self.limit = Some(limit as u32);
34        self
35    }
36
37    /// Set the exclusive start key for pagination
38    pub fn start_after(mut self, pk: &[u8], sk: Option<&[u8]>) -> Self {
39        self.exclusive_start_key = Some(proto::LastKey {
40            partition_key: pk.to_vec(),
41            sort_key: sk.map(|s| s.to_vec()),
42        });
43        self
44    }
45
46    /// Scan a secondary index instead of the base table
47    pub fn index(mut self, index_name: impl Into<String>) -> Self {
48        self.index_name = Some(index_name.into());
49        self
50    }
51
52    /// Configure parallel scan (segment must be < total_segments)
53    pub fn segment(mut self, segment: usize, total_segments: usize) -> Self {
54        self.segment = Some(segment as u32);
55        self.total_segments = Some(total_segments as u32);
56        self
57    }
58
59    /// Execute the scan and get a stream of responses
60    ///
61    /// Note: The server currently returns a single response, but this
62    /// interface is prepared for future streaming support.
63    pub async fn execute(
64        self,
65        client: &mut KeystoneDbClient<Channel>,
66    ) -> Result<RemoteScanResponse> {
67        let request = proto::ScanRequest {
68            filter_expression: None,
69            expression_values: std::collections::HashMap::new(),
70            limit: self.limit,
71            exclusive_start_key: self.exclusive_start_key,
72            index_name: self.index_name,
73            segment: self.segment,
74            total_segments: self.total_segments,
75        };
76
77        let mut stream: Streaming<proto::ScanResponse> = client
78            .scan(request)
79            .await?
80            .into_inner();
81
82        // Collect all items from the stream
83        let mut all_items = Vec::new();
84        let mut total_count = 0;
85        let mut total_scanned_count = 0;
86        let mut last_key = None;
87
88        while let Some(response) = stream.message().await? {
89            let items: Vec<Item> = response
90                .items
91                .into_iter()
92                .map(|proto_item| proto_item_to_ks(proto_item).expect("Invalid item from server"))
93                .collect();
94
95            total_count += response.count as usize;
96            total_scanned_count += response.scanned_count as usize;
97
98            if let Some(key) = response.last_evaluated_key {
99                let (pk, sk) = proto_last_key_to_ks(key);
100                last_key = Some((pk, sk));
101            }
102
103            all_items.extend(items);
104        }
105
106        Ok(RemoteScanResponse {
107            items: all_items,
108            count: total_count,
109            scanned_count: total_scanned_count,
110            last_key,
111        })
112    }
113}
114
115impl Default for RemoteScan {
116    fn default() -> Self {
117        Self::new()
118    }
119}
120
121/// Scan response
122pub struct RemoteScanResponse {
123    /// Items found
124    pub items: Vec<Item>,
125    /// Number of items returned
126    pub count: usize,
127    /// Last evaluated key (for pagination)
128    pub last_key: Option<(Bytes, Option<Bytes>)>,
129    /// Number of items examined
130    pub scanned_count: usize,
131}
kstone_client/scan.rs

kstone_client/
scan.rs
