1use crate::error::LaserError;
2use crate::laser::Laser;
3use crate::query::{
4 AGDX_QUERY_CODE, AggCall, AggFunc, Aggregate, CmpOp, Consistency, Dir, Filter, KeyMatch,
5 QUERY_OP_VERSION, Query, QueryError, QueryExecutionId, QueryExecutionStatus, QueryResult,
6 QueryTarget, RawSql, Row, SnapshotSelector, Sort, SqlDialect, TextQuery, TypedValue,
7 VECTOR_FIELD, VectorQuery, Window,
8};
9use crate::stream::Decoder;
10use crate::types::{ConversationId, MintUlid};
11use laser_wire::framing::encode_named;
12use laser_wire::query::{
13 QueryCancelEnvelope, QueryCancelReply, QueryEnvelope, QueryPageEnvelope, QueryReply,
14 QueryStatusEnvelope, QueryStatusReply,
15};
16use laser_wire::schema::ORIGINAL_PAYLOAD_FIELD_NAME;
17use laser_wire::validate::Validate;
18use serde::de::DeserializeOwned;
19use std::time::{Duration, SystemTime, UNIX_EPOCH};
20
21use crate::retry::{QUERY_DEADLINE as DEFAULT_QUERY_DEADLINE, QUERY_REPLY_GRACE};
22
23impl Laser {
24 pub fn query<'a>(&'a self, index: &'a str) -> QueryRequest<'a> {
43 QueryRequest::new(self, QueryTarget::operational(self.resource_name(index)))
44 }
45
46 pub fn query_target(&self, target: QueryTarget) -> QueryRequest<'_> {
50 let target = match target {
51 QueryTarget::Operational { index } => {
52 QueryTarget::operational(self.resource_name(&index))
53 }
54 other => other,
55 };
56 QueryRequest::new(self, target)
57 }
58
59 pub fn query_lakehouse(
61 &self,
62 destination_id: laser_wire::destination::DestinationId,
63 destination_generation: u64,
64 ) -> QueryRequest<'_> {
65 self.query_target(QueryTarget::Lakehouse {
66 destination_id,
67 destination_generation,
68 snapshot: None,
69 })
70 }
71
72 pub async fn execute_query(&self, query: Query) -> Result<QueryResult, LaserError> {
79 query.validate()?;
80 let execution_id = query.execution_id;
81 let capabilities = self.capabilities().await;
82 if !capabilities.query.available {
83 return Err(LaserError::unsupported(
84 "query",
85 "the query surface is not served by this deployment",
86 ));
87 }
88 if !capabilities.serves_consistency(query.consistency) {
95 return Err(LaserError::unsupported_feature(
96 "query",
97 "consistency",
98 format!(
99 "consistency level {:?} is not served by this deployment",
100 query.consistency
101 ),
102 ));
103 }
104 if query.text.is_some() && !capabilities.query.keyword {
109 return Err(LaserError::unsupported_feature(
110 "query",
111 "keyword",
112 "lexical text search is not advertised by this deployment",
113 ));
114 }
115 if let Some(versions) = capabilities.versions
120 && versions.query != QUERY_OP_VERSION
121 {
122 return Err(QueryError::Version {
123 expected: versions.query,
124 got: QUERY_OP_VERSION,
125 }
126 .into());
127 }
128 let deadline_micros = query.deadline_micros;
129 let request = QueryEnvelope::new(query);
130 let payload = encode_named(&request)
131 .map_err(|error| LaserError::Codec(format!("encode request: {error}")))?;
132 let payload = self
136 .send_query_until(AGDX_QUERY_CODE, payload, deadline_micros)
137 .await?;
138 match crate::error::decode_managed_reply::<QueryReply>(&payload)? {
139 QueryReply::Ok(result) => validate_query_result(*result, execution_id),
140 QueryReply::Err(error) => Err(error.into()),
141 _ => Err(LaserError::Protocol(
142 "query: unknown reply variant".to_owned(),
143 )),
144 }
145 }
146
147 async fn send_query_until(
153 &self,
154 code: u32,
155 payload: Vec<u8>,
156 deadline_micros: u64,
157 ) -> Result<Vec<u8>, LaserError> {
158 tokio::time::timeout(
159 query_reply_wait(deadline_micros, now_micros())?,
160 self.send_raw_with_response(code, payload),
161 )
162 .await
163 .map_err(|_| LaserError::Timeout("no query reply by the query deadline"))?
164 }
165
166 pub async fn query_page(
168 &self,
169 execution_id: QueryExecutionId,
170 cursor: impl Into<String>,
171 deadline_micros: u64,
172 ) -> Result<QueryResult, LaserError> {
173 let capabilities = self.capabilities().await;
174 if !capabilities.query.cursor_paging {
175 return Err(LaserError::unsupported_feature(
176 "query",
177 "cursor_paging",
178 "query cursor paging is not advertised by this deployment",
179 ));
180 }
181 let request = QueryPageEnvelope::new(execution_id, cursor, deadline_micros);
182 request.validate()?;
183 let payload = encode_named(&request)
184 .map_err(|error| LaserError::Codec(format!("encode request: {error}")))?;
185 let payload = self
186 .send_query_until(crate::query::AGDX_QUERY_PAGE_CODE, payload, deadline_micros)
187 .await?;
188 match crate::error::decode_managed_reply::<QueryReply>(&payload)? {
189 QueryReply::Ok(result) => validate_query_result(*result, execution_id),
190 QueryReply::Err(error) => Err(error.into()),
191 _ => Err(LaserError::Protocol(
192 "query page: unknown reply variant".to_owned(),
193 )),
194 }
195 }
196
197 pub async fn cancel_query(
199 &self,
200 execution_id: QueryExecutionId,
201 ) -> Result<QueryExecutionStatus, LaserError> {
202 if !self.capabilities().await.query.cancellation {
203 return Err(LaserError::unsupported_feature(
204 "query",
205 "cancellation",
206 "query cancellation is not advertised by this deployment",
207 ));
208 }
209 let request = QueryCancelEnvelope::new(execution_id);
210 request.validate()?;
211 let payload = encode_named(&request)
212 .map_err(|error| LaserError::Codec(format!("encode request: {error}")))?;
213 let payload = self
214 .send_raw_with_response(crate::query::AGDX_QUERY_CANCEL_CODE, payload)
215 .await?;
216 match crate::error::decode_managed_reply::<QueryCancelReply>(&payload)? {
217 QueryCancelReply::Ok(status) => validate_query_status(status, execution_id),
218 QueryCancelReply::Err(error) => Err(error.into()),
219 _ => Err(LaserError::Protocol(
220 "query cancellation: unknown reply variant".to_owned(),
221 )),
222 }
223 }
224
225 pub async fn query_status(
227 &self,
228 execution_id: QueryExecutionId,
229 ) -> Result<QueryExecutionStatus, LaserError> {
230 if !self.capabilities().await.query.execution_status {
231 return Err(LaserError::unsupported_feature(
232 "query",
233 "execution_status",
234 "query execution status is not advertised by this deployment",
235 ));
236 }
237 let request = QueryStatusEnvelope::new(execution_id);
238 request.validate()?;
239 let payload = encode_named(&request)
240 .map_err(|error| LaserError::Codec(format!("encode request: {error}")))?;
241 let payload = self
242 .send_raw_with_response(crate::query::AGDX_QUERY_STATUS_CODE, payload)
243 .await?;
244 match crate::error::decode_managed_reply::<QueryStatusReply>(&payload)? {
245 QueryStatusReply::Ok(status) => validate_query_status(status, execution_id),
246 QueryStatusReply::Err(error) => Err(error.into()),
247 _ => Err(LaserError::Protocol(
248 "query status: unknown reply variant".to_owned(),
249 )),
250 }
251 }
252}
253
254fn validate_query_result(
255 result: QueryResult,
256 execution_id: QueryExecutionId,
257) -> Result<QueryResult, LaserError> {
258 result.validate()?;
259 if result.context.execution_id != execution_id {
260 return Err(LaserError::Protocol(
261 "query reply execution id does not match the request".to_owned(),
262 ));
263 }
264 Ok(result)
265}
266
267fn validate_query_status(
268 status: QueryExecutionStatus,
269 execution_id: QueryExecutionId,
270) -> Result<QueryExecutionStatus, LaserError> {
271 status.validate()?;
272 if status.execution_id != execution_id {
273 return Err(LaserError::Protocol(
274 "query status execution id does not match the request".to_owned(),
275 ));
276 }
277 Ok(status)
278}
279
280#[must_use = "call .fetch().await (or a fetch_* variant) to run the query"]
284pub struct QueryRequest<'a> {
285 laser: &'a Laser,
286 query: Query,
287 max_rows: Option<usize>,
288}
289
290impl<'a> QueryRequest<'a> {
291 fn new(laser: &'a Laser, target: QueryTarget) -> Self {
292 Self {
293 laser,
294 query: Query::builder()
295 .execution_id(crate::query::QueryExecutionId::mint())
296 .target(target)
297 .deadline_micros(query_deadline(DEFAULT_QUERY_DEADLINE))
298 .build(),
299 max_rows: None,
300 }
301 }
302
303 pub fn deadline(mut self, timeout: Duration) -> Self {
307 self.query.deadline_micros = query_deadline(timeout);
308 self
309 }
310
311 pub fn deadline_micros(mut self, deadline_micros: u64) -> Self {
315 self.query.deadline_micros = deadline_micros;
316 self
317 }
318
319 pub const fn execution_id(&self) -> QueryExecutionId {
321 self.query.execution_id
322 }
323
324 pub async fn status(&self) -> Result<QueryExecutionStatus, LaserError> {
326 self.laser.query_status(self.query.execution_id).await
327 }
328
329 pub async fn cancel(&self) -> Result<QueryExecutionStatus, LaserError> {
331 self.laser.cancel_query(self.query.execution_id).await
332 }
333
334 pub fn at_snapshot(mut self, snapshot_id: i64) -> Result<Self, LaserError> {
336 let QueryTarget::Lakehouse { snapshot, .. } = &mut self.query.target else {
337 return Err(LaserError::Invalid(
338 "snapshot selection requires a lakehouse query target".to_owned(),
339 ));
340 };
341 *snapshot = Some(SnapshotSelector::SnapshotId(snapshot_id));
342 Ok(self)
343 }
344
345 pub fn at_timestamp_micros(mut self, timestamp_micros: i64) -> Result<Self, LaserError> {
347 let QueryTarget::Lakehouse { snapshot, .. } = &mut self.query.target else {
348 return Err(LaserError::Invalid(
349 "timestamp selection requires a lakehouse query target".to_owned(),
350 ));
351 };
352 *snapshot = Some(SnapshotSelector::TimestampMicros(timestamp_micros));
353 Ok(self)
354 }
355
356 pub fn where_eq(mut self, field: impl Into<String>, value: impl Into<TypedValue>) -> Self {
358 self.query.by_key.push(KeyMatch::new(field, value));
359 self
360 }
361
362 pub fn conversation(self, conversation: ConversationId) -> Self {
369 self.where_eq(
370 laser_wire::headers::CONVERSATION_FIELD,
371 conversation.to_string(),
372 )
373 }
374
375 pub fn fork(mut self, fork_id: impl Into<String>) -> Self {
379 self.query.fork = Some(self.laser.resource_name(&fork_id.into()));
380 self
381 }
382
383 pub fn filter_eq(self, field: impl Into<String>, value: impl Into<TypedValue>) -> Self {
385 self.predicate(field, CmpOp::Eq, value)
386 }
387
388 pub fn filter_ne(self, field: impl Into<String>, value: impl Into<TypedValue>) -> Self {
390 self.predicate(field, CmpOp::Ne, value)
391 }
392
393 pub fn filter_gt(self, field: impl Into<String>, value: impl Into<TypedValue>) -> Self {
395 self.predicate(field, CmpOp::Gt, value)
396 }
397
398 pub fn filter_gte(self, field: impl Into<String>, value: impl Into<TypedValue>) -> Self {
400 self.predicate(field, CmpOp::Gte, value)
401 }
402
403 pub fn filter_lt(self, field: impl Into<String>, value: impl Into<TypedValue>) -> Self {
405 self.predicate(field, CmpOp::Lt, value)
406 }
407
408 pub fn filter_lte(self, field: impl Into<String>, value: impl Into<TypedValue>) -> Self {
410 self.predicate(field, CmpOp::Lte, value)
411 }
412
413 pub fn filter_in<T: Into<TypedValue>>(
415 self,
416 field: impl Into<String>,
417 values: impl IntoIterator<Item = T>,
418 ) -> Self {
419 let list = TypedValue::List(values.into_iter().map(Into::into).collect());
420 self.predicate(field, CmpOp::In, list)
421 }
422
423 pub fn filter_contains(self, field: impl Into<String>, value: impl Into<String>) -> Self {
425 self.predicate(field, CmpOp::Contains, TypedValue::String(value.into()))
426 }
427
428 pub fn filter_prefix(self, field: impl Into<String>, value: impl Into<String>) -> Self {
430 self.predicate(field, CmpOp::Prefix, TypedValue::String(value.into()))
431 }
432
433 pub fn message_type(mut self, value: impl Into<String>) -> Self {
435 self.query.message_type = Some(value.into());
436 self
437 }
438
439 pub fn time_range(mut self, start: u64, end: u64) -> Self {
442 self.query.time_range = Some((start, end));
443 self
444 }
445
446 pub fn order_asc(mut self, field: impl Into<String>) -> Self {
448 self.query.order.push(Sort {
449 field: field.into(),
450 dir: Dir::Asc,
451 });
452 self
453 }
454
455 pub fn order_desc(mut self, field: impl Into<String>) -> Self {
457 self.query.order.push(Sort {
458 field: field.into(),
459 dir: Dir::Desc,
460 });
461 self
462 }
463
464 pub fn limit(mut self, n: usize) -> Self {
467 self.query.page.limit = u32::try_from(n).unwrap_or(u32::MAX);
468 self
469 }
470
471 pub fn cursor(mut self, cursor: impl Into<String>) -> Self {
474 self.query.page.cursor = Some(cursor.into());
475 self.query.page.offset = None;
476 self
477 }
478
479 pub fn offset(mut self, n: usize) -> Self {
481 self.query.page.offset = Some(n as u64);
482 self.query.page.cursor = None;
483 self
484 }
485
486 pub fn with_payload(mut self) -> Self {
488 self.query.select.payload = true;
489 self
490 }
491
492 pub fn with_total(mut self) -> Self {
497 self.query.page.want_total = true;
498 self
499 }
500
501 pub fn read_your_writes(mut self) -> Self {
508 self.query.consistency = Consistency::ReadYourWrites;
509 self
510 }
511
512 pub fn consistency(mut self, level: Consistency) -> Self {
515 self.query.consistency = level;
516 self
517 }
518
519 pub fn text(mut self, query: impl Into<String>) -> Self {
524 self.query.text = Some(TextQuery {
525 field: None,
526 query: query.into(),
527 });
528 self
529 }
530
531 pub fn text_in(mut self, field: impl Into<String>, query: impl Into<String>) -> Self {
533 self.query.text = Some(TextQuery {
534 field: Some(field.into()),
535 query: query.into(),
536 });
537 self
538 }
539
540 pub fn select_fields<S: Into<String>>(mut self, fields: impl IntoIterator<Item = S>) -> Self {
542 self.query.select.fields = fields.into_iter().map(Into::into).collect();
543 self
544 }
545
546 pub fn count(self) -> Self {
548 self.push_agg(agg_call(AggFunc::Count, None, None, "count"))
549 }
550
551 pub fn sum(self, field: impl Into<String>) -> Self {
553 self.push_agg(agg_call(AggFunc::Sum, Some(field.into()), None, "sum"))
554 }
555
556 pub fn avg(self, field: impl Into<String>) -> Self {
558 self.push_agg(agg_call(AggFunc::Avg, Some(field.into()), None, "avg"))
559 }
560
561 pub fn min(self, field: impl Into<String>) -> Self {
563 self.push_agg(agg_call(AggFunc::Min, Some(field.into()), None, "min"))
564 }
565
566 pub fn max(self, field: impl Into<String>) -> Self {
568 self.push_agg(agg_call(AggFunc::Max, Some(field.into()), None, "max"))
569 }
570
571 pub fn count_distinct(self, field: impl Into<String>) -> Self {
574 self.push_agg(agg_call(
575 AggFunc::CountDistinct,
576 Some(field.into()),
577 None,
578 "count_distinct",
579 ))
580 }
581
582 pub fn stddev(self, field: impl Into<String>) -> Self {
585 self.push_agg(agg_call(
586 AggFunc::StdDev,
587 Some(field.into()),
588 None,
589 "stddev",
590 ))
591 }
592
593 pub fn percentile(self, field: impl Into<String>, fraction: f64) -> Self {
596 self.push_agg(agg_call(
597 AggFunc::Percentile,
598 Some(field.into()),
599 Some(fraction),
600 "percentile",
601 ))
602 }
603
604 pub fn agg_as(
607 self,
608 func: AggFunc,
609 field: Option<String>,
610 fraction: Option<f64>,
611 alias: impl Into<String>,
612 ) -> Self {
613 let alias = alias.into();
614 self.push_agg(AggCall {
615 func,
616 field,
617 arg: fraction,
618 alias,
619 })
620 }
621
622 pub fn group_by<S: Into<String>>(mut self, fields: impl IntoIterator<Item = S>) -> Self {
624 let group_by: Vec<String> = fields.into_iter().map(Into::into).collect();
625 match self.query.aggregate.as_mut() {
626 Some(aggregate) => aggregate.group_by = group_by,
627 None => {
628 self.query.aggregate = Some(Aggregate {
629 group_by,
630 funcs: Vec::new(),
631 window: None,
632 });
633 }
634 }
635 self
636 }
637
638 pub fn window(mut self, field: impl Into<String>, every_micros: u64) -> Self {
642 let window = Some(Window {
643 field: field.into(),
644 every_micros,
645 });
646 match self.query.aggregate.as_mut() {
647 Some(aggregate) => aggregate.window = window,
648 None => {
649 self.query.aggregate = Some(Aggregate {
650 group_by: Vec::new(),
651 funcs: Vec::new(),
652 window,
653 });
654 }
655 }
656 self
657 }
658
659 pub fn having(mut self, filter: Filter) -> Self {
662 self.query.having = Some(filter);
663 self
664 }
665
666 pub fn distinct(mut self) -> Self {
669 self.query.distinct = true;
670 self
671 }
672
673 pub fn raw_sql(mut self, dialect: SqlDialect, sql: impl Into<String>) -> Self {
677 self.query.raw_sql = Some(RawSql {
678 dialect,
679 sql: sql.into(),
680 params: Vec::new(),
681 });
682 self
683 }
684
685 pub fn raw_sql_with<V: Into<TypedValue>>(
687 mut self,
688 dialect: SqlDialect,
689 sql: impl Into<String>,
690 params: impl IntoIterator<Item = V>,
691 ) -> Self {
692 self.query.raw_sql = Some(RawSql {
693 dialect,
694 sql: sql.into(),
695 params: params.into_iter().map(Into::into).collect(),
696 });
697 self
698 }
699
700 pub fn nearest(self, embedding: Vec<f32>, top_k: usize) -> Self {
704 self.nearest_in(VECTOR_FIELD, embedding, top_k)
705 }
706
707 pub fn nearest_in(
710 mut self,
711 field: impl Into<String>,
712 embedding: Vec<f32>,
713 top_k: usize,
714 ) -> Self {
715 self.query.vector = Some(VectorQuery {
716 field: field.into(),
717 embedding,
718 top_k: u32::try_from(top_k).unwrap_or(u32::MAX),
719 });
720 self
721 }
722
723 pub async fn fetch(self) -> Result<QueryResult, LaserError> {
725 self.laser.execute_query(self.query).await
726 }
727
728 pub async fn fetch_typed<T: DeserializeOwned>(mut self) -> Result<Vec<T>, LaserError> {
732 self.query.select.payload = true;
733 let result = self.laser.execute_query(self.query).await?;
734 result
735 .rows
736 .iter()
737 .map(|row| decode_row::<crate::stream::Json, T>(&result.fields, row))
738 .collect()
739 }
740
741 pub async fn fetch_one<T: DeserializeOwned>(mut self) -> Result<Option<T>, LaserError> {
743 self.query.select.payload = true;
744 self.query.page.limit = 1;
745 let result = self.laser.execute_query(self.query).await?;
746 match result.rows.first() {
747 Some(row) => decode_row::<crate::stream::Json, T>(&result.fields, row).map(Some),
748 None => Ok(None),
749 }
750 }
751
752 pub async fn fetch_typed_with<C, T>(mut self) -> Result<Vec<T>, LaserError>
766 where
767 C: Decoder<T>,
768 {
769 self.query.select.payload = true;
770 let result = self.laser.execute_query(self.query).await?;
771 result
772 .rows
773 .iter()
774 .map(|row| decode_row::<C, T>(&result.fields, row))
775 .collect()
776 }
777
778 pub async fn fetch_one_with<C, T>(mut self) -> Result<Option<T>, LaserError>
781 where
782 C: Decoder<T>,
783 {
784 self.query.select.payload = true;
785 self.query.page.limit = 1;
786 let result = self.laser.execute_query(self.query).await?;
787 match result.rows.first() {
788 Some(row) => decode_row::<C, T>(&result.fields, row).map(Some),
789 None => Ok(None),
790 }
791 }
792
793 fn page_stream(mut self) -> QueryStream<'a> {
801 if self.query.page.limit == 0 {
802 self.query.page.limit = crate::query::DEFAULT_STREAM_PAGE_SIZE as u32;
803 }
804 let single_page = self.query.aggregate.is_some() || self.query.vector.is_some();
805 QueryStream::new(self.laser, self.query, single_page)
806 }
807
808 fn typed_page_stream<T: DeserializeOwned>(mut self) -> TypedQueryStream<'a, T> {
812 self.query.select.payload = true;
813 if self.query.page.limit == 0 {
814 self.query.page.limit = crate::query::DEFAULT_STREAM_PAGE_SIZE as u32;
815 }
816 let single_page = self.query.aggregate.is_some() || self.query.vector.is_some();
817 TypedQueryStream::new(self.laser, self.query, single_page)
818 }
819
820 pub fn max_rows(mut self, n: usize) -> Self {
825 self.max_rows = Some(n);
826 self
827 }
828
829 pub fn rows(self) -> Result<QueryRows<'a>, LaserError> {
844 let Some(max_rows) = self.max_rows else {
845 return Err(LaserError::Invalid(
846 "rows() needs an explicit ceiling: chain .max_rows(n) first".to_owned(),
847 ));
848 };
849 Ok(QueryRows {
850 stream: self.page_stream(),
851 remaining: max_rows,
852 })
853 }
854
855 pub fn rows_typed<T: DeserializeOwned>(self) -> Result<TypedQueryRows<'a, T>, LaserError> {
860 let Some(max_rows) = self.max_rows else {
861 return Err(LaserError::Invalid(
862 "rows_typed() needs an explicit ceiling: chain .max_rows(n) first".to_owned(),
863 ));
864 };
865 Ok(TypedQueryRows {
866 stream: self.typed_page_stream(),
867 remaining: max_rows,
868 })
869 }
870
871 pub async fn fetch_all(self) -> Result<Vec<Row>, LaserError> {
876 let mut stream = self.page_stream();
877 let mut rows = Vec::new();
878 while let Some(row) = stream.next().await? {
879 rows.push(row);
880 }
881 Ok(rows)
882 }
883
884 pub async fn fetch_all_typed<T: DeserializeOwned>(self) -> Result<Vec<T>, LaserError> {
887 let mut stream = self.typed_page_stream::<T>();
888 let mut rows = Vec::new();
889 while let Some(row) = stream.next().await? {
890 rows.push(row);
891 }
892 Ok(rows)
893 }
894
895 pub fn into_query(self) -> Query {
897 self.query
898 }
899
900 fn predicate(self, field: impl Into<String>, op: CmpOp, value: impl Into<TypedValue>) -> Self {
901 self.and_filter(Filter::pred(field, op, value))
902 }
903
904 pub fn filter(self, filter: Filter) -> Self {
909 self.and_filter(filter)
910 }
911
912 fn and_filter(mut self, filter: Filter) -> Self {
913 self.query.filter = Some(match self.query.filter.take() {
914 None => filter,
915 Some(Filter::All(mut existing)) => {
916 existing.push(filter);
917 Filter::All(existing)
918 }
919 Some(other) => Filter::All(vec![other, filter]),
920 });
921 self
922 }
923
924 fn push_agg(mut self, call: AggCall) -> Self {
925 match self.query.aggregate.as_mut() {
926 Some(aggregate) => aggregate.funcs.push(call),
927 None => {
928 self.query.aggregate = Some(Aggregate {
929 group_by: Vec::new(),
930 funcs: vec![call],
931 window: None,
932 });
933 }
934 }
935 self
936 }
937}
938
939fn agg_call(func: AggFunc, field: Option<String>, arg: Option<f64>, alias: &str) -> AggCall {
941 AggCall {
942 func,
943 field,
944 arg,
945 alias: alias.to_owned(),
946 }
947}
948
949impl<'a> From<QueryRequest<'a>> for Query {
950 fn from(request: QueryRequest<'a>) -> Self {
951 request.query
952 }
953}
954
955fn query_deadline(timeout: Duration) -> u64 {
956 let now = SystemTime::now()
957 .duration_since(UNIX_EPOCH)
958 .unwrap_or_default();
959 now.saturating_add(timeout)
960 .as_micros()
961 .min(u128::from(u64::MAX)) as u64
962}
963
964fn now_micros() -> u64 {
965 query_deadline(Duration::ZERO)
966}
967
968fn query_reply_wait(deadline_micros: u64, now_micros: u64) -> Result<Duration, LaserError> {
970 let remaining = Duration::from_micros(deadline_micros.saturating_sub(now_micros));
971 if remaining > DEFAULT_QUERY_DEADLINE {
972 return Err(LaserError::Invalid(format!(
973 "the query deadline is {}s away, the longest allowed is {}s",
974 remaining.as_secs(),
975 DEFAULT_QUERY_DEADLINE.as_secs()
976 )));
977 }
978 Ok(remaining + QUERY_REPLY_GRACE)
979}
980
981fn decode_row<C, T>(fields: &[laser_wire::schema::LogicalField], row: &Row) -> Result<T, LaserError>
982where
983 C: Decoder<T>,
984{
985 let payload_index = fields
986 .iter()
987 .position(|field| field.name == ORIGINAL_PAYLOAD_FIELD_NAME)
988 .ok_or_else(|| {
989 LaserError::from(laser_wire::error::DecodeError::MissingPayload(
990 "query result does not contain the original payload field",
991 ))
992 })?;
993 let value = row.values.get(payload_index).ok_or_else(|| {
994 LaserError::Protocol("query row is shorter than its declared result schema".to_owned())
995 })?;
996 match value {
997 TypedValue::Binary(value) => C::decode(&value.0).map_err(LaserError::from),
998 TypedValue::Null => Err(LaserError::from(
999 laser_wire::error::DecodeError::MissingPayload("query row has no original payload"),
1000 )),
1001 _ => Err(LaserError::Protocol(
1002 "query original payload field is not binary".to_owned(),
1003 )),
1004 }
1005}
1006
1007struct QueryStream<'a> {
1014 laser: &'a Laser,
1015 query: Query,
1016 finished: bool,
1017 single_page: bool,
1020 fields: Vec<laser_wire::schema::LogicalField>,
1021 buffer: std::vec::IntoIter<Row>,
1022}
1023
1024impl<'a> QueryStream<'a> {
1025 fn new(laser: &'a Laser, query: Query, single_page: bool) -> Self {
1026 Self {
1027 laser,
1028 query,
1029 finished: false,
1030 single_page,
1031 fields: Vec::new(),
1032 buffer: Vec::new().into_iter(),
1033 }
1034 }
1035
1036 async fn next(&mut self) -> Result<Option<Row>, LaserError> {
1039 if let Some(row) = self.buffer.next() {
1040 return Ok(Some(row));
1041 }
1042 if self.finished {
1043 return Ok(None);
1044 }
1045 let page = self.laser.execute_query(self.query.clone()).await?;
1046 self.finished = self.single_page
1047 || page_ends_walk(
1048 page.rows.len(),
1049 &page.page,
1050 self.query.page.cursor.as_deref(),
1051 );
1052 self.query.page.cursor = page.page.next_cursor.clone();
1053 self.query.page.offset = None;
1054 self.fields = page.fields;
1055 self.buffer = page.rows.into_iter();
1056 Ok(self.buffer.next())
1057 }
1058}
1059
1060fn page_ends_walk(fetched: usize, page: &laser_wire::query::Page, sent: Option<&str>) -> bool {
1065 fetched == 0
1066 || !page.has_more
1067 || page
1068 .next_cursor
1069 .as_deref()
1070 .is_none_or(|next| Some(next) == sent)
1071}
1072
1073pub struct QueryRows<'a> {
1077 stream: QueryStream<'a>,
1078 remaining: usize,
1079}
1080
1081impl QueryRows<'_> {
1082 pub async fn next(&mut self) -> Result<Option<Row>, LaserError> {
1084 if self.remaining == 0 {
1085 return Ok(None);
1086 }
1087 let row = self.stream.next().await?;
1088 if row.is_some() {
1089 self.remaining -= 1;
1090 }
1091 Ok(row)
1092 }
1093}
1094
1095struct TypedQueryStream<'a, T> {
1099 inner: QueryStream<'a>,
1100 _marker: std::marker::PhantomData<fn() -> T>,
1101}
1102
1103impl<'a, T: DeserializeOwned> TypedQueryStream<'a, T> {
1104 fn new(laser: &'a Laser, query: Query, single_page: bool) -> Self {
1105 Self {
1106 inner: QueryStream::new(laser, query, single_page),
1107 _marker: std::marker::PhantomData,
1108 }
1109 }
1110
1111 async fn next(&mut self) -> Result<Option<T>, LaserError> {
1113 match self.inner.next().await? {
1114 Some(row) => decode_row::<crate::stream::Json, T>(&self.inner.fields, &row).map(Some),
1115 None => Ok(None),
1116 }
1117 }
1118}
1119
1120pub struct TypedQueryRows<'a, T> {
1125 stream: TypedQueryStream<'a, T>,
1126 remaining: usize,
1127}
1128
1129impl<T: DeserializeOwned> TypedQueryRows<'_, T> {
1130 pub async fn next(&mut self) -> Result<Option<T>, LaserError> {
1133 if self.remaining == 0 {
1134 return Ok(None);
1135 }
1136 let row = self.stream.next().await?;
1137 if row.is_some() {
1138 self.remaining -= 1;
1139 }
1140 Ok(row)
1141 }
1142}
1143
1144#[cfg(test)]
1145mod tests {
1146 use super::*;
1147
1148 #[test]
1149 fn given_query_deadlines_when_waiting_then_should_cap_them_and_add_the_grace() {
1150 let now = 1_000_000_000;
1151 assert_eq!(
1152 query_reply_wait(now + 5_000_000, now).unwrap(),
1153 Duration::from_secs(5) + QUERY_REPLY_GRACE
1154 );
1155 assert_eq!(query_reply_wait(now - 1, now).unwrap(), QUERY_REPLY_GRACE);
1156 assert!(query_reply_wait(now + DEFAULT_QUERY_DEADLINE.as_micros() as u64, now).is_ok());
1157 assert!(matches!(
1158 query_reply_wait(now + 60_000_000, now),
1159 Err(LaserError::Invalid(_))
1160 ));
1161 }
1162
1163 #[tokio::test]
1164 async fn given_an_over_cap_limit_when_executed_then_should_be_invalid_before_any_round_trip() {
1165 let laser = Laser::from_client(crate::iggy::prelude::IggyClient::default());
1166 let query = laser
1167 .query("readings")
1168 .limit(crate::query::MAX_PAGE_SIZE + 1)
1169 .into_query();
1170 assert!(matches!(
1171 laser.execute_query(query).await,
1172 Err(LaserError::Invalid(_))
1173 ));
1174 }
1175
1176 #[test]
1177 fn given_a_default_stream_when_querying_then_should_scope_the_index_and_the_fork() {
1178 let laser = Laser::from_client(crate::iggy::prelude::IggyClient::default())
1179 .with_default_stream("acme");
1180 let query = laser.query("readings").fork("experiment").into_query();
1181 assert_eq!(
1182 query.target,
1183 QueryTarget::operational("stream:acme/readings")
1184 );
1185 assert_eq!(query.fork.as_deref(), Some("stream:acme/experiment"));
1186 let explicit = laser
1187 .query_target(QueryTarget::operational("readings"))
1188 .into_query();
1189 assert_eq!(
1190 explicit.target,
1191 QueryTarget::operational("stream:acme/readings")
1192 );
1193 let bare = laser
1194 .with_resource_naming(crate::laser::ResourceNaming::Bare)
1195 .query("readings")
1196 .into_query()
1197 .target;
1198 assert_eq!(bare, QueryTarget::operational("readings"));
1199 }
1200
1201 #[test]
1202 fn given_more_rows_without_a_new_cursor_when_paging_then_should_end_the_walk() {
1203 let page = |has_more: bool, next_cursor: Option<&str>| laser_wire::query::Page {
1204 offset: None,
1205 limit: 2,
1206 total: None,
1207 has_more,
1208 next_cursor: next_cursor.map(str::to_owned),
1209 };
1210 assert!(!page_ends_walk(2, &page(true, Some("b")), Some("a")));
1211 assert!(page_ends_walk(2, &page(true, None), Some("a")));
1212 assert!(page_ends_walk(2, &page(true, None), None));
1213 assert!(page_ends_walk(2, &page(true, Some("a")), Some("a")));
1214 assert!(page_ends_walk(0, &page(true, Some("b")), Some("a")));
1215 assert!(page_ends_walk(2, &page(false, Some("b")), Some("a")));
1216 }
1217
1218 #[test]
1219 fn given_a_cursor_when_set_then_should_replace_the_offset() {
1220 let laser = Laser::from_client(crate::iggy::prelude::IggyClient::default());
1221 let query = laser
1222 .query("readings")
1223 .offset(20)
1224 .cursor("c-2")
1225 .into_query();
1226 assert_eq!(query.page.cursor.as_deref(), Some("c-2"));
1227 assert_eq!(query.page.offset, None);
1228 let query = laser.query("readings").cursor("c-2").offset(5).into_query();
1229 assert_eq!(query.page.cursor, None);
1230 assert_eq!(query.page.offset, Some(5));
1231 }
1232
1233 #[test]
1234 fn given_an_absolute_deadline_when_set_then_should_carry_it_unchanged() {
1235 let laser = Laser::from_client(crate::iggy::prelude::IggyClient::default());
1236 let query = laser
1237 .query("readings")
1238 .deadline_micros(1_700_000_000_000_000)
1239 .into_query();
1240 assert_eq!(query.deadline_micros, 1_700_000_000_000_000);
1241 }
1242}