feat: Re-introduce Streaming Load API (#17921)

* feat: support streaming load API.

Author: youngsofun

Cargo.lock

@@ -21,6 +21,7 @@ use derive_visitor::DriveMut;
 use crate::ast::write_comma_separated_list;
 use crate::ast::write_dot_separated_list;
 use crate::ast::Expr;
+use crate::ast::FileFormatOptions;
 use crate::ast::Hint;
 use crate::ast::Identifier;
 use crate::ast::Query;
@@ -71,9 +72,20 @@ impl Display for InsertStmt {
 
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub enum InsertSource {
-    Values { rows: Vec<Vec<Expr>> },
-    RawValues { rest_str: String, start: usize },
-    Select { query: Box<Query> },
+    Values {
+        rows: Vec<Vec<Expr>>,
+    },
+    RawValues {
+        rest_str: String,
+        start: usize,
+    },
+    Select {
+        query: Box<Query>,
+    },
+    StreamingLoad {
+        format_options: FileFormatOptions,
+        on_error_mode: Option<String>,
+    },
 }
 
 impl Display for InsertSource {
@@ -93,6 +105,17 @@ impl Display for InsertSource {
             }
             InsertSource::RawValues { rest_str, .. } => write!(f, "VALUES {rest_str}"),
             InsertSource::Select { query } => write!(f, "{query}"),
+            InsertSource::StreamingLoad {
+                format_options,
+                on_error_mode,
+            } => {
+                write!(f, " FILE_FORMAT = ({})", format_options)?;
+                write!(
+                    f,
+                    " ON_ERROR = '{}'",
+                    on_error_mode.as_ref().unwrap_or(&"Abort".to_string())
+                )
+            }
         }
     }
 }

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

@@ -21,7 +21,6 @@ use derive_visitor::DriveMut;
 use crate::ast::write_comma_separated_list;
 use crate::ast::write_dot_separated_list;
 use crate::ast::Expr;
-use crate::ast::FileFormatOptions;
 use crate::ast::Hint;
 use crate::ast::Identifier;
 use crate::ast::Query;
