diff --git a/Cargo.lock b/Cargo.lock index f8cf79b66c7..76ec94f73e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2243,6 +2243,7 @@ name = "datanode" version = "0.1.0" dependencies = [ "api", + "async-compat", "async-stream", "async-trait", "axum", @@ -2267,8 +2268,10 @@ dependencies = [ "datafusion-expr", "datatypes", "futures", + "futures-util", "humantime-serde", "hyper", + "lazy_static", "log-store", "meta-client", "meta-srv", @@ -2278,6 +2281,7 @@ dependencies = [ "pin-project", "prost", "query", + "regex", "script", "serde", "serde_json", @@ -2297,6 +2301,7 @@ dependencies = [ "tonic", "tower", "tower-http", + "url", ] [[package]] @@ -3205,9 +3210,9 @@ dependencies = [ [[package]] name = "http" -version = "0.2.8" +version = "0.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75f43d41e26995c17e71ee126451dd3941010b0514a81a9d11f3b341debc2399" +checksum = "bd6effc99afb63425aff9b05836f029929e345a6148a14b7ecd5ab67af944482" dependencies = [ "bytes", "fnv", diff --git a/src/datanode/Cargo.toml b/src/datanode/Cargo.toml index 8e7b6497c25..af37f9a8f3d 100644 --- a/src/datanode/Cargo.toml +++ b/src/datanode/Cargo.toml @@ -32,6 +32,7 @@ datafusion-common.workspace = true datafusion-expr.workspace = true datatypes = { path = "../datatypes" } futures = "0.3" +futures-util.workspace = true hyper = { version = "0.14", features = ["full"] } humantime-serde = "1.1" log-store = { path = "../log-store" } @@ -60,6 +61,10 @@ tokio-stream = { version = "0.1", features = ["net"] } tonic.workspace = true tower = { version = "0.4", features = ["full"] } tower-http = { version = "0.3", features = ["full"] } +regex = "1.6" +lazy_static = "1.4" +async-compat = "0.2" +url = "2.3.1" [dev-dependencies] axum-test-helper = { git = "https://github.com/sunng87/axum-test-helper.git", branch = "patch-1" } diff --git a/src/datanode/src/error.rs b/src/datanode/src/error.rs index c7e583c46d3..2eeabb12e6b 100644 --- a/src/datanode/src/error.rs +++ b/src/datanode/src/error.rs @@ -21,6 +21,7 @@ use datafusion::parquet; use datatypes::prelude::ConcreteDataType; use storage::error::Error as StorageError; use table::error::Error as TableError; +use url::ParseError; use crate::datanode::ObjectStoreConfig; @@ -205,6 +206,30 @@ pub enum Error { #[snafu(display("Invalid SQL, error: {}", msg))] InvalidSql { msg: String }, + #[snafu(display("Invalid url: {}, error :{}", url, source))] + InvalidUrl { url: String, source: ParseError }, + + #[snafu(display("Invalid filepath: {}", path))] + InvalidPath { path: String }, + + #[snafu(display("Invalid connection: {}", msg))] + InvalidConnection { msg: String }, + + #[snafu(display("Unsupported backend protocol: {}", protocol))] + UnsupportedBackendProtocol { protocol: String }, + + #[snafu(display("Failed to regex, source: {}", source))] + BuildRegex { + backtrace: Backtrace, + source: regex::Error, + }, + + #[snafu(display("Failed to parse the data, source: {}", source))] + ParseDataTypes { + #[snafu(backtrace)] + source: common_recordbatch::error::Error, + }, + #[snafu(display("Not support SQL, error: {}", msg))] NotSupportSql { msg: String }, @@ -377,6 +402,22 @@ pub enum Error { source: common_query::error::Error, }, + #[snafu(display( + "File Schema mismatch, expected table schema: {} but found :{}", + table_schema, + file_schema + ))] + InvalidSchema { + table_schema: String, + file_schema: String, + }, + + #[snafu(display("Failed to read parquet file, source: {}", source))] + ReadParquet { + source: parquet::errors::ParquetError, + backtrace: Backtrace, + }, + #[snafu(display("Failed to write parquet file, source: {}", source))] WriteParquet { source: parquet::errors::ParquetError, @@ -389,6 +430,19 @@ pub enum Error { backtrace: Backtrace, }, + #[snafu(display("Failed to build parquet record batch stream, source: {}", source))] + BuildParquetRecordBatchStream { + backtrace: Backtrace, + source: parquet::errors::ParquetError, + }, + + #[snafu(display("Failed to read object in path: {}, source: {}", path, source))] + ReadObject { + path: String, + backtrace: Backtrace, + source: object_store::Error, + }, + #[snafu(display("Failed to write object into path: {}, source: {}", path, source))] WriteObject { path: String, @@ -396,6 +450,13 @@ pub enum Error { source: object_store::Error, }, + #[snafu(display("Failed to lists object in path: {}, source: {}", path, source))] + ListObjects { + path: String, + backtrace: Backtrace, + source: object_store::Error, + }, + #[snafu(display("Unrecognized table option: {}", source))] UnrecognizedTableOption { #[snafu(backtrace)] @@ -456,6 +517,11 @@ impl ErrorExt for Error { ColumnValuesNumberMismatch { .. } | ColumnTypeMismatch { .. } | InvalidSql { .. } + | InvalidUrl { .. } + | InvalidPath { .. } + | InvalidConnection { .. } + | UnsupportedBackendProtocol { .. } + | BuildRegex { .. } | NotSupportSql { .. } | KeyColumnNotFound { .. } | IllegalPrimaryKeysDef { .. } @@ -481,13 +547,19 @@ impl ErrorExt for Error { | RenameTable { .. } | Catalog { .. } | MissingRequiredField { .. } + | BuildParquetRecordBatchStream { .. } + | InvalidSchema { .. } + | ParseDataTypes { .. } | IncorrectInternalState { .. } => StatusCode::Internal, BuildBackend { .. } | InitBackend { .. } + | ReadParquet { .. } | WriteParquet { .. } | PollStream { .. } - | WriteObject { .. } => StatusCode::StorageUnavailable, + | ReadObject { .. } + | WriteObject { .. } + | ListObjects { .. } => StatusCode::StorageUnavailable, OpenLogStore { source } => source.status_code(), StartScriptManager { source } => source.status_code(), OpenStorageEngine { source } => source.status_code(), diff --git a/src/datanode/src/instance/sql.rs b/src/datanode/src/instance/sql.rs index 3fe5b1ada6d..45019d26757 100644 --- a/src/datanode/src/instance/sql.rs +++ b/src/datanode/src/instance/sql.rs @@ -33,7 +33,9 @@ use sql::statements::copy::CopyTable; use sql::statements::statement::Statement; use sql::statements::tql::Tql; use table::engine::TableReference; -use table::requests::{CopyTableRequest, CreateDatabaseRequest, DropTableRequest}; +use table::requests::{ + CopyTableFromRequest, CopyTableRequest, CreateDatabaseRequest, DropTableRequest, +}; use crate::error::{self, BumpTableIdSnafu, ExecuteSqlSnafu, Result, TableIdProviderNotFoundSnafu}; use crate::instance::Instance; @@ -202,7 +204,21 @@ impl Instance { .execute(SqlRequest::CopyTable(req), query_ctx) .await } - CopyTable::From(_) => todo!(), + CopyTable::From(copy_table) => { + let (catalog_name, schema_name, table_name) = + table_idents_to_full_name(©_table.table_name, query_ctx.clone())?; + let req = CopyTableFromRequest { + catalog_name, + schema_name, + table_name, + connection: copy_table.connection, + pattern: copy_table.pattern, + from: copy_table.from, + }; + self.sql_handler + .execute(SqlRequest::CopyTableFrom(req), query_ctx) + .await + } }, QueryStatement::Sql(Statement::Tql(tql)) => self.execute_tql(tql, query_ctx).await, } diff --git a/src/datanode/src/sql.rs b/src/datanode/src/sql.rs index d6975bce049..c0a9f5f4e0c 100644 --- a/src/datanode/src/sql.rs +++ b/src/datanode/src/sql.rs @@ -33,6 +33,7 @@ use crate::instance::sql::table_idents_to_full_name; mod alter; mod copy_table; +mod copy_table_from; mod create; mod delete; mod drop_table; @@ -51,6 +52,7 @@ pub enum SqlRequest { Explain(Box), Delete(Delete), CopyTable(CopyTableRequest), + CopyTableFrom(CopyTableFromRequest), } // Handler to execute SQL except query @@ -92,6 +94,7 @@ impl SqlHandler { SqlRequest::DropTable(req) => self.drop_table(req).await, SqlRequest::Delete(req) => self.delete(query_ctx.clone(), req).await, SqlRequest::CopyTable(req) => self.copy_table(req).await, + SqlRequest::CopyTableFrom(req) => self.copy_table_from(req).await, SqlRequest::ShowDatabases(req) => { show_databases(req, self.catalog_manager.clone()).context(ExecuteSqlSnafu) } diff --git a/src/datanode/src/sql/copy_table_from.rs b/src/datanode/src/sql/copy_table_from.rs new file mode 100644 index 00000000000..fb938c7d112 --- /dev/null +++ b/src/datanode/src/sql/copy_table_from.rs @@ -0,0 +1,451 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; + +use async_compat::CompatExt; +use common_query::Output; +use common_recordbatch::error::DataTypesSnafu; +use datafusion::parquet::arrow::ParquetRecordBatchStreamBuilder; +use datatypes::arrow::record_batch::RecordBatch; +use datatypes::vectors::{Helper, VectorRef}; +use futures::future; +use futures_util::TryStreamExt; +use lazy_static::lazy_static; +use object_store::services::{Fs, S3}; +use object_store::{Object, ObjectStore, ObjectStoreBuilder}; +use regex::Regex; +use snafu::{ensure, ResultExt}; +use table::engine::TableReference; +use table::requests::{CopyTableFromRequest, InsertRequest}; +use tokio::io::BufReader; +use url::{ParseError, Url}; + +use crate::error::{self, Result}; +use crate::sql::SqlHandler; + +lazy_static! { + static ref SCHEMA_PATTERN: Regex = Regex::new(r"^\w*://").unwrap(); +} + +const S3_SCHEMA: &str = "S3"; +const ENDPOINT_URL: &str = "ENDPOINT_URL"; +const ACCESS_KEY_ID: &str = "ACCESS_KEY_ID"; +const SECRET_ACCESS_KEY: &str = "SECRET_ACCESS_KEY"; +const SESSION_TOKEN: &str = "SESSION_TOKEN"; +const REGION: &str = "REGION"; +const ENABLE_VIRTUAL_HOST_STYLE: &str = "ENABLE_VIRTUAL_HOST_STYLE"; + +impl SqlHandler { + pub(crate) async fn copy_table_from(&self, req: CopyTableFromRequest) -> Result { + let table_ref = TableReference { + catalog: &req.catalog_name, + schema: &req.schema_name, + table: &req.table_name, + }; + let table = self.get_table(&table_ref)?; + + let datasource = DataSource::new(&req.from, req.pattern, req.connection)?; + + let objects = datasource.list().await?; + + let mut buf: Vec = Vec::new(); + + for obj in objects.iter() { + let reader = obj.reader().await.context(error::ReadObjectSnafu { + path: &obj.path().to_string(), + })?; + + let buf_reader = BufReader::new(reader.compat()); + + let builder = ParquetRecordBatchStreamBuilder::new(buf_reader) + .await + .context(error::ReadParquetSnafu)?; + + ensure!( + builder.schema() == table.schema().arrow_schema(), + error::InvalidSchemaSnafu { + table_schema: table.schema().arrow_schema().to_string(), + file_schema: (*(builder.schema())).to_string() + } + ); + + let stream = builder + .build() + .context(error::BuildParquetRecordBatchStreamSnafu)?; + + let chunk = stream + .try_collect::>() + .await + .context(error::ReadParquetSnafu)?; + + buf.extend(chunk.into_iter()); + } + + let fields = table + .schema() + .arrow_schema() + .fields() + .iter() + .map(|f| f.name().to_string()) + .collect::>(); + + // Vec + let column_chunks = buf + .into_iter() + .map(|c| Helper::try_into_vectors(c.columns()).context(DataTypesSnafu)) + .collect::>(); + + let mut futs = Vec::with_capacity(column_chunks.len()); + + for column_chunk in column_chunks.into_iter() { + let column_chunk = column_chunk.context(error::ParseDataTypesSnafu)?; + let columns_values = fields + .iter() + .cloned() + .zip(column_chunk.into_iter()) + .collect::>(); + + futs.push(table.insert(InsertRequest { + catalog_name: req.catalog_name.to_string(), + schema_name: req.schema_name.to_string(), + table_name: req.table_name.to_string(), + columns_values, + //TODO: support multi-regions + region_number: 0, + })) + } + + let result = futures::future::try_join_all(futs) + .await + .context(error::InsertSnafu { + table_name: req.table_name.to_string(), + })?; + + Ok(Output::AffectedRows(result.iter().sum())) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +enum Source { + Filename(String), + Dir, +} + +struct DataSource { + object_store: ObjectStore, + source: Source, + path: String, + regex: Option, +} + +impl DataSource { + fn from_path(url: &str, regex: Option) -> Result { + let result = if url.ends_with('/') { + Url::from_directory_path(url) + } else { + Url::from_file_path(url) + }; + + match result { + Ok(url) => { + let path = url.path(); + + let (path, filename) = DataSource::find_dir_and_filename(path); + + let source = if let Some(filename) = filename { + Source::Filename(filename) + } else { + Source::Dir + }; + + let accessor = Fs::default() + .root(&path) + .build() + .context(error::BuildBackendSnafu)?; + + Ok(DataSource { + object_store: ObjectStore::new(accessor).finish(), + source, + path, + regex, + }) + } + Err(()) => error::InvalidPathSnafu { + path: url.to_string(), + } + .fail(), + } + } + + fn build_s3_backend( + host: Option<&str>, + path: &str, + connection: HashMap, + ) -> Result { + let mut builder = S3::default(); + + builder.root(path); + + if let Some(bucket) = host { + builder.bucket(bucket); + } + + if let Some(endpoint) = connection.get(ENDPOINT_URL) { + builder.endpoint(endpoint); + } + + if let Some(region) = connection.get(REGION) { + builder.region(region); + } + + if let Some(key_id) = connection.get(ACCESS_KEY_ID) { + builder.access_key_id(key_id); + } + + if let Some(key) = connection.get(SECRET_ACCESS_KEY) { + builder.secret_access_key(key); + } + + if let Some(session_token) = connection.get(SESSION_TOKEN) { + builder.security_token(session_token); + } + + if let Some(enable_str) = connection.get(ENABLE_VIRTUAL_HOST_STYLE) { + let enable = enable_str.as_str().parse::(); + + match enable { + Ok(true) => { + builder.enable_virtual_host_style(); + } + Err(_) => { + return Err(error::Error::InvalidConnection { + msg: format!("failed to parse the {}", ENABLE_VIRTUAL_HOST_STYLE), + }) + } + _ => (), + } + } + + let accessor = builder.build().context(error::BuildBackendSnafu)?; + + Ok(ObjectStore::new(accessor).finish()) + } + + fn from_url( + url: Url, + regex: Option, + connection: HashMap, + ) -> Result { + let host = url.host_str(); + + let path = url.path(); + + let schema = url.scheme(); + + let (dir, filename) = DataSource::find_dir_and_filename(path); + + let source = if let Some(filename) = filename { + Source::Filename(filename) + } else { + Source::Dir + }; + + let object_store = match schema.to_uppercase().as_str() { + S3_SCHEMA => DataSource::build_s3_backend(host, &dir, connection)?, + _ => { + return error::UnsupportedBackendProtocolSnafu { + protocol: schema.to_string(), + } + .fail() + } + }; + + Ok(DataSource { + object_store, + source, + path: dir, + regex, + }) + } + + pub fn new( + url: &str, + pattern: Option, + connection: HashMap, + ) -> Result { + let regex = if let Some(pattern) = pattern { + let regex = Regex::new(&pattern).context(error::BuildRegexSnafu)?; + Some(regex) + } else { + None + }; + let result = Url::parse(url); + + match result { + Ok(url) => DataSource::from_url(url, regex, connection), + Err(err) => { + if ParseError::RelativeUrlWithoutBase == err { + DataSource::from_path(url, regex) + } else { + Err(error::Error::InvalidUrl { + url: url.to_string(), + source: err, + }) + } + } + } + } + + pub async fn list(&self) -> Result> { + match &self.source { + Source::Dir => { + let streamer = self + .object_store + .object("/") + .list() + .await + .context(error::ListObjectsSnafu { path: &self.path })?; + streamer + .try_filter(|f| { + let res = if let Some(regex) = &self.regex { + regex.is_match(f.name()) + } else { + true + }; + future::ready(res) + }) + .try_collect::>() + .await + .context(error::ListObjectsSnafu { path: &self.path }) + } + Source::Filename(filename) => { + let obj = self.object_store.object(filename); + + Ok(vec![obj]) + } + } + } + + fn find_dir_and_filename(path: &str) -> (String, Option) { + if path.is_empty() { + ("/".to_string(), None) + } else if path.ends_with('/') { + (path.to_string(), None) + } else if let Some(idx) = path.rfind('/') { + ( + path[..idx + 1].to_string(), + Some(path[idx + 1..].to_string()), + ) + } else { + ("/".to_string(), Some(path.to_string())) + } + } +} + +#[cfg(test)] +mod tests { + + use url::Url; + + use super::*; + #[test] + fn test_parse_uri() { + struct Test<'a> { + uri: &'a str, + expected_path: &'a str, + expected_schema: &'a str, + } + + let tests = [ + Test { + uri: "s3://bucket/to/path/", + expected_path: "/to/path/", + expected_schema: "s3", + }, + Test { + uri: "fs:///to/path/", + expected_path: "/to/path/", + expected_schema: "fs", + }, + Test { + uri: "fs:///to/path/file", + expected_path: "/to/path/file", + expected_schema: "fs", + }, + ]; + for test in tests { + let parsed_uri = Url::parse(test.uri).unwrap(); + assert_eq!(parsed_uri.path(), test.expected_path); + assert_eq!(parsed_uri.scheme(), test.expected_schema); + } + } + + #[test] + fn test_parse_path_and_dir() { + let parsed = Url::from_file_path("/to/path/file").unwrap(); + assert_eq!(parsed.path(), "/to/path/file"); + + let parsed = Url::from_directory_path("/to/path/").unwrap(); + assert_eq!(parsed.path(), "/to/path/"); + } + + #[test] + fn test_find_dir_and_filename() { + struct Test<'a> { + path: &'a str, + expected_dir: &'a str, + expected_filename: Option, + } + + let tests = [ + Test { + path: "to/path/", + expected_dir: "to/path/", + expected_filename: None, + }, + Test { + path: "to/path/filename", + expected_dir: "to/path/", + expected_filename: Some("filename".into()), + }, + Test { + path: "/to/path/filename", + expected_dir: "/to/path/", + expected_filename: Some("filename".into()), + }, + Test { + path: "/", + expected_dir: "/", + expected_filename: None, + }, + Test { + path: "filename", + expected_dir: "/", + expected_filename: Some("filename".into()), + }, + Test { + path: "", + expected_dir: "/", + expected_filename: None, + }, + ]; + + for test in tests { + let (path, filename) = DataSource::find_dir_and_filename(test.path); + assert_eq!(test.expected_dir, path); + assert_eq!(test.expected_filename, filename) + } + } +} diff --git a/src/datanode/src/tests/instance_test.rs b/src/datanode/src/tests/instance_test.rs index 5c1f48d3b91..e311b45a6a4 100644 --- a/src/datanode/src/tests/instance_test.rs +++ b/src/datanode/src/tests/instance_test.rs @@ -795,6 +795,91 @@ async fn test_execute_copy_to() { assert!(matches!(output, Output::AffectedRows(2))); } +#[tokio::test(flavor = "multi_thread")] +async fn test_execute_copy_from() { + let instance = setup_test_instance("test_execute_copy_from").await; + + // setups + execute_sql( + &instance, + "create table demo(host string, cpu double, memory double, ts timestamp time index);", + ) + .await; + + let output = execute_sql( + &instance, + r#"insert into demo(host, cpu, memory, ts) values + ('host1', 66.6, 1024, 1655276557000), + ('host2', 88.8, 333.3, 1655276558000) + "#, + ) + .await; + assert!(matches!(output, Output::AffectedRows(2))); + + // export + let data_dir = instance.data_tmp_dir().path(); + + let copy_to_stmt = format!("Copy demo TO '{}/export/demo.parquet'", data_dir.display()); + + let output = execute_sql(&instance, ©_to_stmt).await; + assert!(matches!(output, Output::AffectedRows(2))); + + struct Test<'a> { + sql: &'a str, + table_name: &'a str, + } + let tests = [ + Test { + sql: &format!( + "Copy with_filename FROM '{}/export/demo.parquet_1_2'", + data_dir.display() + ), + table_name: "with_filename", + }, + Test { + sql: &format!("Copy with_path FROM '{}/export/'", data_dir.display()), + table_name: "with_path", + }, + Test { + sql: &format!( + "Copy with_pattern FROM '{}/export/' WITH (PATTERN = 'demo.*')", + data_dir.display() + ), + table_name: "with_pattern", + }, + ]; + + for test in tests { + // import + execute_sql( + &instance, + &format!( + "create table {}(host string, cpu double, memory double, ts timestamp time index);", + test.table_name + ), + ) + .await; + + let output = execute_sql(&instance, test.sql).await; + assert!(matches!(output, Output::AffectedRows(2))); + + let output = execute_sql( + &instance, + &format!("select * from {} order by ts", test.table_name), + ) + .await; + let expected = "\ ++-------+------+--------+---------------------+ +| host | cpu | memory | ts | ++-------+------+--------+---------------------+ +| host1 | 66.6 | 1024.0 | 2022-06-15T07:02:37 | +| host2 | 88.8 | 333.3 | 2022-06-15T07:02:38 | ++-------+------+--------+---------------------+" + .to_string(); + check_output_stream(output, expected).await; + } +} + #[tokio::test(flavor = "multi_thread")] async fn test_create_by_procedure() { common_telemetry::init_default_ut_logging(); diff --git a/src/table/src/requests.rs b/src/table/src/requests.rs index ec77815187f..8cd3472cab3 100644 --- a/src/table/src/requests.rs +++ b/src/table/src/requests.rs @@ -199,6 +199,16 @@ pub struct CopyTableRequest { pub file_name: String, } +#[derive(Debug)] +pub struct CopyTableFromRequest { + pub catalog_name: String, + pub schema_name: String, + pub table_name: String, + pub connection: HashMap, + pub pattern: Option, + pub from: String, +} + #[cfg(test)] mod tests { use super::*;