Skip to main content

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