1use crate::codes::QUERY_OP_VERSION;
2use serde::{Deserialize, Serialize};
3use std::collections::BTreeMap;
4
5#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
7pub struct KeyMatch {
8 pub field: String,
9 pub value: String,
10}
11
12impl KeyMatch {
13 pub fn new(field: impl Into<String>, value: impl Into<String>) -> Self {
15 Self {
16 field: field.into(),
17 value: value.into(),
18 }
19 }
20}
21
22#[derive(Clone, Debug, Default, Serialize, Deserialize)]
25#[cfg_attr(feature = "builders", derive(bon::Builder))]
26pub struct Query {
27 #[cfg_attr(feature = "builders", builder(into))]
29 pub index: String,
30 #[cfg_attr(feature = "builders", builder(default))]
31 #[serde(default, skip_serializing_if = "Vec::is_empty")]
32 pub by_key: Vec<KeyMatch>,
33 #[cfg_attr(feature = "builders", builder(into))]
34 #[serde(default, skip_serializing_if = "Option::is_none")]
35 pub message_type: Option<String>,
36 #[serde(default, skip_serializing_if = "Option::is_none")]
38 pub time_range: Option<(u64, u64)>,
39 #[serde(default, skip_serializing_if = "Option::is_none")]
42 pub filter: Option<Filter>,
43 #[serde(default, skip_serializing_if = "Option::is_none")]
44 pub vector: Option<VectorQuery>,
45 #[serde(default, skip_serializing_if = "Option::is_none")]
49 pub text: Option<TextQuery>,
50 #[cfg_attr(feature = "builders", builder(default))]
51 #[serde(default, skip_serializing_if = "Vec::is_empty")]
52 pub order: Vec<Sort>,
53 #[cfg_attr(feature = "builders", builder(default = 50))]
56 pub limit: usize,
57 #[cfg_attr(feature = "builders", builder(default))]
58 #[serde(default)]
59 pub offset: usize,
60 #[serde(default, skip_serializing_if = "Option::is_none")]
62 pub aggregate: Option<Aggregate>,
63 #[serde(default, skip_serializing_if = "Option::is_none")]
66 pub having: Option<Filter>,
67 #[cfg_attr(feature = "builders", builder(default))]
69 #[serde(default, skip_serializing_if = "is_false")]
70 pub distinct: bool,
71 #[cfg_attr(feature = "builders", builder(default))]
72 #[serde(default)]
73 pub select: Select,
74 #[cfg_attr(feature = "builders", builder(into))]
78 #[serde(default, skip_serializing_if = "Option::is_none")]
79 pub fork: Option<String>,
80 #[serde(default, skip_serializing_if = "Option::is_none")]
82 pub raw_sql: Option<RawSql>,
83 #[cfg_attr(feature = "builders", builder(default))]
86 #[serde(default, skip_serializing_if = "Consistency::is_eventual")]
87 pub consistency: Consistency,
88}
89
90fn is_false(value: &bool) -> bool {
91 !*value
92}
93
94#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
105#[serde(rename_all = "snake_case")]
106#[non_exhaustive]
107pub enum Consistency {
108 #[default]
112 Eventual,
113 ReadYourWrites,
119 Strong,
124}
125
126impl Consistency {
127 pub fn is_eventual(&self) -> bool {
129 matches!(self, Consistency::Eventual)
130 }
131}
132
133#[derive(Clone, Copy, Debug, PartialEq, Eq)]
146pub struct ConsistencyGate {
147 pub applied: u64,
149 pub required: u64,
151}
152
153impl ConsistencyGate {
154 pub fn new(applied: u64, required: u64) -> Self {
157 Self { applied, required }
158 }
159
160 pub fn is_caught_up(&self) -> bool {
162 self.applied >= self.required
163 }
164
165 pub fn check(&self, level: Consistency, what: impl Into<String>) -> Result<(), QueryError> {
170 if level.is_eventual() || self.is_caught_up() {
171 return Ok(());
172 }
173 Err(QueryError::Stale {
174 what: what.into(),
175 applied: self.applied,
176 required: self.required,
177 })
178 }
179}
180
181#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
185#[serde(rename_all = "snake_case")]
186pub enum Filter {
187 All(Vec<Filter>),
188 Any(Vec<Filter>),
189 Not(Box<Filter>),
190 Pred(Predicate),
191}
192
193impl Filter {
194 pub fn all(filters: impl IntoIterator<Item = Filter>) -> Self {
196 Filter::All(filters.into_iter().collect())
197 }
198
199 pub fn any(filters: impl IntoIterator<Item = Filter>) -> Self {
201 Filter::Any(filters.into_iter().collect())
202 }
203
204 pub fn negate(filter: Filter) -> Self {
206 Filter::Not(Box::new(filter))
207 }
208
209 pub fn pred(field: impl Into<String>, op: CmpOp, value: impl Into<Value>) -> Self {
211 Filter::Pred(Predicate {
212 field: field.into(),
213 op,
214 value: value.into(),
215 })
216 }
217}
218
219#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
221pub struct Predicate {
222 pub field: String,
223 pub op: CmpOp,
224 pub value: Value,
225}
226
227#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
230pub struct RawSql {
231 pub sql: String,
232 #[serde(default, skip_serializing_if = "Vec::is_empty")]
233 pub params: Vec<Value>,
234}
235
236#[derive(
238 Clone,
239 Copy,
240 Debug,
241 PartialEq,
242 Eq,
243 Serialize,
244 Deserialize,
245 strum::Display,
246 strum::EnumString,
247 strum::VariantArray,
248)]
249#[serde(rename_all = "snake_case")]
250#[strum(serialize_all = "snake_case")]
251pub enum CmpOp {
252 Eq,
253 Ne,
254 Lt,
255 Lte,
256 Gt,
257 Gte,
258 In,
259 Contains,
260 Prefix,
261}
262
263#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
265pub struct Sort {
266 pub field: String,
267 #[serde(default)]
268 pub dir: Dir,
269}
270
271#[derive(
273 Clone,
274 Copy,
275 Debug,
276 Default,
277 PartialEq,
278 Eq,
279 Serialize,
280 Deserialize,
281 strum::Display,
282 strum::EnumString,
283 strum::VariantArray,
284)]
285#[serde(rename_all = "snake_case")]
286#[strum(serialize_all = "snake_case")]
287pub enum Dir {
288 #[default]
289 Asc,
290 Desc,
291}
292
293#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
300pub struct TextQuery {
301 #[serde(default, skip_serializing_if = "Option::is_none")]
302 pub field: Option<String>,
303 pub query: String,
304}
305
306#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
308pub struct VectorQuery {
309 pub field: String,
310 pub embedding: Vec<f32>,
311 pub top_k: usize,
312}
313
314#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
318pub struct Aggregate {
319 #[serde(default, skip_serializing_if = "Vec::is_empty")]
320 pub group_by: Vec<String>,
321 pub funcs: Vec<AggCall>,
322 #[serde(default, skip_serializing_if = "Option::is_none")]
323 pub window: Option<Window>,
324}
325
326#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
330pub struct AggCall {
331 pub func: AggFunc,
332 #[serde(default, skip_serializing_if = "Option::is_none")]
333 pub field: Option<String>,
334 #[serde(default, skip_serializing_if = "Option::is_none")]
335 pub arg: Option<f64>,
336 pub alias: String,
337}
338
339#[derive(
342 Clone,
343 Copy,
344 Debug,
345 PartialEq,
346 Eq,
347 Serialize,
348 Deserialize,
349 strum::Display,
350 strum::EnumString,
351 strum::VariantArray,
352)]
353#[serde(rename_all = "snake_case")]
354#[strum(serialize_all = "snake_case")]
355pub enum AggFunc {
356 Count,
357 CountDistinct,
358 Sum,
359 Avg,
360 Min,
361 Max,
362 Percentile,
363 StdDev,
364}
365
366#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
368pub struct Window {
369 pub field: String,
370 pub every_micros: u64,
371}
372
373#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
375pub struct Select {
376 #[serde(default, skip_serializing_if = "Vec::is_empty")]
378 pub fields: Vec<String>,
379 #[serde(default)]
381 pub payload: bool,
382}
383
384#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
390#[serde(untagged)]
391pub enum Value {
392 Str(String),
393 Int(i64),
394 Uint(u64),
395 Float(f64),
396 Bool(bool),
397 Null,
398 List(Vec<Value>),
399}
400
401impl From<&str> for Value {
402 fn from(value: &str) -> Self {
403 Self::Str(value.to_owned())
404 }
405}
406
407impl From<String> for Value {
408 fn from(value: String) -> Self {
409 Self::Str(value)
410 }
411}
412
413impl From<&String> for Value {
414 fn from(value: &String) -> Self {
415 Self::Str(value.clone())
416 }
417}
418
419impl From<i64> for Value {
420 fn from(value: i64) -> Self {
421 Self::Int(value)
422 }
423}
424
425impl From<u64> for Value {
426 fn from(value: u64) -> Self {
427 Self::Uint(value)
428 }
429}
430
431impl From<i32> for Value {
432 fn from(value: i32) -> Self {
433 Self::Int(value as i64)
434 }
435}
436
437impl From<u32> for Value {
438 fn from(value: u32) -> Self {
439 Self::Int(value as i64)
440 }
441}
442
443impl From<f64> for Value {
444 fn from(value: f64) -> Self {
445 Self::Float(value)
446 }
447}
448
449impl From<f32> for Value {
450 fn from(value: f32) -> Self {
451 Self::Float(value as f64)
452 }
453}
454
455impl From<bool> for Value {
456 fn from(value: bool) -> Self {
457 Self::Bool(value)
458 }
459}
460
461impl<T: Into<Value>> From<Vec<T>> for Value {
462 fn from(values: Vec<T>) -> Self {
463 Self::List(values.into_iter().map(Into::into).collect())
464 }
465}
466
467impl Value {
468 pub fn from_input(input: &str) -> Self {
479 match input {
480 "null" => return Value::Null,
481 "true" => return Value::Bool(true),
482 "false" => return Value::Bool(false),
483 _ => {}
484 }
485 if let Ok(int) = input.parse::<i64>() {
486 return Value::Int(int);
487 }
488 if let Ok(uint) = input.parse::<u64>() {
489 return Value::Uint(uint);
490 }
491 if !input.is_empty()
494 && input
495 .bytes()
496 .all(|b| b.is_ascii_digit() || b == b'.' || b == b'-' || b == b'+')
497 && let Ok(float) = input.parse::<f64>()
498 {
499 return Value::Float(float);
500 }
501 Value::Str(input.to_owned())
502 }
503}
504
505impl std::fmt::Display for Value {
506 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
510 match self {
511 Value::Str(value) => f.write_str(value),
512 Value::Int(value) => write!(f, "{value}"),
513 Value::Uint(value) => write!(f, "{value}"),
514 Value::Float(value) => write!(f, "{value}"),
515 Value::Bool(value) => write!(f, "{value}"),
516 Value::Null => f.write_str("null"),
517 Value::List(values) => {
518 f.write_str("[")?;
519 for (index, value) in values.iter().enumerate() {
520 if index > 0 {
521 f.write_str(", ")?;
522 }
523 write!(f, "{value}")?;
524 }
525 f.write_str("]")
526 }
527 }
528 }
529}
530
531impl std::str::FromStr for Value {
532 type Err = std::convert::Infallible;
533
534 fn from_str(s: &str) -> Result<Self, Self::Err> {
535 Ok(Value::from_input(s))
536 }
537}
538
539#[derive(Clone, Debug, Default, Serialize, Deserialize)]
541pub struct QueryResult {
542 pub rows: Vec<Row>,
543 #[serde(default)]
545 pub page: Page,
546}
547
548#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
550pub struct Page {
551 pub offset: usize,
553 pub limit: usize,
555 pub total: usize,
557 pub has_more: bool,
559}
560
561impl Page {
562 pub fn total_pages(&self) -> usize {
564 if self.limit == 0 {
565 0
566 } else {
567 self.total.div_ceil(self.limit)
568 }
569 }
570}
571
572#[derive(Clone, Debug, Default, Serialize, Deserialize)]
574pub struct Row {
575 pub headers: BTreeMap<String, String>,
577 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
581 pub metadata: BTreeMap<String, String>,
582 #[serde(default, skip_serializing_if = "Option::is_none")]
585 pub partition: Option<u32>,
586 #[serde(default, skip_serializing_if = "Option::is_none")]
588 pub offset: Option<u64>,
589 #[serde(default, skip_serializing_if = "Option::is_none")]
592 pub stream: Option<u32>,
593 #[serde(default, skip_serializing_if = "Option::is_none")]
594 pub topic: Option<u32>,
595 #[serde(
599 default,
600 skip_serializing_if = "Option::is_none",
601 with = "crate::encoding::opt_bin_bytes"
602 )]
603 pub payload: Option<Vec<u8>>,
604 #[serde(default, skip_serializing_if = "Option::is_none")]
606 pub score: Option<f32>,
607}
608
609#[derive(Clone, Debug, Serialize, Deserialize)]
612#[non_exhaustive]
613pub struct QueryEnvelope {
614 pub v: u32,
615 pub query: Query,
616}
617
618impl QueryEnvelope {
619 pub fn new(query: Query) -> Self {
621 Self {
622 v: QUERY_OP_VERSION,
623 query,
624 }
625 }
626}
627
628#[derive(Clone, Debug, Serialize, Deserialize)]
630#[non_exhaustive]
631pub enum QueryReply {
632 Ok(QueryResult),
633 Err(QueryError),
634}
635
636#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, thiserror::Error)]
638#[non_exhaustive]
639pub enum QueryError {
640 #[error("query not supported: {0}")]
641 Unsupported(String),
642 #[error("unauthorized: {0}")]
643 Unauthorized(String),
644 #[error("index not found: {0}")]
645 IndexNotFound(String),
646 #[error("fork not found: {0}")]
647 ForkNotFound(String),
648 #[error("backend error: {0}")]
649 Backend(String),
650 #[error("result too large: {what} {size} exceeds cap {cap}")]
657 TooLarge {
658 what: String,
659 size: usize,
660 cap: usize,
661 },
662 #[error("unsupported envelope version (expected {expected}, got {got})")]
663 Version { expected: u32, got: u32 },
664 #[error("stale read: {what} applied {applied}, required {required}")]
670 Stale {
671 what: String,
672 applied: u64,
673 required: u64,
674 },
675}
676
677#[cfg(test)]
678mod tests {
679 use super::*;
680
681 #[test]
682 fn given_dsl_enums_when_displayed_then_should_be_snake_case() {
683 assert_eq!(CmpOp::Gte.to_string(), "gte");
684 assert_eq!(CmpOp::Prefix.to_string(), "prefix");
685 assert_eq!("ne".parse::<CmpOp>().expect("ne parses"), CmpOp::Ne);
686 assert_eq!(Dir::Desc.to_string(), "desc");
687 assert_eq!(AggFunc::Count.to_string(), "count");
688 }
689
690 #[test]
691 fn given_a_consistency_gate_when_checked_then_should_fail_not_downgrade() {
692 assert!(
694 ConsistencyGate::new(0, 100)
695 .check(Consistency::Eventual, "orders")
696 .is_ok()
697 );
698 assert!(
700 ConsistencyGate::new(100, 100)
701 .check(Consistency::ReadYourWrites, "orders")
702 .is_ok()
703 );
704 let stale = ConsistencyGate::new(41, 57)
705 .check(Consistency::Strong, "orders")
706 .expect_err("a lagging projector must fail, never downgrade");
707 assert!(matches!(
708 stale,
709 QueryError::Stale {
710 applied: 41,
711 required: 57,
712 ..
713 }
714 ));
715 }
716
717 #[test]
718 fn given_a_page_when_computing_total_pages_then_should_divide_by_limit() {
719 let page = Page {
720 offset: 0,
721 limit: 3,
722 total: 10,
723 has_more: true,
724 };
725 assert_eq!(page.total_pages(), 4);
726 assert_eq!(Page::default().total_pages(), 0);
727 }
728}
729
730#[cfg(all(test, feature = "codecs"))]
731mod serde_tests {
732 use super::*;
733 #[cfg(feature = "builders")]
734 use crate::codes::QUERY_OP_VERSION;
735 use crate::framing::{decode_named, encode_named};
736
737 #[test]
738 fn given_dsl_enums_when_serialized_then_serde_should_match_display() {
739 assert_eq!(
740 serde_json::to_string(&CmpOp::Lte).expect("CmpOp serializes"),
741 "\"lte\""
742 );
743 assert_eq!(
744 serde_json::from_str::<CmpOp>("\"in\"").expect("CmpOp deserializes"),
745 CmpOp::In
746 );
747 assert_eq!(
748 serde_json::to_string(&Dir::Asc).expect("Dir serializes"),
749 "\"asc\""
750 );
751 }
752
753 #[test]
754 #[cfg(feature = "builders")]
755 fn given_a_query_when_round_tripped_through_the_envelope_then_should_be_unchanged() {
756 let query = Query::builder()
757 .index("orders")
758 .by_key(vec![KeyMatch::new("customer_id", "abc")])
759 .filter(Filter::pred("status", CmpOp::Eq, "paid"))
760 .order(vec![Sort {
761 field: "ts".to_owned(),
762 dir: Dir::Desc,
763 }])
764 .limit(20)
765 .build();
766 let request = QueryEnvelope::new(query);
767
768 let json = serde_json::to_string(&request).expect("the request serializes");
769 let back: QueryEnvelope = serde_json::from_str(&json).expect("the request deserializes");
770 assert_eq!(back.v, QUERY_OP_VERSION);
771 assert_eq!(back.query.index, "orders");
772 assert_eq!(back.query.limit, 20);
773 assert_eq!(back.query.by_key, vec![KeyMatch::new("customer_id", "abc")]);
774 let Some(Filter::Pred(predicate)) = &back.query.filter else {
775 panic!("expected a single predicate filter");
776 };
777 assert_eq!(predicate.value, Value::Str("paid".to_owned()));
778 assert_eq!(back.query.order[0].dir, Dir::Desc);
779 }
780
781 #[test]
782 #[cfg(feature = "builders")]
783 fn given_each_consistency_level_when_round_tripped_then_should_preserve_it_and_skip_eventual() {
784 for level in [
785 Consistency::Eventual,
786 Consistency::ReadYourWrites,
787 Consistency::Strong,
788 ] {
789 let query = Query::builder().index("orders").consistency(level).build();
790 let bytes = encode_named(&QueryEnvelope::new(query)).expect("serializes");
791 let back: QueryEnvelope = decode_named(&bytes).expect("deserializes");
792 assert_eq!(back.query.consistency, level);
793 }
794 let default = Query::builder().index("orders").build();
797 assert_eq!(default.consistency, Consistency::Eventual);
798 let json = serde_json::to_string(&default).expect("json");
799 assert!(
800 !json.contains("consistency"),
801 "default Eventual must be omitted: {json}"
802 );
803 }
804
805 #[test]
806 fn given_a_stale_reply_when_round_tripped_then_should_preserve_the_offsets() {
807 let reply = QueryReply::Err(QueryError::Stale {
808 what: "orders".to_owned(),
809 applied: 41,
810 required: 57,
811 });
812 let bytes = encode_named(&reply).expect("serializes");
813 let back: QueryReply = decode_named(&bytes).expect("deserializes");
814 let QueryReply::Err(QueryError::Stale {
815 what,
816 applied,
817 required,
818 }) = back
819 else {
820 panic!("expected a Stale error");
821 };
822 assert_eq!((what.as_str(), applied, required), ("orders", 41, 57));
823 }
824
825 #[test]
826 #[cfg(feature = "builders")]
827 fn given_a_vector_query_when_round_tripped_then_should_preserve_the_embedding() {
828 let query = Query::builder()
829 .index("mem:conv-1")
830 .vector(VectorQuery {
831 field: "embedding".to_owned(),
832 embedding: vec![0.1, 0.2, 0.3],
833 top_k: 5,
834 })
835 .build();
836 let json = serde_json::to_string(&query).expect("the query serializes");
837 let back: Query = serde_json::from_str(&json).expect("the query deserializes");
838 let vector = back.vector.expect("the vector survives the round-trip");
839 assert_eq!(vector.embedding, vec![0.1, 0.2, 0.3]);
840 assert_eq!(vector.top_k, 5);
841 }
842
843 #[test]
844 fn given_a_reply_with_a_payload_row_when_round_tripped_then_should_preserve_the_bytes() {
845 let mut headers = BTreeMap::new();
846 headers.insert("order_id".to_owned(), "123".to_owned());
847 let reply = QueryReply::Ok(QueryResult {
848 rows: vec![Row {
849 headers,
850 metadata: BTreeMap::from([("agdx.ct".to_owned(), "1".to_owned())]),
851 partition: Some(2),
852 offset: Some(17),
853 stream: Some(5),
854 topic: Some(3),
855 payload: Some(b"{\"total\":42}".to_vec()),
856 score: None,
857 }],
858 page: Page {
859 offset: 0,
860 limit: 50,
861 total: 1,
862 has_more: false,
863 },
864 });
865 let bytes = encode_named(&reply).expect("the reply serializes");
866 let back: QueryReply = decode_named(&bytes).expect("the reply deserializes");
867 let QueryReply::Ok(result) = back else {
868 panic!("the reply should decode as Ok");
869 };
870 assert_eq!(result.rows[0].headers["order_id"], "123");
871 assert_eq!(
872 result.rows[0].payload.as_deref(),
873 Some(b"{\"total\":42}".as_ref())
874 );
875 assert_eq!(result.page.total, 1);
876 assert!(!result.page.has_more);
877 }
878}