feat: insert with stage support from and file_format clause. (#18256)

* feat: insert with stage support from and file_format clause.

* update tests.

Author: youngsofun

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

@@ -82,10 +82,13 @@ pub enum InsertSource {
     Select {
         query: Box<Query>,
     },
-    StreamingLoad {
+    LoadFile {
         format_options: FileFormatOptions,
-        on_error_mode: Option<String>,
         value: Option<Vec<Expr>>,
+
+        // '_databend_upload' => read from streaming upload handler body
+        // _ => read from stage
+        location: String,
     },
 }
 
@@ -106,22 +109,19 @@ impl Display for InsertSource {
             }
             InsertSource::RawValues { rest_str, .. } => write!(f, "VALUES {rest_str}"),
             InsertSource::Select { query } => write!(f, "{query}"),
-            InsertSource::StreamingLoad {
+            InsertSource::LoadFile {
                 value,
                 format_options,
-                on_error_mode,
+                location,
             } => {
                 if let Some(value) = value {
                     write!(f, "(")?;
                     write_comma_separated_list(f, value)?;
                     write!(f, ")")?;
                 }
+                write!(f, " FROM @{location}",)?;
                 write!(f, " FILE_FORMAT = ({})", format_options)?;
-                write!(
-                    f,
-                    " ON_ERROR = '{}'",
-                    on_error_mode.as_ref().unwrap_or(&"Abort".to_string())
-                )
+                Ok(())
             }
         }
     }

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

@@ -2965,22 +2965,14 @@ pub fn insert_source_file(i: Input) -> IResult<InsertSource> {
     );
     map(
         rule! {
-           VALUES ~ #value? ~ #file_format_clause ~ (ON_ERROR ~ ^"=" ~ ^#ident)?
+           (VALUES ~ #value?)? ~ FROM ~ #at_string ~  #file_format_clause
         },
-        |(_, value, options, on_error_opt)| InsertSource::StreamingLoad {
-            format_options: options,
-            on_error_mode: on_error_opt.map(|v| v.2.to_string()),
-            value,
+        |(values, _, location, format_options)| InsertSource::LoadFile {
+            value: values.map(|(_, value)| value).unwrap_or_default(),
+            location,
+            format_options,
         },
     )(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
@@ -3004,6 +2996,7 @@ pub fn insert_source_fast_values(i: Input) -> IResult<InsertSource> {
     );
 
     rule!(
+        #insert_source_file |
         #values
         | #query
     )(i)

src/query/service/src/interpreters/interpreter_insert.rs

@@ -259,7 +259,6 @@ impl Interpreter for InsertInterpreter {
                     plan.required_source_schema.clone(),
                     plan.default_exprs.clone(),
                     plan.block_thresholds,
-                    plan.on_error_mode.clone(),
                 )?;
                 if !plan.values_consts.is_empty() {
                     let input_schema = Arc::new(DataSchema::from(&plan.required_source_schema));

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

@@ -68,6 +68,7 @@ use log::warn;
 use parking_lot::RwLock;
 
 use crate::binder::bind_query::MaxColumnPosition;
+use crate::binder::insert::STAGE_PLACEHOLDER;
 use crate::binder::location::parse_uri_location;
 use crate::binder::Binder;
 use crate::plans::CopyIntoTableMode;
@@ -358,19 +359,52 @@ impl Binder {
         values_str: &str,
         write_mode: CopyIntoTableMode,
     ) -> Result<Plan> {
-        let (required_source_schema, const_columns) = if values_str.is_empty() {
-            (required_values_schema.clone(), vec![])
+        let expr_or_placeholders = if values_str.is_empty() {
+            None
         } else {
-            self.prepared_values_str(values_str, &required_values_schema)
-                .await?
+            Some(self.parse_values_str(values_str).await?)
         };
 
+        let (stage_info, files_info, options) = self.bind_attachment(attachment).await?;
+        self.bind_copy_from_upload(
+            bind_context,
+            catalog_name,
+            database_name,
+            table_name,
+            required_values_schema,
+            expr_or_placeholders,
+            stage_info,
+            files_info,
+            options,
+            write_mode,
+        )
+        .await
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    #[async_backtrace::framed]
+    pub(crate) async fn bind_copy_from_upload(
+        &mut self,
+        bind_context: &mut BindContext,
+        catalog_name: String,
+        database_name: String,
+        table_name: String,
+        required_values_schema: TableSchemaRef,
+        expr_or_placeholders: Option<Vec<Expr>>,
+        stage_info: StageInfo,
+        files_info: StageFilesInfo,
+        copy_into_table_options: CopyIntoTableOptions,
+        write_mode: CopyIntoTableMode,
+    ) -> Result<Plan> {
         let catalog = self.ctx.get_catalog(&catalog_name).await?;
         let catalog_info = catalog.info();
-
-        let thread_num = self.ctx.get_settings().get_max_threads()? as usize;
-
-        let (stage_info, files_info, options) = self.bind_attachment(attachment).await?;
+        let settings = self.ctx.get_settings();
+        let (required_source_schema, values_consts) = if let Some(exprs) = expr_or_placeholders {
+            self.prepared_values(exprs, &required_values_schema, settings.clone())
+                .await?
+        } else {
+            (required_values_schema.clone(), vec![])
+        };
 
         // list the files to be copied in binding phase
         // note that, this method(`bind_copy_from_attachment`) are used by
@@ -379,6 +413,7 @@ impl Binder {
         // currently, they do NOT enforce the deduplication detection rules,
         // as the vanilla Copy-Into does.
         // thus, we do not care about the "duplicated_files_detected", just set it to empty vector.
+        let thread_num = settings.get_max_threads()? as usize;
         let files_to_copy = list_stage_files(&stage_info, &files_info, thread_num, None).await?;
         let duplicated_files_detected = vec![];
 
@@ -396,7 +431,7 @@ impl Binder {
             required_values_schema,
             dedup_full_path: false,
             path_prefix: None,
-            values_consts: const_columns,
+            values_consts,
             stage_table_info: StageTableInfo {
                 schema: required_source_schema,
                 files_info,
@@ -405,7 +440,7 @@ impl Binder {
                 duplicated_files_detected,
                 is_select: false,
                 default_exprs: Some(default_values),
-                copy_into_table_options: options,
+                copy_into_table_options,
                 stage_root: "".to_string(),
                 is_variant: false,
             },
@@ -526,17 +561,12 @@ impl Binder {
         Ok(Plan::CopyIntoTable(Box::new(plan)))
     }
 
-    pub(crate) async fn prepared_values_str(
-        &self,
-        values_str: &str,
-        source_schema: &TableSchemaRef,
-    ) -> Result<(TableSchemaRef, Vec<Scalar>)> {
+    pub(crate) async fn parse_values_str(&self, values_str: &str) -> Result<Vec<Expr>> {
         let settings = self.ctx.get_settings();
         let sql_dialect = settings.get_sql_dialect()?;
         let tokens = tokenize_sql(values_str)?;
-        let expr_or_placeholders = parse_values(&tokens, sql_dialect)?;
-        self.prepared_values(expr_or_placeholders, source_schema, settings)
-            .await
+        let values = parse_values(&tokens, sql_dialect)?;
+        Ok(values)
     }
 
     #[async_backtrace::framed]
@@ -660,6 +690,9 @@ pub async fn resolve_stage_location(
 ) -> Result<(StageInfo, String)> {
     // my_named_stage/abc/
     let names: Vec<&str> = location.splitn(2, '/').filter(|v| !v.is_empty()).collect();
+    if names[0] == STAGE_PLACEHOLDER {
+        return Err(ErrorCode::BadArguments("placeholder @_databend_upload should be used in streaming_load handler or replaced in client."));
+    }
 
     let stage = if names[0] == "~" {
         StageInfo::new_user_stage(&ctx.get_current_user()?.name)

src/query/sql/src/planner/binder/ddl/stage.rs

@@ -23,6 +23,7 @@ use databend_common_meta_app::principal::StageInfo;
 use databend_common_storage::init_operator;
 
 use super::super::copy_into_table::resolve_stage_location;
+use crate::binder::insert::STAGE_PLACEHOLDER;
 use crate::binder::location::parse_storage_params_from_uri;
 use crate::binder::Binder;
 use crate::plans::CreateStagePlan;
@@ -59,6 +60,12 @@ impl Binder {
             comments: _,
         } = stmt;
 
+        if stage_name.to_lowercase() == STAGE_PLACEHOLDER {
+            return Err(ErrorCode::InvalidArgument(format!(
+                "Can not create stage with name `{STAGE_PLACEHOLDER}`, which is reserved"
+            )));
+        }
+
         let mut stage_info = match location {
             None => {
                 if stage_name == "~" {

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

@@ -12,14 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::str::FromStr;
 use std::sync::Arc;
 
+use databend_common_ast::ast::CopyIntoTableOptions;
 use databend_common_ast::ast::Identifier;
 use databend_common_ast::ast::InsertSource;
 use databend_common_ast::ast::InsertStmt;
-use databend_common_ast::ast::OnErrorMode;
 use databend_common_ast::ast::Statement;
+use databend_common_catalog::session_type::SessionType;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_expression::DataSchemaRef;
@@ -29,8 +29,10 @@ use databend_common_expression::TableSchema;
 use databend_common_expression::TableSchemaRefExt;
 use databend_common_meta_app::principal::FileFormatOptionsReader;
 use databend_common_meta_app::principal::FileFormatParams;
+use databend_common_storage::StageFilesInfo;
 
 use super::util::TableIdentifier;
+use crate::binder::resolve_stage_location;
 use crate::binder::Binder;
 use crate::normalize_identifier;
 use crate::plans::CopyIntoTableMode;
@@ -42,6 +44,7 @@ use crate::plans::StreamingLoadPlan;
 use crate::BindContext;
 use crate::DefaultExprBinder;
 
+pub const STAGE_PLACEHOLDER: &str = "_databend_load";
 impl Binder {
     pub fn schema_project(
         &self,
@@ -178,44 +181,81 @@ impl Binder {
                 let select_plan = self.bind_statement(bind_context, &statement).await?;
                 Ok(InsertInputSource::SelectPlan(Box::new(select_plan)))
             }
-            InsertSource::StreamingLoad {
+            InsertSource::LoadFile {
                 value,
                 format_options,
-                on_error_mode,
+                location,
             } => {
                 let settings = self.ctx.get_settings();
-                let (required_source_schema, values_consts) = if let Some(value) = value {
-                    self.prepared_values(value, &schema, settings).await?
-                } else {
-                    (schema.clone(), vec![])
-                };
                 let file_format_params = FileFormatParams::try_from_reader(
                     FileFormatOptionsReader::from_ast(&format_options),
                     false,
                 )?;
-                let required_values_schema: DataSchemaRef = Arc::new(schema.clone().into());
-
-                let default_exprs = if file_format_params.need_field_default() {
-                    Some(
-                        DefaultExprBinder::try_new(self.ctx.clone())?
-                            .prepare_default_values(&required_values_schema)?,
-                    )
-                } else {
-                    None
-                };
-                Ok(InsertInputSource::StreamingLoad(StreamingLoadPlan {
-                    file_format: Box::new(file_format_params),
-                    on_error_mode: OnErrorMode::from_str(
-                        &on_error_mode.unwrap_or("abort".to_string()),
-                    )?,
-                    required_values_schema,
-                    values_consts,
-                    block_thresholds: table.get_block_thresholds(),
-                    default_exprs,
-                    // fill it in HTTP handler
-                    receiver: Default::default(),
-                    required_source_schema,
-                }))
+                match location.as_str() {
+                    STAGE_PLACEHOLDER => {
+                        if self.ctx.get_session_type() != SessionType::HTTPStreamingLoad {
+                            return Err(ErrorCode::BadArguments("placeholder @_databend_upload should be used in streaming_load handler or replaced in client."));
+                        }
+                        let (required_source_schema, values_consts) = if let Some(value) = value {
+                            self.prepared_values(value, &schema, settings).await?
+                        } else {
+                            (schema.clone(), vec![])
+                        };
+
+                        let required_values_schema: DataSchemaRef = Arc::new(schema.clone().into());
+
+                        let default_exprs = if file_format_params.need_field_default() {
+                            Some(
+                                DefaultExprBinder::try_new(self.ctx.clone())?
+                                    .prepare_default_values(&required_values_schema)?,
+                            )
+                        } else {
+                            None
+                        };
+                        Ok(InsertInputSource::StreamingLoad(StreamingLoadPlan {
+                            file_format: Box::new(file_format_params),
+                            required_values_schema,
+                            values_consts,
+                            block_thresholds: table.get_block_thresholds(),
+                            default_exprs,
+                            // fill it in HTTP handler
+                            receiver: Default::default(),
+                            required_source_schema,
+                        }))
+                    }
+                    loc => {
+                        let (mut stage_info, path) =
+                            resolve_stage_location(self.ctx.as_ref(), loc).await?;
+                        stage_info.file_format_params = file_format_params;
+
+                        let files_info = StageFilesInfo {
+                            path,
+                            files: None,
+                            pattern: None,
+                        };
+                        let options = CopyIntoTableOptions {
+                            purge: true,
+                            force: true,
+                            ..Default::default()
+                        };
+                        return self
+                            .bind_copy_from_upload(
+                                bind_context,
+                                catalog_name,
+                                database_name,
+                                table_name,
+                                schema,
+                                value,
+                                stage_info,
+                                files_info,
+                                options,
+                                CopyIntoTableMode::Insert {
+                                    overwrite: *overwrite,
+                                },
+                            )
+                            .await;
+                    }
+                }
             }
         };
 

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

@@ -139,7 +139,7 @@ impl Binder {
                 let select_plan = self.bind_statement(bind_context, &statement).await?;
                 Ok(InsertInputSource::SelectPlan(Box::new(select_plan)))
             }
-            InsertSource::StreamingLoad { .. } => Err(ErrorCode::Unimplemented(
+            InsertSource::LoadFile { .. } => Err(ErrorCode::Unimplemented(
                 "Replace with streaming load not supported yet.",
             )),
         };