feat: implement COPY for cluster (#1388)

This commit is contained in:
LFC
2023-04-17 11:04:47 +08:00
committed by GitHub
parent cc7c313937
commit f7b7a9c801
15 changed files with 199 additions and 218 deletions
Generated
+5
View File
@@ -2937,6 +2937,7 @@ name = "frontend"
version = "0.1.1"
dependencies = [
"api",
"async-compat",
"async-stream",
"async-trait",
"catalog",
@@ -2944,6 +2945,7 @@ dependencies = [
"client",
"common-base",
"common-catalog",
"common-datasource",
"common-error",
"common-function",
"common-grpc",
@@ -2965,10 +2967,12 @@ dependencies = [
"meta-srv",
"mito",
"moka",
"object-store",
"openmetrics-parser",
"partition",
"prost",
"query",
"regex",
"rstest",
"rstest_reuse",
"script",
@@ -2978,6 +2982,7 @@ dependencies = [
"session",
"snafu",
"sql",
"storage",
"store-api",
"strfmt",
"substrait 0.1.1",
+3 -128
View File
@@ -14,14 +14,11 @@
use std::any::Any;
use common_datasource::error::Error as DataSourceError;
use common_error::prelude::*;
use common_procedure::ProcedureId;
use datafusion::parquet;
use snafu::Location;
use storage::error::Error as StorageError;
use table::error::Error as TableError;
use url::ParseError;
use crate::datanode::ObjectStoreConfig;
@@ -208,18 +205,6 @@ pub enum Error {
location: Location,
},
#[snafu(display("Failed to build backend, source: {}", source))]
BuildBackend {
#[snafu(backtrace)]
source: DataSourceError,
},
#[snafu(display("Failed to parse url, source: {}", source))]
ParseUrl {
source: DataSourceError,
location: Location,
},
#[snafu(display("Runtime resource error, source: {}", source))]
RuntimeResource {
#[snafu(backtrace)]
@@ -229,36 +214,6 @@ 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 {
location: Location,
source: regex::Error,
},
#[snafu(display("Failed to list objects, source: {}", source))]
ListObjects {
#[snafu(backtrace)]
source: DataSourceError,
},
#[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 },
@@ -409,63 +364,6 @@ pub enum Error {
source: query::error::Error,
},
#[snafu(display("Failed to copy data from table: {}, source: {}", table_name, source))]
CopyTable {
table_name: String,
#[snafu(backtrace)]
source: TableError,
},
#[snafu(display("Failed to execute table scan, source: {}", source))]
TableScanExec {
#[snafu(backtrace)]
source: common_query::error::Error,
},
#[snafu(display(
"File schema mismatch at index {}, expected table schema: {} but found: {}",
index,
table_schema,
file_schema
))]
InvalidSchema {
index: usize,
table_schema: String,
file_schema: String,
},
#[snafu(display("Failed to read parquet file, source: {}", source))]
ReadParquet {
source: parquet::errors::ParquetError,
location: Location,
},
#[snafu(display("Failed to poll stream, source: {}", source))]
PollStream {
source: datafusion_common::DataFusionError,
location: Location,
},
#[snafu(display("Failed to build parquet record batch stream, source: {}", source))]
BuildParquetRecordBatchStream {
location: Location,
source: parquet::errors::ParquetError,
},
#[snafu(display("Failed to read object in path: {}, source: {}", path, source))]
ReadObject {
path: String,
location: Location,
source: object_store::Error,
},
#[snafu(display("Failed to write object into path: {}, source: {}", path, source))]
WriteObject {
path: String,
location: Location,
source: object_store::Error,
},
#[snafu(display("Unrecognized table option: {}", source))]
UnrecognizedTableOption {
#[snafu(backtrace)]
@@ -509,12 +407,6 @@ pub enum Error {
#[snafu(backtrace)]
source: BoxedError,
},
#[snafu(display("Failed to copy table to parquet file, source: {}", source))]
WriteParquet {
#[snafu(backtrace)]
source: storage::error::Error,
},
}
pub type Result<T> = std::result::Result<T, Error>;
@@ -554,11 +446,6 @@ impl ErrorExt for Error {
ColumnValuesNumberMismatch { .. }
| InvalidSql { .. }
| InvalidUrl { .. }
| InvalidPath { .. }
| InvalidConnection { .. }
| UnsupportedBackendProtocol { .. }
| BuildRegex { .. }
| NotSupportSql { .. }
| KeyColumnNotFound { .. }
| IllegalPrimaryKeysDef { .. }
@@ -572,8 +459,7 @@ impl ErrorExt for Error {
| DatabaseNotFound { .. }
| MissingNodeId { .. }
| MissingMetasrvOpts { .. }
| ColumnNoneDefaultValue { .. }
| ParseUrl { .. } => StatusCode::InvalidArguments,
| ColumnNoneDefaultValue { .. } => StatusCode::InvalidArguments,
// TODO(yingwen): Further categorize http error.
StartServer { .. }
@@ -586,22 +472,13 @@ impl ErrorExt for Error {
| RenameTable { .. }
| Catalog { .. }
| MissingRequiredField { .. }
| BuildParquetRecordBatchStream { .. }
| InvalidSchema { .. }
| ParseDataTypes { .. }
| IncorrectInternalState { .. }
| ShutdownServer { .. }
| ShutdownInstance { .. }
| CloseTableEngine { .. } => StatusCode::Internal,
BuildBackend { .. }
| InitBackend { .. }
| ReadParquet { .. }
| WriteParquet { .. }
| PollStream { .. }
| ReadObject { .. }
| WriteObject { .. }
| ListObjects { .. } => StatusCode::StorageUnavailable,
InitBackend { .. } => StatusCode::StorageUnavailable,
OpenLogStore { source } => source.status_code(),
OpenStorageEngine { source } => source.status_code(),
RuntimeResource { .. } => StatusCode::RuntimeResourcesExhausted,
@@ -609,8 +486,6 @@ impl ErrorExt for Error {
TableIdProviderNotFound { .. } => StatusCode::Unsupported,
BumpTableId { source, .. } => source.status_code(),
ColumnDefaultValue { source, .. } => source.status_code(),
CopyTable { source, .. } => source.status_code(),
TableScanExec { source, .. } => source.status_code(),
UnrecognizedTableOption { .. } => StatusCode::InvalidArguments,
RecoverProcedure { source, .. } | SubmitProcedure { source, .. } => {
source.status_code()
+1 -50
View File
@@ -25,10 +25,9 @@ use query::query_engine::SqlStatementExecutor;
use session::context::QueryContextRef;
use snafu::prelude::*;
use sql::ast::ObjectName;
use sql::statements::copy::{CopyTable, CopyTableArgument};
use sql::statements::statement::Statement;
use table::engine::TableReference;
use table::requests::{CopyDirection, CopyTableRequest, CreateDatabaseRequest, DropTableRequest};
use table::requests::{CreateDatabaseRequest, DropTableRequest};
use crate::error::{
self, BumpTableIdSnafu, ExecuteSqlSnafu, ExecuteStatementSnafu, NotSupportSqlSnafu,
@@ -105,54 +104,6 @@ impl Instance {
.execute(SqlRequest::DropTable(req), query_ctx)
.await
}
Statement::Copy(copy_table) => {
let req = match copy_table {
CopyTable::To(copy_table) => {
let CopyTableArgument {
location,
connection,
pattern,
table_name,
..
} = copy_table;
let (catalog_name, schema_name, table_name) =
table_idents_to_full_name(&table_name, query_ctx.clone())?;
CopyTableRequest {
catalog_name,
schema_name,
table_name,
location,
connection,
pattern,
direction: CopyDirection::Export,
}
}
CopyTable::From(copy_table) => {
let CopyTableArgument {
location,
connection,
pattern,
table_name,
..
} = copy_table;
let (catalog_name, schema_name, table_name) =
table_idents_to_full_name(&table_name, query_ctx.clone())?;
CopyTableRequest {
catalog_name,
schema_name,
table_name,
location,
connection,
pattern,
direction: CopyDirection::Import,
}
}
};
self.sql_handler
.execute(SqlRequest::CopyTable(req), query_ctx)
.await
}
_ => NotSupportSqlSnafu {
msg: format!("not supported to execute {stmt:?}"),
}
-7
View File
@@ -32,8 +32,6 @@ use crate::error::{
use crate::instance::sql::table_idents_to_full_name;
mod alter;
mod copy_table_from;
mod copy_table_to;
mod create;
mod drop_table;
mod flush_table;
@@ -46,7 +44,6 @@ pub enum SqlRequest {
Alter(AlterTableRequest),
DropTable(DropTableRequest),
FlushTable(FlushTableRequest),
CopyTable(CopyTableRequest),
}
// Handler to execute SQL except query
@@ -85,10 +82,6 @@ impl SqlHandler {
SqlRequest::CreateDatabase(req) => self.create_database(req, query_ctx.clone()).await,
SqlRequest::Alter(req) => self.alter(req).await,
SqlRequest::DropTable(req) => self.drop_table(req).await,
SqlRequest::CopyTable(req) => match req.direction {
CopyDirection::Export => self.copy_table_to(req).await,
CopyDirection::Import => self.copy_table_from(req).await,
},
SqlRequest::FlushTable(req) => self.flush_table(req).await,
};
if let Err(e) = &result {
+5
View File
@@ -10,6 +10,7 @@ python = ["dep:script"]
[dependencies]
api = { path = "../api" }
async-compat = "0.2"
async-stream.workspace = true
async-trait = "0.1"
catalog = { path = "../catalog" }
@@ -17,6 +18,7 @@ chrono.workspace = true
client = { path = "../client" }
common-base = { path = "../common/base" }
common-catalog = { path = "../common/catalog" }
common-datasource = { path = "../common/datasource" }
common-error = { path = "../common/error" }
common-function = { path = "../common/function" }
common-grpc = { path = "../common/grpc" }
@@ -36,10 +38,12 @@ itertools = "0.10"
meta-client = { path = "../meta-client" }
mito = { path = "../mito", features = ["test"] }
moka = { version = "0.9", features = ["future"] }
object-store = { path = "../object-store" }
openmetrics-parser = "0.4"
partition = { path = "../partition" }
prost.workspace = true
query = { path = "../query" }
regex = "1.6"
script = { path = "../script", features = ["python"], optional = true }
serde = "1.0"
serde_json = "1.0"
@@ -47,6 +51,7 @@ servers = { path = "../servers" }
session = { path = "../session" }
snafu.workspace = true
sql = { path = "../sql" }
storage = { path = "../storage" }
store-api = { path = "../store-api" }
substrait = { path = "../common/substrait" }
table = { path = "../table" }
+109 -11
View File
@@ -15,6 +15,7 @@
use std::any::Any;
use common_error::prelude::*;
use datafusion::parquet;
use snafu::Location;
use store_api::storage::RegionId;
@@ -310,13 +311,8 @@ pub enum Error {
source: common_grpc_expr::error::Error,
},
#[snafu(display(
"Failed to build default value, column: {}, source: {}",
column,
source
))]
ColumnDefaultValue {
column: String,
#[snafu(display("Failed to convert into vectors, source: {}", source))]
IntoVectors {
#[snafu(backtrace)]
source: datatypes::error::Error,
},
@@ -370,6 +366,92 @@ pub enum Error {
#[snafu(backtrace)]
source: script::error::Error,
},
#[snafu(display("Failed to build regex, source: {}", source))]
BuildRegex {
location: Location,
source: regex::Error,
},
#[snafu(display("Failed to copy table: {}, source: {}", table_name, source))]
CopyTable {
table_name: String,
#[snafu(backtrace)]
source: table::error::Error,
},
#[snafu(display(
"Failed to insert value into table: {}, source: {}",
table_name,
source
))]
Insert {
table_name: String,
#[snafu(backtrace)]
source: table::error::Error,
},
#[snafu(display("Failed to execute table scan, source: {}", source))]
TableScanExec {
#[snafu(backtrace)]
source: common_query::error::Error,
},
#[snafu(display("Failed to parse data source url, source: {}", source))]
ParseUrl {
#[snafu(backtrace)]
source: common_datasource::error::Error,
},
#[snafu(display("Failed to build data source backend, source: {}", source))]
BuildBackend {
#[snafu(backtrace)]
source: common_datasource::error::Error,
},
#[snafu(display("Failed to list objects, source: {}", source))]
ListObjects {
#[snafu(backtrace)]
source: common_datasource::error::Error,
},
#[snafu(display("Failed to read object in path: {}, source: {}", path, source))]
ReadObject {
path: String,
location: Location,
source: object_store::Error,
},
#[snafu(display("Failed to read parquet file, source: {}", source))]
ReadParquet {
source: parquet::errors::ParquetError,
location: Location,
},
#[snafu(display("Failed to build parquet record batch stream, source: {}", source))]
BuildParquetRecordBatchStream {
location: Location,
source: parquet::errors::ParquetError,
},
#[snafu(display("Failed to write parquet file, source: {}", source))]
WriteParquet {
#[snafu(backtrace)]
source: storage::error::Error,
},
#[snafu(display(
"Schema datatypes not match at index {}, expected table schema: {}, actual file schema: {}",
index,
table_schema,
file_schema
))]
InvalidSchema {
index: usize,
table_schema: String,
file_schema: String,
location: Location,
},
}
pub type Result<T> = std::result::Result<T, Error>;
@@ -388,7 +470,9 @@ impl ErrorExt for Error {
| Error::MissingInsertValues { .. }
| Error::PrimaryKeyNotFound { .. }
| Error::MissingMetasrvOpts { .. }
| Error::ColumnNoneDefaultValue { .. } => StatusCode::InvalidArguments,
| Error::ColumnNoneDefaultValue { .. }
| Error::BuildRegex { .. }
| Error::InvalidSchema { .. } => StatusCode::InvalidArguments,
Error::NotSupported { .. } => StatusCode::Unsupported,
@@ -401,10 +485,13 @@ impl ErrorExt for Error {
Error::ParseSql { source } => source.status_code(),
Error::Table { source } => source.status_code(),
Error::Table { source }
| Error::CopyTable { source, .. }
| Error::Insert { source, .. } => source.status_code(),
Error::ConvertColumnDefaultConstraint { source, .. }
| Error::CreateTableInfo { source } => source.status_code(),
| Error::CreateTableInfo { source }
| Error::IntoVectors { source } => source.status_code(),
Error::RequestDatanode { source } => source.status_code(),
@@ -447,7 +534,6 @@ impl ErrorExt for Error {
Error::TableAlreadyExist { .. } => StatusCode::TableAlreadyExists,
Error::EncodeSubstraitLogicalPlan { source } => source.status_code(),
Error::InvokeDatanode { source } => source.status_code(),
Error::ColumnDefaultValue { source, .. } => source.status_code(),
Error::External { source } => source.status_code(),
Error::DeserializePartition { source, .. } | Error::FindTableRoute { source, .. } => {
@@ -456,6 +542,18 @@ impl ErrorExt for Error {
Error::UnrecognizedTableOption { .. } => StatusCode::InvalidArguments,
Error::StartScriptManager { source } => source.status_code(),
Error::TableScanExec { source, .. } => source.status_code(),
Error::ReadObject { .. }
| Error::ReadParquet { .. }
| Error::BuildParquetRecordBatchStream { .. } => StatusCode::StorageUnavailable,
Error::ListObjects { source }
| Error::ParseUrl { source }
| Error::BuildBackend { source } => source.status_code(),
Error::WriteParquet { source, .. } => source.status_code(),
}
}
+66 -9
View File
@@ -12,23 +12,31 @@
// See the License for the specific language governing permissions and
// limitations under the License.
mod copy_table_from;
mod copy_table_to;
mod describe;
mod show;
mod tql;
use catalog::CatalogManagerRef;
use common_error::prelude::BoxedError;
use common_query::Output;
use common_recordbatch::RecordBatches;
use datanode::instance::sql::table_idents_to_full_name;
use query::parser::QueryStatement;
use query::query_engine::SqlStatementExecutorRef;
use query::QueryEngineRef;
use session::context::QueryContextRef;
use snafu::{ensure, ResultExt};
use snafu::{ensure, OptionExt, ResultExt};
use sql::statements::copy::{CopyTable, CopyTableArgument};
use sql::statements::statement::Statement;
use table::engine::TableReference;
use table::requests::{CopyDirection, CopyTableRequest};
use table::TableRef;
use crate::error::{
CatalogSnafu, ExecLogicalPlanSnafu, ExecuteStatementSnafu, PlanStatementSnafu, Result,
SchemaNotFoundSnafu,
CatalogSnafu, ExecLogicalPlanSnafu, ExecuteStatementSnafu, ExternalSnafu, PlanStatementSnafu,
Result, SchemaNotFoundSnafu, TableNotFoundSnafu,
};
#[derive(Clone)]
@@ -84,13 +92,20 @@ impl StatementExecutor {
Statement::ShowTables(stmt) => self.show_tables(stmt, query_ctx),
Statement::Copy(stmt) => {
let req = to_copy_table_request(stmt, query_ctx)?;
match req.direction {
CopyDirection::Export => self.copy_table_to(req).await,
CopyDirection::Import => self.copy_table_from(req).await,
}
}
Statement::CreateDatabase(_)
| Statement::CreateTable(_)
| Statement::CreateExternalTable(_)
| Statement::Insert(_)
| Statement::Alter(_)
| Statement::DropTable(_)
| Statement::Copy(_)
| Statement::ShowCreateTable(_) => self
.sql_stmt_executor
.execute_sql(stmt, query_ctx)
@@ -99,11 +114,7 @@ impl StatementExecutor {
}
}
pub(crate) async fn plan_exec(
&self,
stmt: QueryStatement,
query_ctx: QueryContextRef,
) -> Result<Output> {
async fn plan_exec(&self, stmt: QueryStatement, query_ctx: QueryContextRef) -> Result<Output> {
let planner = self.query_engine.planner();
let plan = planner
.plan(stmt, query_ctx.clone())
@@ -129,4 +140,50 @@ impl StatementExecutor {
Ok(Output::RecordBatches(RecordBatches::empty()))
}
async fn get_table(&self, table_ref: &TableReference<'_>) -> Result<TableRef> {
let TableReference {
catalog,
schema,
table,
} = table_ref;
self.catalog_manager
.table(catalog, schema, table)
.await
.context(CatalogSnafu)?
.with_context(|| TableNotFoundSnafu {
table_name: table_ref.to_string(),
})
}
}
fn to_copy_table_request(stmt: CopyTable, query_ctx: QueryContextRef) -> Result<CopyTableRequest> {
let direction = match stmt {
CopyTable::To(_) => CopyDirection::Export,
CopyTable::From(_) => CopyDirection::Import,
};
let CopyTableArgument {
location,
connection,
pattern,
table_name,
..
} = match stmt {
CopyTable::To(arg) => arg,
CopyTable::From(arg) => arg,
};
let (catalog_name, schema_name, table_name) = table_idents_to_full_name(&table_name, query_ctx)
.map_err(BoxedError::new)
.context(ExternalSnafu)?;
Ok(CopyTableRequest {
catalog_name,
schema_name,
table_name,
location,
connection,
pattern,
direction,
})
}
@@ -21,7 +21,6 @@ use common_datasource::lister::{Lister, Source};
use common_datasource::object_store::{build_backend, parse_url};
use common_datasource::util::find_dir_and_filename;
use common_query::Output;
use common_recordbatch::error::DataTypesSnafu;
use datafusion::parquet::arrow::ParquetRecordBatchStreamBuilder;
use datatypes::arrow::datatypes::{DataType, SchemaRef};
use datatypes::vectors::Helper;
@@ -32,10 +31,10 @@ use table::engine::TableReference;
use table::requests::{CopyTableRequest, InsertRequest};
use tokio::io::BufReader;
use crate::error::{self, ParseDataTypesSnafu, Result};
use crate::sql::SqlHandler;
use crate::error::{self, IntoVectorsSnafu, Result};
use crate::statement::StatementExecutor;
impl SqlHandler {
impl StatementExecutor {
pub(crate) async fn copy_table_from(&self, req: CopyTableRequest) -> Result<Output> {
let table_ref = TableReference {
catalog: &req.catalog_name,
@@ -102,9 +101,8 @@ impl SqlHandler {
while let Some(r) = stream.next().await {
let record_batch = r.context(error::ReadParquetSnafu)?;
let vectors = Helper::try_into_vectors(record_batch.columns())
.context(DataTypesSnafu)
.context(ParseDataTypesSnafu)?;
let vectors =
Helper::try_into_vectors(record_batch.columns()).context(IntoVectorsSnafu)?;
pending_mem_size += vectors.iter().map(|v| v.memory_size()).sum::<usize>();
@@ -12,7 +12,6 @@
// See the License for the specific language governing permissions and
// limitations under the License.
use common_datasource;
use common_datasource::object_store::{build_backend, parse_url};
use common_query::physical_plan::SessionContext;
use common_query::Output;
@@ -23,9 +22,9 @@ use table::engine::TableReference;
use table::requests::CopyTableRequest;
use crate::error::{self, Result, WriteParquetSnafu};
use crate::sql::SqlHandler;
use crate::statement::StatementExecutor;
impl SqlHandler {
impl StatementExecutor {
pub(crate) async fn copy_table_to(&self, req: CopyTableRequest) -> Result<Output> {
let table_ref = TableReference {
catalog: &req.catalog_name,
+2 -2
View File
@@ -765,7 +765,7 @@ async fn test_delete(instance: Arc<dyn MockInstance>) {
check_output_stream(output, expect).await;
}
#[apply(standalone_instance_case)]
#[apply(both_instances_cases)]
async fn test_execute_copy_to_s3(instance: Arc<dyn MockInstance>) {
if let Ok(bucket) = env::var("GT_S3_BUCKET") {
if !bucket.is_empty() {
@@ -803,7 +803,7 @@ async fn test_execute_copy_to_s3(instance: Arc<dyn MockInstance>) {
}
}
#[apply(standalone_instance_case)]
#[apply(both_instances_cases)]
async fn test_execute_copy_from_s3(instance: Arc<dyn MockInstance>) {
logging::init_default_ut_logging();
if let Ok(bucket) = env::var("GT_S3_BUCKET") {
+1 -1
View File
@@ -54,7 +54,7 @@ async fn create_insert_query_assert(
let query_output = instance
.statement_executor()
.plan_exec(QueryStatement::Promql(eval_stmt), QueryContext::arc())
.execute_stmt(QueryStatement::Promql(eval_stmt), QueryContext::arc())
.await
.unwrap();
check_unordered_output_stream(query_output, expected).await;