databend_common_ast/ast/statements/
insert.rs

1// Copyright 2021 Datafuse Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt::Display;
16use std::fmt::Formatter;
17
18use derive_visitor::Drive;
19use derive_visitor::DriveMut;
20
21use crate::ast::write_comma_separated_list;
22use crate::ast::write_dot_separated_list;
23use crate::ast::Expr;
24use crate::ast::FileFormatOptions;
25use crate::ast::Hint;
26use crate::ast::Identifier;
27use crate::ast::Query;
28use crate::ast::With;
29
30#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
31pub struct InsertStmt {
32    pub hints: Option<Hint>,
33    // With clause, common table expression
34    pub with: Option<With>,
35    pub catalog: Option<Identifier>,
36    pub database: Option<Identifier>,
37    pub table: Identifier,
38    pub columns: Vec<Identifier>,
39    pub source: InsertSource,
40    pub overwrite: bool,
41}
42
43impl Display for InsertStmt {
44    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
45        if let Some(cte) = &self.with {
46            write!(f, "WITH {} ", cte)?;
47        }
48        write!(f, "INSERT ")?;
49        if let Some(hints) = &self.hints {
50            write!(f, "{} ", hints)?;
51        }
52        if self.overwrite {
53            write!(f, "OVERWRITE ")?;
54        } else {
55            write!(f, "INTO ")?;
56        }
57        write_dot_separated_list(
58            f,
59            self.catalog
60                .iter()
61                .chain(&self.database)
62                .chain(Some(&self.table)),
63        )?;
64        if !self.columns.is_empty() {
65            write!(f, " (")?;
66            write_comma_separated_list(f, &self.columns)?;
67            write!(f, ")")?;
68        }
69        write!(f, " {}", self.source)
70    }
71}
72
73#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
74pub enum InsertSource {
75    Values {
76        rows: Vec<Vec<Expr>>,
77    },
78    RawValues {
79        rest_str: String,
80        start: usize,
81    },
82    Select {
83        query: Box<Query>,
84    },
85    LoadFile {
86        format_options: FileFormatOptions,
87        value: Option<Vec<Expr>>,
88
89        // '_databend_upload' => read from streaming upload handler body
90        // _ => read from stage
91        location: String,
92    },
93}
94
95impl Display for InsertSource {
96    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
97        match self {
98            InsertSource::Values { rows } => {
99                write!(f, "VALUES ")?;
100                for (i, row) in rows.iter().enumerate() {
101                    if i > 0 {
102                        write!(f, ", ")?;
103                    }
104                    write!(f, "(")?;
105                    write_comma_separated_list(f, row)?;
106                    write!(f, ")")?;
107                }
108                Ok(())
109            }
110            InsertSource::RawValues { rest_str, .. } => write!(f, "VALUES {rest_str}"),
111            InsertSource::Select { query } => write!(f, "{query}"),
112            InsertSource::LoadFile {
113                value,
114                format_options,
115                location,
116            } => {
117                if let Some(value) = value {
118                    write!(f, "(")?;
119                    write_comma_separated_list(f, value)?;
120                    write!(f, ")")?;
121                }
122                write!(f, " FROM @{location}",)?;
123                write!(f, " FILE_FORMAT = ({})", format_options)?;
124                Ok(())
125            }
126        }
127    }
128}