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 #[cfg_attr(feature = "builders", builder(default))]
93 #[serde(default, skip_serializing_if = "is_false")]
94 pub want_total: bool,
95}
96
97fn is_false(value: &bool) -> bool {
98 !*value
99}
100
101#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
112#[serde(rename_all = "snake_case")]
113#[non_exhaustive]
114pub enum Consistency {
115 #[default]
119 Eventual,
120 ReadYourWrites,
126 Strong,
131}
132
133impl Consistency {
134 pub fn is_eventual(&self) -> bool {
136 matches!(self, Consistency::Eventual)
137 }
138}
139
140#[derive(Clone, Copy, Debug, PartialEq, Eq)]
153pub struct ConsistencyGate {
154 pub applied: u64,
156 pub required: u64,
158}
159
160impl ConsistencyGate {
161 pub fn new(applied: u64, required: u64) -> Self {
164 Self { applied, required }
165 }
166
167 pub fn is_caught_up(&self) -> bool {
169 self.applied >= self.required
170 }
171
172 pub fn check(&self, level: Consistency, what: impl Into<String>) -> Result<(), QueryError> {
177 if level.is_eventual() || self.is_caught_up() {
178 return Ok(());
179 }
180 Err(QueryError::Stale {
181 what: what.into(),
182 applied: self.applied,
183 required: self.required,
184 })
185 }
186}
187
188#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
192#[serde(rename_all = "snake_case")]
193pub enum Filter {
194 All(Vec<Filter>),
195 Any(Vec<Filter>),
196 Not(Box<Filter>),
197 Pred(Predicate),
198}
199
200impl Filter {
201 pub fn all(filters: impl IntoIterator<Item = Filter>) -> Self {
203 Filter::All(filters.into_iter().collect())
204 }
205
206 pub fn any(filters: impl IntoIterator<Item = Filter>) -> Self {
208 Filter::Any(filters.into_iter().collect())
209 }
210
211 pub fn negate(filter: Filter) -> Self {
213 Filter::Not(Box::new(filter))
214 }
215
216 pub fn pred(field: impl Into<String>, op: CmpOp, value: impl Into<Value>) -> Self {
218 Filter::Pred(Predicate {
219 field: field.into(),
220 op,
221 value: value.into(),
222 })
223 }
224}
225
226#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
228pub struct Predicate {
229 pub field: String,
230 pub op: CmpOp,
231 pub value: Value,
232}
233
234#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
237pub struct RawSql {
238 pub sql: String,
239 #[serde(default, skip_serializing_if = "Vec::is_empty")]
240 pub params: Vec<Value>,
241}
242
243#[derive(
245 Clone,
246 Copy,
247 Debug,
248 PartialEq,
249 Eq,
250 Serialize,
251 Deserialize,
252 strum::Display,
253 strum::EnumString,
254 strum::VariantArray,
255)]
256#[serde(rename_all = "snake_case")]
257#[strum(serialize_all = "snake_case")]
258pub enum CmpOp {
259 Eq,
260 Ne,
261 Lt,
262 Lte,
263 Gt,
264 Gte,
265 In,
266 Contains,
267 Prefix,
268}
269
270#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
272pub struct Sort {
273 pub field: String,
274 #[serde(default)]
275 pub dir: Dir,
276}
277
278#[derive(
280 Clone,
281 Copy,
282 Debug,
283 Default,
284 PartialEq,
285 Eq,
286 Serialize,
287 Deserialize,
288 strum::Display,
289 strum::EnumString,
290 strum::VariantArray,
291)]
292#[serde(rename_all = "snake_case")]
293#[strum(serialize_all = "snake_case")]
294pub enum Dir {
295 #[default]
296 Asc,
297 Desc,
298}
299
300#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
307pub struct TextQuery {
308 #[serde(default, skip_serializing_if = "Option::is_none")]
309 pub field: Option<String>,
310 pub query: String,
311}
312
313#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
315pub struct VectorQuery {
316 pub field: String,
317 pub embedding: Vec<f32>,
318 pub top_k: usize,
319}
320
321#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
325pub struct Aggregate {
326 #[serde(default, skip_serializing_if = "Vec::is_empty")]
327 pub group_by: Vec<String>,
328 pub funcs: Vec<AggCall>,
329 #[serde(default, skip_serializing_if = "Option::is_none")]
330 pub window: Option<Window>,
331}
332
333#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
337pub struct AggCall {
338 pub func: AggFunc,
339 #[serde(default, skip_serializing_if = "Option::is_none")]
340 pub field: Option<String>,
341 #[serde(default, skip_serializing_if = "Option::is_none")]
342 pub arg: Option<f64>,
343 pub alias: String,
344}
345
346#[derive(
349 Clone,
350 Copy,
351 Debug,
352 PartialEq,
353 Eq,
354 Serialize,
355 Deserialize,
356 strum::Display,
357 strum::EnumString,
358 strum::VariantArray,
359)]
360#[serde(rename_all = "snake_case")]
361#[strum(serialize_all = "snake_case")]
362pub enum AggFunc {
363 Count,
364 CountDistinct,
365 Sum,
366 Avg,
367 Min,
368 Max,
369 Percentile,
370 StdDev,
371}
372
373#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
375pub struct Window {
376 pub field: String,
377 pub every_micros: u64,
378}
379
380#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
382pub struct Select {
383 #[serde(default, skip_serializing_if = "Vec::is_empty")]
385 pub fields: Vec<String>,
386 #[serde(default)]
388 pub payload: bool,
389}
390
391#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
397#[serde(untagged)]
398pub enum Value {
399 Str(String),
400 Int(i64),
401 Uint(u64),
402 Float(f64),
403 Bool(bool),
404 Null,
405 List(Vec<Value>),
406}
407
408impl From<&str> for Value {
409 fn from(value: &str) -> Self {
410 Self::Str(value.to_owned())
411 }
412}
413
414impl From<String> for Value {
415 fn from(value: String) -> Self {
416 Self::Str(value)
417 }
418}
419
420impl From<&String> for Value {
421 fn from(value: &String) -> Self {
422 Self::Str(value.clone())
423 }
424}
425
426impl From<i64> for Value {
427 fn from(value: i64) -> Self {
428 Self::Int(value)
429 }
430}
431
432impl From<u64> for Value {
433 fn from(value: u64) -> Self {
434 Self::Uint(value)
435 }
436}
437
438impl From<i32> for Value {
439 fn from(value: i32) -> Self {
440 Self::Int(value as i64)
441 }
442}
443
444impl From<u32> for Value {
445 fn from(value: u32) -> Self {
446 Self::Int(value as i64)
447 }
448}
449
450impl From<f64> for Value {
451 fn from(value: f64) -> Self {
452 Self::Float(value)
453 }
454}
455
456impl From<f32> for Value {
457 fn from(value: f32) -> Self {
458 Self::Float(value as f64)
459 }
460}
461
462impl From<bool> for Value {
463 fn from(value: bool) -> Self {
464 Self::Bool(value)
465 }
466}
467
468impl<T: Into<Value>> From<Vec<T>> for Value {
469 fn from(values: Vec<T>) -> Self {
470 Self::List(values.into_iter().map(Into::into).collect())
471 }
472}
473
474impl Value {
475 pub fn from_input(input: &str) -> Self {
486 match input {
487 "null" => return Value::Null,
488 "true" => return Value::Bool(true),
489 "false" => return Value::Bool(false),
490 _ => {}
491 }
492 if let Ok(int) = input.parse::<i64>() {
493 return Value::Int(int);
494 }
495 if let Ok(uint) = input.parse::<u64>() {
496 return Value::Uint(uint);
497 }
498 if !input.is_empty()
501 && input
502 .bytes()
503 .all(|b| b.is_ascii_digit() || b == b'.' || b == b'-' || b == b'+')
504 && let Ok(float) = input.parse::<f64>()
505 {
506 return Value::Float(float);
507 }
508 Value::Str(input.to_owned())
509 }
510}
511
512impl std::fmt::Display for Value {
513 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
517 match self {
518 Value::Str(value) => f.write_str(value),
519 Value::Int(value) => write!(f, "{value}"),
520 Value::Uint(value) => write!(f, "{value}"),
521 Value::Float(value) => write!(f, "{value}"),
522 Value::Bool(value) => write!(f, "{value}"),
523 Value::Null => f.write_str("null"),
524 Value::List(values) => {
525 f.write_str("[")?;
526 for (index, value) in values.iter().enumerate() {
527 if index > 0 {
528 f.write_str(", ")?;
529 }
530 write!(f, "{value}")?;
531 }
532 f.write_str("]")
533 }
534 }
535 }
536}
537
538impl std::str::FromStr for Value {
539 type Err = std::convert::Infallible;
540
541 fn from_str(s: &str) -> Result<Self, Self::Err> {
542 Ok(Value::from_input(s))
543 }
544}
545
546#[derive(Clone, Debug, Default, Serialize, Deserialize)]
548pub struct QueryResult {
549 pub rows: Vec<Row>,
550 #[serde(default)]
552 pub page: Page,
553}
554
555#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
564pub struct Page {
565 pub offset: usize,
567 pub limit: usize,
569 #[serde(default, skip_serializing_if = "Option::is_none")]
571 pub total: Option<u64>,
572 pub has_more: bool,
574}
575
576impl Page {
577 pub fn at_least(&self, rows_on_page: usize) -> usize {
581 self.offset + rows_on_page
582 }
583
584 pub fn total_pages(&self) -> Option<u64> {
588 match (self.total, self.limit) {
589 (Some(total), limit) if limit > 0 => Some(total.div_ceil(limit as u64)),
590 _ => None,
591 }
592 }
593}
594
595#[derive(Clone, Debug, Default, Serialize, Deserialize)]
597pub struct Row {
598 pub headers: BTreeMap<String, String>,
600 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
604 pub metadata: BTreeMap<String, String>,
605 #[serde(default, skip_serializing_if = "Option::is_none")]
608 pub partition: Option<u32>,
609 #[serde(default, skip_serializing_if = "Option::is_none")]
611 pub offset: Option<u64>,
612 #[serde(default, skip_serializing_if = "Option::is_none")]
615 pub stream: Option<u32>,
616 #[serde(default, skip_serializing_if = "Option::is_none")]
617 pub topic: Option<u32>,
618 #[serde(
622 default,
623 skip_serializing_if = "Option::is_none",
624 with = "crate::encoding::opt_bin_bytes"
625 )]
626 pub payload: Option<Vec<u8>>,
627 #[serde(default, skip_serializing_if = "Option::is_none")]
629 pub score: Option<f32>,
630}
631
632#[derive(Clone, Debug, Serialize, Deserialize)]
635#[non_exhaustive]
636pub struct QueryEnvelope {
637 pub v: u32,
638 pub query: Query,
639}
640
641impl QueryEnvelope {
642 pub fn new(query: Query) -> Self {
644 Self {
645 v: QUERY_OP_VERSION,
646 query,
647 }
648 }
649}
650
651#[derive(Clone, Debug, Serialize, Deserialize)]
653#[non_exhaustive]
654pub enum QueryReply {
655 Ok(QueryResult),
656 Err(QueryError),
657}
658
659#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, thiserror::Error)]
661#[non_exhaustive]
662pub enum QueryError {
663 #[error("query not supported: {0}")]
664 Unsupported(String),
665 #[error("unauthorized: {0}")]
666 Unauthorized(String),
667 #[error("index not found: {0}")]
668 IndexNotFound(String),
669 #[error("fork not found: {0}")]
670 ForkNotFound(String),
671 #[error("backend error: {0}")]
672 Backend(String),
673 #[error("result too large: {what} {size} exceeds cap {cap}")]
680 TooLarge {
681 what: String,
682 size: usize,
683 cap: usize,
684 },
685 #[error("unsupported envelope version (expected {expected}, got {got})")]
686 Version { expected: u32, got: u32 },
687 #[error("stale read: {what} applied {applied}, required {required}")]
693 Stale {
694 what: String,
695 applied: u64,
696 required: u64,
697 },
698}
699
700#[cfg(test)]
701mod tests {
702 use super::*;
703
704 #[test]
705 fn given_dsl_enums_when_displayed_then_should_be_snake_case() {
706 assert_eq!(CmpOp::Gte.to_string(), "gte");
707 assert_eq!(CmpOp::Prefix.to_string(), "prefix");
708 assert_eq!("ne".parse::<CmpOp>().expect("ne parses"), CmpOp::Ne);
709 assert_eq!(Dir::Desc.to_string(), "desc");
710 assert_eq!(AggFunc::Count.to_string(), "count");
711 }
712
713 #[test]
714 fn given_a_consistency_gate_when_checked_then_should_fail_not_downgrade() {
715 assert!(
717 ConsistencyGate::new(0, 100)
718 .check(Consistency::Eventual, "orders")
719 .is_ok()
720 );
721 assert!(
723 ConsistencyGate::new(100, 100)
724 .check(Consistency::ReadYourWrites, "orders")
725 .is_ok()
726 );
727 let stale = ConsistencyGate::new(41, 57)
728 .check(Consistency::Strong, "orders")
729 .expect_err("a lagging projector must fail, never downgrade");
730 assert!(matches!(
731 stale,
732 QueryError::Stale {
733 applied: 41,
734 required: 57,
735 ..
736 }
737 ));
738 }
739
740 #[test]
741 fn given_a_page_when_computing_total_pages_then_should_divide_by_limit() {
742 let page = Page {
743 offset: 0,
744 limit: 3,
745 total: Some(10),
746 has_more: true,
747 };
748 assert_eq!(page.total_pages(), Some(4));
749 assert_eq!(page.at_least(3), 3, "offset 0 plus this page's rows");
750 assert_eq!(Page::default().total_pages(), None);
752 }
753}
754
755#[cfg(all(test, feature = "codecs"))]
756mod serde_tests {
757 use super::*;
758 #[cfg(feature = "builders")]
759 use crate::codes::QUERY_OP_VERSION;
760 use crate::framing::{decode_named, encode_named};
761
762 #[test]
763 fn given_dsl_enums_when_serialized_then_serde_should_match_display() {
764 assert_eq!(
765 serde_json::to_string(&CmpOp::Lte).expect("CmpOp serializes"),
766 "\"lte\""
767 );
768 assert_eq!(
769 serde_json::from_str::<CmpOp>("\"in\"").expect("CmpOp deserializes"),
770 CmpOp::In
771 );
772 assert_eq!(
773 serde_json::to_string(&Dir::Asc).expect("Dir serializes"),
774 "\"asc\""
775 );
776 }
777
778 #[test]
779 #[cfg(feature = "builders")]
780 fn given_a_query_when_round_tripped_through_the_envelope_then_should_be_unchanged() {
781 let query = Query::builder()
782 .index("orders")
783 .by_key(vec![KeyMatch::new("customer_id", "abc")])
784 .filter(Filter::pred("status", CmpOp::Eq, "paid"))
785 .order(vec![Sort {
786 field: "ts".to_owned(),
787 dir: Dir::Desc,
788 }])
789 .limit(20)
790 .build();
791 let request = QueryEnvelope::new(query);
792
793 let json = serde_json::to_string(&request).expect("the request serializes");
794 let back: QueryEnvelope = serde_json::from_str(&json).expect("the request deserializes");
795 assert_eq!(back.v, QUERY_OP_VERSION);
796 assert_eq!(back.query.index, "orders");
797 assert_eq!(back.query.limit, 20);
798 assert_eq!(back.query.by_key, vec![KeyMatch::new("customer_id", "abc")]);
799 let Some(Filter::Pred(predicate)) = &back.query.filter else {
800 panic!("expected a single predicate filter");
801 };
802 assert_eq!(predicate.value, Value::Str("paid".to_owned()));
803 assert_eq!(back.query.order[0].dir, Dir::Desc);
804 }
805
806 #[test]
807 #[cfg(feature = "builders")]
808 fn given_each_consistency_level_when_round_tripped_then_should_preserve_it_and_skip_eventual() {
809 for level in [
810 Consistency::Eventual,
811 Consistency::ReadYourWrites,
812 Consistency::Strong,
813 ] {
814 let query = Query::builder().index("orders").consistency(level).build();
815 let bytes = encode_named(&QueryEnvelope::new(query)).expect("serializes");
816 let back: QueryEnvelope = decode_named(&bytes).expect("deserializes");
817 assert_eq!(back.query.consistency, level);
818 }
819 let default = Query::builder().index("orders").build();
822 assert_eq!(default.consistency, Consistency::Eventual);
823 let json = serde_json::to_string(&default).expect("json");
824 assert!(
825 !json.contains("consistency"),
826 "default Eventual must be omitted: {json}"
827 );
828 }
829
830 #[test]
831 fn given_a_stale_reply_when_round_tripped_then_should_preserve_the_offsets() {
832 let reply = QueryReply::Err(QueryError::Stale {
833 what: "orders".to_owned(),
834 applied: 41,
835 required: 57,
836 });
837 let bytes = encode_named(&reply).expect("serializes");
838 let back: QueryReply = decode_named(&bytes).expect("deserializes");
839 let QueryReply::Err(QueryError::Stale {
840 what,
841 applied,
842 required,
843 }) = back
844 else {
845 panic!("expected a Stale error");
846 };
847 assert_eq!((what.as_str(), applied, required), ("orders", 41, 57));
848 }
849
850 #[test]
851 #[cfg(feature = "builders")]
852 fn given_a_vector_query_when_round_tripped_then_should_preserve_the_embedding() {
853 let query = Query::builder()
854 .index("mem:conv-1")
855 .vector(VectorQuery {
856 field: "embedding".to_owned(),
857 embedding: vec![0.1, 0.2, 0.3],
858 top_k: 5,
859 })
860 .build();
861 let json = serde_json::to_string(&query).expect("the query serializes");
862 let back: Query = serde_json::from_str(&json).expect("the query deserializes");
863 let vector = back.vector.expect("the vector survives the round-trip");
864 assert_eq!(vector.embedding, vec![0.1, 0.2, 0.3]);
865 assert_eq!(vector.top_k, 5);
866 }
867
868 #[test]
869 fn given_a_reply_with_a_payload_row_when_round_tripped_then_should_preserve_the_bytes() {
870 let mut headers = BTreeMap::new();
871 headers.insert("order_id".to_owned(), "123".to_owned());
872 let reply = QueryReply::Ok(QueryResult {
873 rows: vec![Row {
874 headers,
875 metadata: BTreeMap::from([("agdx.ct".to_owned(), "1".to_owned())]),
876 partition: Some(2),
877 offset: Some(17),
878 stream: Some(5),
879 topic: Some(3),
880 payload: Some(b"{\"total\":42}".to_vec()),
881 score: None,
882 }],
883 page: Page {
884 offset: 0,
885 limit: 50,
886 total: Some(1),
887 has_more: false,
888 },
889 });
890 let bytes = encode_named(&reply).expect("the reply serializes");
891 let back: QueryReply = decode_named(&bytes).expect("the reply deserializes");
892 let QueryReply::Ok(result) = back else {
893 panic!("the reply should decode as Ok");
894 };
895 assert_eq!(result.rows[0].headers["order_id"], "123");
896 assert_eq!(
897 result.rows[0].payload.as_deref(),
898 Some(b"{\"total\":42}".as_ref())
899 );
900 assert_eq!(result.page.total, Some(1));
901 assert!(!result.page.has_more);
902 }
903}