mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-10-26 13:46:23 +00:00
feat: implement COPY for cluster (#1388)
This commit is contained in:
Generated
+5
@@ -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
@@ -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()
|
||||
|
||||
@@ -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:?}"),
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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(),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
})
|
||||
}
|
||||
|
||||
+5
-7
@@ -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,
|
||||
@@ -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") {
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user