sql/statements/
insert.rs

1// Copyright 2023 Greptime Team
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 serde::Serialize;
16use sqlparser::ast::{
17    Insert as SpInsert, ObjectName, ObjectNamePart, Parens, Query, SetExpr, Statement, TableObject,
18    UnaryOperator, ValueWithSpan, Values,
19};
20use sqlparser::parser::ParserError;
21use sqlparser_derive::{Visit, VisitMut};
22
23use crate::ast::{Expr, Value};
24use crate::error::{Result, UnsupportedSnafu};
25use crate::statements::query::Query as GtQuery;
26
27#[derive(Debug, Clone, PartialEq, Eq, Visit, VisitMut, Serialize)]
28pub struct Insert {
29    // Can only be sqlparser::ast::Statement::Insert variant
30    pub inner: Statement,
31}
32
33macro_rules! parse_fail {
34    ($expr: expr) => {
35        return crate::error::ParseSqlValueSnafu {
36            msg: format!("{:?}", $expr),
37        }
38        .fail();
39    };
40}
41
42impl Insert {
43    pub fn table_name(&self) -> Result<&ObjectName> {
44        match &self.inner {
45            Statement::Insert(insert) => {
46                let TableObject::TableName(name) = &insert.table else {
47                    return UnsupportedSnafu {
48                        keyword: "TABLE FUNCTION".to_string(),
49                    }
50                    .fail();
51                };
52                Ok(name)
53            }
54            _ => unreachable!(),
55        }
56    }
57
58    pub fn columns(&self) -> Vec<&String> {
59        match &self.inner {
60            Statement::Insert(insert) => insert
61                .columns
62                .iter()
63                .filter_map(single_part_column_ident)
64                .map(|ident| &ident.value)
65                .collect(),
66            _ => unreachable!(),
67        }
68    }
69
70    /// Extracts the literal insert statement body if possible
71    pub fn values_body(&self) -> Result<Vec<Vec<Value>>> {
72        match &self.inner {
73            Statement::Insert(SpInsert {
74                source:
75                    Some(box Query {
76                        body: box SetExpr::Values(Values { rows, .. }),
77                        ..
78                    }),
79                ..
80            }) => sql_exprs_to_values(rows),
81            _ => unreachable!(),
82        }
83    }
84
85    /// Returns true when the insert statement can extract literal values.
86    /// The rules is the same as function `values_body()`.
87    pub fn can_extract_values(&self) -> bool {
88        match &self.inner {
89            Statement::Insert(SpInsert {
90                source:
91                    Some(box Query {
92                        body: box SetExpr::Values(Values { rows, .. }),
93                        ..
94                    }),
95                ..
96            }) => rows.iter().all(|es| {
97                es.iter().all(|expr| match expr {
98                    Expr::Value(_) => true,
99                    Expr::Identifier(ident) => {
100                        if ident.quote_style.is_none() {
101                            ident.value.to_lowercase() == "default"
102                        } else {
103                            ident.quote_style == Some('"')
104                        }
105                    }
106                    Expr::UnaryOp { op, expr } => {
107                        matches!(op, UnaryOperator::Minus | UnaryOperator::Plus)
108                            && matches!(
109                                &**expr,
110                                Expr::Value(ValueWithSpan {
111                                    value: Value::Number(_, _),
112                                    ..
113                                })
114                            )
115                    }
116                    _ => false,
117                })
118            }),
119            _ => false,
120        }
121    }
122
123    /// Returns true when the insert source is a query rather than `VALUES`.
124    pub fn has_non_values_query_source(&self) -> bool {
125        match &self.inner {
126            Statement::Insert(SpInsert {
127                source: Some(box query),
128                ..
129            }) => !matches!(&*query.body, SetExpr::Values(_)),
130            _ => false,
131        }
132    }
133
134    pub fn query_body(&self) -> Result<Option<GtQuery>> {
135        Ok(match &self.inner {
136            Statement::Insert(SpInsert {
137                source: Some(box query),
138                ..
139            }) => Some(query.clone().try_into()?),
140            _ => None,
141        })
142    }
143}
144
145fn sql_exprs_to_values(exprs: &[Parens<Vec<Expr>>]) -> Result<Vec<Vec<Value>>> {
146    let mut values = Vec::with_capacity(exprs.len());
147    for es in exprs.iter() {
148        let mut vs = Vec::with_capacity(es.len());
149        for expr in es.iter() {
150            vs.push(match expr {
151                Expr::Value(v) => v.value.clone(),
152                Expr::Identifier(ident) => {
153                    if ident.quote_style.is_none() {
154                        // Special processing for `default` value
155                        if ident.value.to_lowercase() == "default" {
156                            Value::Placeholder(ident.value.clone())
157                        } else {
158                            parse_fail!(expr);
159                        }
160                    } else {
161                        // Identifiers with double quotes, we treat them as strings.
162                        if ident.quote_style == Some('"') {
163                            Value::SingleQuotedString(ident.value.clone())
164                        } else {
165                            parse_fail!(expr);
166                        }
167                    }
168                }
169                Expr::UnaryOp { op, expr }
170                    if matches!(op, UnaryOperator::Minus | UnaryOperator::Plus) =>
171                {
172                    if let Expr::Value(ValueWithSpan {
173                        value: Value::Number(s, b),
174                        ..
175                    }) = &**expr
176                    {
177                        match op {
178                            UnaryOperator::Minus => Value::Number(format!("-{s}"), *b),
179                            UnaryOperator::Plus => Value::Number(s.clone(), *b),
180                            _ => unreachable!(),
181                        }
182                    } else {
183                        parse_fail!(expr);
184                    }
185                }
186                _ => {
187                    parse_fail!(expr);
188                }
189            });
190        }
191        values.push(vs);
192    }
193    Ok(values)
194}
195
196fn single_part_column_ident(name: &ObjectName) -> Option<&sqlparser::ast::Ident> {
197    let [ObjectNamePart::Identifier(ident)] = name.0.as_slice() else {
198        return None;
199    };
200    Some(ident)
201}
202
203impl TryFrom<Statement> for Insert {
204    type Error = ParserError;
205
206    fn try_from(value: Statement) -> std::result::Result<Self, Self::Error> {
207        let Statement::Insert(insert) = &value else {
208            return Err(ParserError::ParserError(format!(
209                "Not expected to be {value}"
210            )));
211        };
212
213        if let Some(column) = insert
214            .columns
215            .iter()
216            .find(|column| single_part_column_ident(column).is_none())
217        {
218            return Err(ParserError::ParserError(format!(
219                "Expected a single-part insert column name, found {column}"
220            )));
221        }
222
223        Ok(Insert { inner: value })
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230    use crate::dialect::GreptimeDbDialect;
231    use crate::parser::{ParseOptions, ParserContext};
232    use crate::statements::statement::Statement;
233
234    #[test]
235    fn test_insert_value_with_unary_op() {
236        // insert "-1"
237        let sql = "INSERT INTO my_table VALUES(-1)";
238        let stmt =
239            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
240                .unwrap()
241                .remove(0);
242        match stmt {
243            Statement::Insert(insert) => {
244                let values = insert.values_body().unwrap();
245                assert_eq!(values, vec![vec![Value::Number("-1".to_string(), false)]]);
246            }
247            _ => unreachable!(),
248        }
249
250        // insert "+1"
251        let sql = "INSERT INTO my_table VALUES(+1)";
252        let stmt =
253            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
254                .unwrap()
255                .remove(0);
256        match stmt {
257            Statement::Insert(insert) => {
258                let values = insert.values_body().unwrap();
259                assert_eq!(values, vec![vec![Value::Number("1".to_string(), false)]]);
260            }
261            _ => unreachable!(),
262        }
263    }
264
265    #[test]
266    fn test_insert_column_names_are_single_identifiers() {
267        let stmt = ParserContext::create_with_dialect(
268            "INSERT INTO my_table (host, \"value\") VALUES (1, 2)",
269            &GreptimeDbDialect {},
270            ParseOptions::default(),
271        )
272        .unwrap()
273        .remove(0);
274        let Statement::Insert(insert) = stmt else {
275            unreachable!()
276        };
277        assert_eq!(insert.columns(), vec!["host", "value"]);
278
279        let result = ParserContext::create_with_dialect(
280            "INSERT INTO my_table (metric.host) VALUES (1)",
281            &GreptimeDbDialect {},
282            ParseOptions::default(),
283        );
284        let error = result.unwrap_err().to_string();
285        assert!(
286            error.contains("Expected a single-part insert column name, found metric.host"),
287            "unexpected error: {error}"
288        );
289    }
290
291    #[test]
292    fn test_insert_value_with_default() {
293        // insert "default"
294        let sql = "INSERT INTO my_table VALUES(default)";
295        let stmt =
296            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
297                .unwrap()
298                .remove(0);
299        match stmt {
300            Statement::Insert(insert) => {
301                let values = insert.values_body().unwrap();
302                assert_eq!(values, vec![vec![Value::Placeholder("default".to_owned())]]);
303            }
304            _ => unreachable!(),
305        }
306    }
307
308    #[test]
309    fn test_insert_value_with_default_uppercase() {
310        // insert "DEFAULT"
311        let sql = "INSERT INTO my_table VALUES(DEFAULT)";
312        let stmt =
313            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
314                .unwrap()
315                .remove(0);
316        match stmt {
317            Statement::Insert(insert) => {
318                let values = insert.values_body().unwrap();
319                assert_eq!(values, vec![vec![Value::Placeholder("DEFAULT".to_owned())]]);
320            }
321            _ => unreachable!(),
322        }
323    }
324
325    #[test]
326    fn test_insert_value_with_quoted_string() {
327        // insert 'default'
328        let sql = "INSERT INTO my_table VALUES('default')";
329        let stmt =
330            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
331                .unwrap()
332                .remove(0);
333        match stmt {
334            Statement::Insert(insert) => {
335                let values = insert.values_body().unwrap();
336                assert_eq!(
337                    values,
338                    vec![vec![Value::SingleQuotedString("default".to_owned())]]
339                );
340            }
341            _ => unreachable!(),
342        }
343
344        // insert "default". Treating double-quoted identifiers as strings.
345        let sql = "INSERT INTO my_table VALUES(\"default\")";
346        let stmt =
347            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
348                .unwrap()
349                .remove(0);
350        match stmt {
351            Statement::Insert(insert) => {
352                let values = insert.values_body().unwrap();
353                assert_eq!(
354                    values,
355                    vec![vec![Value::SingleQuotedString("default".to_owned())]]
356                );
357            }
358            _ => unreachable!(),
359        }
360
361        let sql = "INSERT INTO my_table VALUES(`default`)";
362        let stmt =
363            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
364                .unwrap()
365                .remove(0);
366        match stmt {
367            Statement::Insert(insert) => {
368                assert!(insert.values_body().is_err());
369            }
370            _ => unreachable!(),
371        }
372    }
373
374    #[test]
375    fn test_insert_select() {
376        let sql = "INSERT INTO my_table select * from other_table";
377        let stmt =
378            ParserContext::create_with_dialect(sql, &GreptimeDbDialect {}, ParseOptions::default())
379                .unwrap()
380                .remove(0);
381        match stmt {
382            Statement::Insert(insert) => {
383                let q = insert.query_body().unwrap().unwrap();
384                assert!(insert.has_non_values_query_source());
385                assert!(matches!(
386                    q.inner,
387                    Query {
388                        body: box SetExpr::Select { .. },
389                        ..
390                    }
391                ));
392            }
393            _ => unreachable!(),
394        }
395    }
396
397    #[test]
398    fn test_has_non_values_query_source() {
399        let cases = [
400            ("INSERT INTO my_table SELECT * FROM other_table", true),
401            (
402                "INSERT INTO my_table WITH cte AS (SELECT * FROM other_table) SELECT * FROM cte",
403                true,
404            ),
405            (
406                "INSERT INTO my_table SELECT * FROM t1 UNION ALL SELECT * FROM t2",
407                true,
408            ),
409            ("INSERT INTO my_table VALUES(1)", false),
410            ("INSERT INTO my_table VALUES(now())", false),
411            ("INSERT INTO my_table VALUES(1 + 1)", false),
412        ];
413
414        for (sql, expected) in cases {
415            let stmt = ParserContext::create_with_dialect(
416                sql,
417                &GreptimeDbDialect {},
418                ParseOptions::default(),
419            )
420            .unwrap()
421            .remove(0);
422            match stmt {
423                Statement::Insert(insert) => {
424                    assert_eq!(insert.has_non_values_query_source(), expected, "{sql}");
425                }
426                _ => unreachable!(),
427            }
428        }
429    }
430}
sql/statements/insert.rs

sql/statements/
insert.rs
