fix(query): fix replace into partial parse insert source (#18085)

* fix(query): fix replace into partial parse insert source

* fix

* add tests

* fix tests

Author: b41sh

src/query/ast/src/ast/statements/replace.rs

@@ -31,6 +31,7 @@ pub struct ReplaceStmt {
     pub catalog: Option<Identifier>,
     pub database: Option<Identifier>,
     pub table: Identifier,
+    pub is_conflict: bool,
     pub on_conflict_columns: Vec<Identifier>,
     pub columns: Vec<Identifier>,
     pub source: InsertSource,
@@ -59,7 +60,10 @@ impl Display for ReplaceStmt {
         }
 
         // on_conflict_columns must be non-empty
-        write!(f, " ON CONFLICT")?;
+        write!(f, " ON")?;
+        if self.is_conflict {
+            write!(f, " CONFLICT")?;
+        }
         write!(f, " (")?;
         write_comma_separated_list(f, &self.on_conflict_columns)?;
         write!(f, ")")?;

src/query/ast/src/parser/statement.rs

@@ -2955,7 +2955,7 @@ pub fn replace_stmt(allow_raw: bool) -> impl FnMut(Input) -> IResult<Statement>
                 (catalog, database, table),
                 opt_columns,
                 _,
-                _,
+                opt_conflict,
                 _,
                 on_conflict_columns,
                 _,
@@ -2967,6 +2967,7 @@ pub fn replace_stmt(allow_raw: bool) -> impl FnMut(Input) -> IResult<Statement>
                     catalog,
                     database,
                     table,
+                    is_conflict: opt_conflict.is_some(),
                     on_conflict_columns,
                     columns: opt_columns
                         .map(|(_, columns, _)| columns)

src/query/ast/tests/it/testdata/stmt.txt

@@ -777,7 +777,7 @@ Some(
 ---------- Input ----------
 replace into test on(c) select sum(c) as c from source group by v;
 ---------- Output ---------
-REPLACE INTO test ON CONFLICT (c) SELECT sum(c) AS c FROM source GROUP BY v
+REPLACE INTO test ON (c) SELECT sum(c) AS c FROM source GROUP BY v
 ---------- AST ------------
 Replace(
     ReplaceStmt {
@@ -792,6 +792,7 @@ Replace(
             quote: None,
             ident_type: None,
         },
+        is_conflict: false,
         on_conflict_columns: [
             Identifier {
                 span: Some(
@@ -1028,7 +1029,7 @@ Explain {
 ---------- Input ----------
 explain replace into test on(c) select sum(c) as c from source group by v;
 ---------- Output ---------
-EXPLAIN REPLACE INTO test ON CONFLICT (c) SELECT sum(c) AS c FROM source GROUP BY v
+EXPLAIN REPLACE INTO test ON (c) SELECT sum(c) AS c FROM source GROUP BY v
 ---------- AST ------------
 Explain {
     kind: Plan,
@@ -1049,6 +1050,7 @@ Explain {
                 quote: None,
                 ident_type: None,
             },
+            is_conflict: false,
             on_conflict_columns: [
                 Identifier {
                     span: Some(

src/query/sql/src/planner/binder/replace.rs

@@ -40,6 +40,7 @@ impl Binder {
             catalog,
             database,
             table,
+            is_conflict: _,
             on_conflict_columns,
             columns,
             source,

src/query/sql/src/planner/planner.rs

@@ -19,6 +19,7 @@ use databend_common_ast::ast::Expr;
 use databend_common_ast::ast::InsertSource;
 use databend_common_ast::ast::InsertStmt;
 use databend_common_ast::ast::Literal;
+use databend_common_ast::ast::ReplaceStmt;
 use databend_common_ast::ast::Statement;
 use databend_common_ast::parser::parse_raw_insert_stmt;
 use databend_common_ast::parser::parse_raw_replace_stmt;
@@ -50,7 +51,6 @@ use crate::NameResolutionContext;
 use crate::VariableNormalizer;
 
 const PROBE_INSERT_INITIAL_TOKENS: usize = 128;
-const PROBE_INSERT_MAX_TOKENS: usize = 128 * 8;
 
 pub struct Planner {
     pub(crate) ctx: Arc<dyn TableContext>,
@@ -180,40 +180,39 @@ impl Planner {
             let mut maybe_partial_insert = false;
 
             if is_insert_or_replace_stmt && matches!(tokenizer.peek(), Some(Ok(_))) {
-                if let Ok(PlanExtras {
-                    statement:
-                        Statement::Insert(InsertStmt {
-                            source: InsertSource::Select { .. },
-                            ..
-                        }),
-                    ..
-                }) = &res
-                {
-                    maybe_partial_insert = true;
+                match res {
+                    Ok(PlanExtras {
+                        statement:
+                            Statement::Insert(InsertStmt {
+                                source: InsertSource::Select { .. },
+                                ..
+                            }),
+                        ..
+                    })
+                    | Ok(PlanExtras {
+                        statement:
+                            Statement::Replace(ReplaceStmt {
+                                source: InsertSource::Select { .. },
+                                ..
+                            }),
+                        ..
+                    }) => {
+                        maybe_partial_insert = true;
+                    }
+                    _ => {}
                 }
             }
 
             if (maybe_partial_insert || res.is_err()) && matches!(tokenizer.peek(), Some(Ok(_))) {
                 // Remove the EOI.
                 tokens.pop();
-                // Tokenize more and try again.
-                if !maybe_partial_insert && tokens.len() < PROBE_INSERT_MAX_TOKENS {
-                    let iter = (&mut tokenizer)
-                        .take(tokens.len() * 2)
-                        .take_while(|token| token.is_ok())
-                        .map(|token| token.unwrap())
-                        // Make sure the tokens stream is always ended with EOI.
-                        .chain(std::iter::once(Token::new_eoi(&final_sql)));
-                    tokens.extend(iter);
-                } else {
-                    // Take the whole tokenizer
-                    let iter = (&mut tokenizer)
-                        .take_while(|token| token.is_ok())
-                        .map(|token| token.unwrap())
-                        // Make sure the tokens stream is always ended with EOI.
-                        .chain(std::iter::once(Token::new_eoi(&final_sql)));
-                    tokens.extend(iter);
-                };
+                // Take the whole tokenizer
+                let iter = (&mut tokenizer)
+                    .take_while(|token| token.is_ok())
+                    .map(|token| token.unwrap())
+                    // Make sure the tokens stream is always ended with EOI.
+                    .chain(std::iter::once(Token::new_eoi(&final_sql)));
+                tokens.extend(iter);
             } else {
                 return res;
             }

src/tests/sqlsmith/src/sql_gen/dml.rs

@@ -195,6 +195,7 @@ impl<'a, R: Rng + 'a> SqlGenerator<'a, R> {
             catalog: None,
             database: table.db_name.clone(),
             table: table.name.clone(),
+            is_conflict: true,
             on_conflict_columns,
             columns,
             source,

tests/sqllogictests/suites/base/09_fuse_engine/09_0023_replace_into.test

@@ -514,7 +514,7 @@ select * from t order by c;
 3
 
 statement ok
-create or replace table test (c int);
+create or replace table test (c int64);
 
 statement ok
 create or replace table source (c int, v string);
@@ -536,5 +536,9 @@ select sum(c) from test;
 statement ok
 explain replace into test on(c) select sum(c) as c from source group by v;
 
+# partial parse insert source
 statement ok
-DROP DATABASE db_09_0023
+replace into test (c) on conflict (c) select sum(c) as c from source where c >= 0 and c >= 1 and c >= 2 and c >= 3 and c >= 4 and c >= 5 and c >= 6 and c >= 7 and c >= 8 and c >= 9 and c >= 10 and c >= 11 and c >= 12 and c >= 13 and c >= 14 and c >= 15 and c >= 16 and c >= 17 and c >= 18 and c >= 19 and c >= 20 and c >= 21 and c >=22 and c >= 23 and c >= 24 and true and true and now() >= '2025-01-01' and date(now()) < date(add_hours(now(),24));
+
+statement ok
+DROP DATABASE db_09_0023