refactor: combine Copy To and Copy From (#1197)

* refactor: combine Copy To and Copy From

* Apply suggestions from code review

Co-authored-by: LFC <bayinamine@gmail.com>

* Apply suggestions from code review

Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com>

---------

Co-authored-by: LFC <bayinamine@gmail.com>
Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com>
This commit is contained in:
Weny Xu
2023-03-20 19:23:25 +08:00
committed by GitHub
parent ad886f5b3e
commit e19c8fa2b6
19 changed files with 643 additions and 508 deletions

View File

@@ -18,7 +18,7 @@ use sqlparser::keywords::Keyword;
use crate::error::{self, Result};
use crate::parser::ParserContext;
use crate::statements::copy::{CopyTable, CopyTableFrom, CopyTableTo, Format};
use crate::statements::copy::{CopyTable, CopyTableArgument, Format};
use crate::statements::statement::Statement;
// COPY tbl TO 'output.parquet';
@@ -40,24 +40,24 @@ impl<'a> ParserContext<'a> {
})?;
if self.parser.parse_keyword(Keyword::TO) {
self.parse_copy_table_to(table_name)
Ok(CopyTable::To(self.parse_copy_table_to(table_name)?))
} else {
self.parser
.expect_keyword(Keyword::FROM)
.context(error::SyntaxSnafu { sql: self.sql })?;
self.parse_copy_table_from(table_name)
Ok(CopyTable::From(self.parse_copy_table_from(table_name)?))
}
}
fn parse_copy_table_from(&mut self, table_name: ObjectName) -> Result<CopyTable> {
let uri = self
.parser
.parse_literal_string()
.with_context(|_| error::UnexpectedSnafu {
sql: self.sql,
expected: "a uri",
actual: self.peek_token_as_string(),
})?;
fn parse_copy_table_from(&mut self, table_name: ObjectName) -> Result<CopyTableArgument> {
let location =
self.parser
.parse_literal_string()
.with_context(|_| error::UnexpectedSnafu {
sql: self.sql,
expected: "a uri",
actual: self.peek_token_as_string(),
})?;
let options = self
.parser
@@ -99,14 +99,17 @@ impl<'a> ParserContext<'a> {
}
})
.collect();
Ok(CopyTable::From(CopyTableFrom::new(
table_name, uri, format, pattern, connection,
)))
Ok(CopyTableArgument {
table_name,
format,
pattern,
connection,
location,
})
}
fn parse_copy_table_to(&mut self, table_name: ObjectName) -> Result<CopyTable> {
let file_name =
fn parse_copy_table_to(&mut self, table_name: ObjectName) -> Result<CopyTableArgument> {
let location =
self.parser
.parse_literal_string()
.with_context(|_| error::UnexpectedSnafu {
@@ -146,9 +149,13 @@ impl<'a> ParserContext<'a> {
})
.collect();
Ok(CopyTable::To(CopyTableTo::new(
table_name, file_name, format, connection,
)))
Ok(CopyTableArgument {
table_name,
format,
connection,
pattern: None,
location,
})
}
fn parse_option_string(value: Value) -> Option<String> {
@@ -197,7 +204,7 @@ mod tests {
assert_eq!("schema0", schema);
assert_eq!("tbl", table);
let file_name = copy_table.file_name;
let file_name = copy_table.location;
assert_eq!("tbl_file.parquet", file_name);
let format = copy_table.format;
@@ -240,7 +247,7 @@ mod tests {
assert_eq!("schema0", schema);
assert_eq!("tbl", table);
let file_name = copy_table.from;
let file_name = copy_table.location;
assert_eq!("tbl_file.parquet", file_name);
let format = copy_table.format;

View File

@@ -20,60 +20,18 @@ use crate::error::{self, Result};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CopyTable {
To(CopyTableTo),
From(CopyTableFrom),
To(CopyTableArgument),
From(CopyTableArgument),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CopyTableTo {
pub table_name: ObjectName,
pub file_name: String,
pub format: Format,
pub connection: HashMap<String, String>,
}
impl CopyTableTo {
pub(crate) fn new(
table_name: ObjectName,
file_name: String,
format: Format,
connection: HashMap<String, String>,
) -> Self {
Self {
table_name,
file_name,
format,
connection,
}
}
}
// TODO: To combine struct CopyTableFrom and CopyTableTo
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CopyTableFrom {
pub struct CopyTableArgument {
pub table_name: ObjectName,
pub format: Format,
pub connection: HashMap<String, String>,
pub pattern: Option<String>,
pub from: String,
}
impl CopyTableFrom {
pub(crate) fn new(
table_name: ObjectName,
from: String,
format: Format,
pattern: Option<String>,
connection: HashMap<String, String>,
) -> Self {
CopyTableFrom {
table_name,
format,
connection,
pattern,
from,
}
}
/// Copy tbl [To|From] 'location'.
pub location: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]