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 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 pub fn partition_by(mut self, field: impl Into<String>) -> Self {
609 self.partition_by = Some(field.into());
610 self
611 }
612
613 pub fn stream(mut self, chunk_size: usize) -> Self {
616 self.stream_config = Some(StreamConfig { chunk_size });
617 self
618 }
619
620 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
629pub 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}