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

datafusion_sql/unparser/
ast.rs
