Skip to main content

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::Expr;
22use crate::ast::FileFormatOptions;
23use crate::ast::Hint;
24use crate::ast::Identifier;
25use crate::ast::Query;
26use crate::ast::TableRef;
27use crate::ast::With;
28use crate::ast::write_comma_separated_list;
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 table: TableRef,
36    pub columns: Vec<Identifier>,
37    pub source: InsertSource,
38    pub overwrite: bool,
39}
40
41impl Display for InsertStmt {
42    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
43        if let Some(cte) = &self.with {
44            write!(f, "WITH {} ", cte)?;
45        }
46        write!(f, "INSERT ")?;
47        if let Some(hints) = &self.hints {
48            write!(f, "{} ", hints)?;
49        }
50        if self.overwrite {
51            write!(f, "OVERWRITE ")?;
52        } else {
53            write!(f, "INTO ")?;
54        }
55        write!(f, "{}", self.table)?;
56        if !self.columns.is_empty() {
57            write!(f, " (")?;
58            write_comma_separated_list(f, &self.columns)?;
59            write!(f, ")")?;
60        }
61        write!(f, " {}", self.source)
62    }
63}
64
65#[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
66pub enum InsertSource {
67    Values {
68        rows: Vec<Vec<Expr>>,
69    },
70    RawValues {
71        rest_str: String,
72        start: usize,
73    },
74    Select {
75        query: Box<Query>,
76    },
77    LoadFile {
78        format_options: FileFormatOptions,
79        value: Option<Vec<Expr>>,
80
81        // '_databend_upload' => read from streaming upload handler body
82        // _ => read from stage
83        location: String,
84    },
85}
86
87impl Display for InsertSource {
88    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
89        match self {
90            InsertSource::Values { rows } => {
91                write!(f, "VALUES ")?;
92                for (i, row) in rows.iter().enumerate() {
93                    if i > 0 {
94                        write!(f, ", ")?;
95                    }
96                    write!(f, "(")?;
97                    write_comma_separated_list(f, row)?;
98                    write!(f, ")")?;
99                }
100                Ok(())
101            }
102            InsertSource::RawValues { rest_str, .. } => write!(f, "VALUES {rest_str}"),
103            InsertSource::Select { query } => write!(f, "{query}"),
104            InsertSource::LoadFile {
105                value,
106                format_options,
107                location,
108            } => {
109                if let Some(value) = value {
110                    write!(f, "VALUES (")?;
111                    write_comma_separated_list(f, value)?;
112                    write!(f, ")")?;
113                }
114                write!(f, " FROM @{location}",)?;
115                write!(f, " FILE_FORMAT = ({})", format_options)?;
116                Ok(())
117            }
118        }
119    }
120}