refactor: replace Copy Format with datasource Format (#1435)

* refactor: replace Copy Format with datasource Format

* chore: apply suggestions from CR

* chore: apply suggestions from CR
This commit is contained in:
Weny Xu
2023-04-23 17:31:54 +09:00
committed by GitHub
parent c5dba29f9e
commit d374859e24
10 changed files with 148 additions and 106 deletions

View File

@@ -137,9 +137,6 @@ pub enum Error {
target_unit: TimeUnit,
},
#[snafu(display("Unsupported format option: {}", name))]
UnsupportedCopyFormatOption { name: String },
#[snafu(display("Unable to convert statement {} to DataFusion statement", statement))]
ConvertToDfStatement {
statement: String,
@@ -178,8 +175,7 @@ impl ErrorExt for Error {
| ColumnTypeMismatch { .. }
| InvalidTableName { .. }
| InvalidSqlValue { .. }
| TimestampOverflow { .. }
| UnsupportedCopyFormatOption { .. } => StatusCode::InvalidArguments,
| TimestampOverflow { .. } => StatusCode::InvalidArguments,
UnsupportedAlterTableStatement { .. } => StatusCode::InvalidSyntax,
SerializeColumnDefaultConstraint { source, .. } => source.status_code(),

View File

@@ -18,7 +18,7 @@ use sqlparser::keywords::Keyword;
use crate::error::{self, Result};
use crate::parser::ParserContext;
use crate::statements::copy::{CopyTable, CopyTableArgument, Format};
use crate::statements::copy::{CopyTable, CopyTableArgument};
use crate::statements::statement::Statement;
use crate::util::parse_option_string;
@@ -65,25 +65,12 @@ impl<'a> ParserContext<'a> {
.parse_options(Keyword::WITH)
.context(error::SyntaxSnafu { sql: self.sql })?;
// default format is parquet
let mut format = Format::Parquet;
let mut pattern = None;
for option in options {
match option.name.value.to_ascii_uppercase().as_str() {
"FORMAT" => {
if let Some(fmt_str) = parse_option_string(option.value) {
format = Format::try_from(fmt_str)?;
}
}
"PATTERN" => {
if let Some(v) = parse_option_string(option.value) {
pattern = Some(v);
}
}
//TODO: throws warnings?
_ => (),
}
}
let with = options
.into_iter()
.filter_map(|option| {
parse_option_string(option.value).map(|v| (option.name.to_string(), v))
})
.collect();
let connection_options = self
.parser
@@ -93,17 +80,12 @@ impl<'a> ParserContext<'a> {
let connection = connection_options
.into_iter()
.filter_map(|option| {
if let Some(v) = parse_option_string(option.value) {
Some((option.name.value.to_uppercase(), v))
} else {
None
}
parse_option_string(option.value).map(|v| (option.name.to_string(), v))
})
.collect();
Ok(CopyTableArgument {
table_name,
format,
pattern,
with,
connection,
location,
})
@@ -124,15 +106,12 @@ impl<'a> ParserContext<'a> {
.parse_options(Keyword::WITH)
.context(error::SyntaxSnafu { sql: self.sql })?;
// default format is parquet
let mut format = Format::Parquet;
for option in options {
if option.name.value.eq_ignore_ascii_case("FORMAT") {
if let Some(fmt_str) = parse_option_string(option.value) {
format = Format::try_from(fmt_str)?;
}
}
}
let with = options
.into_iter()
.filter_map(|option| {
parse_option_string(option.value).map(|v| (option.name.to_string(), v))
})
.collect();
let connection_options = self
.parser
@@ -142,19 +121,14 @@ impl<'a> ParserContext<'a> {
let connection = connection_options
.into_iter()
.filter_map(|option| {
if let Some(v) = parse_option_string(option.value) {
Some((option.name.value.to_uppercase(), v))
} else {
None
}
parse_option_string(option.value).map(|v| (option.name.to_string(), v))
})
.collect();
Ok(CopyTableArgument {
table_name,
format,
with,
connection,
pattern: None,
location,
})
}
@@ -198,11 +172,11 @@ mod tests {
assert_eq!("schema0", schema);
assert_eq!("tbl", table);
let file_name = copy_table.location;
let file_name = &copy_table.location;
assert_eq!("tbl_file.parquet", file_name);
let format = copy_table.format;
assert_eq!(Format::Parquet, format);
let format = copy_table.format().unwrap();
assert_eq!("parquet", format.to_lowercase());
}
_ => unreachable!(),
}
@@ -241,11 +215,11 @@ mod tests {
assert_eq!("schema0", schema);
assert_eq!("tbl", table);
let file_name = copy_table.location;
let file_name = &copy_table.location;
assert_eq!("tbl_file.parquet", file_name);
let format = copy_table.format;
assert_eq!(Format::Parquet, format);
let format = copy_table.format().unwrap();
assert_eq!("parquet", format.to_lowercase());
}
_ => unreachable!(),
}
@@ -283,7 +257,7 @@ mod tests {
match statement {
Statement::Copy(CopyTable::From(copy_table)) => {
if let Some(expected_pattern) = test.expected_pattern {
assert_eq!(copy_table.pattern.clone().unwrap(), expected_pattern);
assert_eq!(copy_table.pattern().unwrap(), expected_pattern);
}
assert_eq!(copy_table.connection.clone(), test.expected_connection);
}
@@ -329,23 +303,4 @@ mod tests {
}
}
}
#[test]
fn test_parse_copy_table_with_unsupopoted_format() {
let results = [
"COPY catalog0.schema0.tbl TO 'tbl_file.parquet' WITH (FORMAT = 'unknow_format')",
"COPY catalog0.schema0.tbl FROM 'tbl_file.parquet' WITH (FORMAT = 'unknow_format')",
]
.iter()
.map(|sql| ParserContext::create_with_dialect(sql, &GenericDialect {}))
.collect::<Vec<_>>();
for result in results {
assert!(result.is_err());
assert_matches!(
result.err().unwrap(),
error::Error::UnsupportedCopyFormatOption { .. }
);
}
}
}

View File

@@ -16,8 +16,6 @@ use std::collections::HashMap;
use sqlparser::ast::ObjectName;
use crate::error::{self, Result};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CopyTable {
To(CopyTableArgument),
@@ -27,25 +25,24 @@ pub enum CopyTable {
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CopyTableArgument {
pub table_name: ObjectName,
pub format: Format,
pub with: HashMap<String, String>,
pub connection: HashMap<String, String>,
pub pattern: Option<String>,
/// Copy tbl [To|From] 'location'.
pub location: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Format {
Parquet,
}
#[cfg(test)]
impl CopyTableArgument {
const FORMAT: &str = "FORMAT";
impl TryFrom<String> for Format {
type Error = error::Error;
pub fn format(&self) -> Option<String> {
self.with
.get(Self::FORMAT)
.cloned()
.or_else(|| Some("PARQUET".to_string()))
}
fn try_from(name: String) -> Result<Self> {
if name.eq_ignore_ascii_case("PARQUET") {
return Ok(Format::Parquet);
}
error::UnsupportedCopyFormatOptionSnafu { name }.fail()
pub fn pattern(&self) -> Option<String> {
self.with.get("PATTERN").cloned()
}
}