Support CREATE TABLE x AS TABLE y (#704)

* first commit

* fix style and edit test

* fix test?

* remove unnecessary logic

* refactor implementation

* codestyle

* add schema support

* codestyle and lint

* Apply suggestions from code review

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

* PartialOrd and Ord

* clean up parser logic

* codestyle and lint

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>
This commit is contained in:
Sarah Yurick 2022-12-01 06:56:14 -08:00 committed by GitHub
parent 8e1c90c0d8
commit 528b3f2234
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 131 additions and 2 deletions

View file

@ -33,7 +33,7 @@ pub use self::operator::{BinaryOperator, UnaryOperator};
pub use self::query::{
Cte, ExcludeSelectItem, Fetch, Join, JoinConstraint, JoinOperator, LateralView, LockType,
Offset, OffsetRows, OrderByExpr, Query, Select, SelectInto, SelectItem, SetExpr, SetOperator,
SetQuantifier, TableAlias, TableFactor, TableWithJoins, Top, Values, With,
SetQuantifier, Table, TableAlias, TableFactor, TableWithJoins, Top, Values, With,
};
pub use self::value::{escape_quoted_string, DateTimeField, TrimWhereField, Value};

View file

@ -84,7 +84,7 @@ pub enum SetExpr {
},
Values(Values),
Insert(Statement),
// TODO: ANSI SQL supports `TABLE` here.
Table(Box<Table>),
}
impl fmt::Display for SetExpr {
@ -94,6 +94,7 @@ impl fmt::Display for SetExpr {
SetExpr::Query(q) => write!(f, "({})", q),
SetExpr::Values(v) => write!(f, "{}", v),
SetExpr::Insert(v) => write!(f, "{}", v),
SetExpr::Table(t) => write!(f, "{}", t),
SetExpr::SetOperation {
left,
right,
@ -152,6 +153,31 @@ impl fmt::Display for SetQuantifier {
}
}
}
#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
/// A [`TABLE` command]( https://www.postgresql.org/docs/current/sql-select.html#SQL-TABLE)
pub struct Table {
pub table_name: Option<String>,
pub schema_name: Option<String>,
}
impl fmt::Display for Table {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
if let Some(ref schema_name) = self.schema_name {
write!(
f,
"TABLE {}.{}",
schema_name,
self.table_name.as_ref().unwrap(),
)?;
} else {
write!(f, "TABLE {}", self.table_name.as_ref().unwrap(),)?;
}
Ok(())
}
}
/// A restricted variant of `SELECT` (without CTEs/`ORDER BY`), which may
/// appear either as the only body item of a `Query`, or as an operand
/// to a set operation like `UNION`.

View file

@ -4368,6 +4368,14 @@ impl<'a> Parser<'a> {
SetExpr::Query(Box::new(subquery))
} else if self.parse_keyword(Keyword::VALUES) {
SetExpr::Values(self.parse_values()?)
} else if self.parse_keyword(Keyword::TABLE) {
let token1 = self.peek_token();
let token2 = self.peek_nth_token(1);
let token3 = self.peek_nth_token(2);
self.next_token();
self.next_token();
self.next_token();
SetExpr::Table(Box::new(self.parse_as_table(token1, token2, token3)?))
} else {
return self.expected(
"SELECT, VALUES, or a subquery in the query body",
@ -4566,6 +4574,52 @@ impl<'a> Parser<'a> {
})
}
/// Parse `CREATE TABLE x AS TABLE y`
pub fn parse_as_table(
&self,
token1: Token,
token2: Token,
token3: Token,
) -> Result<Table, ParserError> {
let table_name;
let schema_name;
if token2 == Token::Period {
match token1 {
Token::Word(w) => {
schema_name = w.value;
}
_ => {
return self.expected("Schema name", token1);
}
}
match token3 {
Token::Word(w) => {
table_name = w.value;
}
_ => {
return self.expected("Table name", token3);
}
}
Ok(Table {
table_name: Some(table_name),
schema_name: Some(schema_name),
})
} else {
match token1 {
Token::Word(w) => {
table_name = w.value;
}
_ => {
return self.expected("Table name", token1);
}
}
Ok(Table {
table_name: Some(table_name),
schema_name: None,
})
}
}
pub fn parse_set(&mut self) -> Result<Statement, ParserError> {
let modifier =
self.parse_one_of_keywords(&[Keyword::SESSION, Keyword::LOCAL, Keyword::HIVEVAR]);

View file

@ -2273,6 +2273,55 @@ fn parse_create_table_as() {
}
}
#[test]
fn parse_create_table_as_table() {
let sql1 = "CREATE TABLE new_table AS TABLE old_table";
let expected_query1 = Box::new(Query {
with: None,
body: Box::new(SetExpr::Table(Box::new(Table {
table_name: Some("old_table".to_string()),
schema_name: None,
}))),
order_by: vec![],
limit: None,
offset: None,
fetch: None,
lock: None,
});
match verified_stmt(sql1) {
Statement::CreateTable { query, name, .. } => {
assert_eq!(name, ObjectName(vec![Ident::new("new_table")]));
assert_eq!(query.unwrap(), expected_query1);
}
_ => unreachable!(),
}
let sql2 = "CREATE TABLE new_table AS TABLE schema_name.old_table";
let expected_query2 = Box::new(Query {
with: None,
body: Box::new(SetExpr::Table(Box::new(Table {
table_name: Some("old_table".to_string()),
schema_name: Some("schema_name".to_string()),
}))),
order_by: vec![],
limit: None,
offset: None,
fetch: None,
lock: None,
});
match verified_stmt(sql2) {
Statement::CreateTable { query, name, .. } => {
assert_eq!(name, ObjectName(vec![Ident::new("new_table")]));
assert_eq!(query.unwrap(), expected_query2);
}
_ => unreachable!(),
}
}
#[test]
fn parse_create_table_on_cluster() {
// Using single-quote literal to define current cluster