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

datafusion_sql/unparser/
ast.rs
