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