@@ -168,12 +167,6 @@ impl Display for MergeIntoStmt {
 
 #[derive(Debug, Clone, PartialEq, Drive, DriveMut)]
 pub enum MutationSource {
-    StreamingV2 {
-        settings: FileFormatOptions,
-        on_error_mode: Option<String>,
-        start: usize,
-    },
-
     Select {
         query: Box<Query>,
         source_alias: TableAlias,
@@ -190,12 +183,6 @@ pub enum MutationSource {
 impl MutationSource {
     pub fn transform_table_reference(&self) -> TableReference {
         match self {
-            Self::StreamingV2 {
-                settings: _,
-                on_error_mode: _,
-                start: _,
-            } => unimplemented!(),
-
             Self::Select {
                 query,
                 source_alias,
@@ -232,19 +219,6 @@ impl MutationSource {
 impl Display for MutationSource {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
         match self {
-            MutationSource::StreamingV2 {
-                settings,
-                on_error_mode,
-                start: _,
-            } => {
-                write!(f, " FILE_FORMAT = ({})", settings)?;
-                write!(
-                    f,
-                    " ON_ERROR = '{}'",
-                    on_error_mode.as_ref().unwrap_or(&"Abort".to_string())
-                )
-            }
-
             MutationSource::Select {
                 query,
                 source_alias,

src/query/ast/src/ast/statements/merge_into.rs

@@ -2988,6 +2988,15 @@ pub fn insert_source(i: Input) -> IResult<InsertSource> {
 //
 // This is a hack to parse large insert statements.
 pub fn raw_insert_source(i: Input) -> IResult<InsertSource> {
+    let streaming = map(
+        rule! {
+           #file_format_clause ~ (ON_ERROR ~ ^"=" ~ ^#ident)?
+        },
+        |(options, on_error_opt)| InsertSource::StreamingLoad {
+            format_options: options,
+            on_error_mode: on_error_opt.map(|v| v.2.to_string()),
+        },
+    );
     let values = map(
         rule! {
             VALUES ~ #rest_str
@@ -3006,21 +3015,11 @@ pub fn raw_insert_source(i: Input) -> IResult<InsertSource> {
     rule!(
         #values
         | #query
+        | #streaming
     )(i)
 }
 
 pub fn mutation_source(i: Input) -> IResult<MutationSource> {
-    let streaming_v2 = map(
-        rule! {
-           #file_format_clause  ~ (ON_ERROR ~ ^"=" ~ ^#ident)? ~  #rest_str
-        },
-        |(options, on_error_opt, (_, start))| MutationSource::StreamingV2 {
-            settings: options,
-            on_error_mode: on_error_opt.map(|v| v.2.to_string()),
-            start,
-        },
-    );
-
     let query = map(rule! {#query ~ #table_alias}, |(query, source_alias)| {
         MutationSource::Select {
             query: Box::new(query),
@@ -3040,8 +3039,7 @@ pub fn mutation_source(i: Input) -> IResult<MutationSource> {
     );
 
     rule!(
-          #streaming_v2
-        | #query
+        #query
         | #source_table
     )(i)
 }

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

@@ -694,6 +694,7 @@ impl PrivilegeAccess {
                 self.check(ctx, plan).await?;
             }
             InsertInputSource::Values(_) => {}
+            InsertInputSource::StreamingLoad { .. } => {}
         }
         Ok(())
     }

src/query/service/src/interpreters/access/privilege_access.rs

@@ -33,6 +33,7 @@ use databend_common_sql::plans::InsertInputSource;
 use databend_common_sql::plans::InsertValue;
 use databend_common_sql::plans::Plan;
 use databend_common_sql::NameResolutionContext;
+use databend_common_storages_stage::build_streaming_load_pipeline;
 use log::info;
 
 use crate::interpreters::common::check_deduplicate_label;
@@ -247,6 +248,25 @@ impl Interpreter for InsertInterpreter {
 
                 return Ok(build_res);
             }
+            InsertInputSource::StreamingLoad {
+                file_format,
+                on_error_mode,
+                schema,
+                default_exprs,
+                block_thresholds,
+                receiver,
+            } => {
+                build_streaming_load_pipeline(
+                    self.ctx.clone(),
+                    &mut build_res.main_pipeline,
+                    file_format,
+                    receiver.clone(),
+                    schema.clone(),
+                    default_exprs.clone(),
+                    *block_thresholds,
+                    on_error_mode.clone(),
+                )?;
+            }
         };
 
         PipelineBuilder::build_append2table_with_commit_pipeline(

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

@@ -432,6 +432,9 @@ impl ReplaceInterpreter {
                 }
                 _ => unreachable!("plan in InsertInputSource::Stag must be CopyIntoTable"),
             },
+            InsertInputSource::StreamingLoad { .. } => {
+                unreachable!("replace with streaming not supported yet")
+            }
         }
     }
 

src/query/service/src/interpreters/interpreter_replace.rs

@@ -96,6 +96,7 @@ pub enum EndpointKind {
     SystemInfo,
     Catalog,
     Metadata,
+    StreamingLoad,
 }
 
 impl EndpointKind {
@@ -122,6 +123,7 @@ impl EndpointKind {
             | EndpointKind::Logout
             | EndpointKind::SystemInfo
             | EndpointKind::HeartBeat
+            | EndpointKind::StreamingLoad
             | EndpointKind::UploadToStage
             | EndpointKind::Metadata
             | EndpointKind::Catalog => {

src/query/service/src/servers/http/middleware/session.rs

@@ -78,6 +78,7 @@ use crate::servers::http::v1::query::blocks_serializer::BlocksSerializer;
 use crate::servers::http::v1::query::Progresses;
 use crate::servers::http::v1::refresh_handler;
 use crate::servers::http::v1::roles::list_roles_handler;
+use crate::servers::http::v1::streaming_load_handler;
 use crate::servers::http::v1::upload_to_stage;
 use crate::servers::http::v1::users::create_user_handler;
 use crate::servers::http::v1::users::list_users_handler;
@@ -703,6 +704,11 @@ pub fn query_route() -> Route {
             get(list_users_handler).post(create_user_handler),
             EndpointKind::Metadata,
         ),
+        (
+            "/streaming_load",
+            put(streaming_load_handler),
+            EndpointKind::StreamingLoad,
+        ),
         ("/roles", get(list_roles_handler), EndpointKind::Metadata),
     ];
 
@@ -742,14 +748,14 @@ fn query_id_to_trace_id(query_id: &str) -> TraceId {
 /// The HTTP query endpoints are expected to be responses within 60 seconds.
 /// If it exceeds far from 60 seconds, there might be something wrong, we should
 /// log it.
-struct SlowRequestLogTracker {
+pub(crate) struct SlowRequestLogTracker {
     started_at: std::time::Instant,
     method: String,
     uri: String,
 }
 
 impl SlowRequestLogTracker {
-    fn new(ctx: &HttpQueryContext) -> Self {
+    pub(crate) fn new(ctx: &HttpQueryContext) -> Self {
         Self {
             started_at: std::time::Instant::now(),
             method: ctx.http_method.clone(),
@@ -776,7 +782,11 @@ impl Drop for SlowRequestLogTracker {
 
 // get_http_tracing_span always return a valid span for tracing
 // it will try to decode w3 traceparent and if empty or failed, it will create a new root span and throw a warning
-fn get_http_tracing_span(name: &'static str, ctx: &HttpQueryContext, query_id: &str) -> Span {
+pub(crate) fn get_http_tracing_span(
+    name: &'static str,
+    ctx: &HttpQueryContext,
+    query_id: &str,
+) -> Span {
     if let Some(parent) = ctx.trace_parent.as_ref() {
         let trace = parent.as_str();
         match SpanContext::decode_w3c_traceparent(trace) {

src/query/service/src/servers/http/v1/http_query_handlers.rs

@@ -19,6 +19,7 @@ mod query;
 pub mod roles;
 mod session;
 mod stage;
+mod streaming_load;
 pub mod users;
 mod verify;
 
@@ -48,6 +49,8 @@ pub use session::ClientSessionManager;
 pub(crate) use session::SessionClaim;
 pub use stage::upload_to_stage;
 pub use stage::UploadToStageResponse;
+pub use streaming_load::streaming_load_handler;
+pub use streaming_load::LoadResponse;
 pub use users::create_user_handler;
 pub use users::list_users_handler;
 pub use verify::verify_handler;