Skip to content

Commit

Permalink
Support CREATE TABLE x AS TABLE y (#704)
Browse files Browse the repository at this point in the history
* 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>
  • Loading branch information
sarahyurick and alamb committed Dec 1, 2022
1 parent 8e1c90c commit 528b3f2
Show file tree
Hide file tree
Showing 4 changed files with 131 additions and 2 deletions.
2 changes: 1 addition & 1 deletion src/ast/mod.rs
Expand Up @@ -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};

Expand Down
28 changes: 27 additions & 1 deletion src/ast/query.rs
Expand Up @@ -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 {
Expand All @@ -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,
Expand Down Expand Up @@ -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`.
Expand Down
54 changes: 54 additions & 0 deletions src/parser.rs
Expand Up @@ -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",
Expand Down Expand Up @@ -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]);
Expand Down
49 changes: 49 additions & 0 deletions tests/sqlparser_common.rs
Expand Up @@ -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
Expand Down

0 comments on commit 528b3f2

Please sign in to comment.