Merge pull request #9247 from youngsofun/stage

BohuTANG · web-flow · commit d6eca280e2fe · 2022-12-14T22:23:39.000+08:00
feat(stage): support select from URI.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/common/storage/src/lib.rs b/src/common/storage/src/lib.rs
@@ -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;
diff --git a/src/query/ast/src/ast/query.rs b/src/query/ast/src/ast/query.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>>,
     },
diff --git a/src/query/ast/src/ast/statements/copy.rs b/src/query/ast/src/ast/statements/copy.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),
diff --git a/src/query/ast/src/parser/query.rs b/src/query/ast/src/parser/query.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,
diff --git a/src/query/sql/Cargo.toml b/src/query/sql/Cargo.toml
@@ -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"
diff --git a/src/query/sql/src/planner/binder/copy.rs b/src/query/sql/src/planner/binder/copy.rs
@@ -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;
diff --git a/src/query/sql/src/planner/binder/ddl/stage.rs b/src/query/sql/src/planner/binder/ddl/stage.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;
diff --git a/src/query/sql/src/planner/binder/ddl/table.rs b/src/query/sql/src/planner/binder/ddl/table.rs
@@ -40,6 +40,7 @@ use common_ast::ast::Statement;
 use common_ast::ast::TableReference;
 use common_ast::ast::TruncateTableStmt;
 use common_ast::ast::UndropTableStmt;
+use common_ast::ast::UriLocation;
 use common_ast::parser::parse_sql;
 use common_ast::parser::tokenize_sql;
 use common_ast::walk_expr_mut;
@@ -54,15 +55,14 @@ use common_datavalues::TypeFactory;
 use common_datavalues::Vu8;
 use common_exception::ErrorCode;
 use common_exception::Result;
-use common_storage::parse_uri_location;
 use common_storage::DataOperator;
-use common_storage::UriLocation;
 use common_storages_table_meta::table::is_reserved_opt_key;
 use common_storages_table_meta::table::OPT_KEY_DATABASE_ID;
 use common_storages_view::view_table::QUERY;
 use common_storages_view::view_table::VIEW_ENGINE;
 use tracing::debug;
 
+use crate::binder::location::parse_uri_location;
 use crate::binder::scalar::ScalarBinder;
 use crate::binder::Binder;
 use crate::binder::Visibility;
diff --git a/src/query/sql/src/planner/binder/location.rs b/src/query/sql/src/planner/binder/location.rs
@@ -12,35 +12,27 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::BTreeMap;
 use std::io::Error;
 use std::io::ErrorKind;
 use std::io::Result;
 
 use anyhow::anyhow;
+use common_ast::ast::UriLocation;
+use common_storage::StorageAzblobConfig;
+use common_storage::StorageFsConfig;
+use common_storage::StorageFtpConfig;
+use common_storage::StorageGcsConfig;
+use common_storage::StorageHttpConfig;
+use common_storage::StorageIpfsConfig;
+use common_storage::StorageOssConfig;
+use common_storage::StorageParams;
+use common_storage::StorageS3Config;
+use common_storage::STORAGE_GCS_DEFAULT_ENDPOINT;
+use common_storage::STORAGE_IPFS_DEFAULT_ENDPOINT;
+use common_storage::STORAGE_S3_DEFAULT_ENDPOINT;
 use opendal::Scheme;
 use percent_encoding::percent_decode_str;
 
-use crate::config::StorageHttpConfig;
-use crate::config::StorageIpfsConfig;
-use crate::config::StorageOssConfig;
-use crate::config::STORAGE_IPFS_DEFAULT_ENDPOINT;
-use crate::config::STORAGE_S3_DEFAULT_ENDPOINT;
-use crate::StorageAzblobConfig;
-use crate::StorageFsConfig;
-use crate::StorageParams;
-use crate::StorageS3Config;
-use crate::STORAGE_GCS_DEFAULT_ENDPOINT;
-
-#[derive(Clone, Debug)]
-pub struct UriLocation {
-    pub protocol: String,
-    pub name: String,
-    pub path: String,
-    /// connection should carry all connection related options for storage.
-    pub connection: BTreeMap<String, String>,
-}
-
 /// secure_omission will fix omitted endpoint url schemes into 'https://'
 #[inline]
 fn secure_omission(endpoint: String) -> String {
@@ -85,7 +77,7 @@ pub fn parse_uri_location(l: &UriLocation) -> Result<(StorageParams, String)> {
                 root: root.to_string(),
             })
         }
