Skip to main content

google_cloud_spanner/
key.rs

1// Copyright 2026 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::value::Value;
16
17/// Represents a primary key or index key for Cloud Spanner.
18///
19/// # Example
20/// ```
21/// use google_cloud_spanner::key;
22///
23/// let key = key![1_i64, "Alice"];
24/// ```
25///
26/// Keys are constructed by appending values in the correct column order.
27#[macro_export]
28macro_rules! key {
29    ($($val:expr),* $(,)?) => {
30        $crate::key::Key::new(vec![
31            $($crate::value::ToValue::to_value(&$val)),*
32        ])
33    };
34}
35
36/// A primary key or index key in Spanner.
37///
38/// A `Key` consists of an ordered collection of one or more Spanner [`Value`]s.
39/// The elements must match the correct column order as defined by the table or
40/// index schema.
41///
42/// For a simpler way to construct a `Key` instance, use the [`key!`] macro.
43#[derive(Clone, Debug, PartialEq)]
44pub struct Key {
45    pub(crate) values: Vec<Value>,
46}
47
48impl Key {
49    /// Creates a new Key from a vector of values.
50    pub fn new(values: Vec<Value>) -> Self {
51        Key { values }
52    }
53
54    pub(crate) fn into_values(self) -> Vec<serde_json::Value> {
55        self.values
56            .into_iter()
57            .map(|v| v.into_serde_value())
58            .collect()
59    }
60}
61
62impl Default for Key {
63    fn default() -> Self {
64        Self::new(vec![])
65    }
66}
67
68/// Defines whether a boundary of a key range is open (exclusive) or closed (inclusive).
69#[derive(Clone, Debug, PartialEq)]
70pub(crate) enum Endpoint {
71    Closed(Key),
72    Open(Key),
73}
74
75/// Represents a contiguous range of keys.
76///
77/// # Example
78/// ```
79/// use google_cloud_spanner::key::KeyRange;
80/// use google_cloud_spanner::key;
81///
82/// let range = KeyRange::closed_open(key![1_i64], key![10_i64]);
83/// ```
84#[derive(Clone, Debug, PartialEq)]
85pub struct KeyRange {
86    pub(crate) start: Endpoint,
87    pub(crate) end: Endpoint,
88}
89
90impl KeyRange {
91    /// Creates a key for the range [start, end).
92    ///
93    /// # Example
94    /// ```
95    /// use google_cloud_spanner::key::KeyRange;
96    /// use google_cloud_spanner::key;
97    ///
98    /// // Creates a key for the range [1, 10)
99    /// let range = KeyRange::closed_open(key![1_i64], key![10_i64]);
100    /// ```
101    pub fn closed_open(start: Key, end: Key) -> Self {
102        KeyRange {
103            start: Endpoint::Closed(start),
104            end: Endpoint::Open(end),
105        }
106    }
107
108    /// Creates a key for the range [start, end].
109    ///
110    /// # Example
111    /// ```
112    /// use google_cloud_spanner::key::KeyRange;
113    /// use google_cloud_spanner::key;
114    ///
115    /// // Creates a key for the range [1, 10]
116    /// let range = KeyRange::closed_closed(key![1_i64], key![10_i64]);
117    /// ```
118    pub fn closed_closed(start: Key, end: Key) -> Self {
119        KeyRange {
120            start: Endpoint::Closed(start),
121            end: Endpoint::Closed(end),
122        }
123    }
124
125    /// Creates a key for the range (start, end].
126    ///
127    /// # Example
128    /// ```
129    /// use google_cloud_spanner::key::KeyRange;
130    /// use google_cloud_spanner::key;
131    ///
132    /// // Creates a key for the range (1, 10]
133    /// let range = KeyRange::open_closed(key![1_i64], key![10_i64]);
134    /// ```
135    pub fn open_closed(start: Key, end: Key) -> Self {
136        KeyRange {
137            start: Endpoint::Open(start),
138            end: Endpoint::Closed(end),
139        }
140    }
141
142    /// Creates a key for the range (start, end).
143    ///
144    /// # Example
145    /// ```
146    /// use google_cloud_spanner::key::KeyRange;
147    /// use google_cloud_spanner::key;
148    ///
149    /// // Creates a key for the range (1, 10)
150    /// let range = KeyRange::open_open(key![1_i64], key![10_i64]);
151    /// ```
152    pub fn open_open(start: Key, end: Key) -> Self {
153        KeyRange {
154            start: Endpoint::Open(start),
155            end: Endpoint::Open(end),
156        }
157    }
158
159    #[allow(dead_code)]
160    pub(crate) fn into_proto(self) -> crate::model::KeyRange {
161        let mut proto = crate::model::KeyRange::new();
162
163        proto = match self.start {
164            Endpoint::Closed(start) => proto.set_start_closed(start.into_values()),
165            Endpoint::Open(start) => proto.set_start_open(start.into_values()),
166        };
167
168        match self.end {
169            Endpoint::Closed(end) => proto.set_end_closed(end.into_values()),
170            Endpoint::Open(end) => proto.set_end_open(end.into_values()),
171        }
172    }
173}
174
175/// A collection of Spanner keys and key ranges.
176///
177/// # Example
178/// ```
179/// use google_cloud_spanner::key::KeySet;
180/// use google_cloud_spanner::key::KeyRange;
181/// use google_cloud_spanner::key;
182///
183/// let keyset = KeySet::builder()
184///     .add_key(key![1_i64])
185///     .add_range(KeyRange::closed_open(key![10_i64], key![100_i64]))
186///     .build();
187/// ```
188///
189/// Defines a collection of Cloud Spanner keys and/or key ranges. All the keys are expected to be in
190/// the same table or index. The keys need not be sorted in any particular way.
191///
192/// `KeySet`s are used for delete mutations and reads.
193///
194/// If the same key is specified multiple times in the set (for example if two ranges, two keys,
195/// or a key and a range overlap), the Cloud Spanner backend behaves as if the key were only
196/// specified once. `KeySet` instances are immutable.
197#[derive(Clone, Debug, PartialEq)]
198pub struct KeySet {
199    pub(crate) keys: Vec<Key>,
200    pub(crate) ranges: Vec<KeyRange>,
201    pub(crate) all: bool,
202}
203
204impl KeySet {
205    /// Creates a builder for `KeySet`.
206    pub fn builder() -> KeySetBuilder {
207        KeySetBuilder::new()
208    }
209
210    /// Creates a `KeySet` that matches all rows.
211    pub fn all() -> Self {
212        KeySet {
213            keys: vec![],
214            ranges: vec![],
215            all: true,
216        }
217    }
218
219    #[allow(dead_code)]
220    pub(crate) fn into_proto(self) -> crate::model::KeySet {
221        let mut proto = crate::model::KeySet::new();
222        if self.all {
223            proto = proto.set_all(true);
224        }
225
226        // Convert keys
227        let keys_proto: Vec<Vec<serde_json::Value>> =
228            self.keys.into_iter().map(|k| k.into_values()).collect();
229        if !keys_proto.is_empty() {
230            proto = proto.set_keys(keys_proto);
231        }
232
233        // Convert ranges
234        let ranges_proto: Vec<crate::model::KeyRange> =
235            self.ranges.into_iter().map(|r| r.into_proto()).collect();
236        if !ranges_proto.is_empty() {
237            proto = proto.set_ranges(ranges_proto);
238        }
239
240        proto
241    }
242}
243
244/// A builder for constructing a `KeySet`.
245#[derive(Clone, Debug)]
246pub struct KeySetBuilder {
247    keys: Vec<Key>,
248    ranges: Vec<KeyRange>,
249}
250
251impl KeySetBuilder {
252    /// Creates a new, empty builder.
253    pub fn new() -> Self {
254        KeySetBuilder {
255            keys: vec![],
256            ranges: vec![],
257        }
258    }
259
260    /// Adds a key to the key set.
261    pub fn add_key(mut self, key: Key) -> Self {
262        self.keys.push(key);
263        self
264    }
265
266    /// Adds a range to the key set.
267    pub fn add_range(mut self, range: KeyRange) -> Self {
268        self.ranges.push(range);
269        self
270    }
271
272    /// Builds the `KeySet`.
273    pub fn build(self) -> KeySet {
274        KeySet {
275            keys: self.keys,
276            ranges: self.ranges,
277            all: false,
278        }
279    }
280}
281
282impl Default for KeySetBuilder {
283    fn default() -> Self {
284        Self::new()
285    }
286}
287
288impl Default for KeySet {
289    fn default() -> Self {
290        KeySet::builder().build()
291    }
292}
293
294impl From<Key> for KeySet {
295    fn from(key: Key) -> Self {
296        KeySet::builder().add_key(key).build()
297    }
298}
299
300impl From<KeyRange> for KeySet {
301    fn from(range: KeyRange) -> Self {
302        KeySet::builder().add_range(range).build()
303    }
304}
305
306#[cfg(test)]
307mod tests {
308    use super::*;
309
310    #[test]
311    fn auto_traits() {
312        static_assertions::assert_impl_all!(Endpoint: Send, Sync, Clone, std::fmt::Debug);
313        static_assertions::assert_impl_all!(Key: Send, Sync, Clone, std::fmt::Debug);
314        static_assertions::assert_impl_all!(KeyRange: Send, Sync, Clone, std::fmt::Debug);
315        static_assertions::assert_impl_all!(KeySet: Send, Sync, Clone, std::fmt::Debug);
316        static_assertions::assert_impl_all!(KeySetBuilder: Send, Sync, Clone, std::fmt::Debug);
317    }
318
319    #[test]
320    fn key_builder() {
321        let key = key![1_i64, "test"];
322        assert_eq!(key.values.len(), 2);
323    }
324
325    #[test]
326    fn key_default() {
327        let key = Key::default();
328        assert_eq!(key.values.len(), 0);
329    }
330
331    #[test]
332    fn key_into_values() {
333        let key = key![1_i64];
334        let values = key.into_values();
335        assert_eq!(values.len(), 1);
336        assert_eq!(values[0], serde_json::json!("1"));
337    }
338
339    #[test]
340    fn keyrange_factories() {
341        let start = key![1_i64];
342        let end = key![10_i64];
343
344        let r1 = KeyRange::closed_open(start.clone(), end.clone());
345        assert_eq!(r1.start, Endpoint::Closed(start.clone()));
346        assert_eq!(r1.end, Endpoint::Open(end.clone()));
347
348        let r2 = KeyRange::closed_closed(start.clone(), end.clone());
349        assert_eq!(r2.start, Endpoint::Closed(start.clone()));
350        assert_eq!(r2.end, Endpoint::Closed(end.clone()));
351
352        let r3 = KeyRange::open_closed(start.clone(), end.clone());
353        assert_eq!(r3.start, Endpoint::Open(start.clone()));
354        assert_eq!(r3.end, Endpoint::Closed(end.clone()));
355
356        let r4 = KeyRange::open_open(start.clone(), end.clone());
357        assert_eq!(r4.start, Endpoint::Open(start.clone()));
358        assert_eq!(r4.end, Endpoint::Open(end.clone()));
359    }
360
361    #[test]
362    fn keyrange_into_proto() {
363        let r1 = KeyRange::closed_open(key![1_i64], key![10_i64]);
364        let proto1 = r1.into_proto();
365        assert_eq!(proto1.start_closed().unwrap().len(), 1);
366        assert_eq!(proto1.end_open().unwrap().len(), 1);
367
368        let r2 = KeyRange::open_closed(key![1_i64], key![10_i64]);
369        let proto2 = r2.into_proto();
370        assert_eq!(proto2.start_open().unwrap().len(), 1);
371        assert_eq!(proto2.end_closed().unwrap().len(), 1);
372    }
373
374    #[test]
375    fn keyset_builder() {
376        let key = key![1_i64];
377        let range = KeyRange::closed_open(key![2_i64], key![5_i64]);
378        let keyset = KeySet::builder().add_key(key).add_range(range).build();
379
380        assert_eq!(keyset.keys.len(), 1);
381        assert_eq!(keyset.ranges.len(), 1);
382        assert!(!keyset.all);
383    }
384
385    #[test]
386    fn keyset_builder_default() {
387        let builder = KeySetBuilder::default();
388        let keyset = builder.build();
389        assert_eq!(keyset.keys.len(), 0);
390        assert_eq!(keyset.ranges.len(), 0);
391    }
392
393    #[test]
394    fn keyset_default() {
395        let keyset = KeySet::default();
396        assert_eq!(keyset.keys.len(), 0);
397        assert_eq!(keyset.ranges.len(), 0);
398        assert!(!keyset.all);
399    }
400
401    #[test]
402    fn keyset_from_key() {
403        let keyset: KeySet = key![1_i64].into();
404        assert_eq!(keyset.keys.len(), 1);
405        assert_eq!(keyset.ranges.len(), 0);
406    }
407
408    #[test]
409    fn keyset_from_keyrange() {
410        let range = KeyRange::closed_open(key![2_i64], key![5_i64]);
411        let keyset: KeySet = range.into();
412        assert_eq!(keyset.keys.len(), 0);
413        assert_eq!(keyset.ranges.len(), 1);
414    }
415
416    #[test]
417    fn keyset_into_proto() {
418        let keyset = KeySet::all();
419        let proto = keyset.into_proto();
420        assert!(proto.all);
421
422        let keyset2 = KeySet::builder()
423            .add_key(key![1_i64])
424            .add_range(KeyRange::closed_open(key![2_i64], key![5_i64]))
425            .build();
426        let proto2 = keyset2.into_proto();
427        assert_eq!(proto2.keys.len(), 1);
428        assert_eq!(proto2.ranges.len(), 1);
429    }
430}