PostgreSQL: Add support for TABLE query command

src/ast/query.rs

@@ -295,23 +295,34 @@ 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)
+/// A [`TABLE` command](https://www.postgresql.org/docs/current/sql-select.html#SQL-TABLE)
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 /// A (possibly schema-qualified) table reference used in `FROM` clauses.
 pub struct Table {
+    /// `ONLY` modifier before the table name.
+    pub only: bool,
     /// Optional table name (absent for e.g. `TABLE` command without argument).
     pub table_name: Option<String>,
     /// Optional schema/catalog name qualifying the table.
     pub schema_name: Option<String>,
+    /// Trailing `*` modifier after the table name.
+    pub with_asterisk: bool,
 }
 
 impl fmt::Display for Table {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         if let Some(ref table_name) = self.table_name {
+            write!(f, "TABLE ")?;
+            if self.only {
+                write!(f, "ONLY ")?;
+            }
             if let Some(ref schema_name) = self.schema_name {
-                write!(f, "TABLE {}.{}", schema_name, table_name,)?;
+                write!(f, "{}.{}", schema_name, table_name,)?;
             } else {
-                write!(f, "TABLE {}", table_name)?;
+                write!(f, "{}", table_name)?;
+            }
+            if self.with_asterisk {
+                write!(f, " *")?;
             }
         } else {
             write!(f, "TABLE")?;

src/parser/mod.rs

@@ -641,7 +641,11 @@ impl<'a> Parser<'a> {
                     self.prev_token();
                     self.parse_raise_stmt().map(Into::into)
                 }
-                Keyword::SELECT | Keyword::WITH | Keyword::VALUES | Keyword::FROM => {
+                Keyword::SELECT
+                | Keyword::TABLE
+                | Keyword::WITH
+                | Keyword::VALUES
+                | Keyword::FROM => {
                     self.prev_token();
                     self.parse_query().map(Into::into)
                 }
@@ -14496,7 +14500,15 @@ impl<'a> Parser<'a> {
             }
         }
 
-        match self.maybe_parse(|parser| parser.parse_statement())? {
+        let statement = if self.dialect.describe_requires_table_keyword()
+            && self.peek_keyword(Keyword::TABLE)
+        {
+            None
+        } else {
+            self.maybe_parse(|parser| parser.parse_statement())?
+        };
+
+        match statement {
             Some(Statement::Explain { .. }) | Some(Statement::ExplainTable { .. }) => Err(
                 ParserError::ParserError("Explain must be root of the plan".to_string()),
             ),
@@ -15625,13 +15637,12 @@ impl<'a> Parser<'a> {
 
     /// Parse `CREATE TABLE x AS TABLE y`
     pub fn parse_as_table(&mut self) -> Result<Table, ParserError> {
+        let only = self.parse_keyword(Keyword::ONLY);
         let token1 = self.next_token();
-        let token2 = self.next_token();
-        let token3 = self.next_token();
 
         let table_name;
         let schema_name;
-        if token2 == Token::Period {
+        if self.consume_token(&Token::Period) {
             match token1.token {
                 Token::Word(w) => {
                     schema_name = w.value;
@@ -15640,17 +15651,20 @@ impl<'a> Parser<'a> {
                     return self.expected("Schema name", token1);
                 }
             }
-            match token3.token {
+            let token2 = self.next_token();
+            match token2.token {
                 Token::Word(w) => {
                     table_name = w.value;
                 }
                 _ => {
-                    return self.expected("Table name", token3);
+                    return self.expected("Table name", token2);
                 }
             }
             Ok(Table {
+                only,
                 table_name: Some(table_name),
                 schema_name: Some(schema_name),
+                with_asterisk: self.consume_token(&Token::Mul),
             })
         } else {
             match token1.token {
@@ -15662,8 +15676,10 @@ impl<'a> Parser<'a> {
                 }
             }
             Ok(Table {
+                only,
                 table_name: Some(table_name),
                 schema_name: None,
+                with_asterisk: self.consume_token(&Token::Mul),
             })
         }
     }

tests/sqlparser_common.rs

@@ -4848,8 +4848,10 @@ fn parse_create_table_as_table() {
     let expected_query1 = Box::new(Query {
         with: None,
         body: Box::new(SetExpr::Table(Box::new(Table {
+            only: false,
             table_name: Some("old_table".to_string()),
             schema_name: None,
+            with_asterisk: false,
         }))),
         order_by: None,
         limit_clause: None,
@@ -4874,8 +4876,10 @@ fn parse_create_table_as_table() {
     let expected_query2 = Box::new(Query {
         with: None,
         body: Box::new(SetExpr::Table(Box::new(Table {
+            only: false,
             table_name: Some("old_table".to_string()),
             schema_name: Some("schema_name".to_string()),
+            with_asterisk: false,
         }))),
         order_by: None,
         limit_clause: None,
@@ -4894,6 +4898,62 @@ fn parse_create_table_as_table() {
         }
         _ => unreachable!(),
     }
+
+    let sql3 = "CREATE TABLE new_table AS TABLE ONLY old_table";
+
+    let expected_query3 = Box::new(Query {
+        with: None,
+        body: Box::new(SetExpr::Table(Box::new(Table {
+            only: true,
+            table_name: Some("old_table".to_string()),
+            schema_name: None,
+            with_asterisk: false,
+        }))),
+        order_by: None,
+        limit_clause: None,
+        fetch: None,
+        locks: vec![],
+        for_clause: None,
+        settings: None,
+        format_clause: None,
+        pipe_operators: vec![],
+    });
+
+    match verified_stmt(sql3) {
+        Statement::CreateTable(CreateTable { query, name, .. }) => {
+            assert_eq!(name, ObjectName::from(vec![Ident::new("new_table")]));
+            assert_eq!(query.unwrap(), expected_query3);
+        }
+        _ => unreachable!(),
+    }
+
+    let sql4 = "CREATE TABLE new_table AS TABLE old_table *";
+
+    let expected_query4 = Box::new(Query {
+        with: None,
+        body: Box::new(SetExpr::Table(Box::new(Table {
+            only: false,
+            table_name: Some("old_table".to_string()),
+            schema_name: None,
+            with_asterisk: true,
+        }))),
+        order_by: None,
+        limit_clause: None,
+        fetch: None,
+        locks: vec![],
+        for_clause: None,
+        settings: None,
+        format_clause: None,
+        pipe_operators: vec![],
+    });
+
+    match verified_stmt(sql4) {
+        Statement::CreateTable(CreateTable { query, name, .. }) => {
+            assert_eq!(name, ObjectName::from(vec![Ident::new("new_table")]));
+            assert_eq!(query.unwrap(), expected_query4);
+        }
+        _ => unreachable!(),
+    }
 }
 
 #[test]

tests/sqlparser_postgres.rs

@@ -9953,3 +9953,10 @@ fn parse_insert_by_name_keywords_as_table_and_alias() {
         statement => panic!("Expected INSERT statement, got: {statement:?}"),
     }
 }
+
+#[test]
+fn parse_table_command() {
+    pg_and_generic().verified_stmt("TABLE customers ORDER BY contact_name LIMIT 1");
+    pg_and_generic().verified_stmt("TABLE ONLY customers");
+    pg_and_generic().verified_stmt("TABLE customers *");
+}