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

sql/statements/
insert.rs
