teaql_core/
query.rs

1use std::collections::BTreeMap;
2
3use crate::{Expr, Value};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6pub enum SortDirection {
7    Asc,
8    Desc,
9}
10
11#[cfg(test)]
12mod hard_limit_tests {
13    use super::*;
14
15    #[test]
16    fn list_limit_defaults_rejects_and_allows_explicit_override() {
17        assert_eq!(
18            SelectQuery::new("Order")
19                .prepare_for_list()
20                .unwrap()
21                .slice
22                .unwrap()
23                .limit,
24            Some(10_000)
25        );
26        assert!(
27            SelectQuery::new("Order")
28                .limit(10_001)
29                .prepare_for_list()
30                .is_err()
31        );
32        assert!(
33            SelectQuery::new("Order")
34                .limit(10_001)
35                .hard_limit(20_000)
36                .prepare_for_list()
37                .is_ok()
38        );
39    }
40
41    #[test]
42    fn continuous_page_fetch_is_explicit_and_validated() {
43        assert!(SelectQuery::new("Order").continuous_page_fetch.is_none());
44        let query =
45            SelectQuery::new("Order").optimize_for_continuous_page_fetch_with("recent-orders", 30);
46        let options = query.continuous_page_fetch.unwrap();
47        assert_eq!(options.namespace, "recent-orders");
48        assert_eq!(options.ttl_seconds, 30);
49    }
50
51    #[test]
52    #[should_panic(expected = "continuous page namespace must not be empty")]
53    fn continuous_page_fetch_rejects_empty_namespace() {
54        let _ = SelectQuery::new("Order").optimize_for_continuous_page_fetch_with(" ", 30);
55    }
56}
57
58#[derive(Debug, Clone, PartialEq)]
59pub struct NamedExpr {
60    pub alias: String,
61    pub expr: Expr,
62}
63
64impl NamedExpr {
65    pub fn new(alias: impl Into<String>, expr: Expr) -> Self {
66        Self {
67            alias: alias.into(),
68            expr,
69        }
70    }
71}
72
73#[derive(Debug, Clone, PartialEq)]
74pub struct OrderBy {
75    pub field: String,
76    pub expr: Option<Expr>,
77    pub direction: SortDirection,
78}
79
80impl OrderBy {
81    pub fn new(field: impl Into<String>, direction: SortDirection) -> Self {
82        Self {
83            field: field.into(),
84            expr: None,
85            direction,
86        }
87    }
88
89    pub fn expr(expr: Expr, direction: SortDirection) -> Self {
90        Self {
91            field: String::new(),
92            expr: Some(expr),
93            direction,
94        }
95    }
96
97    pub fn asc(field: impl Into<String>) -> Self {
98        Self::new(field, SortDirection::Asc)
99    }
100
101    pub fn desc(field: impl Into<String>) -> Self {
102        Self::new(field, SortDirection::Desc)
103    }
104
105    pub fn asc_expr(expr: Expr) -> Self {
106        Self::expr(expr, SortDirection::Asc)
107    }
108
109    pub fn desc_expr(expr: Expr) -> Self {
110        Self::expr(expr, SortDirection::Desc)
111    }
112
113    pub fn asc_gbk(field: impl Into<String>) -> Self {
114        Self::asc_expr(Expr::gbk(Expr::column(field)))
115    }
116
117    pub fn desc_gbk(field: impl Into<String>) -> Self {
118        Self::desc_expr(Expr::gbk(Expr::column(field)))
119    }
120}
121
122#[derive(Debug, Clone, Copy, PartialEq, Eq)]
123pub enum AggregateFunction {
124    Count,
125    Sum,
126    Avg,
127    Min,
128    Max,
129    Stddev,
130    StddevPop,
131    VarSamp,
132    VarPop,
133    BitAnd,
134    BitOr,
135    BitXor,
136}
137
138#[derive(Debug, Clone, PartialEq, Eq)]
139pub struct Aggregate {
140    pub function: AggregateFunction,
141    pub field: String,
142    pub alias: String,
143}
144
145impl Aggregate {
146    pub fn new(
147        function: AggregateFunction,
148        field: impl Into<String>,
149        alias: impl Into<String>,
150    ) -> Self {
151        Self {
152            function,
153            field: field.into(),
154            alias: alias.into(),
155        }
156    }
157
158    pub fn count(alias: impl Into<String>) -> Self {
159        Self::new(AggregateFunction::Count, "*", alias)
160    }
161
162    pub fn count_field(field: impl Into<String>, alias: impl Into<String>) -> Self {
163        Self::new(AggregateFunction::Count, field, alias)
164    }
165
166    pub fn sum(field: impl Into<String>, alias: impl Into<String>) -> Self {
167        Self::new(AggregateFunction::Sum, field, alias)
168    }
169
170    pub fn avg(field: impl Into<String>, alias: impl Into<String>) -> Self {
171        Self::new(AggregateFunction::Avg, field, alias)
172    }
173
174    pub fn min(field: impl Into<String>, alias: impl Into<String>) -> Self {
175        Self::new(AggregateFunction::Min, field, alias)
176    }
177
178    pub fn max(field: impl Into<String>, alias: impl Into<String>) -> Self {
179        Self::new(AggregateFunction::Max, field, alias)
180    }
181
182    pub fn stddev(field: impl Into<String>, alias: impl Into<String>) -> Self {
183        Self::new(AggregateFunction::Stddev, field, alias)
184    }
185
186    pub fn stddev_pop(field: impl Into<String>, alias: impl Into<String>) -> Self {
187        Self::new(AggregateFunction::StddevPop, field, alias)
188    }
189
190    pub fn var_samp(field: impl Into<String>, alias: impl Into<String>) -> Self {
191        Self::new(AggregateFunction::VarSamp, field, alias)
192    }
193
194    pub fn var_pop(field: impl Into<String>, alias: impl Into<String>) -> Self {
195        Self::new(AggregateFunction::VarPop, field, alias)
196    }
197
198    pub fn bit_and(field: impl Into<String>, alias: impl Into<String>) -> Self {
199        Self::new(AggregateFunction::BitAnd, field, alias)
200    }
201
202    pub fn bit_or(field: impl Into<String>, alias: impl Into<String>) -> Self {
203        Self::new(AggregateFunction::BitOr, field, alias)
204    }
205
206    pub fn bit_xor(field: impl Into<String>, alias: impl Into<String>) -> Self {
207        Self::new(AggregateFunction::BitXor, field, alias)
208    }
209}
210
211#[derive(Debug, Clone, Copy, PartialEq, Eq)]
212pub struct Slice {
213    pub limit: Option<u64>,
214    pub offset: u64,
215}
216
217#[derive(Debug, Clone, PartialEq)]
218pub struct RelationLoad {
219    pub name: String,
220    pub query: Option<Box<SelectQuery>>,
221}
222
223impl RelationLoad {
224    pub fn new(name: impl Into<String>) -> Self {
225        Self {
226            name: name.into(),
227            query: None,
228        }
229    }
230
231    pub fn with_query(name: impl Into<String>, query: SelectQuery) -> Self {
232        Self {
233            name: name.into(),
234            query: Some(Box::new(query)),
235        }
236    }
237}
238
239#[derive(Debug, Clone, PartialEq)]
240pub struct RelationAggregate {
241    pub relation_name: String,
242    pub alias: String,
243    pub query: SelectQuery,
244    pub single_result: bool,
245}
246
247impl RelationAggregate {
248    pub fn new(
249        relation_name: impl Into<String>,
250        alias: impl Into<String>,
251        query: SelectQuery,
252        single_result: bool,
253    ) -> Self {
254        Self {
255            relation_name: relation_name.into(),
256            alias: alias.into(),
257            query,
258            single_result,
259        }
260    }
261}
262
263#[derive(Debug, Clone, PartialEq, Eq)]
264pub struct RawSqlProjection {
265    pub property_name: String,
266    pub raw_sql_segment: String,
267}
268
269impl RawSqlProjection {
270    pub fn new(property_name: impl Into<String>, raw_sql_segment: impl Into<String>) -> Self {
271        Self {
272            property_name: property_name.into(),
273            raw_sql_segment: raw_sql_segment.into(),
274        }
275    }
276}
277
278#[derive(Debug, Clone, PartialEq)]
279pub struct ObjectGroupBy {
280    pub property_name: String,
281    pub storage_field: String,
282    pub query: SelectQuery,
283}
284
285impl ObjectGroupBy {
286    pub fn new(
287        property_name: impl Into<String>,
288        storage_field: impl Into<String>,
289        query: SelectQuery,
290    ) -> Self {
291        Self {
292            property_name: property_name.into(),
293            storage_field: storage_field.into(),
294            query,
295        }
296    }
297}
298
299#[derive(Debug, Clone, Copy, PartialEq, Eq)]
300pub struct AggregationCacheOptions {
301    pub enabled: bool,
302    pub cache_expired_millis: u64,
303    pub propagate: bool,
304    pub propagate_cache_expired_millis: u64,
305}
306
307impl AggregationCacheOptions {
308    pub fn enabled(cache_expired_millis: u64) -> Self {
309        Self {
310            enabled: true,
311            cache_expired_millis,
312            propagate: false,
313            propagate_cache_expired_millis: 0,
314        }
315    }
316
317    pub fn propagate(mut self, cache_expired_millis: u64) -> Self {
318        self.propagate = true;
319        self.propagate_cache_expired_millis = cache_expired_millis;
320        self
321    }
322}
323
324#[derive(Debug, Clone, PartialEq)]
325pub struct StreamConfig {
326    pub chunk_size: usize,
327}
328
329#[derive(Debug, Clone, PartialEq, Eq)]
330pub struct ContinuousPageFetchOptions {
331    pub namespace: String,
332    pub ttl_seconds: u64,
333}
334
335impl ContinuousPageFetchOptions {
336    pub const DEFAULT_TTL_SECONDS: u64 = 600;
337
338    pub fn new(namespace: impl Into<String>, ttl_seconds: u64) -> Self {
339        let namespace = namespace.into();
340        assert!(
341            !namespace.trim().is_empty(),
342            "continuous page namespace must not be empty"
343        );
344        assert!(
345            ttl_seconds > 0,
346            "continuous page ttl_seconds must be positive"
347        );
348        Self {
349            namespace,
350            ttl_seconds,
351        }
352    }
353}
354
355impl Default for StreamConfig {
356    fn default() -> Self {
357        Self { chunk_size: 1000 }
358    }
359}
360
361#[derive(Debug, Clone, PartialEq)]
362pub struct SelectQuery {
363    /// Safety ceiling for a fully materialized outer query.
364    pub hard_limit: u64,
365    pub entity: String,
366    pub projection: Vec<String>,
367    pub expr_projection: Vec<NamedExpr>,
368    pub search_with_text: Option<String>,
369    pub filter: Option<Expr>,
370    pub having: Option<Expr>,
371    pub order_by: Vec<OrderBy>,
372    pub slice: Option<Slice>,
373    /// Apply `slice` independently inside each value of this property.
374    pub partition_by: Option<String>,
375    pub aggregates: Vec<Aggregate>,
376    pub group_by: Vec<String>,
377    pub relations: Vec<RelationLoad>,
378    pub aggregation_cache: Option<AggregationCacheOptions>,
379    pub comment: Option<String>,
380    pub trace_chain: Vec<crate::TraceNode>,
381    pub raw_sql: Option<String>,
382    pub raw_sql_search_criteria: Vec<String>,
383    pub dynamic_properties: Vec<RawSqlProjection>,
384    pub raw_projections: Vec<RawSqlProjection>,
385    pub object_group_bys: Vec<ObjectGroupBy>,
386    pub child_enhancements: Vec<SelectQuery>,
387    pub stream_config: Option<StreamConfig>,
388    /// Explicit, process-local hint for transparent seek pagination of outer list queries.
389    pub continuous_page_fetch: Option<ContinuousPageFetchOptions>,
390}
391
392impl SelectQuery {
393    pub fn new(entity: impl Into<String>) -> Self {
394        Self {
395            hard_limit: 10_000,
396            entity: entity.into(),
397            projection: Vec::new(),
398            expr_projection: Vec::new(),
399            search_with_text: None,
400            filter: None,
401            having: None,
402            order_by: Vec::new(),
403            slice: None,
404            partition_by: None,
405            aggregates: Vec::new(),
406            group_by: Vec::new(),
407            relations: Vec::new(),
408            aggregation_cache: None,
409            comment: None,
410            trace_chain: Vec::new(),
411            raw_sql: None,
412            raw_sql_search_criteria: Vec::new(),
413            dynamic_properties: Vec::new(),
414            raw_projections: Vec::new(),
415            object_group_bys: Vec::new(),
416            child_enhancements: Vec::new(),
417            stream_config: None,
418            continuous_page_fetch: None,
419        }
420    }
421
422    pub fn project(mut self, field: impl Into<String>) -> Self {
423        self.projection.push(field.into());
424        self
425    }
426
427    pub fn projects(mut self, fields: impl IntoIterator<Item = impl Into<String>>) -> Self {
428        self.projection.extend(fields.into_iter().map(Into::into));
429        self
430    }
431
432    pub fn project_expr(mut self, alias: impl Into<String>, expr: Expr) -> Self {
433        self.expr_projection.push(NamedExpr::new(alias, expr));
434        self
435    }
436
437    pub fn project_raw(
438        mut self,
439        alias: impl Into<String>,
440        raw_sql_segment: impl Into<String>,
441    ) -> Self {
442        self.raw_projections
443            .push(RawSqlProjection::new(alias, raw_sql_segment));
444        self
445    }
446
447    pub fn dynamic_property_raw(
448        mut self,
449        alias: impl Into<String>,
450        raw_sql_segment: impl Into<String>,
451    ) -> Self {
452        self.dynamic_properties
453            .push(RawSqlProjection::new(alias, raw_sql_segment));
454        self
455    }
456
457    pub fn search_with_text(mut self, text: impl Into<String>) -> Self {
458        self.search_with_text = Some(text.into());
459        self
460    }
461
462    pub fn filter(mut self, filter: Expr) -> Self {
463        self.filter = Some(filter);
464        self
465    }
466
467    pub fn and_filter(mut self, filter: Expr) -> Self {
468        self.filter = Some(match self.filter.take() {
469            Some(existing) => existing.and_expr(filter),
470            None => filter,
471        });
472        self
473    }
474
475    pub fn or_filter(mut self, filter: Expr) -> Self {
476        self.filter = Some(match self.filter.take() {
477            Some(existing) => existing.or_expr(filter),
478            None => filter,
479        });
480        self
481    }
482
483    pub fn having(mut self, having: Expr) -> Self {
484        self.having = Some(having);
485        self
486    }
487
488    pub fn and_having(mut self, having: Expr) -> Self {
489        self.having = Some(match self.having.take() {
490            Some(existing) => existing.and_expr(having),
491            None => having,
492        });
493        self
494    }
495
496    pub fn or_having(mut self, having: Expr) -> Self {
497        self.having = Some(match self.having.take() {
498            Some(existing) => existing.or_expr(having),
499            None => having,
500        });
501        self
502    }
503
504    pub fn order_by(mut self, order: OrderBy) -> Self {
505        self.order_by.push(order);
506        self
507    }
508
509    pub fn order_asc(self, field: impl Into<String>) -> Self {
510        self.order_by(OrderBy::asc(field))
511    }
512
513    pub fn order_desc(self, field: impl Into<String>) -> Self {
514        self.order_by(OrderBy::desc(field))
515    }
516
517    pub fn order_expr_asc(self, expr: Expr) -> Self {
518        self.order_by(OrderBy::asc_expr(expr))
519    }
520
521    pub fn order_expr_desc(self, expr: Expr) -> Self {
522        self.order_by(OrderBy::desc_expr(expr))
523    }
524
525    pub fn order_gbk_asc(self, field: impl Into<String>) -> Self {
526        self.order_by(OrderBy::asc_gbk(field))
527    }
528
529    pub fn order_gbk_desc(self, field: impl Into<String>) -> Self {
530        self.order_by(OrderBy::desc_gbk(field))
531    }
532
533    pub fn group_by(mut self, field: impl Into<String>) -> Self {
534        self.group_by.push(field.into());
535        self
536    }
537
538    pub fn aggregate(mut self, aggregate: Aggregate) -> Self {
539        self.aggregates.push(aggregate);
540        self
541    }
542
543    pub fn count(self, alias: impl Into<String>) -> Self {
544        self.aggregate(Aggregate::count(alias))
545    }
546
547    pub fn count_field(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
548        self.aggregate(Aggregate::count_field(field, alias))
549    }
550
551    pub fn sum(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
552        self.aggregate(Aggregate::sum(field, alias))
553    }
554
555    pub fn avg(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
556        self.aggregate(Aggregate::avg(field, alias))
557    }
558
559    pub fn min(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
560        self.aggregate(Aggregate::min(field, alias))
561    }
562
563    pub fn max(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
564        self.aggregate(Aggregate::max(field, alias))
565    }
566
567    pub fn stddev(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
568        self.aggregate(Aggregate::stddev(field, alias))
569    }
570
571    pub fn stddev_pop(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
572        self.aggregate(Aggregate::stddev_pop(field, alias))
573    }
574
575    pub fn var_samp(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
576        self.aggregate(Aggregate::var_samp(field, alias))
577    }
578
579    pub fn var_pop(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
580        self.aggregate(Aggregate::var_pop(field, alias))
581    }
582
583    pub fn bit_and(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
584        self.aggregate(Aggregate::bit_and(field, alias))
585    }
586
587    pub fn bit_or(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
588        self.aggregate(Aggregate::bit_or(field, alias))
589    }
590
591    pub fn bit_xor(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
592        self.aggregate(Aggregate::bit_xor(field, alias))
593    }
594
595    pub fn enable_aggregation_cache(self) -> Self {
596        self.enable_aggregation_cache_for(0)
597    }
598
599    pub fn enable_aggregation_cache_for(mut self, cache_expired_millis: u64) -> Self {
600        self.aggregation_cache = Some(AggregationCacheOptions::enabled(cache_expired_millis));
601        self
602    }
603
604    pub fn propagate_aggregation_cache(mut self, cache_expired_millis: u64) -> Self {
605        self.aggregation_cache = Some(
606            self.aggregation_cache
607                .unwrap_or_else(|| AggregationCacheOptions::enabled(0))
608                .propagate(cache_expired_millis),
609        );
610        self
611    }
612
613    pub fn comment(mut self, comment: impl Into<String>) -> Self {
614        let comment_str = comment.into();
615        self.comment = Some(comment_str.clone());
616        self.trace_chain.push(crate::TraceNode {
617            entity_type: self.entity.clone(),
618            entity_id: None,
619            comment: comment_str,
620        });
621        self
622    }
623
624    pub fn raw_sql(mut self, raw_sql: impl Into<String>) -> Self {
625        self.raw_sql = Some(raw_sql.into());
626        self
627    }
628
629    pub fn raw_sql_search_criteria(mut self, raw_sql: impl Into<String>) -> Self {
630        self.raw_sql_search_criteria.push(raw_sql.into());
631        self
632    }
633
634    pub fn object_group_by(
635        mut self,
636        property_name: impl Into<String>,
637        storage_field: impl Into<String>,
638        query: SelectQuery,
639    ) -> Self {
640        self.object_group_bys
641            .push(ObjectGroupBy::new(property_name, storage_field, query));
642        self
643    }
644
645    pub fn child_enhancement(mut self, query: SelectQuery) -> Self {
646        self.child_enhancements.push(query);
647        self
648    }
649
650    pub fn relation(mut self, name: impl Into<String>) -> Self {
651        self.relations.push(RelationLoad::new(name));
652        self
653    }
654
655    pub fn relation_query(mut self, name: impl Into<String>, query: SelectQuery) -> Self {
656        self.relations.push(RelationLoad::with_query(name, query));
657        self
658    }
659
660    pub fn limit(mut self, limit: u64) -> Self {
661        let slice = self.slice.get_or_insert(Slice {
662            limit: None,
663            offset: 0,
664        });
665        slice.limit = Some(limit);
666        self
667    }
668
669    /// Override the outer materialized-list ceiling. Most callers should keep 10,000.
670    pub fn hard_limit(mut self, hard_limit: u64) -> Self {
671        assert!(hard_limit > 0, "hard_limit must be positive");
672        self.hard_limit = hard_limit;
673        self
674    }
675
676    /// Apply and validate list-materialization limits. This is intentionally not
677    /// used by streaming execution.
678    pub fn prepare_for_list(mut self) -> Result<Self, String> {
679        self.apply_list_limit(self.hard_limit, true)?;
680        Ok(self)
681    }
682
683    fn apply_list_limit(&mut self, ceiling: u64, outer: bool) -> Result<(), String> {
684        let slice = self.slice.get_or_insert(Slice {
685            limit: None,
686            offset: 0,
687        });
688        match slice.limit {
689            Some(limit) if limit > ceiling => {
690                return Err(format!(
691                    "QUERY_HARD_LIMIT_EXCEEDED: requested limit {limit} exceeds hard limit {ceiling}"
692                ));
693            }
694            None => slice.limit = Some(ceiling),
695            _ => {}
696        }
697        for relation in &mut self.relations {
698            if let Some(query) = relation.query.as_mut() {
699                query.apply_list_limit(10_000, false)?;
700            }
701        }
702        for query in &mut self.child_enhancements {
703            query.apply_list_limit(10_000, false)?;
704        }
705        let _ = outer;
706        Ok(())
707    }
708
709    pub fn offset(mut self, offset: u64) -> Self {
710        let slice = self.slice.get_or_insert(Slice {
711            limit: None,
712            offset: 0,
713        });
714        slice.offset = offset;
715        self
716    }
717
718    pub fn page(self, offset: u64, limit: u64) -> Self {
719        self.offset(offset).limit(limit)
720    }
721
722    pub fn optimize_for_continuous_page_fetch(mut self) -> Self {
723        self.continuous_page_fetch = Some(ContinuousPageFetchOptions::new(
724            "default",
725            ContinuousPageFetchOptions::DEFAULT_TTL_SECONDS,
726        ));
727        self
728    }
729
730    pub fn optimize_for_continuous_page_fetch_with(
731        mut self,
732        namespace: impl Into<String>,
733        ttl_seconds: u64,
734    ) -> Self {
735        self.continuous_page_fetch = Some(ContinuousPageFetchOptions::new(namespace, ttl_seconds));
736        self
737    }
738
739    /// Scope pagination to each distinct value of `field`.
740    ///
741    /// Relation loading sets this automatically. Most application queries
742    /// should use a generated relation selector instead of calling this
743    /// method directly.
744    pub fn partition_by(mut self, field: impl Into<String>) -> Self {
745        self.partition_by = Some(field.into());
746        self
747    }
748
749    /// Enable streaming mode with the given chunk size.
750    /// When streaming, rows are fetched and enhanced in batches rather than all at once.
751    pub fn stream(mut self, chunk_size: usize) -> Self {
752        self.stream_config = Some(StreamConfig { chunk_size });
753        self
754    }
755
756    /// Enable streaming mode with default chunk size (1000).
757    pub fn stream_default(mut self) -> Self {
758        self.stream_config = Some(StreamConfig::default());
759        self
760    }
761}
762
763pub type Record = BTreeMap<String, Value>;
764
765/// Internal projection used to implement per-parent pagination for relation
766/// loads. Runtime relation attachment removes it before exposing child rows.
767pub const PARTITION_RANK_PROPERTY: &str = "__teaql_partition_rank";
768
769pub fn record_to_json_value(record: &Record) -> serde_json::Value {
770    serde_json::Value::Object(
771        record
772            .iter()
773            .map(|(key, value)| (key.clone(), value.to_json_value()))
774            .collect(),
775    )
776}
teaql_core/query.rs

teaql_core/
query.rs
