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

datafusion_sql/unparser/
ast.rs
