datafusion_sql/unparser/
ast.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use core::fmt;
19use std::ops::ControlFlow;
20
21use sqlparser::ast::helpers::attached_token::AttachedToken;
22use sqlparser::ast::{
23    self, visit_expressions_mut, LimitClause, OrderByKind, SelectFlavor,
24};
25
26#[derive(Clone)]
27pub struct QueryBuilder {
28    with: Option<ast::With>,
29    body: Option<Box<ast::SetExpr>>,
30    order_by_kind: Option<OrderByKind>,
31    limit: Option<ast::Expr>,
32    limit_by: Vec<ast::Expr>,
33    offset: Option<ast::Offset>,
34    fetch: Option<ast::Fetch>,
35    locks: Vec<ast::LockClause>,
36    for_clause: Option<ast::ForClause>,
37    // If true, we need to unparse LogicalPlan::Union as a SQL `UNION` rather than a `UNION ALL`.
38    distinct_union: bool,
39}
40
41#[allow(dead_code)]
42impl QueryBuilder {
43    pub fn with(&mut self, value: Option<ast::With>) -> &mut Self {
44        self.with = value;
45        self
46    }
47    pub fn body(&mut self, value: Box<ast::SetExpr>) -> &mut Self {
48        self.body = Some(value);
49        self
50    }
51    pub fn take_body(&mut self) -> Option<Box<ast::SetExpr>> {
52        self.body.take()
53    }
54    pub fn order_by(&mut self, value: OrderByKind) -> &mut Self {
55        self.order_by_kind = Some(value);
56        self
57    }
58    pub fn limit(&mut self, value: Option<ast::Expr>) -> &mut Self {
59        self.limit = value;
60        self
61    }
62    pub fn limit_by(&mut self, value: Vec<ast::Expr>) -> &mut Self {
63        self.limit_by = value;
64        self
65    }
66    pub fn offset(&mut self, value: Option<ast::Offset>) -> &mut Self {
67        self.offset = value;
68        self
69    }
70    pub fn fetch(&mut self, value: Option<ast::Fetch>) -> &mut Self {
71        self.fetch = value;
72        self
73    }
74    pub fn locks(&mut self, value: Vec<ast::LockClause>) -> &mut Self {
75        self.locks = value;
76        self
77    }
78    pub fn for_clause(&mut self, value: Option<ast::ForClause>) -> &mut Self {
79        self.for_clause = value;
80        self
81    }
82    pub fn distinct_union(&mut self) -> &mut Self {
83        self.distinct_union = true;
84        self
85    }
86    pub fn is_distinct_union(&self) -> bool {
87        self.distinct_union
88    }
89    pub fn build(&self) -> Result<ast::Query, BuilderError> {
90        let order_by = self
91            .order_by_kind
92            .as_ref()
93            .map(|order_by_kind| ast::OrderBy {
94                kind: order_by_kind.clone(),
95                interpolate: None,
96            });
97
98        Ok(ast::Query {
99            with: self.with.clone(),
100            body: match self.body {
101                Some(ref value) => value.clone(),
102                None => return Err(Into::into(UninitializedFieldError::from("body"))),
103            },
104            order_by,
105            limit_clause: Some(LimitClause::LimitOffset {
106                limit: self.limit.clone(),
107                offset: self.offset.clone(),
108                limit_by: self.limit_by.clone(),
109            }),
110            fetch: self.fetch.clone(),
111            locks: self.locks.clone(),
112            for_clause: self.for_clause.clone(),
113            settings: None,
114            format_clause: None,
115            pipe_operators: vec![],
116        })
117    }
118    fn create_empty() -> Self {
119        Self {
120            with: Default::default(),
121            body: Default::default(),
122            order_by_kind: Default::default(),
123            limit: Default::default(),
124            limit_by: Default::default(),
125            offset: Default::default(),
126            fetch: Default::default(),
127            locks: Default::default(),
128            for_clause: Default::default(),
129            distinct_union: false,
130        }
131    }
132}
133impl Default for QueryBuilder {
134    fn default() -> Self {
135        Self::create_empty()
136    }
137}
138
139#[derive(Clone)]
140pub struct SelectBuilder {
141    distinct: Option<ast::Distinct>,
142    top: Option<ast::Top>,
143    projection: Vec<ast::SelectItem>,
144    into: Option<ast::SelectInto>,
145    from: Vec<TableWithJoinsBuilder>,
146    lateral_views: Vec<ast::LateralView>,
147    selection: Option<ast::Expr>,
148    group_by: Option<ast::GroupByExpr>,
149    cluster_by: Vec<ast::Expr>,
150    distribute_by: Vec<ast::Expr>,
151    sort_by: Vec<ast::OrderByExpr>,
152    having: Option<ast::Expr>,
153    named_window: Vec<ast::NamedWindowDefinition>,
154    qualify: Option<ast::Expr>,
155    value_table_mode: Option<ast::ValueTableMode>,
156    flavor: Option<SelectFlavor>,
157}
158
159#[allow(dead_code)]
160impl SelectBuilder {
161    pub fn distinct(&mut self, value: Option<ast::Distinct>) -> &mut Self {
162        self.distinct = value;
163        self
164    }
165    pub fn top(&mut self, value: Option<ast::Top>) -> &mut Self {
166        self.top = value;
167        self
168    }
169    pub fn projection(&mut self, value: Vec<ast::SelectItem>) -> &mut Self {
170        self.projection = value;
171        self
172    }
173    pub fn pop_projections(&mut self) -> Vec<ast::SelectItem> {
174        let ret = self.projection.clone();
175        self.projection.clear();
176        ret
177    }
178    pub fn already_projected(&self) -> bool {
179        !self.projection.is_empty()
180    }
181    pub fn into(&mut self, value: Option<ast::SelectInto>) -> &mut Self {
182        self.into = value;
183        self
184    }
185    pub fn from(&mut self, value: Vec<TableWithJoinsBuilder>) -> &mut Self {
186        self.from = value;
187        self
188    }
189    pub fn push_from(&mut self, value: TableWithJoinsBuilder) -> &mut Self {
190        self.from.push(value);
191        self
192    }
193    pub fn pop_from(&mut self) -> Option<TableWithJoinsBuilder> {
194        self.from.pop()
195    }
196    pub fn lateral_views(&mut self, value: Vec<ast::LateralView>) -> &mut Self {
197        self.lateral_views = value;
198        self
199    }
200
201    /// Replaces the selection with a new value.
202    ///
203    /// This function is used to replace a specific expression within the selection.
204    /// Unlike the `selection` method which combines existing and new selections with AND,
205    /// this method searches for and replaces occurrences of a specific expression.
206    ///
207    /// This method is primarily used to modify LEFT MARK JOIN expressions.
208    /// When processing a LEFT MARK JOIN, we need to replace the placeholder expression
209    /// with the actual join condition in the selection clause.
210    ///
211    /// # Arguments
212    ///
213    /// * `existing_expr` - The expression to replace
214    /// * `value` - The new expression to set as the selection
215    pub fn replace_mark(
216        &mut self,
217        existing_expr: &ast::Expr,
218        value: &ast::Expr,
219    ) -> &mut Self {
220        if let Some(selection) = &mut self.selection {
221            let _ = visit_expressions_mut(selection, |expr| {
222                if expr == existing_expr {
223                    *expr = value.clone();
224                }
225                ControlFlow::<()>::Continue(())
226            });
227        }
228        self
229    }
230
231    pub fn selection(&mut self, value: Option<ast::Expr>) -> &mut Self {
232        // With filter pushdown optimization, the LogicalPlan can have filters defined as part of `TableScan` and `Filter` nodes.
233        // To avoid overwriting one of the filters, we combine the existing filter with the additional filter.
234        // Example:                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                       |
235        // |  Projection: customer.c_phone AS cntrycode, customer.c_acctbal                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                      |
236        // |   Filter: CAST(customer.c_acctbal AS Decimal128(38, 6)) > (<subquery>)                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                              |
237        // |     Subquery:
238        // |     ..                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                 |
239        // |     TableScan: customer, full_filters=[customer.c_mktsegment = Utf8("BUILDING")]
240        match (&self.selection, value) {
241            (Some(existing_selection), Some(new_selection)) => {
242                self.selection = Some(ast::Expr::BinaryOp {
243                    left: Box::new(existing_selection.clone()),
244                    op: ast::BinaryOperator::And,
245                    right: Box::new(new_selection),
246                });
247            }
248            (None, Some(new_selection)) => {
249                self.selection = Some(new_selection);
250            }
251            (_, None) => (),
252        }
253
254        self
255    }
256    pub fn group_by(&mut self, value: ast::GroupByExpr) -> &mut Self {
257        self.group_by = Some(value);
258        self
259    }
260    pub fn cluster_by(&mut self, value: Vec<ast::Expr>) -> &mut Self {
261        self.cluster_by = value;
262        self
263    }
264    pub fn distribute_by(&mut self, value: Vec<ast::Expr>) -> &mut Self {
265        self.distribute_by = value;
266        self
267    }
268    pub fn sort_by(&mut self, value: Vec<ast::OrderByExpr>) -> &mut Self {
269        self.sort_by = value;
270        self
271    }
272    pub fn having(&mut self, value: Option<ast::Expr>) -> &mut Self {
273        self.having = value;
274        self
275    }
276    pub fn named_window(&mut self, value: Vec<ast::NamedWindowDefinition>) -> &mut Self {
277        self.named_window = value;
278        self
279    }
280    pub fn qualify(&mut self, value: Option<ast::Expr>) -> &mut Self {
281        self.qualify = value;
282        self
283    }
284    pub fn value_table_mode(&mut self, value: Option<ast::ValueTableMode>) -> &mut Self {
285        self.value_table_mode = value;
286        self
287    }
288    pub fn build(&self) -> Result<ast::Select, BuilderError> {
289        Ok(ast::Select {
290            distinct: self.distinct.clone(),
291            top_before_distinct: false,
292            top: self.top.clone(),
293            projection: self.projection.clone(),
294            into: self.into.clone(),
295            from: self
296                .from
297                .iter()
298                .filter_map(|b| b.build().transpose())
299                .collect::<Result<Vec<_>, BuilderError>>()?,
300            lateral_views: self.lateral_views.clone(),
301            selection: self.selection.clone(),
302            group_by: match self.group_by {
303                Some(ref value) => value.clone(),
304                None => {
305                    return Err(Into::into(UninitializedFieldError::from("group_by")))
306                }
307            },
308            cluster_by: self.cluster_by.clone(),
309            distribute_by: self.distribute_by.clone(),
310            sort_by: self.sort_by.clone(),
311            having: self.having.clone(),
312            named_window: self.named_window.clone(),
313            qualify: self.qualify.clone(),
314            value_table_mode: self.value_table_mode,
315            connect_by: None,
316            window_before_qualify: false,
317            prewhere: None,
318            select_token: AttachedToken::empty(),
319            flavor: match self.flavor {
320                Some(ref value) => value.clone(),
321                None => return Err(Into::into(UninitializedFieldError::from("flavor"))),
322            },
323            exclude: None,
324        })
325    }
326    fn create_empty() -> Self {
327        Self {
328            distinct: Default::default(),
329            top: Default::default(),
330            projection: Default::default(),
331            into: Default::default(),
332            from: Default::default(),
333            lateral_views: Default::default(),
334            selection: Default::default(),
335            group_by: Some(ast::GroupByExpr::Expressions(Vec::new(), Vec::new())),
336            cluster_by: Default::default(),
337            distribute_by: Default::default(),
338            sort_by: Default::default(),
339            having: Default::default(),
340            named_window: Default::default(),
341            qualify: Default::default(),
342            value_table_mode: Default::default(),
343            flavor: Some(SelectFlavor::Standard),
344        }
345    }
346}
347impl Default for SelectBuilder {
348    fn default() -> Self {
349        Self::create_empty()
350    }
351}
352
353#[derive(Clone)]
354pub struct TableWithJoinsBuilder {
355    relation: Option<RelationBuilder>,
356    joins: Vec<ast::Join>,
357}
358
359#[allow(dead_code)]
360impl TableWithJoinsBuilder {
361    pub fn relation(&mut self, value: RelationBuilder) -> &mut Self {
362        self.relation = Some(value);
363        self
364    }
365
366    pub fn joins(&mut self, value: Vec<ast::Join>) -> &mut Self {
367        self.joins = value;
368        self
369    }
370    pub fn push_join(&mut self, value: ast::Join) -> &mut Self {
371        self.joins.push(value);
372        self
373    }
374
375    pub fn build(&self) -> Result<Option<ast::TableWithJoins>, BuilderError> {
376        match self.relation {
377            Some(ref value) => match value.build()? {
378                Some(relation) => Ok(Some(ast::TableWithJoins {
379                    relation,
380                    joins: self.joins.clone(),
381                })),
382                None => Ok(None),
383            },
384            None => Err(Into::into(UninitializedFieldError::from("relation"))),
385        }
386    }
387    fn create_empty() -> Self {
388        Self {
389            relation: Default::default(),
390            joins: Default::default(),
391        }
392    }
393}
394impl Default for TableWithJoinsBuilder {
395    fn default() -> Self {
396        Self::create_empty()
397    }
398}
399
400#[derive(Clone)]
401pub struct RelationBuilder {
402    relation: Option<TableFactorBuilder>,
403}
404
405#[allow(dead_code)]
406#[derive(Clone)]
407#[allow(clippy::large_enum_variant)]
408enum TableFactorBuilder {
409    Table(TableRelationBuilder),
410    Derived(DerivedRelationBuilder),
411    Unnest(UnnestRelationBuilder),
412    Empty,
413}
414
415#[allow(dead_code)]
416impl RelationBuilder {
417    pub fn has_relation(&self) -> bool {
418        self.relation.is_some()
419    }
420    pub fn table(&mut self, value: TableRelationBuilder) -> &mut Self {
421        self.relation = Some(TableFactorBuilder::Table(value));
422        self
423    }
424    pub fn derived(&mut self, value: DerivedRelationBuilder) -> &mut Self {
425        self.relation = Some(TableFactorBuilder::Derived(value));
426        self
427    }
428
429    pub fn unnest(&mut self, value: UnnestRelationBuilder) -> &mut Self {
430        self.relation = Some(TableFactorBuilder::Unnest(value));
431        self
432    }
433
434    pub fn empty(&mut self) -> &mut Self {
435        self.relation = Some(TableFactorBuilder::Empty);
436        self
437    }
438    pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
439        let new = self;
440        match new.relation {
441            Some(TableFactorBuilder::Table(ref mut rel_builder)) => {
442                rel_builder.alias = value;
443            }
444            Some(TableFactorBuilder::Derived(ref mut rel_builder)) => {
445                rel_builder.alias = value;
446            }
447            Some(TableFactorBuilder::Unnest(ref mut rel_builder)) => {
448                rel_builder.alias = value;
449            }
450            Some(TableFactorBuilder::Empty) => (),
451            None => (),
452        }
453        new
454    }
455    pub fn build(&self) -> Result<Option<ast::TableFactor>, BuilderError> {
456        Ok(match self.relation {
457            Some(TableFactorBuilder::Table(ref value)) => Some(value.build()?),
458            Some(TableFactorBuilder::Derived(ref value)) => Some(value.build()?),
459            Some(TableFactorBuilder::Unnest(ref value)) => Some(value.build()?),
460            Some(TableFactorBuilder::Empty) => None,
461            None => return Err(Into::into(UninitializedFieldError::from("relation"))),
462        })
463    }
464    fn create_empty() -> Self {
465        Self {
466            relation: Default::default(),
467        }
468    }
469}
470impl Default for RelationBuilder {
471    fn default() -> Self {
472        Self::create_empty()
473    }
474}
475
476#[derive(Clone)]
477pub struct TableRelationBuilder {
478    name: Option<ast::ObjectName>,
479    alias: Option<ast::TableAlias>,
480    args: Option<Vec<ast::FunctionArg>>,
481    with_hints: Vec<ast::Expr>,
482    version: Option<ast::TableVersion>,
483    partitions: Vec<ast::Ident>,
484    index_hints: Vec<ast::TableIndexHints>,
485}
486
487#[allow(dead_code)]
488impl TableRelationBuilder {
489    pub fn name(&mut self, value: ast::ObjectName) -> &mut Self {
490        self.name = Some(value);
491        self
492    }
493    pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
494        self.alias = value;
495        self
496    }
497    pub fn args(&mut self, value: Option<Vec<ast::FunctionArg>>) -> &mut Self {
498        self.args = value;
499        self
500    }
501    pub fn with_hints(&mut self, value: Vec<ast::Expr>) -> &mut Self {
502        self.with_hints = value;
503        self
504    }
505    pub fn version(&mut self, value: Option<ast::TableVersion>) -> &mut Self {
506        self.version = value;
507        self
508    }
509    pub fn partitions(&mut self, value: Vec<ast::Ident>) -> &mut Self {
510        self.partitions = value;
511        self
512    }
513    pub fn index_hints(&mut self, value: Vec<ast::TableIndexHints>) -> &mut Self {
514        self.index_hints = value;
515        self
516    }
517    pub fn build(&self) -> Result<ast::TableFactor, BuilderError> {
518        Ok(ast::TableFactor::Table {
519            name: match self.name {
520                Some(ref value) => value.clone(),
521                None => return Err(Into::into(UninitializedFieldError::from("name"))),
522            },
523            alias: self.alias.clone(),
524            args: self.args.clone().map(|args| ast::TableFunctionArgs {
525                args,
526                settings: None,
527            }),
528            with_hints: self.with_hints.clone(),
529            version: self.version.clone(),
530            partitions: self.partitions.clone(),
531            with_ordinality: false,
532            json_path: None,
533            sample: None,
534            index_hints: self.index_hints.clone(),
535        })
536    }
537    fn create_empty() -> Self {
538        Self {
539            name: Default::default(),
540            alias: Default::default(),
541            args: Default::default(),
542            with_hints: Default::default(),
543            version: Default::default(),
544            partitions: Default::default(),
545            index_hints: Default::default(),
546        }
547    }
548}
549impl Default for TableRelationBuilder {
550    fn default() -> Self {
551        Self::create_empty()
552    }
553}
554#[derive(Clone)]
555pub struct DerivedRelationBuilder {
556    lateral: Option<bool>,
557    subquery: Option<Box<ast::Query>>,
558    alias: Option<ast::TableAlias>,
559}
560
561#[allow(dead_code)]
562impl DerivedRelationBuilder {
563    pub fn lateral(&mut self, value: bool) -> &mut Self {
564        self.lateral = Some(value);
565        self
566    }
567    pub fn subquery(&mut self, value: Box<ast::Query>) -> &mut Self {
568        self.subquery = Some(value);
569        self
570    }
571    pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
572        self.alias = value;
573        self
574    }
575    fn build(&self) -> Result<ast::TableFactor, BuilderError> {
576        Ok(ast::TableFactor::Derived {
577            lateral: match self.lateral {
578                Some(ref value) => *value,
579                None => return Err(Into::into(UninitializedFieldError::from("lateral"))),
580            },
581            subquery: match self.subquery {
582                Some(ref value) => value.clone(),
583                None => {
584                    return Err(Into::into(UninitializedFieldError::from("subquery")))
585                }
586            },
587            alias: self.alias.clone(),
588        })
589    }
590    fn create_empty() -> Self {
591        Self {
592            lateral: Default::default(),
593            subquery: Default::default(),
594            alias: Default::default(),
595        }
596    }
597}
598impl Default for DerivedRelationBuilder {
599    fn default() -> Self {
600        Self::create_empty()
601    }
602}
603
604#[derive(Clone)]
605pub struct UnnestRelationBuilder {
606    pub alias: Option<ast::TableAlias>,
607    pub array_exprs: Vec<ast::Expr>,
608    with_offset: bool,
609    with_offset_alias: Option<ast::Ident>,
610    with_ordinality: bool,
611}
612
613#[allow(dead_code)]
614impl UnnestRelationBuilder {
615    pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
616        self.alias = value;
617        self
618    }
619    pub fn array_exprs(&mut self, value: Vec<ast::Expr>) -> &mut Self {
620        self.array_exprs = value;
621        self
622    }
623
624    pub fn with_offset(&mut self, value: bool) -> &mut Self {
625        self.with_offset = value;
626        self
627    }
628
629    pub fn with_offset_alias(&mut self, value: Option<ast::Ident>) -> &mut Self {
630        self.with_offset_alias = value;
631        self
632    }
633
634    pub fn with_ordinality(&mut self, value: bool) -> &mut Self {
635        self.with_ordinality = value;
636        self
637    }
638
639    pub fn build(&self) -> Result<ast::TableFactor, BuilderError> {
640        Ok(ast::TableFactor::UNNEST {
641            alias: self.alias.clone(),
642            array_exprs: self.array_exprs.clone(),
643            with_offset: self.with_offset,
644            with_offset_alias: self.with_offset_alias.clone(),
645            with_ordinality: self.with_ordinality,
646        })
647    }
648
649    fn create_empty() -> Self {
650        Self {
651            alias: Default::default(),
652            array_exprs: Default::default(),
653            with_offset: Default::default(),
654            with_offset_alias: Default::default(),
655            with_ordinality: Default::default(),
656        }
657    }
658}
659
660impl Default for UnnestRelationBuilder {
661    fn default() -> Self {
662        Self::create_empty()
663    }
664}
665
666/// Runtime error when a `build()` method is called and one or more required fields
667/// do not have a value.
668#[derive(Debug, Clone)]
669pub struct UninitializedFieldError(&'static str);
670
671impl UninitializedFieldError {
672    /// Create a new `UninitializedFieldError` for the specified field name.
673    pub fn new(field_name: &'static str) -> Self {
674        UninitializedFieldError(field_name)
675    }
676
677    /// Get the name of the first-declared field that wasn't initialized
678    pub fn field_name(&self) -> &'static str {
679        self.0
680    }
681}
682
683impl fmt::Display for UninitializedFieldError {
684    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
685        write!(f, "Field not initialized: {}", self.0)
686    }
687}
688
689impl From<&'static str> for UninitializedFieldError {
690    fn from(field_name: &'static str) -> Self {
691        Self::new(field_name)
692    }
693}
694impl std::error::Error for UninitializedFieldError {}
695
696#[derive(Debug)]
697pub enum BuilderError {
698    UninitializedField(&'static str),
699    ValidationError(String),
700}
701impl From<UninitializedFieldError> for BuilderError {
702    fn from(s: UninitializedFieldError) -> Self {
703        Self::UninitializedField(s.field_name())
704    }
705}
706impl From<String> for BuilderError {
707    fn from(s: String) -> Self {
708        Self::ValidationError(s)
709    }
710}
711impl fmt::Display for BuilderError {
712    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
713        match self {
714            Self::UninitializedField(ref field) => {
715                write!(f, "`{field}` must be initialized")
716            }
717            Self::ValidationError(ref error) => write!(f, "{error}"),
718        }
719    }
720}
721impl std::error::Error for BuilderError {}
datafusion_sql/unparser/ast.rs

datafusion_sql/unparser/
ast.rs
