Merge branch 'main' into fix-ci

Author: mergify[bot]

Cargo.lock

@@ -55,10 +55,6 @@ pub use operator::init_operator;
 pub use operator::CacheOperator;
 pub use operator::DataOperator;
 
-mod location;
-pub use location::parse_uri_location;
-pub use location::UriLocation;
-
 mod metrics;
 pub use metrics::StorageMetrics;
 pub use metrics::StorageMetricsLayer;

src/common/storage/src/lib.rs

@@ -18,8 +18,8 @@ use std::fmt::Formatter;
 use crate::ast::write_comma_separated_list;
 use crate::ast::write_period_separated_list;
 use crate::ast::Expr;
+use crate::ast::FileLocation;
 use crate::ast::Identifier;
-use crate::ast::StageLocation;
 use crate::parser::token::Token;
 
 /// Root node of a query tree
@@ -178,7 +178,7 @@ pub enum TableReference<'a> {
     },
     Stage {
         span: &'a [Token<'a>],
-        location: StageLocation,
+        location: FileLocation,
         files: Vec<String>,
         alias: Option<TableAlias<'a>>,
     },

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

@@ -213,6 +213,25 @@ impl Display for StageLocation {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum FileLocation {
+    Stage(StageLocation),
+    Uri(UriLocation),
+}
+
+impl Display for FileLocation {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            FileLocation::Uri(loc) => {
+                write!(f, "{}", loc)
+            }
+            FileLocation::Stage(loc) => {
+                write!(f, "{}", loc)
+            }
+        }
+    }
+}
+
 pub enum CopyOption {
     Files(Vec<String>),
     Pattern(String),

src/query/ast/src/ast/statements/copy.rs

@@ -22,6 +22,7 @@ use pratt::PrattParser;
 use pratt::Precedence;
 
 use super::stage::stage_location;
+use super::stage::uri_location;
 use crate::ast::*;
 use crate::input::Input;
 use crate::input::WithSpan;
@@ -259,7 +260,7 @@ pub enum TableReferenceElement<'a> {
     JoinCondition(JoinCondition<'a>),
     Group(TableReference<'a>),
     Stage {
-        location: StageLocation,
+        location: FileLocation,
         files: Vec<String>,
         alias: Option<TableAlias<'a>>,
     },
@@ -325,9 +326,27 @@ pub fn table_reference_element(i: Input) -> IResult<WithSpan<TableReferenceEleme
         |(_, table_ref, _)| TableReferenceElement::Group(table_ref),
     );
 
+    let stage_location = |i| {
+        map_res(
+            rule! {
+                #stage_location
+            },
+            |v| Ok(FileLocation::Stage(v)),
+        )(i)
+    };
+
+    let uri_location = |i| {
+        map_res(
+            rule! {
+                #uri_location
+            },
+            |v| Ok(FileLocation::Uri(v)),
+        )(i)
+    };
+
     let aliased_stage = map(
         rule! {
-            #stage_location ~ ( FILES ~ "=" ~ "(" ~ #comma_separated_list0(literal_string) ~ ")")? ~ #table_alias?
+            (#stage_location | #uri_location) ~ ( FILES ~ "=" ~ "(" ~ #comma_separated_list0(literal_string) ~ ")")? ~ #table_alias?
         },
         |(location, files, alias)| TableReferenceElement::Stage {
             location,

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

@@ -469,4 +469,9 @@ impl ProgressReporter for ContextProgressReporter {
             convert_byte_size((progress.bytes as f64) / (seconds)),
         )
     }
+
+    fn affected_rows(&self) -> u64 {
+        let progress = self.context.get_scan_progress_value();
+        progress.rows as u64
+    }
 }

src/query/service/src/servers/mysql/mysql_interactive_worker.rs

@@ -38,6 +38,7 @@ use tracing::error;
 /// "Read x rows, y MiB in z sec., A million rows/sec., B MiB/sec."
 pub trait ProgressReporter {
     fn progress_info(&self) -> String;
+    fn affected_rows(&self) -> u64;
 }
 
 pub struct QueryResult {
@@ -122,7 +123,16 @@ impl<'a, W: AsyncWrite + Send + Unpin> DFQueryResultWriter<'a, W> {
                 }
             }
 
-            dataset_writer.completed(OkResponse::default()).await?;
+            let affected_rows = query_result
+                .extra_info
+                .map(|r| r.affected_rows())
+                .unwrap_or_default();
+            dataset_writer
+                .completed(OkResponse {
+                    affected_rows,
+                    ..Default::default()
+                })
+                .await?;
             return Ok(());
         }
 

src/query/service/src/servers/mysql/writers/query_result_writer.rs

@@ -48,6 +48,7 @@ common-users = { path = "../users" }
 
 # Crates.io dependencies
 ahash = { version = "0.8.2", features = ["no-rng"] }
+anyhow = { workspace = true }
 async-channel = "1.7.1"
 async-recursion = "1.0.0"
 async-stream = "0.3.3"
@@ -61,6 +62,7 @@ chrono-tz = { workspace = true }
 dashmap = "5.4"
 futures = "0.3.24"
 futures-util = "0.3.24"
+globiter = "0.1"
 headers = "0.3.8"
 http = "0.2.8"
 itertools = "0.10.5"
@@ -73,6 +75,7 @@ opendal = { version = "0.22", features = ["layers-tracing", "layers-metrics", "c
 opensrv-mysql = { git = "https://github.com/datafuselabs/opensrv", rev = "5e37788" }
 openssl = { version = "0.10.41", features = ["vendored"] }
 parking_lot = "0.12.1"
+percent-encoding = "2"
 petgraph = "0.6.2"
 poem = { version = "1", features = ["rustls", "multipart", "compression"] }
 primitive-types = "0.12.0"

src/query/sql/Cargo.toml

@@ -20,6 +20,7 @@ use common_ast::ast::CopyStmt;
 use common_ast::ast::CopyUnit;
 use common_ast::ast::Query;
 use common_ast::ast::Statement;
+use common_ast::ast::UriLocation;
 use common_ast::parser::parse_sql;
 use common_ast::parser::tokenize_sql;
 use common_ast::Backtrace;
@@ -35,11 +36,10 @@ use common_exception::Result;
 use common_meta_types::FileFormatOptions;
 use common_meta_types::StageFileFormatType;
 use common_meta_types::UserStageInfo;
-use common_storage::parse_uri_location;
-use common_storage::UriLocation;
 use common_users::UserApiProvider;
 use tracing::debug;
 
+use crate::binder::location::parse_uri_location;
 use crate::binder::Binder;
 use crate::normalize_identifier;
 use crate::plans::CopyPlanV2;

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

@@ -15,15 +15,15 @@
 use std::str::FromStr;
 
 use common_ast::ast::CreateStageStmt;
+use common_ast::ast::UriLocation;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_meta_types::OnErrorMode;
 use common_meta_types::UserStageInfo;
-use common_storage::parse_uri_location;
-use common_storage::UriLocation;
 
 use super::super::copy::parse_copy_file_format_options;
 use super::super::copy::parse_stage_location;
+use crate::binder::location::parse_uri_location;
 use crate::binder::Binder;
 use crate::plans::CreateStagePlan;
 use crate::plans::ListPlan;