-        Scheme::Ftp => StorageParams::Ftp(crate::StorageFtpConfig {
+        Scheme::Ftp => StorageParams::Ftp(StorageFtpConfig {
             endpoint: if !l.protocol.is_empty() {
                 format!("{}://{}", l.protocol, l.name)
             } else {
@@ -102,7 +94,7 @@ pub fn parse_uri_location(l: &UriLocation) -> Result<(StorageParams, String)> {
                 .get("endpoint_url")
                 .cloned()
                 .unwrap_or_else(|| STORAGE_GCS_DEFAULT_ENDPOINT.to_string());
-            StorageParams::Gcs(crate::StorageGcsConfig {
+            StorageParams::Gcs(StorageGcsConfig {
                 endpoint_url: secure_omission(endpoint),
                 bucket: l.name.clone(),
                 root: l.path.clone(),
diff --git a/src/query/sql/src/planner/binder/mod.rs b/src/query/sql/src/planner/binder/mod.rs
@@ -27,6 +27,7 @@ mod insert;
 mod join;
 mod kill;
 mod limit;
+mod location;
 mod presign;
 mod project;
 mod scalar;
@@ -43,5 +44,6 @@ pub use aggregate::AggregateInfo;
 pub use bind_context::*;
 pub use binder::Binder;
 pub use builders::*;
+pub use location::parse_uri_location;
 pub use scalar::ScalarBinder;
 pub use scalar_common::*;
diff --git a/src/query/sql/src/planner/binder/table.rs b/src/query/sql/src/planner/binder/table.rs
@@ -17,6 +17,7 @@ use std::default::Default;
 use std::path::Path;
 use std::sync::Arc;
 
+use common_ast::ast::FileLocation;
 use common_ast::ast::Indirection;
 use common_ast::ast::SelectStmt;
 use common_ast::ast::SelectTarget;
@@ -36,19 +37,21 @@ use common_catalog::table::ColumnStatistics;
 use common_catalog::table::NavigationPoint;
 use common_catalog::table::Table;
 use common_catalog::table_function::TableFunction;
+use common_config::GlobalConfig;
 use common_datavalues::prelude::*;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_meta_types::FileFormatOptions;
 use common_meta_types::StageFileCompression;
 use common_meta_types::StageFileFormatType;
+use common_meta_types::UserStageInfo;
 use common_pipeline_sources::processors::sources::input_formats::InputContext;
 use common_storages_stage::get_first_file;
 use common_storages_stage::StageTable;
 use common_storages_view::view_table::QUERY;
-use opendal::layers::SubdirLayer;
 
 use crate::binder::copy::parse_stage_location_v2;
+use crate::binder::location::parse_uri_location;
 use crate::binder::scalar::ScalarBinder;
 use crate::binder::Binder;
 use crate::binder::ColumnBinding;
@@ -296,10 +299,25 @@ impl<'a> Binder {
                     None
                 };
 
-                let (mut user_stage_info, path) =
-                    parse_stage_location_v2(&self.ctx, &location.name, &location.path).await?;
-                let op = self.ctx.get_data_operator()?.operator();
-                let op = op.layer(SubdirLayer::new(&user_stage_info.stage_prefix()));
+                let (mut user_stage_info, path) = match location {
+                    FileLocation::Stage(location) => {
+                        parse_stage_location_v2(&self.ctx, &location.name, &location.path).await?
+                    }
+                    FileLocation::Uri(location) => {
+                        let (storage_params, path) = parse_uri_location(location)?;
+                        if !storage_params.is_secure()
+                            && !GlobalConfig::instance().storage.allow_insecure
+                        {
+                            return Err(ErrorCode::StorageInsecure(
+                                "copy from insecure storage is not allowed",
+                            ));
+                        }
+                        let stage_info = UserStageInfo::new_external_stage(storage_params, &path);
+                        (stage_info, path)
+                    }
+                };
+
+                let op = StageTable::get_op(&self.ctx, &user_stage_info)?;
 
                 let first_file = if files.is_empty() {
                     let file = get_first_file(&op, &path).await?;
diff --git a/src/query/sql/src/planner/mod.rs b/src/query/sql/src/planner/mod.rs
@@ -19,7 +19,7 @@ mod metadata;
 mod planner;
 mod semantic;
 
-pub(crate) mod binder;
+pub mod binder;
 pub mod optimizer;
 pub mod plans;
 
diff --git a/src/query/sql/tests/location.rs b/src/query/sql/tests/location.rs
@@ -15,7 +15,8 @@
 use std::collections::BTreeMap;
 use std::io::Result;
 
-use common_storage::parse_uri_location;
+use common_ast::ast::UriLocation;
+use common_sql::planner::binder::parse_uri_location;
 use common_storage::StorageFsConfig;
 use common_storage::StorageFtpConfig;
 use common_storage::StorageGcsConfig;
@@ -24,7 +25,6 @@ use common_storage::StorageIpfsConfig;
 use common_storage::StorageOssConfig;
 use common_storage::StorageParams;
 use common_storage::StorageS3Config;
-use common_storage::UriLocation;
 use common_storage::STORAGE_GCS_DEFAULT_ENDPOINT;
 use common_storage::STORAGE_IPFS_DEFAULT_ENDPOINT;
 use common_storage::STORAGE_S3_DEFAULT_ENDPOINT;
diff --git a/tests/suites/0_stateless/20+_others/20_0011_select_stage.result b/tests/suites/0_stateless/20+_others/20_0011_select_stage.result
diff --git a/tests/suites/0_stateless/20+_others/20_0011_select_stage.sh b/tests/suites/0_stateless/20+_others/20_0011_select_stage.sh
diff --git a/tests/suites/1_stateful/00_copy/00_0004_select_stage.result b/tests/suites/1_stateful/00_copy/00_0004_select_stage.result
diff --git a/tests/suites/1_stateful/00_copy/00_0004_select_stage.sh b/tests/suites/1_stateful/00_copy/00_0004_select_stage.sh
