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    pub aggregates: Vec<Aggregate>,
299    pub group_by: Vec<String>,
300    pub relations: Vec<RelationLoad>,
301    pub aggregation_cache: Option<AggregationCacheOptions>,
302    pub comment: Option<String>,
303    pub trace_chain: Vec<crate::TraceNode>,
304    pub raw_sql: Option<String>,
305    pub raw_sql_search_criteria: Vec<String>,
306    pub dynamic_properties: Vec<RawSqlProjection>,
307    pub raw_projections: Vec<RawSqlProjection>,
308    pub object_group_bys: Vec<ObjectGroupBy>,
309    pub child_enhancements: Vec<SelectQuery>,
310    pub stream_config: Option<StreamConfig>,
311}
312
313impl SelectQuery {
314    pub fn new(entity: impl Into<String>) -> Self {
315        Self {
316            entity: entity.into(),
317            projection: Vec::new(),
318            expr_projection: Vec::new(),
319            search_with_text: None,
320            filter: None,
321            having: None,
322            order_by: Vec::new(),
323            slice: None,
324            aggregates: Vec::new(),
325            group_by: Vec::new(),
326            relations: Vec::new(),
327            aggregation_cache: None,
328            comment: None,
329            trace_chain: Vec::new(),
330            raw_sql: None,
331            raw_sql_search_criteria: Vec::new(),
332            dynamic_properties: Vec::new(),
333            raw_projections: Vec::new(),
334            object_group_bys: Vec::new(),
335            child_enhancements: Vec::new(),
336            stream_config: None,
337        }
338    }
339
340    pub fn project(mut self, field: impl Into<String>) -> Self {
341        self.projection.push(field.into());
342        self
343    }
344
345    pub fn projects(mut self, fields: impl IntoIterator<Item = impl Into<String>>) -> Self {
346        self.projection.extend(fields.into_iter().map(Into::into));
347        self
348    }
349
350    pub fn project_expr(mut self, alias: impl Into<String>, expr: Expr) -> Self {
351        self.expr_projection.push(NamedExpr::new(alias, expr));
352        self
353    }
354
355    pub fn project_raw(
356        mut self,
357        alias: impl Into<String>,
358        raw_sql_segment: impl Into<String>,
359    ) -> Self {
360        self.raw_projections
361            .push(RawSqlProjection::new(alias, raw_sql_segment));
362        self
363    }
364
365    pub fn dynamic_property_raw(
366        mut self,
367        alias: impl Into<String>,
368        raw_sql_segment: impl Into<String>,
369    ) -> Self {
370        self.dynamic_properties
371            .push(RawSqlProjection::new(alias, raw_sql_segment));
372        self
373    }
374
375    pub fn search_with_text(mut self, text: impl Into<String>) -> Self {
376        self.search_with_text = Some(text.into());
377        self
378    }
379
380    pub fn filter(mut self, filter: Expr) -> Self {
381        self.filter = Some(filter);
382        self
383    }
384
385    pub fn and_filter(mut self, filter: Expr) -> Self {
386        self.filter = Some(match self.filter.take() {
387            Some(existing) => existing.and_expr(filter),
388            None => filter,
389        });
390        self
391    }
392
393    pub fn or_filter(mut self, filter: Expr) -> Self {
394        self.filter = Some(match self.filter.take() {
395            Some(existing) => existing.or_expr(filter),
396            None => filter,
397        });
398        self
399    }
400
401    pub fn having(mut self, having: Expr) -> Self {
402        self.having = Some(having);
403        self
404    }
405
406    pub fn and_having(mut self, having: Expr) -> Self {
407        self.having = Some(match self.having.take() {
408            Some(existing) => existing.and_expr(having),
409            None => having,
410        });
411        self
412    }
413
414    pub fn or_having(mut self, having: Expr) -> Self {
415        self.having = Some(match self.having.take() {
416            Some(existing) => existing.or_expr(having),
417            None => having,
418        });
419        self
420    }
421
422    pub fn order_by(mut self, order: OrderBy) -> Self {
423        self.order_by.push(order);
424        self
425    }
426
427    pub fn order_asc(self, field: impl Into<String>) -> Self {
428        self.order_by(OrderBy::asc(field))
429    }
430
431    pub fn order_desc(self, field: impl Into<String>) -> Self {
432        self.order_by(OrderBy::desc(field))
433    }
434
435    pub fn order_expr_asc(self, expr: Expr) -> Self {
436        self.order_by(OrderBy::asc_expr(expr))
437    }
438
439    pub fn order_expr_desc(self, expr: Expr) -> Self {
440        self.order_by(OrderBy::desc_expr(expr))
441    }
442
443    pub fn order_gbk_asc(self, field: impl Into<String>) -> Self {
444        self.order_by(OrderBy::asc_gbk(field))
445    }
446
447    pub fn order_gbk_desc(self, field: impl Into<String>) -> Self {
448        self.order_by(OrderBy::desc_gbk(field))
449    }
450
451    pub fn group_by(mut self, field: impl Into<String>) -> Self {
452        self.group_by.push(field.into());
453        self
454    }
455
456    pub fn aggregate(mut self, aggregate: Aggregate) -> Self {
457        self.aggregates.push(aggregate);
458        self
459    }
460
461    pub fn count(self, alias: impl Into<String>) -> Self {
462        self.aggregate(Aggregate::count(alias))
463    }
464
465    pub fn count_field(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
466        self.aggregate(Aggregate::count_field(field, alias))
467    }
468
469    pub fn sum(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
470        self.aggregate(Aggregate::sum(field, alias))
471    }
472
473    pub fn avg(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
474        self.aggregate(Aggregate::avg(field, alias))
475    }
476
477    pub fn min(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
478        self.aggregate(Aggregate::min(field, alias))
479    }
480
481    pub fn max(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
482        self.aggregate(Aggregate::max(field, alias))
483    }
484
485    pub fn stddev(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
486        self.aggregate(Aggregate::stddev(field, alias))
487    }
488
489    pub fn stddev_pop(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
490        self.aggregate(Aggregate::stddev_pop(field, alias))
491    }
492
493    pub fn var_samp(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
494        self.aggregate(Aggregate::var_samp(field, alias))
495    }
496
497    pub fn var_pop(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
498        self.aggregate(Aggregate::var_pop(field, alias))
499    }
500
501    pub fn bit_and(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
502        self.aggregate(Aggregate::bit_and(field, alias))
503    }
504
505    pub fn bit_or(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
506        self.aggregate(Aggregate::bit_or(field, alias))
507    }
508
509    pub fn bit_xor(self, field: impl Into<String>, alias: impl Into<String>) -> Self {
510        self.aggregate(Aggregate::bit_xor(field, alias))
511    }
512
513    pub fn enable_aggregation_cache(self) -> Self {
514        self.enable_aggregation_cache_for(0)
515    }
516
517    pub fn enable_aggregation_cache_for(mut self, cache_expired_millis: u64) -> Self {
518        self.aggregation_cache = Some(AggregationCacheOptions::enabled(cache_expired_millis));
519        self
520    }
521
522    pub fn propagate_aggregation_cache(mut self, cache_expired_millis: u64) -> Self {
523        self.aggregation_cache = Some(
524            self.aggregation_cache
525                .unwrap_or_else(|| AggregationCacheOptions::enabled(0))
526                .propagate(cache_expired_millis),
527        );
528        self
529    }
530
531    pub fn comment(mut self, comment: impl Into<String>) -> Self {
532        let comment_str = comment.into();
533        self.comment = Some(comment_str.clone());
534        self.trace_chain.push(crate::TraceNode {
535            entity_type: self.entity.clone(),
536            entity_id: None,
537            comment: comment_str,
538        });
539        self
540    }
541
542    pub fn raw_sql(mut self, raw_sql: impl Into<String>) -> Self {
543        self.raw_sql = Some(raw_sql.into());
544        self
545    }
546
547    pub fn raw_sql_search_criteria(mut self, raw_sql: impl Into<String>) -> Self {
548        self.raw_sql_search_criteria.push(raw_sql.into());
549        self
550    }
551
552    pub fn object_group_by(
553        mut self,
554        property_name: impl Into<String>,
555        storage_field: impl Into<String>,
556        query: SelectQuery,
557    ) -> Self {
558        self.object_group_bys
559            .push(ObjectGroupBy::new(property_name, storage_field, query));
560        self
561    }
562
563    pub fn child_enhancement(mut self, query: SelectQuery) -> Self {
564        self.child_enhancements.push(query);
565        self
566    }
567
568    pub fn relation(mut self, name: impl Into<String>) -> Self {
569        self.relations.push(RelationLoad::new(name));
570        self
571    }
572
573    pub fn relation_query(mut self, name: impl Into<String>, query: SelectQuery) -> Self {
574        self.relations.push(RelationLoad::with_query(name, query));
575        self
576    }
577
578    pub fn limit(mut self, limit: u64) -> Self {
579        let slice = self.slice.get_or_insert(Slice {
580            limit: None,
581            offset: 0,
582        });
583        slice.limit = Some(limit);
584        self
585    }
586
587    pub fn offset(mut self, offset: u64) -> Self {
588        let slice = self.slice.get_or_insert(Slice {
589            limit: None,
590            offset: 0,
591        });
592        slice.offset = offset;
593        self
594    }
595
596    pub fn page(self, offset: u64, limit: u64) -> Self {
597        self.offset(offset).limit(limit)
598    }
599
600    /// Enable streaming mode with the given chunk size.
601    /// When streaming, rows are fetched and enhanced in batches rather than all at once.
602    pub fn stream(mut self, chunk_size: usize) -> Self {
603        self.stream_config = Some(StreamConfig { chunk_size });
604        self
605    }
606
607    /// Enable streaming mode with default chunk size (1000).
608    pub fn stream_default(mut self) -> Self {
609        self.stream_config = Some(StreamConfig::default());
610        self
611    }
612}
613
614pub type Record = BTreeMap<String, Value>;
615
616pub fn record_to_json_value(record: &Record) -> serde_json::Value {
617    serde_json::Value::Object(
618        record
619            .iter()
620            .map(|(key, value)| (key.clone(), value.to_json_value()))
621            .collect(),
622    )
623}