krabka_client_streams/runtime/iqv2/
query.rs

1//! `IQv2` query objects. Each builder lowers to a `store::iq::Iq2Query` with no
2//! serde, and the store supplies the actual serdes.
3
4use std::{any::Any, marker::PhantomData};
5
6use crate::store::{
7    iq::{Iq2Query, StoreKind},
8    versioned::VersionedRecord,
9};
10
11mod sealed {
12    pub trait Sealed {}
13}
14
15/// A typed `IQv2` query. `Result` is the type that `KafkaStreams::query` returns
16/// per partition. This trait is sealed, so only the in-crate query types
17/// implement it.
18pub trait Query: sealed::Sealed {
19    /// What a successful `QueryResult` carries.
20    type Result: 'static;
21    #[doc(hidden)]
22    fn store_kind(&self) -> StoreKind;
23    #[doc(hidden)]
24    fn lower(self) -> Iq2Query;
25}
26
27/// Single-key lookup. Result: `Option<V>`.
28pub struct KeyQuery<K, V> {
29    key: K,
30    _v: PhantomData<fn() -> V>,
31}
32impl<K, V> KeyQuery<K, V> {
33    #[must_use]
34    pub fn with_key(key: K) -> Self {
35        Self {
36            key,
37            _v: PhantomData,
38        }
39    }
40}
41impl<K: Send + Sync + 'static, V: 'static> sealed::Sealed for KeyQuery<K, V> {}
42impl<K: Send + Sync + 'static, V: 'static> Query for KeyQuery<K, V> {
43    type Result = Option<V>;
44    fn store_kind(&self) -> StoreKind {
45        StoreKind::KeyValue
46    }
47    fn lower(self) -> Iq2Query {
48        Iq2Query::Key {
49            key: Box::new(self.key),
50        }
51    }
52}
53
54/// Key-range scan. A `None` bound is unbounded on that side. Result:
55/// `Vec<(K,V)>`.
56pub struct RangeQuery<K, V> {
57    lo: Option<K>,
58    hi: Option<K>,
59    descending: bool,
60    _v: PhantomData<fn() -> V>,
61}
62impl<K, V> RangeQuery<K, V> {
63    #[must_use]
64    pub fn with_range(lo: K, hi: K) -> Self {
65        Self {
66            lo: Some(lo),
67            hi: Some(hi),
68            descending: false,
69            _v: PhantomData,
70        }
71    }
72    #[must_use]
73    pub fn with_lower_bound(lo: K) -> Self {
74        Self {
75            lo: Some(lo),
76            hi: None,
77            descending: false,
78            _v: PhantomData,
79        }
80    }
81    #[must_use]
82    pub fn with_upper_bound(hi: K) -> Self {
83        Self {
84            lo: None,
85            hi: Some(hi),
86            descending: false,
87            _v: PhantomData,
88        }
89    }
90    #[must_use]
91    pub fn with_no_bounds() -> Self {
92        Self {
93            lo: None,
94            hi: None,
95            descending: false,
96            _v: PhantomData,
97        }
98    }
99    #[must_use]
100    pub fn with_ascending_keys(mut self) -> Self {
101        self.descending = false;
102        self
103    }
104    #[must_use]
105    pub fn with_descending_keys(mut self) -> Self {
106        self.descending = true;
107        self
108    }
109}
110impl<K: Send + Sync + 'static, V: 'static> sealed::Sealed for RangeQuery<K, V> {}
111impl<K: Send + Sync + 'static, V: 'static> Query for RangeQuery<K, V> {
112    type Result = Vec<(K, V)>;
113    fn store_kind(&self) -> StoreKind {
114        StoreKind::KeyValue
115    }
116    fn lower(self) -> Iq2Query {
117        let bx = |k: K| -> Box<dyn Any + Send + Sync> { Box::new(k) };
118        Iq2Query::Range {
119            lo: self.lo.map(bx),
120            hi: self.hi.map(bx),
121            descending: self.descending,
122        }
123    }
124}
125
126/// One key, window starts in `[from_time, to_time]`. Result: `Vec<(i64, V)>`.
127pub struct WindowKeyQuery<K, V> {
128    key: K,
129    from_ts: i64,
130    to_ts: i64,
131    _v: PhantomData<fn() -> V>,
132}
133impl<K, V> WindowKeyQuery<K, V> {
134    #[must_use]
135    pub fn with_key(key: K) -> Self {
136        Self {
137            key,
138            from_ts: i64::MIN,
139            to_ts: i64::MAX,
140            _v: PhantomData,
141        }
142    }
143    #[must_use]
144    pub fn from_time(mut self, t: i64) -> Self {
145        self.from_ts = t;
146        self
147    }
148    #[must_use]
149    pub fn to_time(mut self, t: i64) -> Self {
150        self.to_ts = t;
151        self
152    }
153}
154impl<K: Send + Sync + 'static, V: 'static> sealed::Sealed for WindowKeyQuery<K, V> {}
155impl<K: Send + Sync + 'static, V: 'static> Query for WindowKeyQuery<K, V> {
156    type Result = Vec<(i64, V)>;
157    fn store_kind(&self) -> StoreKind {
158        StoreKind::Window
159    }
160    fn lower(self) -> Iq2Query {
161        Iq2Query::WindowKey {
162            key: Box::new(self.key),
163            from_ts: self.from_ts,
164            to_ts: self.to_ts,
165        }
166    }
167}
168
169/// Key range by window-start range. Result: `Vec<((K, i64), V)>`.
170pub struct WindowRangeQuery<K, V> {
171    lo: Option<K>,
172    hi: Option<K>,
173    from_ts: i64,
174    to_ts: i64,
175    _v: PhantomData<fn() -> V>,
176}
177impl<K, V> WindowRangeQuery<K, V> {
178    #[must_use]
179    pub fn with_key_range(lo: K, hi: K) -> Self {
180        Self {
181            lo: Some(lo),
182            hi: Some(hi),
183            from_ts: i64::MIN,
184            to_ts: i64::MAX,
185            _v: PhantomData,
186        }
187    }
188    #[must_use]
189    pub fn with_all_keys() -> Self {
190        Self {
191            lo: None,
192            hi: None,
193            from_ts: i64::MIN,
194            to_ts: i64::MAX,
195            _v: PhantomData,
196        }
197    }
198    #[must_use]
199    pub fn from_time(mut self, t: i64) -> Self {
200        self.from_ts = t;
201        self
202    }
203    #[must_use]
204    pub fn to_time(mut self, t: i64) -> Self {
205        self.to_ts = t;
206        self
207    }
208}
209impl<K: Send + Sync + 'static, V: 'static> sealed::Sealed for WindowRangeQuery<K, V> {}
210impl<K: Send + Sync + 'static, V: 'static> Query for WindowRangeQuery<K, V> {
211    type Result = Vec<((K, i64), V)>;
212    fn store_kind(&self) -> StoreKind {
213        StoreKind::Window
214    }
215    fn lower(self) -> Iq2Query {
216        let bx = |k: K| -> Box<dyn Any + Send + Sync> { Box::new(k) };
217        Iq2Query::WindowRange {
218            lo: self.lo.map(bx),
219            hi: self.hi.map(bx),
220            from_ts: self.from_ts,
221            to_ts: self.to_ts,
222        }
223    }
224}
225
226/// Single versioned-key lookup (KIP-960). `as_of = None` gives the latest live
227/// version. Result: `Option<VersionedRecord<V>>`.
228pub struct VersionedKeyQuery<K, V> {
229    key: K,
230    as_of: Option<i64>,
231    _v: PhantomData<fn() -> V>,
232}
233impl<K, V> VersionedKeyQuery<K, V> {
234    #[must_use]
235    pub fn with_key(key: K) -> Self {
236        Self {
237            key,
238            as_of: None,
239            _v: PhantomData,
240        }
241    }
242    #[must_use]
243    pub fn as_of(mut self, timestamp: i64) -> Self {
244        self.as_of = Some(timestamp);
245        self
246    }
247}
248impl<K: Send + Sync + 'static, V: 'static> sealed::Sealed for VersionedKeyQuery<K, V> {}
249impl<K: Send + Sync + 'static, V: 'static> Query for VersionedKeyQuery<K, V> {
250    type Result = Option<VersionedRecord<V>>;
251    fn store_kind(&self) -> StoreKind {
252        StoreKind::Versioned
253    }
254    fn lower(self) -> Iq2Query {
255        Iq2Query::VersionedKey {
256            key: Box::new(self.key),
257            as_of: self.as_of,
258        }
259    }
260}
261
262/// All versions of a key whose validity overlaps `[from_time, to_time]`
263/// (KIP-968). A `None` bound is unbounded on that side. The order is ascending by
264/// `valid_from`, unless you call `with_descending_timestamps()`. Result:
265/// `Vec<VersionedRecord<V>>`.
266pub struct MultiVersionedKeyQuery<K, V> {
267    key: K,
268    from_ts: Option<i64>,
269    to_ts: Option<i64>,
270    descending: bool,
271    _v: PhantomData<fn() -> V>,
272}
273impl<K, V> MultiVersionedKeyQuery<K, V> {
274    #[must_use]
275    pub fn with_key(key: K) -> Self {
276        Self {
277            key,
278            from_ts: None,
279            to_ts: None,
280            descending: false,
281            _v: PhantomData,
282        }
283    }
284    #[must_use]
285    pub fn from_time(mut self, t: i64) -> Self {
286        self.from_ts = Some(t);
287        self
288    }
289    #[must_use]
290    pub fn to_time(mut self, t: i64) -> Self {
291        self.to_ts = Some(t);
292        self
293    }
294    #[must_use]
295    pub fn with_ascending_timestamps(mut self) -> Self {
296        self.descending = false;
297        self
298    }
299    #[must_use]
300    pub fn with_descending_timestamps(mut self) -> Self {
301        self.descending = true;
302        self
303    }
304}
305impl<K: Send + Sync + 'static, V: 'static> sealed::Sealed for MultiVersionedKeyQuery<K, V> {}
306impl<K: Send + Sync + 'static, V: 'static> Query for MultiVersionedKeyQuery<K, V> {
307    type Result = Vec<VersionedRecord<V>>;
308    fn store_kind(&self) -> StoreKind {
309        StoreKind::Versioned
310    }
311    fn lower(self) -> Iq2Query {
312        Iq2Query::MultiVersionedKey {
313            key: Box::new(self.key),
314            from_ts: self.from_ts,
315            to_ts: self.to_ts,
316            descending: self.descending,
317        }
318    }
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324
325    #[test]
326    fn lowering_picks_the_right_variant_and_kind() {
327        let kq = KeyQuery::<String, i64>::with_key("a".into());
328        assert_eq!(kq.store_kind(), StoreKind::KeyValue);
329        assert!(matches!(kq.lower(), Iq2Query::Key { .. }));
330
331        let rq = RangeQuery::<String, i64>::with_lower_bound("a".into()).with_descending_keys();
332        assert!(matches!(
333            rq.lower(),
334            Iq2Query::Range {
335                lo: Some(_),
336                hi: None,
337                descending: true
338            }
339        ));
340
341        let wk = WindowKeyQuery::<String, i64>::with_key("a".into())
342            .from_time(0)
343            .to_time(9);
344        assert_eq!(wk.store_kind(), StoreKind::Window);
345        assert!(matches!(
346            wk.lower(),
347            Iq2Query::WindowKey {
348                from_ts: 0,
349                to_ts: 9,
350                ..
351            }
352        ));
353
354        let wr = WindowRangeQuery::<String, i64>::with_all_keys();
355        assert!(matches!(
356            wr.lower(),
357            Iq2Query::WindowRange {
358                lo: None,
359                hi: None,
360                ..
361            }
362        ));
363    }
364
365    #[test]
366    fn versioned_queries_lower_correctly() {
367        let vk = VersionedKeyQuery::<String, i64>::with_key("k".into()).as_of(250);
368        assert_eq!(vk.store_kind(), StoreKind::Versioned);
369        assert!(matches!(
370            vk.lower(),
371            Iq2Query::VersionedKey {
372                as_of: Some(250),
373                ..
374            }
375        ));
376
377        let vk_latest = VersionedKeyQuery::<String, i64>::with_key("k".into());
378        assert!(matches!(
379            vk_latest.lower(),
380            Iq2Query::VersionedKey { as_of: None, .. }
381        ));
382
383        let mv = MultiVersionedKeyQuery::<String, i64>::with_key("k".into())
384            .from_time(150)
385            .to_time(250)
386            .with_descending_timestamps();
387        assert_eq!(mv.store_kind(), StoreKind::Versioned);
388        assert!(matches!(
389            mv.lower(),
390            Iq2Query::MultiVersionedKey {
391                from_ts: Some(150),
392                to_ts: Some(250),
393                descending: true,
394                ..
395            }
396        ));
397
398        let mv_all = MultiVersionedKeyQuery::<String, i64>::with_key("k".into());
399        assert!(matches!(
400            mv_all.lower(),
401            Iq2Query::MultiVersionedKey {
402                from_ts: None,
403                to_ts: None,
404                descending: false,
405                ..
406            }
407        ));
408
409        // `with_ascending_timestamps()` is the default; calling it explicitly
410        // (after a descending toggle) must restore ascending.
411        let mv_asc = MultiVersionedKeyQuery::<String, i64>::with_key("k".into())
412            .with_descending_timestamps()
413            .with_ascending_timestamps();
414        assert!(matches!(
415            mv_asc.lower(),
416            Iq2Query::MultiVersionedKey {
417                descending: false,
418                ..
419            }
420        ));
421    }
422
423    #[test]
424    fn range_query_bound_variants_lower_correctly() {
425        let both = RangeQuery::<String, i64>::with_range("a".into(), "b".into());
426        assert_eq!(both.store_kind(), StoreKind::KeyValue);
427        assert!(matches!(
428            both.lower(),
429            Iq2Query::Range {
430                lo: Some(_),
431                hi: Some(_),
432                descending: false,
433            }
434        ));
435
436        let upper = RangeQuery::<String, i64>::with_upper_bound("b".into());
437        assert!(matches!(
438            upper.lower(),
439            Iq2Query::Range {
440                lo: None,
441                hi: Some(_),
442                descending: false,
443            }
444        ));
445
446        let none = RangeQuery::<String, i64>::with_no_bounds().with_ascending_keys();
447        assert!(matches!(
448            none.lower(),
449            Iq2Query::Range {
450                lo: None,
451                hi: None,
452                descending: false,
453            }
454        ));
455    }
456
457    #[test]
458    fn window_key_query_default_and_explicit_bounds() {
459        let dflt = WindowKeyQuery::<String, i64>::with_key("a".into());
460        assert!(matches!(
461            dflt.lower(),
462            Iq2Query::WindowKey {
463                from_ts: i64::MIN,
464                to_ts: i64::MAX,
465                ..
466            }
467        ));
468
469        let bounded = WindowKeyQuery::<String, i64>::with_key("a".into())
470            .from_time(5)
471            .to_time(9);
472        assert!(matches!(
473            bounded.lower(),
474            Iq2Query::WindowKey {
475                from_ts: 5,
476                to_ts: 9,
477                ..
478            }
479        ));
480    }
481
482    #[test]
483    fn window_range_query_key_range_and_all_keys() {
484        let ranged = WindowRangeQuery::<String, i64>::with_key_range("a".into(), "b".into())
485            .from_time(1)
486            .to_time(2);
487        assert_eq!(ranged.store_kind(), StoreKind::Window);
488        assert!(matches!(
489            ranged.lower(),
490            Iq2Query::WindowRange {
491                lo: Some(_),
492                hi: Some(_),
493                from_ts: 1,
494                to_ts: 2,
495            }
496        ));
497
498        let all = WindowRangeQuery::<String, i64>::with_all_keys();
499        assert!(matches!(
500            all.lower(),
501            Iq2Query::WindowRange {
502                lo: None,
503                hi: None,
504                ..
505            }
506        ));
507    }
508}
krabka_client_streams/runtime/iqv2/query.rs

krabka_client_streams/runtime/iqv2/
query.rs
