feat: adjust syntax of streaming load insert SQL. (#18075)

Author: youngsofun

Cargo.lock

@@ -18,6 +18,7 @@ crate-type = ["cdylib"]
 arrow = { workspace = true, features = ["pyarrow"] }
 arrow-schema = { workspace = true }
 ctor = { workspace = true }
+databend-common-catalog = { workspace = true }
 databend-common-config = { workspace = true }
 databend-common-exception = { workspace = true }
 databend-common-expression = { workspace = true }

src/bendpy/Cargo.toml

@@ -14,6 +14,7 @@
 
 use std::sync::Arc;
 
+use databend_common_catalog::session_type::SessionType;
 use databend_common_config::GlobalConfig;
 use databend_common_exception::Result;
 use databend_common_meta_app::principal::GrantObject;
@@ -24,7 +25,6 @@ use databend_common_users::UserApiProvider;
 use databend_query::sessions::QueryContext;
 use databend_query::sessions::Session;
 use databend_query::sessions::SessionManager;
-use databend_query::sessions::SessionType;
 use databend_query::sql::Planner;
 use pyo3::prelude::*;
 

src/bendpy/src/context.rs

@@ -8,6 +8,7 @@ edition = "2021"
 
 [dependencies]
 databend-common-base = { workspace = true }
+databend-common-catalog = { workspace = true }
 databend-common-config = { workspace = true }
 databend-common-license = { workspace = true }
 databend-common-meta-client = { workspace = true }

src/bendsave/Cargo.toml

@@ -24,6 +24,7 @@ use bytes::Bytes;
 use bytes::BytesMut;
 use databend_common_base::base::GlobalInstance;
 use databend_common_base::runtime::GlobalIORuntime;
+use databend_common_catalog::session_type::SessionType;
 use databend_common_config::Config;
 use databend_common_config::GlobalConfig;
 use databend_common_config::InnerConfig;
@@ -38,7 +39,6 @@ use databend_common_users::builtin::BuiltIn;
 use databend_common_users::UserApiProvider;
 use databend_enterprise_query::license::RealLicenseManager;
 use databend_query::sessions::SessionManager;
-use databend_query::sessions::SessionType;
 use futures::TryStream;
 use futures::TryStreamExt;
 use log::debug;

src/bendsave/src/storage.rs

@@ -85,6 +85,7 @@ pub enum InsertSource {
     StreamingLoad {
         format_options: FileFormatOptions,
         on_error_mode: Option<String>,
+        value: Option<Vec<Expr>>,
     },
 }
 
@@ -106,9 +107,15 @@ impl Display for InsertSource {
             InsertSource::RawValues { rest_str, .. } => write!(f, "VALUES {rest_str}"),
             InsertSource::Select { query } => write!(f, "{query}"),
             InsertSource::StreamingLoad {
+                value,
                 format_options,
                 on_error_mode,
             } => {
+                if let Some(value) = value {
+                    write!(f, "(")?;
+                    write_comma_separated_list(f, value)?;
+                    write!(f, ")")?;
+                }
                 write!(f, " FILE_FORMAT = ({})", format_options)?;
                 write!(
                     f,

src/query/ast/src/ast/statements/insert.rs

@@ -81,13 +81,17 @@ pub fn parse_values(tokens: &[Token], dialect: Dialect) -> Result<Vec<Expr>> {
     run_parser(tokens, dialect, ParseMode::Default, false, values)
 }
 
-pub fn parse_raw_insert_stmt(tokens: &[Token], dialect: Dialect) -> Result<Statement> {
+pub fn parse_raw_insert_stmt(
+    tokens: &[Token],
+    dialect: Dialect,
+    in_streaming_load: bool,
+) -> Result<Statement> {
     run_parser(
         tokens,
         dialect,
         ParseMode::Default,
         false,
-        insert_stmt(true),
+        insert_stmt(true, in_streaming_load),
     )
 }
 

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

@@ -2605,7 +2605,7 @@ pub fn statement_body(i: Input) -> IResult<Statement> {
         rule!(
             #conditional_multi_table_insert() : "`INSERT [OVERWRITE] {FIRST|ALL} { WHEN <condition> THEN intoClause [ ... ] } [ ... ] [ ELSE intoClause ] <subquery>`"
             | #unconditional_multi_table_insert() : "`INSERT [OVERWRITE] ALL intoClause [ ... ] <subquery>`"
-            | #insert_stmt(false) : "`INSERT INTO [TABLE] <table> [(<column>, ...)] (FORMAT <format> | VALUES <values> | <query>)`"
+            | #insert_stmt(false, false) : "`INSERT INTO [TABLE] <table> [(<column>, ...)] (VALUES <values> | <query>)`"
             | #replace_stmt(false) : "`REPLACE INTO [TABLE] <table> [(<column>, ...)] (FORMAT <format> | VALUES <values> | <query>)`"
             | #merge : "`MERGE INTO <target_table> USING <source> ON <join_expr> { matchedClause | notMatchedClause } [ ... ]`"
             | #delete : "`DELETE FROM <table> [WHERE ...]`"
@@ -2798,10 +2798,15 @@ pub fn parse_create_option(
     }
 }
 
-pub fn insert_stmt(allow_raw: bool) -> impl FnMut(Input) -> IResult<Statement> {
+pub fn insert_stmt(
+    allow_raw: bool,
+    in_streaming_load: bool,
+) -> impl FnMut(Input) -> IResult<Statement> {
     move |i| {
-        let insert_source_parser = if allow_raw {
-            raw_insert_source
+        let insert_source_parser = if in_streaming_load {
+            insert_source_file
+        } else if allow_raw {
+            insert_source_fast_values
         } else {
             insert_source
         };
@@ -2935,7 +2940,7 @@ fn else_clause(i: Input) -> IResult<ElseClause> {
 pub fn replace_stmt(allow_raw: bool) -> impl FnMut(Input) -> IResult<Statement> {
     move |i| {
         let insert_source_parser = if allow_raw {
-            raw_insert_source
+            insert_source_fast_values
         } else {
             insert_source
         };
@@ -3005,20 +3010,38 @@ pub fn insert_source(i: Input) -> IResult<InsertSource> {
     )(i)
 }
 
-// `INSERT INTO ... VALUES` statement will
-// stop the parser immediately and return the rest tokens in `InsertSource`.
-//
-// This is a hack to parse large insert statements.
-pub fn raw_insert_source(i: Input) -> IResult<InsertSource> {
-    let streaming = map(
+pub fn insert_source_file(i: Input) -> IResult<InsertSource> {
+    let value = map(
         rule! {
-           #file_format_clause ~ (ON_ERROR ~ ^"=" ~ ^#ident)?
+            "(" ~ #comma_separated_list1(expr) ~ ")"
         },
-        |(options, on_error_opt)| InsertSource::StreamingLoad {
+        |(_, values, _)| values,
+    );
+    map(
+        rule! {
+           VALUES ~ #value? ~ #file_format_clause ~ (ON_ERROR ~ ^"=" ~ ^#ident)?
+        },
+        |(_, value, options, on_error_opt)| InsertSource::StreamingLoad {
             format_options: options,
             on_error_mode: on_error_opt.map(|v| v.2.to_string()),
+            value,
         },
-    );
+    )(i)
+    // TODO: support query later
+    // let query = map(query, |query| InsertSource::Select {
+    //     query: Box::new(query),
+    // });
+    // rule!(
+    //     #file
+    //     | #query
+    // )(i)
+}
+
+// `INSERT INTO ... VALUES` statement will
+// stop the parser immediately and return the rest tokens in `InsertSource`.
+//
+// This is a hack to parse large insert statements.
+pub fn insert_source_fast_values(i: Input) -> IResult<InsertSource> {
     let values = map(
         rule! {
             VALUES ~ #rest_str
@@ -3037,7 +3060,6 @@ pub fn raw_insert_source(i: Input) -> IResult<InsertSource> {
     rule!(
         #values
         | #query
-        | #streaming
     )(i)
 }
 

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

@@ -1102,7 +1102,7 @@ fn test_raw_insert_stmt() {
     ];
 
     for case in cases {
-        run_parser(file, insert_stmt(true), case);
+        run_parser(file, insert_stmt(true, false), case);
     }
 }
 

src/query/ast/tests/it/parser.rs

@@ -133,7 +133,7 @@ error:
   | ^^^^^^
   | |
   | INSERT statement must be followed by 'overwrite' or 'into'
-  | while parsing `INSERT INTO [TABLE] <table> [(<column>, ...)] (FORMAT <format> | VALUES <values> | <query>)`
+  | while parsing `INSERT INTO [TABLE] <table> [(<column>, ...)] (VALUES <values> | <query>)`
 
 
 ---------- Input ----------
@@ -167,7 +167,7 @@ error:
 1 | insert into t format
   | ------        ^^^^^^ unexpected `format`, expecting `FROM`, `ORDER`, `LIMIT`, `OFFSET`, `IGNORE_RESULT`, `WITH`, `VALUES`, `EXCEPT`, `SELECT`, `INTERSECT`, `(`, `UNION`, or `.`
   | |              
-  | while parsing `INSERT INTO [TABLE] <table> [(<column>, ...)] (FORMAT <format> | VALUES <values> | <query>)`
+  | while parsing `INSERT INTO [TABLE] <table> [(<column>, ...)] (VALUES <values> | <query>)`
 
 
 ---------- Input ----------