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

teaql_core/
query.rs
