1use crate::agent::wire_id;
2use crate::codes::QUERY_OP_VERSION;
3use crate::destination::BackendResourceId;
4use crate::error::InvalidError;
5use crate::limits::{
6 MAX_PAGE_SIZE, MAX_QUERY_FIELDS, MAX_QUERY_NAME_BYTES, MAX_QUERY_OFFSET, MAX_QUERY_PARAMETERS,
7 MAX_QUERY_PREDICATES, MAX_RAW_SQL_BYTES,
8};
9use crate::schema::{LogicalField, TypedValue, validate_result_fields};
10use crate::validate::Validate;
11use serde::{Deserialize, Serialize};
12use std::collections::BTreeSet;
13
14wire_id!(
15 QueryExecutionId
17);
18
19#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
21#[serde(tag = "kind", rename_all = "snake_case")]
22#[non_exhaustive]
23pub enum QueryTarget {
24 Operational { index: String },
25}
26
27impl QueryTarget {
28 pub fn operational(index: impl Into<String>) -> Self {
29 Self::Operational {
30 index: index.into(),
31 }
32 }
33}
34
35#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
37#[serde(rename_all = "snake_case")]
38#[non_exhaustive]
39pub enum SqlDialect {
40 Postgres,
41 MySql,
42 Sqlite,
43}
44
45#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
47pub struct KeyMatch {
48 pub field: String,
49 pub value: TypedValue,
50}
51
52impl KeyMatch {
53 pub fn new(field: impl Into<String>, value: impl Into<TypedValue>) -> Self {
55 Self {
56 field: field.into(),
57 value: value.into(),
58 }
59 }
60}
61
62#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
65#[cfg_attr(feature = "builders", derive(bon::Builder))]
66pub struct Query {
67 pub execution_id: QueryExecutionId,
68 pub target: QueryTarget,
69 pub deadline_micros: u64,
71 #[cfg_attr(feature = "builders", builder(default))]
72 #[serde(default, skip_serializing_if = "Vec::is_empty")]
73 pub by_key: Vec<KeyMatch>,
74 #[cfg_attr(feature = "builders", builder(into))]
75 #[serde(default, skip_serializing_if = "Option::is_none")]
76 pub message_type: Option<String>,
77 #[serde(default, skip_serializing_if = "Option::is_none")]
80 pub time_range: Option<(u64, u64)>,
81 #[serde(default, skip_serializing_if = "Option::is_none")]
84 pub filter: Option<Filter>,
85 #[serde(default, skip_serializing_if = "Option::is_none")]
86 pub vector: Option<VectorQuery>,
87 #[serde(default, skip_serializing_if = "Option::is_none")]
89 pub text: Option<TextQuery>,
90 #[cfg_attr(feature = "builders", builder(default))]
91 #[serde(default, skip_serializing_if = "Vec::is_empty")]
92 pub order: Vec<Sort>,
93 #[cfg_attr(feature = "builders", builder(default))]
94 #[serde(default)]
95 pub page: QueryPageRequest,
96 #[serde(default, skip_serializing_if = "Option::is_none")]
98 pub aggregate: Option<Aggregate>,
99 #[serde(default, skip_serializing_if = "Option::is_none")]
102 pub having: Option<Filter>,
103 #[cfg_attr(feature = "builders", builder(default))]
105 #[serde(default, skip_serializing_if = "is_false")]
106 pub distinct: bool,
107 #[cfg_attr(feature = "builders", builder(default))]
108 #[serde(default)]
109 pub select: Select,
110 #[cfg_attr(feature = "builders", builder(into))]
114 #[serde(default, skip_serializing_if = "Option::is_none")]
115 pub fork: Option<String>,
116 #[serde(default, skip_serializing_if = "Option::is_none")]
118 pub raw_sql: Option<RawSql>,
119 #[cfg_attr(feature = "builders", builder(default))]
122 #[serde(default, skip_serializing_if = "Consistency::is_eventual")]
123 pub consistency: Consistency,
124}
125
126impl Query {
127 pub fn new(execution_id: QueryExecutionId, target: QueryTarget, deadline_micros: u64) -> Self {
129 Self {
130 execution_id,
131 target,
132 deadline_micros,
133 by_key: Vec::new(),
134 message_type: None,
135 time_range: None,
136 filter: None,
137 vector: None,
138 text: None,
139 order: Vec::new(),
140 page: QueryPageRequest::default(),
141 aggregate: None,
142 having: None,
143 distinct: false,
144 select: Select::default(),
145 fork: None,
146 raw_sql: None,
147 consistency: Consistency::Eventual,
148 }
149 }
150
151 pub fn operational(
153 execution_id: QueryExecutionId,
154 index: impl Into<String>,
155 deadline_micros: u64,
156 ) -> Self {
157 Self::new(
158 execution_id,
159 QueryTarget::operational(index),
160 deadline_micros,
161 )
162 }
163}
164
165#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
166pub struct QueryPageRequest {
167 pub limit: u32,
168 #[serde(default, skip_serializing_if = "Option::is_none")]
169 pub offset: Option<u64>,
170 #[serde(default, skip_serializing_if = "is_false")]
171 pub want_total: bool,
172}
173
174impl Default for QueryPageRequest {
175 fn default() -> Self {
176 Self {
177 limit: 50,
178 offset: Some(0),
179 want_total: false,
180 }
181 }
182}
183
184impl Validate for Query {
185 fn validate(&self) -> Result<(), InvalidError> {
186 if self.execution_id.as_u128() == 0 {
187 return Err(InvalidError::new("query execution id must be nonzero"));
188 }
189 if self.deadline_micros == 0 {
190 return Err(InvalidError::new("query deadline must be nonzero"));
191 }
192 match &self.target {
193 QueryTarget::Operational { index } => validate_query_name("index", index)?,
194 }
195 if self.page.limit == 0 || self.page.limit > MAX_PAGE_SIZE as u32 {
196 return Err(InvalidError::new(format!(
197 "query limit {} is outside 1..={MAX_PAGE_SIZE}",
198 self.page.limit
199 )));
200 }
201 for (label, count) in [
202 ("exact matches", self.by_key.len()),
203 ("sort fields", self.order.len()),
204 ("selected fields", self.select.fields.len()),
205 ] {
206 if count > MAX_QUERY_FIELDS {
207 return Err(InvalidError::new(format!(
208 "query {label} count {count} exceeds cap {MAX_QUERY_FIELDS}"
209 )));
210 }
211 }
212 if let Some(offset) = self.page.offset
213 && offset > MAX_QUERY_OFFSET
214 {
215 return Err(InvalidError::new(format!(
216 "query offset {offset} exceeds cap {MAX_QUERY_OFFSET}"
217 )));
218 }
219 for key_match in &self.by_key {
220 validate_query_name("key-match field", &key_match.field)?;
221 key_match.value.validate()?;
222 }
223 if let Some(message_type) = &self.message_type {
224 validate_query_name("message type", message_type)?;
225 }
226 if let Some((start, end)) = self.time_range
227 && start > end
228 {
229 return Err(InvalidError::new(
230 "query time range start must not be after its end, and both bounds are inclusive",
231 ));
232 }
233 for sort in &self.order {
234 validate_query_name("sort field", &sort.field)?;
235 }
236 for field in &self.select.fields {
237 validate_query_name("selected field", field)?;
238 }
239 if let Some(text) = &self.text {
240 text.validate()?;
241 }
242 if let Some(vector) = &self.vector {
243 vector.validate()?;
244 }
245 validate_filter(self.filter.as_ref())?;
246 validate_filter(self.having.as_ref())?;
247 if self.having.is_some() && self.aggregate.is_none() {
248 return Err(InvalidError::new(
249 "query having filter requires an aggregate",
250 ));
251 }
252 if let Some(aggregate) = &self.aggregate {
253 aggregate.validate()?;
254 }
255 if self.aggregate.is_some() && (!self.select.fields.is_empty() || self.select.payload) {
256 return Err(InvalidError::new(
257 "aggregate queries cannot request row selection",
258 ));
259 }
260 let shapes = [
261 self.text.is_some(),
262 self.vector.is_some(),
263 self.aggregate.is_some(),
264 self.distinct,
265 ];
266 if shapes.into_iter().filter(|set| *set).count() > 1 {
267 return Err(InvalidError::new(
268 "a query takes at most one of text, vector, aggregate and distinct",
269 ));
270 }
271 if self.text.is_some() && (!self.order.is_empty() || self.page.want_total) {
272 return Err(InvalidError::new(
273 "a text query ranks by relevance and takes no order or total",
274 ));
275 }
276 if self.vector.is_some()
277 && (!self.order.is_empty()
278 || self.page.want_total
279 || self.page.offset.is_some_and(|offset| offset > 0))
280 {
281 return Err(InvalidError::new(
282 "a vector query returns its nearest rows and takes no order, offset or total",
283 ));
284 }
285 if let Some(raw_sql) = &self.raw_sql {
286 raw_sql.validate()?;
287 if !self.by_key.is_empty()
288 || self.message_type.is_some()
289 || self.time_range.is_some()
290 || self.filter.is_some()
291 || self.vector.is_some()
292 || self.text.is_some()
293 || !self.order.is_empty()
294 || self.aggregate.is_some()
295 || self.having.is_some()
296 || self.distinct
297 || !self.select.fields.is_empty()
298 || self.select.payload
299 {
300 return Err(InvalidError::new(
301 "raw SQL cannot be combined with the structured query expression",
302 ));
303 }
304 }
305 Ok(())
306 }
307}
308
309impl Validate for Aggregate {
310 fn validate(&self) -> Result<(), InvalidError> {
311 if self.funcs.is_empty() || self.funcs.len() > MAX_QUERY_PARAMETERS {
312 return Err(InvalidError::new(format!(
313 "aggregate function count must be in 1..={MAX_QUERY_PARAMETERS}"
314 )));
315 }
316 for field in &self.group_by {
317 validate_query_name("group field", field)?;
318 }
319 if self.group_by.len() > MAX_QUERY_FIELDS {
320 return Err(InvalidError::new(format!(
321 "aggregate group field count exceeds cap {MAX_QUERY_FIELDS}"
322 )));
323 }
324 let mut aliases = BTreeSet::new();
325 for call in &self.funcs {
326 validate_query_name("aggregate alias", &call.alias)?;
327 if !aliases.insert(&call.alias) {
328 return Err(InvalidError::new(format!(
329 "aggregate repeats alias `{}`",
330 call.alias
331 )));
332 }
333 if let Some(field) = &call.field {
334 validate_query_name("aggregate field", field)?;
335 }
336 match call.func {
337 AggFunc::Count if call.arg.is_none() => {}
338 AggFunc::Percentile
339 if call.field.is_some()
340 && call.arg.is_some_and(|value| {
341 value.is_finite() && (0.0..=1.0).contains(&value)
342 }) => {}
343 AggFunc::Percentile => {
344 return Err(InvalidError::new(
345 "percentile requires a field and a finite fraction in 0..=1",
346 ));
347 }
348 _ if call.field.is_some() && call.arg.is_none() => {}
349 _ => {
350 return Err(InvalidError::new(
351 "aggregate field and argument do not match the selected function",
352 ));
353 }
354 }
355 }
356 if let Some(window) = &self.window {
357 validate_query_name("window field", &window.field)?;
358 if window.every_micros == 0 {
359 return Err(InvalidError::new("window duration must be nonzero"));
360 }
361 }
362 Ok(())
363 }
364}
365
366impl Validate for RawSql {
367 fn validate(&self) -> Result<(), InvalidError> {
368 if self.sql.trim().is_empty() || self.sql.len() > MAX_RAW_SQL_BYTES {
369 return Err(InvalidError::new(format!(
370 "raw SQL must contain 1..={MAX_RAW_SQL_BYTES} bytes"
371 )));
372 }
373 if self
374 .sql
375 .chars()
376 .any(|character| character.is_control() && !matches!(character, '\n' | '\r' | '\t'))
377 {
378 return Err(InvalidError::new("raw SQL contains a control character"));
379 }
380 if self.params.len() > MAX_QUERY_PARAMETERS {
381 return Err(InvalidError::new(format!(
382 "raw SQL has {} parameters, exceeds cap {MAX_QUERY_PARAMETERS}",
383 self.params.len()
384 )));
385 }
386 for param in &self.params {
387 param.validate()?;
388 }
389 Ok(())
390 }
391}
392
393fn validate_query_name(label: &str, value: &str) -> Result<(), InvalidError> {
394 if value.is_empty() || value.len() > MAX_QUERY_NAME_BYTES {
395 return Err(InvalidError::new(format!(
396 "{label} must contain 1..={MAX_QUERY_NAME_BYTES} bytes"
397 )));
398 }
399 if value.trim() != value || value.chars().any(char::is_control) {
400 return Err(InvalidError::new(format!(
401 "{label} contains surrounding whitespace or a control character"
402 )));
403 }
404 Ok(())
405}
406
407fn validate_filter(filter: Option<&Filter>) -> Result<(), InvalidError> {
408 fn visit(filter: &Filter, depth: usize, count: &mut usize) -> Result<(), InvalidError> {
409 *count += 1;
410 if *count > MAX_QUERY_PREDICATES {
411 return Err(InvalidError::new(format!(
412 "query filter node count exceeds cap {MAX_QUERY_PREDICATES}"
413 )));
414 }
415 if depth > 64 {
416 return Err(InvalidError::new("query filter depth exceeds cap 64"));
417 }
418 match filter {
419 Filter::All(filters) | Filter::Any(filters) => {
420 if filters.is_empty() {
421 return Err(InvalidError::new(
422 "query filter conjunction must not be empty",
423 ));
424 }
425 for filter in filters {
426 visit(filter, depth + 1, count)?;
427 }
428 Ok(())
429 }
430 Filter::Not(filter) => visit(filter, depth + 1, count),
431 Filter::Pred(predicate) => {
432 validate_query_name("predicate field", &predicate.field)?;
433 predicate.value.validate()?;
434 if predicate.op == CmpOp::In
435 && !matches!(&predicate.value, TypedValue::List(values) if !values.is_empty())
436 {
437 return Err(InvalidError::new(
438 "query IN predicate requires a nonempty typed list",
439 ));
440 }
441 if predicate.op == CmpOp::Prefix
442 && !matches!(&predicate.value, TypedValue::String(_))
443 {
444 return Err(InvalidError::new(
445 "query prefix predicate requires a string value",
446 ));
447 }
448 Ok(())
449 }
450 }
451 }
452
453 let mut count = 0;
454 filter.map_or(Ok(()), |filter| visit(filter, 1, &mut count))
455}
456
457fn is_false(value: &bool) -> bool {
458 !*value
459}
460
461#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
472#[serde(rename_all = "snake_case")]
473#[non_exhaustive]
474pub enum Consistency {
475 #[default]
479 Eventual,
480 ReadYourWrites,
487}
488
489impl Consistency {
490 pub fn is_eventual(&self) -> bool {
492 matches!(self, Consistency::Eventual)
493 }
494}
495
496#[derive(Clone, Copy, Debug, PartialEq, Eq)]
505pub struct ConsistencyGate {
506 pub applied: u64,
508 pub required: u64,
510}
511
512impl ConsistencyGate {
513 pub fn new(applied: u64, required: u64) -> Self {
516 Self { applied, required }
517 }
518
519 pub fn is_caught_up(&self) -> bool {
521 self.applied >= self.required
522 }
523
524 pub fn check(&self, level: Consistency, what: impl Into<String>) -> Result<(), QueryError> {
529 if level.is_eventual() || self.is_caught_up() {
530 return Ok(());
531 }
532 Err(QueryError::Stale {
533 what: what.into(),
534 applied: self.applied,
535 required: self.required,
536 })
537 }
538}
539
540#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
544#[serde(rename_all = "snake_case")]
545pub enum Filter {
546 All(Vec<Filter>),
547 Any(Vec<Filter>),
548 Not(Box<Filter>),
549 Pred(Predicate),
550}
551
552impl Filter {
553 pub fn all(filters: impl IntoIterator<Item = Filter>) -> Self {
555 Filter::All(filters.into_iter().collect())
556 }
557
558 pub fn any(filters: impl IntoIterator<Item = Filter>) -> Self {
560 Filter::Any(filters.into_iter().collect())
561 }
562
563 pub fn negate(filter: Filter) -> Self {
565 Filter::Not(Box::new(filter))
566 }
567
568 pub fn pred(field: impl Into<String>, op: CmpOp, value: impl Into<TypedValue>) -> Self {
570 Filter::Pred(Predicate {
571 field: field.into(),
572 op,
573 value: value.into(),
574 })
575 }
576}
577
578#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
580pub struct Predicate {
581 pub field: String,
582 pub op: CmpOp,
583 pub value: TypedValue,
584}
585
586#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
589pub struct RawSql {
590 pub dialect: SqlDialect,
591 pub sql: String,
592 #[serde(default, skip_serializing_if = "Vec::is_empty")]
593 pub params: Vec<TypedValue>,
594}
595
596#[derive(
598 Clone,
599 Copy,
600 Debug,
601 PartialEq,
602 Eq,
603 Serialize,
604 Deserialize,
605 strum::Display,
606 strum::EnumString,
607 strum::VariantArray,
608)]
609#[serde(rename_all = "snake_case")]
610#[strum(serialize_all = "snake_case")]
611pub enum CmpOp {
612 Eq,
613 Ne,
614 Lt,
615 Lte,
616 Gt,
617 Gte,
618 In,
619 Contains,
620 Prefix,
621}
622
623#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
625pub struct Sort {
626 pub field: String,
627 #[serde(default)]
628 pub dir: Dir,
629}
630
631#[derive(
633 Clone,
634 Copy,
635 Debug,
636 Default,
637 PartialEq,
638 Eq,
639 Serialize,
640 Deserialize,
641 strum::Display,
642 strum::EnumString,
643 strum::VariantArray,
644)]
645#[serde(rename_all = "snake_case")]
646#[strum(serialize_all = "snake_case")]
647pub enum Dir {
648 #[default]
649 Asc,
650 Desc,
651}
652
653#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
659pub struct TextQuery {
660 #[serde(default, skip_serializing_if = "Option::is_none")]
661 pub field: Option<String>,
662 pub query: String,
663}
664
665#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
667pub struct VectorQuery {
668 pub field: String,
669 pub embedding: Vec<f32>,
670 pub top_k: u32,
671}
672
673#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
677pub struct Aggregate {
678 #[serde(default, skip_serializing_if = "Vec::is_empty")]
679 pub group_by: Vec<String>,
680 pub funcs: Vec<AggCall>,
681 #[serde(default, skip_serializing_if = "Option::is_none")]
682 pub window: Option<Window>,
683}
684
685#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
689pub struct AggCall {
690 pub func: AggFunc,
691 #[serde(default, skip_serializing_if = "Option::is_none")]
692 pub field: Option<String>,
693 #[serde(default, skip_serializing_if = "Option::is_none")]
694 pub arg: Option<f64>,
695 pub alias: String,
696}
697
698#[derive(
701 Clone,
702 Copy,
703 Debug,
704 PartialEq,
705 Eq,
706 Serialize,
707 Deserialize,
708 strum::Display,
709 strum::EnumString,
710 strum::VariantArray,
711)]
712#[serde(rename_all = "snake_case")]
713#[strum(serialize_all = "snake_case")]
714pub enum AggFunc {
715 Count,
716 CountDistinct,
717 Sum,
718 Avg,
719 Min,
720 Max,
721 Percentile,
722 StdDev,
723}
724
725#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
727pub struct Window {
728 pub field: String,
729 pub every_micros: u64,
730}
731
732#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
734pub struct Select {
735 #[serde(default, skip_serializing_if = "Vec::is_empty")]
737 pub fields: Vec<String>,
738 #[serde(default)]
740 pub payload: bool,
741}
742
743#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
753#[serde(untagged)]
754pub enum Value {
755 Str(String),
756 Int(i64),
757 Uint(u64),
758 Float(f64),
759 Bool(bool),
760 Null,
761 List(Vec<Value>),
762}
763
764impl From<&str> for Value {
765 fn from(value: &str) -> Self {
766 Self::Str(value.to_owned())
767 }
768}
769
770impl From<String> for Value {
771 fn from(value: String) -> Self {
772 Self::Str(value)
773 }
774}
775
776impl From<&String> for Value {
777 fn from(value: &String) -> Self {
778 Self::Str(value.clone())
779 }
780}
781
782impl From<i64> for Value {
783 fn from(value: i64) -> Self {
784 Self::Int(value)
785 }
786}
787
788impl From<u64> for Value {
789 fn from(value: u64) -> Self {
790 Self::Uint(value)
791 }
792}
793
794impl From<i32> for Value {
795 fn from(value: i32) -> Self {
796 Self::Int(value as i64)
797 }
798}
799
800impl From<u32> for Value {
801 fn from(value: u32) -> Self {
802 Self::Int(value as i64)
803 }
804}
805
806impl From<f64> for Value {
807 fn from(value: f64) -> Self {
808 Self::Float(value)
809 }
810}
811
812impl From<f32> for Value {
813 fn from(value: f32) -> Self {
814 Self::Float(value as f64)
815 }
816}
817
818impl From<bool> for Value {
819 fn from(value: bool) -> Self {
820 Self::Bool(value)
821 }
822}
823
824impl<T: Into<Value>> From<Vec<T>> for Value {
825 fn from(values: Vec<T>) -> Self {
826 Self::List(values.into_iter().map(Into::into).collect())
827 }
828}
829
830impl Value {
831 pub fn from_input(input: &str) -> Self {
842 match input {
843 "null" => return Value::Null,
844 "true" => return Value::Bool(true),
845 "false" => return Value::Bool(false),
846 _ => {}
847 }
848 if let Ok(int) = input.parse::<i64>() {
849 return Value::Int(int);
850 }
851 if let Ok(uint) = input.parse::<u64>() {
852 return Value::Uint(uint);
853 }
854 if !input.is_empty()
857 && input
858 .bytes()
859 .all(|b| b.is_ascii_digit() || b == b'.' || b == b'-' || b == b'+')
860 && let Ok(float) = input.parse::<f64>()
861 {
862 return Value::Float(float);
863 }
864 Value::Str(input.to_owned())
865 }
866}
867
868impl std::fmt::Display for Value {
869 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
873 match self {
874 Value::Str(value) => f.write_str(value),
875 Value::Int(value) => write!(f, "{value}"),
876 Value::Uint(value) => write!(f, "{value}"),
877 Value::Float(value) => write!(f, "{value}"),
878 Value::Bool(value) => write!(f, "{value}"),
879 Value::Null => f.write_str("null"),
880 Value::List(values) => {
881 f.write_str("[")?;
882 for (index, value) in values.iter().enumerate() {
883 if index > 0 {
884 f.write_str(", ")?;
885 }
886 write!(f, "{value}")?;
887 }
888 f.write_str("]")
889 }
890 }
891 }
892}
893
894impl std::str::FromStr for Value {
895 type Err = std::convert::Infallible;
896
897 fn from_str(s: &str) -> Result<Self, Self::Err> {
898 Ok(Value::from_input(s))
899 }
900}
901
902#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
904pub struct QueryResult {
905 pub fields: Vec<LogicalField>,
906 pub rows: Vec<Row>,
907 #[serde(default)]
908 pub page: Page,
909 pub context: QueryContext,
910}
911
912#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
914pub struct Row {
915 pub values: Vec<TypedValue>,
916 #[serde(default, skip_serializing_if = "Option::is_none")]
919 pub score: Option<f32>,
920}
921
922impl Validate for QueryResult {
923 fn validate(&self) -> Result<(), InvalidError> {
924 validate_result_fields(&self.fields)?;
925 for row in &self.rows {
926 if row.values.len() != self.fields.len() {
927 return Err(InvalidError::new(
928 "query row value count does not match the result schema",
929 ));
930 }
931 for (value, field) in row.values.iter().zip(&self.fields) {
932 value.validate_against(&field.field_type, field.required)?;
933 }
934 if row.score.is_some_and(|score| !score.is_finite()) {
935 return Err(InvalidError::new("query row score must be finite"));
936 }
937 }
938 if self.context.execution_id.as_u128() == 0 {
939 return Err(InvalidError::new(
940 "query result execution id must be nonzero",
941 ));
942 }
943 if self.context.row_count != self.rows.len() as u64 {
944 return Err(InvalidError::new(
945 "query result row count does not match the page rows",
946 ));
947 }
948 self.context.validate()?;
949 self.page.validate()
950 }
951}
952
953impl Validate for QueryContext {
954 fn validate(&self) -> Result<(), InvalidError> {
955 if self.execution_id.as_u128() == 0 {
956 return Err(InvalidError::new(
957 "query context execution id must be nonzero",
958 ));
959 }
960 validate_query_name("query engine name", &self.engine.name)?;
961 validate_query_name("query engine version", &self.engine.version)?;
962 if self.delivered_consistency < self.requested_consistency {
963 return Err(InvalidError::new(
964 "delivered consistency is weaker than requested consistency",
965 ));
966 }
967 match &self.resolved_target {
968 ResolvedQueryTarget::Operational {
969 index,
970 backend_resource_id,
971 backend_generation,
972 runtime_configuration_revision,
973 } => {
974 validate_query_name("resolved index", index)?;
975 if backend_resource_id.as_u128() == 0
976 || *backend_generation == 0
977 || *runtime_configuration_revision == 0
978 {
979 return Err(InvalidError::new(
980 "operational query context is missing backend evidence",
981 ));
982 }
983 }
984 }
985 Ok(())
986 }
987}
988
989impl QueryResult {
990 pub fn field_index(&self, name: &str) -> Option<usize> {
991 self.fields.iter().position(|field| field.name == name)
992 }
993
994 pub fn value<'a>(&self, row: &'a Row, field: &str) -> Option<&'a TypedValue> {
995 self.field_index(field)
996 .and_then(|index| row.values.get(index))
997 }
998
999 pub fn value_text(&self, row: &Row, field: &str) -> Option<String> {
1000 self.value(row, field).map(TypedValue::diagnostic_text)
1001 }
1002
1003 pub fn value_u64(&self, row: &Row, field: &str) -> Option<u64> {
1004 self.value(row, field).and_then(TypedValue::as_u64)
1005 }
1006
1007 pub fn value_i64(&self, row: &Row, field: &str) -> Option<i64> {
1008 self.value(row, field).and_then(TypedValue::as_i64)
1009 }
1010}
1011
1012#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1014pub struct QueryContext {
1015 pub execution_id: QueryExecutionId,
1016 pub engine: QueryEngine,
1017 pub resolved_target: ResolvedQueryTarget,
1018 pub requested_consistency: Consistency,
1019 pub delivered_consistency: Consistency,
1020 pub truncated: bool,
1021 pub elapsed_micros: u64,
1022 pub scanned_bytes: u64,
1023 pub produced_bytes: u64,
1024 pub row_count: u64,
1025 #[serde(default, skip_serializing_if = "Vec::is_empty")]
1030 pub frontier: Vec<QueryFrontier>,
1031}
1032
1033#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1035pub struct QueryFrontier {
1036 pub stream: String,
1037 pub topic: String,
1038 pub partition_id: u32,
1039 #[serde(default, skip_serializing_if = "Option::is_none")]
1043 pub topic_created_at_micros: Option<u64>,
1044 pub applied_offset: u64,
1046}
1047
1048#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
1049pub struct QueryEngine {
1050 pub name: String,
1051 pub version: String,
1052 #[serde(default, skip_serializing_if = "Option::is_none")]
1053 pub dialect: Option<SqlDialect>,
1054}
1055
1056#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1057#[serde(tag = "kind", rename_all = "snake_case")]
1058#[non_exhaustive]
1059pub enum ResolvedQueryTarget {
1060 Operational {
1061 index: String,
1062 backend_resource_id: BackendResourceId,
1063 backend_generation: u64,
1064 runtime_configuration_revision: u64,
1065 },
1066}
1067
1068#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
1077pub struct Page {
1078 #[serde(default, skip_serializing_if = "Option::is_none")]
1079 pub offset: Option<u64>,
1080 pub limit: u32,
1082 #[serde(default, skip_serializing_if = "Option::is_none")]
1084 pub total: Option<u64>,
1085 pub has_more: bool,
1087}
1088
1089impl Page {
1090 pub fn at_least(&self, rows_on_page: usize) -> Option<u64> {
1094 self.offset
1095 .and_then(|offset| offset.checked_add(rows_on_page as u64))
1096 }
1097
1098 pub fn total_pages(&self) -> Option<u64> {
1102 match (self.total, self.limit) {
1103 (Some(total), limit) if limit > 0 => Some(total.div_ceil(u64::from(limit))),
1104 _ => None,
1105 }
1106 }
1107}
1108
1109impl Validate for Page {
1110 fn validate(&self) -> Result<(), InvalidError> {
1111 if self.limit == 0 || self.limit > MAX_PAGE_SIZE as u32 {
1112 return Err(InvalidError::new(format!(
1113 "query result page limit exceeds cap {MAX_PAGE_SIZE}"
1114 )));
1115 }
1116 Ok(())
1117 }
1118}
1119
1120#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1123#[non_exhaustive]
1124pub struct QueryEnvelope {
1125 pub v: u32,
1126 pub query: Query,
1127}
1128
1129impl QueryEnvelope {
1130 pub fn new(query: Query) -> Self {
1132 Self {
1133 v: QUERY_OP_VERSION,
1134 query,
1135 }
1136 }
1137}
1138
1139impl Validate for QueryEnvelope {
1140 fn validate(&self) -> Result<(), InvalidError> {
1141 if self.v != QUERY_OP_VERSION {
1142 return Err(InvalidError::new(format!(
1143 "query version must be {QUERY_OP_VERSION}, got {}",
1144 self.v
1145 )));
1146 }
1147 self.query.validate()
1148 }
1149}
1150
1151#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
1153#[non_exhaustive]
1154pub enum QueryReply {
1155 Ok(Box<QueryResult>),
1156 Err(QueryError),
1157}
1158
1159#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, thiserror::Error)]
1161#[non_exhaustive]
1162pub enum QueryError {
1163 #[error("query not supported: {0}")]
1164 Unsupported(String),
1165 #[error("invalid query: {0}")]
1169 Invalid(String),
1170 #[error("unauthorized: {0}")]
1171 Unauthorized(String),
1172 #[error("index not found: {0}")]
1173 IndexNotFound(String),
1174 #[error("fork not found: {0}")]
1175 ForkNotFound(String),
1176 #[error("writer schema {0} is not registered")]
1177 SchemaNotFound(u32),
1178 #[error("backend error: {0}")]
1179 Backend(String),
1180 #[error("temporarily unavailable: {0}")]
1185 Unavailable(String),
1186 #[error("result too large: {what} {size} exceeds cap {cap}")]
1193 TooLarge { what: String, size: u64, cap: u64 },
1194 #[error("unsupported envelope version (expected {expected}, got {got})")]
1195 Version { expected: u32, got: u32 },
1196 #[error("stale read: {what} applied {applied}, required {required}")]
1202 Stale {
1203 what: String,
1204 applied: u64,
1205 required: u64,
1206 },
1207 #[error("query {execution_id} exceeded its deadline")]
1208 DeadlineExceeded { execution_id: QueryExecutionId },
1209 #[error("stale {what} generation: requested {requested}, observed {observed}")]
1210 StaleGeneration {
1211 what: String,
1212 requested: u64,
1213 observed: u64,
1214 },
1215}
1216
1217#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
1219#[serde(rename_all = "snake_case")]
1220#[non_exhaustive]
1221pub enum QueryErrorCode {
1222 Unsupported,
1223 Invalid,
1224 Unauthorized,
1225 IndexNotFound,
1226 ForkNotFound,
1227 SchemaNotFound,
1228 Backend,
1229 Unavailable,
1230 TooLarge,
1231 Version,
1232 Stale,
1233 DeadlineExceeded,
1234 StaleGeneration,
1235}
1236
1237impl QueryError {
1238 pub fn code(&self) -> QueryErrorCode {
1239 match self {
1240 Self::Unsupported(_) => QueryErrorCode::Unsupported,
1241 Self::Invalid(_) => QueryErrorCode::Invalid,
1242 Self::Unauthorized(_) => QueryErrorCode::Unauthorized,
1243 Self::IndexNotFound(_) => QueryErrorCode::IndexNotFound,
1244 Self::ForkNotFound(_) => QueryErrorCode::ForkNotFound,
1245 Self::SchemaNotFound(_) => QueryErrorCode::SchemaNotFound,
1246 Self::Backend(_) => QueryErrorCode::Backend,
1247 Self::Unavailable(_) => QueryErrorCode::Unavailable,
1248 Self::TooLarge { .. } => QueryErrorCode::TooLarge,
1249 Self::Version { .. } => QueryErrorCode::Version,
1250 Self::Stale { .. } => QueryErrorCode::Stale,
1251 Self::DeadlineExceeded { .. } => QueryErrorCode::DeadlineExceeded,
1252 Self::StaleGeneration { .. } => QueryErrorCode::StaleGeneration,
1253 }
1254 }
1255}
1256
1257#[cfg(test)]
1258mod tests {
1259 use super::*;
1260 use crate::schema::LogicalType;
1261
1262 #[test]
1263 fn given_dsl_enums_when_displayed_then_should_be_snake_case() {
1264 assert_eq!(CmpOp::Gte.to_string(), "gte");
1265 assert_eq!(CmpOp::Prefix.to_string(), "prefix");
1266 assert_eq!("ne".parse::<CmpOp>().expect("ne parses"), CmpOp::Ne);
1267 assert_eq!(Dir::Desc.to_string(), "desc");
1268 assert_eq!(AggFunc::Count.to_string(), "count");
1269 }
1270
1271 #[test]
1272 fn given_a_consistency_gate_when_checked_then_should_fail_not_downgrade() {
1273 assert!(
1275 ConsistencyGate::new(0, 100)
1276 .check(Consistency::Eventual, "readings")
1277 .is_ok()
1278 );
1279 assert!(
1281 ConsistencyGate::new(100, 100)
1282 .check(Consistency::ReadYourWrites, "readings")
1283 .is_ok()
1284 );
1285 let stale = ConsistencyGate::new(41, 57)
1286 .check(Consistency::ReadYourWrites, "readings")
1287 .expect_err("a lagging projector must fail, never downgrade");
1288 assert!(matches!(
1289 stale,
1290 QueryError::Stale {
1291 applied: 41,
1292 required: 57,
1293 ..
1294 }
1295 ));
1296 }
1297
1298 #[test]
1299 fn given_a_query_with_more_than_one_shape_when_validated_then_should_refuse_it() {
1300 let vector = VectorQuery {
1301 field: "embedding".to_owned(),
1302 embedding: vec![0.1, 0.2],
1303 top_k: 10,
1304 };
1305 let text = TextQuery {
1306 field: None,
1307 query: "incident".to_owned(),
1308 };
1309 let aggregate = Aggregate {
1310 group_by: Vec::new(),
1311 funcs: vec![AggCall {
1312 func: AggFunc::Count,
1313 field: None,
1314 arg: None,
1315 alias: "n".to_owned(),
1316 }],
1317 window: None,
1318 };
1319 let base = || Query::operational(QueryExecutionId::from_u128(1), "docs", 1_000_000);
1320 let shapes: [(&str, Query); 4] = [
1321 (
1322 "vector and aggregate",
1323 Query {
1324 vector: Some(vector.clone()),
1325 aggregate: Some(aggregate.clone()),
1326 ..base()
1327 },
1328 ),
1329 (
1330 "text and vector",
1331 Query {
1332 text: Some(text.clone()),
1333 vector: Some(vector.clone()),
1334 ..base()
1335 },
1336 ),
1337 (
1338 "text and distinct",
1339 Query {
1340 text: Some(text),
1341 distinct: true,
1342 ..base()
1343 },
1344 ),
1345 (
1346 "aggregate and distinct",
1347 Query {
1348 aggregate: Some(aggregate),
1349 distinct: true,
1350 ..base()
1351 },
1352 ),
1353 ];
1354 for (label, query) in shapes {
1355 let error = query
1356 .validate()
1357 .expect_err("a query with two shapes must be refused");
1358 assert!(error.0.contains("one of"), "{label}: {error:?}");
1359 }
1360 assert!(
1361 Query {
1362 vector: Some(vector),
1363 ..base()
1364 }
1365 .validate()
1366 .is_ok()
1367 );
1368 }
1369
1370 #[test]
1371 fn given_a_ranked_query_with_paging_it_cannot_honour_when_validated_then_should_refuse_it() {
1372 let base = || Query::operational(QueryExecutionId::from_u128(1), "docs", 1_000_000);
1373 let text = || {
1374 Some(TextQuery {
1375 field: None,
1376 query: "incident".to_owned(),
1377 })
1378 };
1379 let vector = || {
1380 Some(VectorQuery {
1381 field: "embedding".to_owned(),
1382 embedding: vec![0.1, 0.2],
1383 top_k: 10,
1384 })
1385 };
1386 let order = || {
1387 vec![Sort {
1388 field: "ts".to_owned(),
1389 dir: Dir::Desc,
1390 }]
1391 };
1392 let total = QueryPageRequest {
1393 want_total: true,
1394 ..QueryPageRequest::default()
1395 };
1396 let refused: [(&str, Query); 5] = [
1397 (
1398 "text with an order",
1399 Query {
1400 text: text(),
1401 order: order(),
1402 ..base()
1403 },
1404 ),
1405 (
1406 "text with a total",
1407 Query {
1408 text: text(),
1409 page: total.clone(),
1410 ..base()
1411 },
1412 ),
1413 (
1414 "vector with an order",
1415 Query {
1416 vector: vector(),
1417 order: order(),
1418 ..base()
1419 },
1420 ),
1421 (
1422 "vector with a total",
1423 Query {
1424 vector: vector(),
1425 page: total,
1426 ..base()
1427 },
1428 ),
1429 (
1430 "vector with an offset",
1431 Query {
1432 vector: vector(),
1433 page: QueryPageRequest {
1434 offset: Some(10),
1435 ..QueryPageRequest::default()
1436 },
1437 ..base()
1438 },
1439 ),
1440 ];
1441 for (label, query) in refused {
1442 let error = query
1443 .validate()
1444 .expect_err("a ranked query cannot honour this paging");
1445 assert!(error.0.contains("takes no"), "{label}: {error:?}");
1446 }
1447 let next_text_page = Query {
1448 text: text(),
1449 page: QueryPageRequest {
1450 offset: Some(50),
1451 ..QueryPageRequest::default()
1452 },
1453 ..base()
1454 };
1455 assert!(next_text_page.validate().is_ok());
1456 }
1457
1458 #[test]
1459 fn given_a_page_when_computing_total_pages_then_should_divide_by_limit() {
1460 let page = Page {
1461 offset: Some(0),
1462 limit: 3,
1463 total: Some(10),
1464 has_more: true,
1465 };
1466 assert_eq!(page.total_pages(), Some(4));
1467 assert_eq!(page.at_least(3), Some(3), "offset 0 plus this page's rows");
1468 assert_eq!(Page::default().total_pages(), None);
1470 }
1471
1472 #[test]
1473 fn given_query_bounds_when_validated_then_should_accept_each_cap_and_refuse_one_past_it() {
1474 let paged = |limit: u32, offset: u64| Query {
1475 page: QueryPageRequest {
1476 limit,
1477 offset: Some(offset),
1478 want_total: false,
1479 },
1480 ..query()
1481 };
1482 let selecting = |fields: usize| Query {
1483 select: Select {
1484 fields: vec!["host".to_owned(); fields],
1485 payload: false,
1486 },
1487 ..query()
1488 };
1489 let ranged = |start: u64, end: u64| Query {
1490 time_range: Some((start, end)),
1491 ..query()
1492 };
1493 for accepted in [
1494 paged(MAX_PAGE_SIZE as u32, MAX_QUERY_OFFSET),
1495 selecting(MAX_QUERY_FIELDS),
1496 ranged(1, 2),
1497 ranged(2, 2),
1498 ] {
1499 accepted.validate().expect("a query at its caps");
1500 }
1501 let refused = [
1502 (
1503 "a zero execution id",
1504 Query {
1505 execution_id: QueryExecutionId::from_u128(0),
1506 ..query()
1507 },
1508 ),
1509 (
1510 "a zero deadline",
1511 Query {
1512 deadline_micros: 0,
1513 ..query()
1514 },
1515 ),
1516 ("a zero page limit", paged(0, 0)),
1517 (
1518 "a page limit past the cap",
1519 paged(MAX_PAGE_SIZE as u32 + 1, 0),
1520 ),
1521 ("an offset past the cap", paged(1, MAX_QUERY_OFFSET + 1)),
1522 (
1523 "selected fields past the cap",
1524 selecting(MAX_QUERY_FIELDS + 1),
1525 ),
1526 ("a range that ends before it starts", ranged(3, 2)),
1527 ];
1528 for (case, rejected) in refused {
1529 assert!(rejected.validate().is_err(), "{case} must be refused");
1530 }
1531 }
1532
1533 #[test]
1534 fn given_aggregate_clauses_when_combined_then_should_refuse_lone_having_and_row_selection() {
1535 let aggregate = Aggregate {
1536 group_by: Vec::new(),
1537 funcs: vec![AggCall {
1538 func: AggFunc::Count,
1539 field: None,
1540 arg: None,
1541 alias: "rows".to_owned(),
1542 }],
1543 window: None,
1544 };
1545 Query {
1546 aggregate: Some(aggregate.clone()),
1547 ..query()
1548 }
1549 .validate()
1550 .expect("an aggregate alone");
1551 let refused = [
1552 (
1553 "a having filter without an aggregate",
1554 Query {
1555 having: Some(Filter::pred("rows", CmpOp::Gt, 1i64)),
1556 ..query()
1557 },
1558 ),
1559 (
1560 "an aggregate with selected fields",
1561 Query {
1562 aggregate: Some(aggregate),
1563 select: Select {
1564 fields: vec!["host".to_owned()],
1565 payload: false,
1566 },
1567 ..query()
1568 },
1569 ),
1570 ];
1571 for (case, rejected) in refused {
1572 assert!(rejected.validate().is_err(), "{case} must be refused");
1573 }
1574 }
1575
1576 #[test]
1577 fn given_raw_sql_with_any_structured_clause_when_validated_then_should_refuse_it() {
1578 let raw_sql = Query {
1579 raw_sql: Some(RawSql {
1580 dialect: SqlDialect::Postgres,
1581 sql: "SELECT 1".to_owned(),
1582 params: Vec::new(),
1583 }),
1584 ..query()
1585 };
1586 raw_sql.validate().expect("raw SQL alone");
1587 let aggregate = Aggregate {
1588 group_by: Vec::new(),
1589 funcs: vec![AggCall {
1590 func: AggFunc::Count,
1591 field: None,
1592 arg: None,
1593 alias: "rows".to_owned(),
1594 }],
1595 window: None,
1596 };
1597 let refused = [
1599 (
1600 "an exact match",
1601 Query {
1602 by_key: vec![KeyMatch::new("host", "a")],
1603 ..raw_sql.clone()
1604 },
1605 ),
1606 (
1607 "a message type",
1608 Query {
1609 message_type: Some("reading".to_owned()),
1610 ..raw_sql.clone()
1611 },
1612 ),
1613 (
1614 "a time range",
1615 Query {
1616 time_range: Some((1, 2)),
1617 ..raw_sql.clone()
1618 },
1619 ),
1620 (
1621 "a filter",
1622 Query {
1623 filter: Some(Filter::pred("host", CmpOp::Eq, "a")),
1624 ..raw_sql.clone()
1625 },
1626 ),
1627 (
1628 "a vector search",
1629 Query {
1630 vector: Some(VectorQuery {
1631 field: "embedding".to_owned(),
1632 embedding: vec![0.5],
1633 top_k: 1,
1634 }),
1635 ..raw_sql.clone()
1636 },
1637 ),
1638 (
1639 "a text search",
1640 Query {
1641 text: Some(TextQuery {
1642 field: None,
1643 query: "incident".to_owned(),
1644 }),
1645 ..raw_sql.clone()
1646 },
1647 ),
1648 (
1649 "a sort",
1650 Query {
1651 order: vec![Sort {
1652 field: "ts".to_owned(),
1653 dir: Dir::Desc,
1654 }],
1655 ..raw_sql.clone()
1656 },
1657 ),
1658 (
1659 "an aggregate",
1660 Query {
1661 aggregate: Some(aggregate.clone()),
1662 ..raw_sql.clone()
1663 },
1664 ),
1665 (
1666 "a having filter",
1667 Query {
1668 aggregate: Some(aggregate),
1669 having: Some(Filter::pred("rows", CmpOp::Gt, 1i64)),
1670 ..raw_sql.clone()
1671 },
1672 ),
1673 (
1674 "distinct rows",
1675 Query {
1676 distinct: true,
1677 ..raw_sql.clone()
1678 },
1679 ),
1680 (
1681 "selected fields",
1682 Query {
1683 select: Select {
1684 fields: vec!["host".to_owned()],
1685 payload: false,
1686 },
1687 ..raw_sql.clone()
1688 },
1689 ),
1690 (
1691 "the payload",
1692 Query {
1693 select: Select {
1694 fields: Vec::new(),
1695 payload: true,
1696 },
1697 ..raw_sql
1698 },
1699 ),
1700 ];
1701 for (clause, rejected) in refused {
1702 assert!(
1703 rejected.validate().is_err(),
1704 "raw SQL with {clause} must be refused"
1705 );
1706 }
1707 }
1708
1709 #[test]
1710 fn given_a_query_envelope_when_validated_then_should_require_the_current_version() {
1711 let envelope = QueryEnvelope::new(query());
1712 envelope.validate().expect("the current version");
1713 let newer = QueryEnvelope {
1714 v: QUERY_OP_VERSION + 1,
1715 ..envelope
1716 };
1717 assert!(newer.validate().is_err());
1718 }
1719
1720 #[test]
1721 fn given_a_query_name_when_validated_then_should_refuse_empty_oversized_or_padded_text() {
1722 let indexed = |index: &str| Query {
1723 target: QueryTarget::operational(index),
1724 ..query()
1725 };
1726 indexed(&"i".repeat(MAX_QUERY_NAME_BYTES))
1727 .validate()
1728 .expect("a name at the cap");
1729 for index in [
1730 String::new(),
1731 "i".repeat(MAX_QUERY_NAME_BYTES + 1),
1732 " readings".to_owned(),
1733 ] {
1734 assert!(
1735 indexed(&index).validate().is_err(),
1736 "{index:?} must be refused"
1737 );
1738 }
1739 }
1740
1741 #[test]
1742 fn given_in_and_prefix_filters_when_validated_then_should_need_a_list_and_a_string() {
1743 let filtered = |filter: Filter| Query {
1744 filter: Some(filter),
1745 ..query()
1746 };
1747 filtered(Filter::all([
1748 Filter::pred("status", CmpOp::Eq, 5i64),
1749 Filter::pred("host", CmpOp::In, vec!["a", "b"]),
1750 Filter::pred("host", CmpOp::Prefix, "node-"),
1751 ]))
1752 .validate()
1753 .expect("well-typed operators");
1754 let refused = [
1755 (
1756 "an empty IN list",
1757 Filter::pred("host", CmpOp::In, Vec::<&str>::new()),
1758 ),
1759 (
1760 "a numeric prefix",
1761 Filter::pred("host", CmpOp::Prefix, 5i64),
1762 ),
1763 ];
1764 for (case, filter) in refused {
1765 assert!(
1766 filtered(filter).validate().is_err(),
1767 "{case} must be refused"
1768 );
1769 }
1770 }
1771
1772 #[test]
1773 fn given_a_filter_at_its_node_and_depth_caps_when_validated_then_should_refuse_one_more() {
1774 let filtered = |filter: Filter| Query {
1775 filter: Some(filter),
1776 ..query()
1777 };
1778 let leaf = Filter::pred("host", CmpOp::Eq, "a");
1779 let wide = |leaves: usize| Filter::All(vec![leaf.clone(); leaves]);
1780 let nested = |depth: usize, wrap: fn(Filter) -> Filter| {
1781 (1..depth).fold(leaf.clone(), |filter, _| wrap(filter))
1782 };
1783 let conjunction: fn(Filter) -> Filter = |filter| Filter::All(vec![filter]);
1784 for accepted in [
1786 wide(MAX_QUERY_PREDICATES - 1),
1787 nested(64, Filter::negate),
1788 nested(64, conjunction),
1789 ] {
1790 filtered(accepted).validate().expect("a filter at its caps");
1791 }
1792 for refused in [
1793 wide(MAX_QUERY_PREDICATES),
1794 nested(65, Filter::negate),
1795 nested(65, conjunction),
1796 ] {
1797 assert!(filtered(refused).validate().is_err(), "one past a cap");
1798 }
1799 }
1800
1801 #[test]
1802 fn given_aggregate_calls_when_validated_then_should_match_each_function_to_its_inputs() {
1803 let aggregate = |func: AggFunc, field: Option<&str>, arg: Option<f64>| Aggregate {
1804 group_by: Vec::new(),
1805 funcs: vec![AggCall {
1806 func,
1807 field: field.map(str::to_owned),
1808 arg,
1809 alias: "value".to_owned(),
1810 }],
1811 window: None,
1812 };
1813 for accepted in [
1814 aggregate(AggFunc::Count, None, None),
1815 aggregate(AggFunc::Percentile, Some("latency"), Some(0.95)),
1816 aggregate(AggFunc::Sum, Some("bytes"), None),
1817 ] {
1818 accepted.validate().expect("a well-formed call");
1819 }
1820 let refused = [
1821 (
1822 "a count with an argument",
1823 aggregate(AggFunc::Count, None, Some(0.5)),
1824 ),
1825 (
1826 "a percentile past one",
1827 aggregate(AggFunc::Percentile, Some("latency"), Some(1.5)),
1828 ),
1829 (
1830 "a percentile without a fraction",
1831 aggregate(AggFunc::Percentile, Some("latency"), None),
1832 ),
1833 ("a sum without a field", aggregate(AggFunc::Sum, None, None)),
1834 ];
1835 for (case, rejected) in refused {
1836 assert!(rejected.validate().is_err(), "{case} must be refused");
1837 }
1838 }
1839
1840 #[test]
1841 fn given_an_aggregate_when_validated_then_should_bound_its_size_and_need_unique_aliases() {
1842 let counts = |calls: usize| {
1843 (0..calls)
1844 .map(|index| AggCall {
1845 func: AggFunc::Count,
1846 field: None,
1847 arg: None,
1848 alias: format!("count_{index}"),
1849 })
1850 .collect::<Vec<_>>()
1851 };
1852 let shaped = |funcs: Vec<AggCall>, groups: usize, every_micros: u64| Aggregate {
1853 group_by: vec!["host".to_owned(); groups],
1854 funcs,
1855 window: Some(Window {
1856 field: "ts".to_owned(),
1857 every_micros,
1858 }),
1859 };
1860 shaped(counts(MAX_QUERY_PARAMETERS), MAX_QUERY_FIELDS, 60_000_000)
1861 .validate()
1862 .expect("an aggregate at its caps");
1863 let refused = [
1864 ("no calls", shaped(Vec::new(), 0, 60_000_000)),
1865 (
1866 "calls past the cap",
1867 shaped(counts(MAX_QUERY_PARAMETERS + 1), 0, 60_000_000),
1868 ),
1869 (
1870 "groups past the cap",
1871 shaped(counts(1), MAX_QUERY_FIELDS + 1, 60_000_000),
1872 ),
1873 (
1874 "a repeated alias",
1875 shaped([counts(1), counts(1)].concat(), 0, 60_000_000),
1876 ),
1877 ("a zero window", shaped(counts(1), 0, 0)),
1878 ];
1879 for (case, rejected) in refused {
1880 assert!(rejected.validate().is_err(), "{case} must be refused");
1881 }
1882 }
1883
1884 #[test]
1885 fn given_raw_sql_when_validated_then_should_bound_its_text_and_parameters() {
1886 let raw = |sql: String, params: usize| RawSql {
1887 dialect: SqlDialect::Postgres,
1888 sql,
1889 params: vec![TypedValue::Long(1); params],
1890 };
1891 for accepted in [
1892 raw("x".repeat(MAX_RAW_SQL_BYTES), MAX_QUERY_PARAMETERS),
1893 raw("SELECT\n\t1\r".to_owned(), 0),
1894 ] {
1895 accepted.validate().expect("raw SQL within its caps");
1896 }
1897 let refused = [
1898 ("blank text", raw(" ".to_owned(), 0)),
1899 (
1900 "text past the cap",
1901 raw("x".repeat(MAX_RAW_SQL_BYTES + 1), 0),
1902 ),
1903 ("a control character", raw("SELECT 1\u{7}".to_owned(), 0)),
1904 (
1905 "parameters past the cap",
1906 raw("SELECT 1".to_owned(), MAX_QUERY_PARAMETERS + 1),
1907 ),
1908 ];
1909 for (case, rejected) in refused {
1910 assert!(rejected.validate().is_err(), "{case} must be refused");
1911 }
1912 }
1913
1914 #[test]
1915 fn given_distinct_and_want_total_when_round_tripped_then_should_omit_them_only_when_false() {
1916 let flagged = Query {
1917 distinct: true,
1918 page: QueryPageRequest {
1919 want_total: true,
1920 ..QueryPageRequest::default()
1921 },
1922 ..query()
1923 };
1924 let json = serde_json::to_value(&flagged).expect("encode a flagged query");
1925 let decoded: Query = serde_json::from_value(json).expect("decode a flagged query");
1926 assert_eq!(decoded, flagged);
1927 let plain = serde_json::to_value(query()).expect("encode a plain query");
1928 assert_eq!(
1929 (plain.get("distinct"), plain["page"].get("want_total")),
1930 (None, None)
1931 );
1932 }
1933
1934 #[test]
1935 fn given_typed_input_when_parsed_then_should_narrow_keywords_and_signed_decimals() {
1936 assert_eq!(
1937 ["null", "true", "false", "1.5", "-1.5", "+1.5"].map(Value::from_input),
1938 [
1939 Value::Null,
1940 Value::Bool(true),
1941 Value::Bool(false),
1942 Value::Float(1.5),
1943 Value::Float(-1.5),
1944 Value::Float(1.5),
1945 ]
1946 );
1947 }
1948
1949 #[test]
1950 fn given_a_list_value_when_displayed_then_should_separate_its_bare_elements() {
1951 let list = Value::List(vec![Value::Int(1), Value::from("a"), Value::Null]);
1952 assert_eq!(list.to_string(), "[1, a, null]");
1953 }
1954
1955 #[test]
1956 fn given_a_query_result_when_validated_then_should_match_rows_to_the_schema_and_the_count() {
1957 let valid = result();
1958 valid.validate().expect("a well-formed result");
1959 let refused = [
1960 (
1961 "a row missing values",
1962 QueryResult {
1963 rows: vec![Row {
1964 values: Vec::new(),
1965 score: None,
1966 }],
1967 ..result()
1968 },
1969 ),
1970 (
1971 "a score that is not a number",
1972 QueryResult {
1973 rows: vec![Row {
1974 score: Some(f32::NAN),
1975 ..valid.rows[0].clone()
1976 }],
1977 ..result()
1978 },
1979 ),
1980 (
1981 "a row count that disagrees with the rows",
1982 QueryResult {
1983 context: QueryContext {
1984 row_count: 2,
1985 ..context()
1986 },
1987 ..result()
1988 },
1989 ),
1990 ];
1991 for (case, rejected) in refused {
1992 assert!(rejected.validate().is_err(), "{case} must be refused");
1993 }
1994 }
1995
1996 #[test]
1997 fn given_a_query_context_when_validated_then_should_need_evidence_and_enough_consistency() {
1998 for accepted in [
1999 context(),
2000 QueryContext {
2001 delivered_consistency: Consistency::ReadYourWrites,
2002 ..context()
2003 },
2004 ] {
2005 accepted
2006 .validate()
2007 .expect("a delivery at least as strong as requested");
2008 }
2009 let evidence = |resource: u128, generation: u64, revision: u64| QueryContext {
2010 resolved_target: ResolvedQueryTarget::Operational {
2011 index: "readings".to_owned(),
2012 backend_resource_id: BackendResourceId::from_u128(resource),
2013 backend_generation: generation,
2014 runtime_configuration_revision: revision,
2015 },
2016 ..context()
2017 };
2018 let refused = [
2019 (
2020 "a zero execution id",
2021 QueryContext {
2022 execution_id: QueryExecutionId::from_u128(0),
2023 ..context()
2024 },
2025 ),
2026 (
2027 "a weaker delivery than requested",
2028 QueryContext {
2029 requested_consistency: Consistency::ReadYourWrites,
2030 ..context()
2031 },
2032 ),
2033 ("a zero backend resource id", evidence(0, 1, 1)),
2034 ("a zero backend generation", evidence(2, 0, 1)),
2035 ("a zero configuration revision", evidence(2, 1, 0)),
2036 ];
2037 for (case, rejected) in refused {
2038 assert!(rejected.validate().is_err(), "{case} must be refused");
2039 }
2040 }
2041
2042 #[test]
2043 fn given_a_result_row_when_read_by_field_name_then_should_return_its_typed_values() {
2044 let reply = result();
2045 let row = &reply.rows[0];
2046 assert_eq!(
2047 (
2048 reply.field_index("delta"),
2049 reply.value(row, "count"),
2050 reply.value_text(row, "reading_id"),
2051 reply.value_u64(row, "count"),
2052 reply.value_i64(row, "delta"),
2053 ),
2054 (
2055 Some(2),
2056 Some(&TypedValue::Long(7)),
2057 Some("r-1".to_owned()),
2058 Some(7),
2059 Some(-3),
2060 )
2061 );
2062 }
2063
2064 #[test]
2065 fn given_a_page_limit_when_validated_then_should_accept_one_to_the_cap() {
2066 let page = |limit: u32| Page {
2067 offset: Some(0),
2068 limit,
2069 total: None,
2070 has_more: false,
2071 };
2072 page(MAX_PAGE_SIZE as u32)
2073 .validate()
2074 .expect("a limit at the cap");
2075 for limit in [0, MAX_PAGE_SIZE as u32 + 1] {
2076 assert!(
2077 page(limit).validate().is_err(),
2078 "limit {limit} must be refused"
2079 );
2080 }
2081 }
2082
2083 #[test]
2084 fn given_a_total_without_a_page_size_when_counting_pages_then_should_have_none() {
2085 let page = Page {
2086 offset: Some(0),
2087 limit: 0,
2088 total: Some(10),
2089 has_more: false,
2090 };
2091 assert_eq!(page.total_pages(), None);
2092 }
2093
2094 fn query() -> Query {
2095 Query::operational(QueryExecutionId::from_u128(1), "readings", 1_000_000)
2096 }
2097
2098 fn result() -> QueryResult {
2099 let field = |id: u32, name: &str, field_type: LogicalType| LogicalField {
2100 id,
2101 name: name.to_owned(),
2102 required: true,
2103 field_type,
2104 doc: None,
2105 };
2106 QueryResult {
2107 fields: vec![
2108 field(1, "reading_id", LogicalType::String),
2109 field(2, "count", LogicalType::Long),
2110 field(3, "delta", LogicalType::Long),
2111 ],
2112 rows: vec![Row {
2113 values: vec![
2114 TypedValue::from("r-1"),
2115 TypedValue::Long(7),
2116 TypedValue::Long(-3),
2117 ],
2118 score: Some(0.25),
2119 }],
2120 page: Page {
2121 offset: Some(0),
2122 limit: 50,
2123 total: None,
2124 has_more: false,
2125 },
2126 context: context(),
2127 }
2128 }
2129
2130 fn context() -> QueryContext {
2131 QueryContext {
2132 execution_id: QueryExecutionId::from_u128(1),
2133 engine: QueryEngine {
2134 name: "postgres".to_owned(),
2135 version: "17.0".to_owned(),
2136 dialect: Some(SqlDialect::Postgres),
2137 },
2138 resolved_target: ResolvedQueryTarget::Operational {
2139 index: "readings".to_owned(),
2140 backend_resource_id: BackendResourceId::from_u128(2),
2141 backend_generation: 1,
2142 runtime_configuration_revision: 1,
2143 },
2144 requested_consistency: Consistency::Eventual,
2145 delivered_consistency: Consistency::Eventual,
2146 truncated: false,
2147 elapsed_micros: 10,
2148 scanned_bytes: 128,
2149 produced_bytes: 32,
2150 row_count: 1,
2151 frontier: Vec::new(),
2152 }
2153 }
2154}
2155
2156#[cfg(all(test, feature = "codecs"))]
2157mod serde_tests {
2158 use super::*;
2159 use crate::codes::QUERY_OP_VERSION;
2160 use crate::framing::{decode_named, encode_named};
2161
2162 fn query() -> Query {
2163 Query {
2164 execution_id: QueryExecutionId::from_u128(1),
2165 target: QueryTarget::operational("readings"),
2166 deadline_micros: 1_000_000,
2167 by_key: Vec::new(),
2168 message_type: None,
2169 time_range: None,
2170 filter: None,
2171 vector: None,
2172 text: None,
2173 order: Vec::new(),
2174 page: QueryPageRequest::default(),
2175 aggregate: None,
2176 having: None,
2177 distinct: false,
2178 select: Select::default(),
2179 fork: None,
2180 raw_sql: None,
2181 consistency: Consistency::Eventual,
2182 }
2183 }
2184
2185 #[test]
2186 fn given_dsl_enums_when_serialized_then_serde_should_match_display() {
2187 assert_eq!(
2188 serde_json::to_string(&CmpOp::Lte).expect("CmpOp serializes"),
2189 "\"lte\""
2190 );
2191 assert_eq!(
2192 serde_json::from_str::<CmpOp>("\"in\"").expect("CmpOp deserializes"),
2193 CmpOp::In
2194 );
2195 assert_eq!(
2196 serde_json::to_string(&Dir::Asc).expect("Dir serializes"),
2197 "\"asc\""
2198 );
2199 }
2200
2201 #[test]
2202 fn given_a_query_when_round_tripped_through_the_envelope_then_should_be_unchanged() {
2203 let mut query = query();
2204 query.by_key = vec![KeyMatch::new("host_id", "abc")];
2205 query.filter = Some(Filter::pred("status", CmpOp::Eq, "degraded"));
2206 query.order = vec![Sort {
2207 field: "ts".to_owned(),
2208 dir: Dir::Desc,
2209 }];
2210 query.page.limit = 20;
2211 let request = QueryEnvelope::new(query);
2212
2213 let json = serde_json::to_string(&request).expect("the request serializes");
2214 let back: QueryEnvelope = serde_json::from_str(&json).expect("the request deserializes");
2215 assert_eq!(back.v, QUERY_OP_VERSION);
2216 assert_eq!(back.query.target, QueryTarget::operational("readings"));
2217 assert_eq!(back.query.page.limit, 20);
2218 assert_eq!(back.query.by_key, vec![KeyMatch::new("host_id", "abc")]);
2219 let Some(Filter::Pred(predicate)) = &back.query.filter else {
2220 panic!("expected a single predicate filter");
2221 };
2222 assert_eq!(predicate.value, TypedValue::String("degraded".to_owned()));
2223 assert_eq!(back.query.order[0].dir, Dir::Desc);
2224 }
2225
2226 #[test]
2227 fn given_each_consistency_level_when_round_tripped_then_should_preserve_it_and_skip_eventual() {
2228 for level in [Consistency::Eventual, Consistency::ReadYourWrites] {
2229 let mut query = query();
2230 query.consistency = level;
2231 let bytes = encode_named(&QueryEnvelope::new(query)).expect("serializes");
2232 let back: QueryEnvelope = decode_named(&bytes).expect("deserializes");
2233 assert_eq!(back.query.consistency, level);
2234 }
2235 let default = query();
2238 assert_eq!(default.consistency, Consistency::Eventual);
2239 let json = serde_json::to_string(&default).expect("json");
2240 assert!(
2241 !json.contains("consistency"),
2242 "default Eventual must be omitted: {json}"
2243 );
2244 }
2245
2246 #[test]
2247 fn given_a_stale_reply_when_round_tripped_then_should_preserve_the_offsets() {
2248 let reply = QueryReply::Err(QueryError::Stale {
2249 what: "readings".to_owned(),
2250 applied: 41,
2251 required: 57,
2252 });
2253 let bytes = encode_named(&reply).expect("serializes");
2254 let back: QueryReply = decode_named(&bytes).expect("deserializes");
2255 let QueryReply::Err(QueryError::Stale {
2256 what,
2257 applied,
2258 required,
2259 }) = back
2260 else {
2261 panic!("expected a Stale error");
2262 };
2263 assert_eq!((what.as_str(), applied, required), ("readings", 41, 57));
2264 }
2265
2266 #[test]
2267 fn given_a_vector_query_when_round_tripped_then_should_preserve_the_embedding() {
2268 let mut query = query();
2269 query.target = QueryTarget::operational("mem:conv-1");
2270 query.vector = Some(VectorQuery {
2271 field: "embedding".to_owned(),
2272 embedding: vec![0.1, 0.2, 0.3],
2273 top_k: 5,
2274 });
2275 let json = serde_json::to_string(&query).expect("the query serializes");
2276 let back: Query = serde_json::from_str(&json).expect("the query deserializes");
2277 let vector = back.vector.expect("the vector survives the round-trip");
2278 assert_eq!(vector.embedding, vec![0.1, 0.2, 0.3]);
2279 assert_eq!(vector.top_k, 5);
2280 }
2281
2282 #[test]
2283 fn given_a_typed_reply_when_round_tripped_then_should_preserve_the_values() {
2284 let reply = QueryReply::Ok(Box::new(QueryResult {
2285 fields: vec![LogicalField {
2286 id: 1,
2287 name: "reading_id".to_owned(),
2288 required: true,
2289 field_type: crate::schema::LogicalType::String,
2290 doc: None,
2291 }],
2292 rows: vec![Row {
2293 values: vec![TypedValue::String("123".to_owned())],
2294 score: Some(0.25),
2295 }],
2296 page: Page {
2297 offset: Some(0),
2298 limit: 50,
2299 total: Some(1),
2300 has_more: false,
2301 },
2302 context: QueryContext {
2303 execution_id: QueryExecutionId::from_u128(1),
2304 engine: QueryEngine {
2305 name: "postgres".to_owned(),
2306 version: "17.0".to_owned(),
2307 dialect: Some(SqlDialect::Postgres),
2308 },
2309 resolved_target: ResolvedQueryTarget::Operational {
2310 index: "readings".to_owned(),
2311 backend_resource_id: BackendResourceId::from_u128(2),
2312 backend_generation: 1,
2313 runtime_configuration_revision: 1,
2314 },
2315 requested_consistency: Consistency::Eventual,
2316 delivered_consistency: Consistency::Eventual,
2317 truncated: false,
2318 elapsed_micros: 10,
2319 scanned_bytes: 128,
2320 produced_bytes: 32,
2321 row_count: 1,
2322 frontier: Vec::new(),
2323 },
2324 }));
2325 let bytes = encode_named(&reply).expect("the reply serializes");
2326 let back: QueryReply = decode_named(&bytes).expect("the reply deserializes");
2327 let QueryReply::Ok(result) = back else {
2328 panic!("the reply should decode as Ok");
2329 };
2330 assert_eq!(
2331 result.rows[0].values,
2332 vec![TypedValue::String("123".to_owned())]
2333 );
2334 assert_eq!(result.page.total, Some(1));
2335 assert!(!result.page.has_more);
2336 }
2337}