krabka_client_streams/store/
iq.rs

1//! Byte-level read surface for Interactive Queries.
2//!
3//! The supervisor calls these methods through `&dyn StateStore::as_iq()` to serve
4//! `KafkaStreams::*_store` queries, and it never learns `K` or `V`. The typed
5//! view owns the serialization and deserialization. All reads take `&self`, and
6//! only key and value **bytes** cross this trait.
7
8use std::any::Any;
9
10use async_trait::async_trait;
11use bytes::Bytes;
12
13/// Which kind of store a query targets. Public so it can appear in `IqError`.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum StoreKind {
16    KeyValue,
17    Window,
18    Session,
19    Versioned,
20}
21
22/// A typed `IQv2` query lowered to the store boundary.
23///
24/// Keys travel as `Box<dyn Any + Send + Sync>`, which holds the raw `K`. The
25/// concrete store downcasts to its own `K`, serializes with its own key serde,
26/// runs the op, and returns the typed result boxed as `Box<dyn Any + Send>`.
27/// That result is an `Option<V>`, a `Vec<(K,V)>`, or a similar type.
28///
29/// Time bounds are plain `i64`, and the ordering and bound choices are flags. No
30/// serde and no `K` or `V` appear here. That is the whole point of the
31/// byte-level boundary.
32pub enum Iq2Query {
33    /// `KeyQuery`: single key. Result: `Option<V>`.
34    Key { key: Box<dyn Any + Send + Sync> },
35    /// `RangeQuery`: a `None` bound is unbounded on that side. Result:
36    /// `Vec<(K,V)>`.
37    Range {
38        lo: Option<Box<dyn Any + Send + Sync>>,
39        hi: Option<Box<dyn Any + Send + Sync>>,
40        descending: bool,
41    },
42    /// `WindowKeyQuery`: one key, window starts in `[from_ts, to_ts]`.
43    /// Result: `Vec<(i64 /*windowStart*/, V)>`, ascending by start.
44    WindowKey {
45        key: Box<dyn Any + Send + Sync>,
46        from_ts: i64,
47        to_ts: i64,
48    },
49    /// `WindowRangeQuery`: key range by window-start range. A `None` bound is
50    /// unbounded on that side. Result: `Vec<((K, i64 /*windowStart*/), V)>`,
51    /// ascending by (key bytes, windowStart).
52    WindowRange {
53        lo: Option<Box<dyn Any + Send + Sync>>,
54        hi: Option<Box<dyn Any + Send + Sync>>,
55        from_ts: i64,
56        to_ts: i64,
57    },
58    /// `VersionedKeyQuery` (KIP-960): one key. `as_of = None` gives the latest
59    /// live version, and `Some(t)` gives the version valid at `t`. Result:
60    /// `Option<VersionedRecord<V>>`.
61    VersionedKey {
62        key: Box<dyn Any + Send + Sync>,
63        as_of: Option<i64>,
64    },
65    /// `MultiVersionedKeyQuery` (KIP-968): one key, and every version whose
66    /// validity `[valid_from, valid_to)` overlaps `[from_ts, to_ts]`. A `None`
67    /// bound is unbounded on that side. The order is ascending by `valid_from`,
68    /// unless `descending` is set. Result: `Vec<VersionedRecord<V>>`.
69    MultiVersionedKey {
70        key: Box<dyn Any + Send + Sync>,
71        from_ts: Option<i64>,
72        to_ts: Option<i64>,
73        descending: bool,
74    },
75}
76
77/// Why a store could not run an `IQv2` query. The runtime maps these variants
78/// into the public `FailureReason`, together with its own conditions:
79/// rebalancing, not-up-to-bound, and not-active.
80#[derive(Debug, Clone, PartialEq, Eq)]
81pub enum Iq2Failure {
82    /// This store kind has no handler for the requested query variant.
83    UnknownQueryType,
84    /// A key `Box<dyn Any>` did not downcast to this store's `K`.
85    KeyTypeMismatch,
86}
87
88/// Byte-level IQ reads. The three materialized `*Bytes` stores implement this
89/// trait. The default methods return empty results, so a non-matching store kind
90/// never produces wrong data. The supervisor's `kind()` check catches that case
91/// earlier.
92#[doc(hidden)]
93#[async_trait]
94pub trait IqQueryable: Send + Sync {
95    fn kind(&self) -> StoreKind;
96
97    async fn iq_kv_get(&self, _key: &[u8]) -> Option<Bytes> {
98        None
99    }
100    /// Inclusive `[lo, hi]` in memcmp order.
101    async fn iq_kv_range(&self, _lo: &[u8], _hi: &[u8]) -> Vec<(Bytes, Bytes)> {
102        Vec::new()
103    }
104    async fn iq_kv_all(&self) -> Vec<(Bytes, Bytes)> {
105        Vec::new()
106    }
107    async fn iq_kv_approx_count(&self) -> u64 {
108        0
109    }
110
111    async fn iq_window_fetch_single(&self, _key: &[u8], _window_start: i64) -> Option<Bytes> {
112        None
113    }
114    /// Ascending by window start, inclusive `[time_from, time_to]`.
115    async fn iq_window_fetch(
116        &self,
117        _key: &[u8],
118        _time_from: i64,
119        _time_to: i64,
120    ) -> Vec<(i64, Bytes)> {
121        Vec::new()
122    }
123
124    /// All sessions for `key`, store order. Tuple is `(start, end)`.
125    async fn iq_session_fetch_key(&self, _key: &[u8]) -> Vec<((i64, i64), Bytes)> {
126        Vec::new()
127    }
128
129    /// Latest live version: `(validFrom, validTo=None, valueBytes)`.
130    async fn iq_versioned_get(&self, _key: &[u8]) -> Option<(i64, Option<i64>, Bytes)> {
131        None
132    }
133    /// Version valid at `as_of`: `(validFrom, validTo, valueBytes)`.
134    async fn iq_versioned_get_as_of(
135        &self,
136        _key: &[u8],
137        _as_of: i64,
138    ) -> Option<(i64, Option<i64>, Bytes)> {
139        None
140    }
141
142    /// `IQv2` entry point. The store downcasts the keys, serializes and
143    /// deserializes with its own serdes, runs the op, and returns the typed
144    /// result boxed. The default is that this store kind handles no `IQv2` query
145    /// variant.
146    async fn iq2_execute(&self, _query: &Iq2Query) -> Result<Box<dyn Any + Send>, Iq2Failure> {
147        Err(Iq2Failure::UnknownQueryType)
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use krabka_units::prelude::*;
154
155    use super::*;
156    use crate::{
157        processor::serde::{I64Serde, Serde, StringSerde},
158        store::{
159            api::{KeyValueStore, StateStore},
160            kv::KeyValueBytesStore,
161            session::{SessionBytesStore, SessionStore},
162            window::{WindowBytesStore, WindowStore},
163        },
164    };
165
166    #[tokio::test]
167    async fn kv_get_range_all_count_inclusive() {
168        let mut s = KeyValueBytesStore::<String, i64>::in_memory(
169            "c".into(),
170            Box::new(StringSerde),
171            Box::new(I64Serde),
172            "c-changelog".into(),
173        );
174        for (k, v) in [("a", 1), ("b", 2), ("c", 3)] {
175            s.put(k.into(), v).await;
176        }
177        let q: &dyn IqQueryable = s.as_iq().unwrap();
178        assert_eq!(q.iq_kv_get(b"b").await, Some(I64Serde.serialize("t", &2)));
179        assert_eq!(q.iq_kv_get(b"z").await, None);
180        let r = q.iq_kv_range(b"a", b"b").await; // inclusive => a,b
181        assert_eq!(r.len(), 2);
182        assert_eq!(r[0].0, bytes::Bytes::from_static(b"a"));
183        assert_eq!(r[1].0, bytes::Bytes::from_static(b"b"));
184        assert!(q.iq_kv_range(b"c", b"a").await.is_empty()); // lo>hi => empty
185        assert_eq!(q.iq_kv_all().await.len(), 3);
186        assert_eq!(q.iq_kv_approx_count().await, 3);
187    }
188
189    #[tokio::test]
190    async fn window_fetch_point_and_range() {
191        let mut s = WindowBytesStore::<String, i64>::in_memory(
192            "w".into(),
193            Box::new(StringSerde),
194            Box::new(I64Serde),
195            "w-changelog".into(),
196            secs(1),
197        );
198        s.put("k".into(), 0, 10, 5).await;
199        s.put("k".into(), 1000, 20, 1005).await;
200        let q: &dyn IqQueryable = s.as_iq().unwrap();
201        assert_eq!(
202            q.iq_window_fetch_single(b"k", 0).await,
203            Some(I64Serde.serialize("t", &10))
204        );
205        assert_eq!(q.iq_window_fetch_single(b"k", 500).await, None);
206        let r = q.iq_window_fetch(b"k", 0, 1000).await;
207        assert_eq!(r.iter().map(|(t, _)| *t).collect::<Vec<_>>(), vec![0, 1000]);
208    }
209
210    #[tokio::test]
211    async fn versioned_get_latest_and_as_of() {
212        use crate::store::versioned::{VersionedBytesStore, VersionedKeyValueStore};
213        let mut s = VersionedBytesStore::<String, i64>::in_memory(
214            "v".into(),
215            secs(1_000),
216            Box::new(StringSerde),
217            Box::new(I64Serde),
218            "v-changelog".into(),
219        );
220        s.put("k".into(), Some(10), 100).await;
221        s.put("k".into(), Some(20), 200).await;
222        let q: &dyn IqQueryable = s.as_iq().unwrap();
223        assert_eq!(q.kind(), StoreKind::Versioned);
224        let (vf, vt, raw) = q.iq_versioned_get(b"k").await.unwrap();
225        assert_eq!((vf, vt), (200, None));
226        assert_eq!(raw, I64Serde.serialize("t", &20));
227        let (vf2, vt2, raw2) = q.iq_versioned_get_as_of(b"k", 150).await.unwrap();
228        assert_eq!((vf2, vt2), (100, Some(200)));
229        assert_eq!(raw2, I64Serde.serialize("t", &10));
230        assert_eq!(q.iq_versioned_get_as_of(b"k", 50).await, None);
231    }
232
233    #[tokio::test]
234    async fn session_fetch_key_carries_start_end() {
235        let mut s = SessionBytesStore::<String, i64>::in_memory(
236            "s".into(),
237            Box::new(StringSerde),
238            Box::new(I64Serde),
239            "s-changelog".into(),
240        );
241        s.put("k".into(), 0, 10, 1).await;
242        s.put("k".into(), 20, 30, 2).await;
243        let q: &dyn IqQueryable = s.as_iq().unwrap();
244        let r = q.iq_session_fetch_key(b"k").await;
245        let windows: Vec<(i64, i64)> = r.iter().map(|((st, en), _)| (*st, *en)).collect();
246        assert!(windows.contains(&(0, 10)) && windows.contains(&(20, 30)));
247    }
248
249    #[tokio::test]
250    async fn iq2_execute_default_is_unknown_query_type() {
251        use super::{Iq2Failure, Iq2Query};
252        // A session store has no `IQv2` handler — default impl must reject.
253        let s = SessionBytesStore::<String, i64>::in_memory(
254            "s".into(),
255            Box::new(StringSerde),
256            Box::new(I64Serde),
257            "s-changelog".into(),
258        );
259        let q: &dyn IqQueryable = s.as_iq().unwrap();
260        let query = Iq2Query::Key {
261            key: Box::new("k".to_string()),
262        };
263        assert_eq!(
264            q.iq2_execute(&query).await.err(),
265            Some(Iq2Failure::UnknownQueryType)
266        );
267
268        // Versioned variants also hit the default (a session store has no handler).
269        let mv = Iq2Query::MultiVersionedKey {
270            key: Box::new("k".to_string()),
271            from_ts: None,
272            to_ts: None,
273            descending: false,
274        };
275        assert_eq!(
276            q.iq2_execute(&mv).await.err(),
277            Some(Iq2Failure::UnknownQueryType)
278        );
279    }
280}
krabka_client_streams/store/iq.rs

krabka_client_streams/store/
iq.rs
