PostgreSQL: Support ONLY and * table inheritance markers (#2770)

Author: LucaCappelletti94

fuzz/fuzz_seeds/postgres_table_inheritance_markers.sql

@@ -0,0 +1,3 @@
+SELECT * FROM ONLY (a) JOIN b * AS x ON a.id = x.id;
+UPDATE ONLY t SET id = 1;
+DELETE FROM t *;

src/ast/mod.rs

@@ -112,8 +112,8 @@ pub use self::query::{
     SelectFlavor, SelectInto, SelectItem, SelectItemQualifiedWildcardKind, SelectModifiers,
     SetExpr, SetOperator, SetQuantifier, Setting, SymbolDefinition, Table, TableAlias,
     TableAliasColumnDef, TableFactor, TableFunctionArgs, TableIndexDirective,
-    TableIndexHintForClause, TableIndexHintType, TableIndexHints, TableIndexType, TableSample,
-    TableSampleBucket, TableSampleKind, TableSampleMethod, TableSampleModifier,
+    TableIndexHintForClause, TableIndexHintType, TableIndexHints, TableIndexType, TableInheritance,
+    TableSample, TableSampleBucket, TableSampleKind, TableSampleMethod, TableSampleModifier,
     TableSampleQuantity, TableSampleSeed, TableSampleSeedModifier, TableSampleUnit, TableVersion,
     TableWithJoins, Top, TopQuantity, UpdateTableFromKind, ValueTableMode, Values,
     WildcardAdditionalOptions, With, WithFill, XmlNamespaceDefinition, XmlPassingArgument,

src/ast/query.rs

@@ -1477,6 +1477,22 @@ impl fmt::Display for TableIndexDirective {
     }
 }
 
+/// PostgreSQL table inheritance marker, see [`TableFactor::Table::inheritance`].
+///
+/// <https://www.postgresql.org/docs/current/sql-select.html#SQL-FROM>
+#[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub enum TableInheritance {
+    /// `ONLY name`, or `ONLY (name)` when `parenthesized`, excluding descendant tables.
+    Only {
+        /// Whether the name was written in parentheses.
+        parenthesized: bool,
+    },
+    /// `name *`, explicitly including descendant tables.
+    IncludeDescendants,
+}
+
 /// A table name or a parenthesized subquery with an optional alias
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
@@ -1516,6 +1532,8 @@ pub enum TableFactor {
         sample: Option<TableSampleKind>,
         /// Optional table-level index-selection hints or requirements.
         index_directives: Vec<TableIndexDirective>,
+        /// PostgreSQL `ONLY` or `*` inheritance marker.
+        inheritance: Option<TableInheritance>,
     },
     /// A derived table (a parenthesized subquery), optionally `LATERAL`.
     Derived {
@@ -2250,8 +2268,18 @@ impl fmt::Display for TableFactor {
                 json_path,
                 sample,
                 index_directives,
+                inheritance,
             } => {
-                name.fmt(f)?;
+                match inheritance {
+                    Some(TableInheritance::Only {
+                        parenthesized: true,
+                    }) => write!(f, "ONLY ({name})")?,
+                    Some(TableInheritance::Only {
+                        parenthesized: false,
+                    }) => write!(f, "ONLY {name}")?,
+                    Some(TableInheritance::IncludeDescendants) => write!(f, "{name} *")?,
+                    None => name.fmt(f)?,
+                }
                 if let Some(json_path) = json_path {
                     json_path.fmt(f)?;
                 }

src/ast/spans.rs

@@ -2059,6 +2059,7 @@ impl Spanned for TableFactor {
                 json_path: _,
                 sample: _,
                 index_directives: _,
+                inheritance: _,
             } => union_spans(
                 name.0
                     .iter()

src/dialect/generic.rs

@@ -333,6 +333,10 @@ impl Dialect for GenericDialect {
         true
     }
 
+    fn supports_table_inheritance_markers(&self) -> bool {
+        true
+    }
+
     fn supports_select_format(&self) -> bool {
         true
     }

src/dialect/mod.rs

@@ -2174,6 +2174,20 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if this dialect supports the `ONLY` and `*` table
+    /// inheritance markers on a table factor.
+    ///
+    /// Example:
+    /// ```sql
+    /// SELECT * FROM ONLY parent;
+    /// SELECT * FROM parent *;
+    /// ```
+    ///
+    /// [PostgreSQL](https://www.postgresql.org/docs/current/sql-select.html#SQL-FROM)
+    fn supports_table_inheritance_markers(&self) -> bool {
+        false
+    }
+
     /// Returns true if this dialect supports the `FORMAT` clause in `SELECT` statements.
     ///
     /// Example:

src/dialect/postgresql.rs

@@ -613,4 +613,9 @@ impl Dialect for PostgreSqlDialect {
     fn supports_attach_partition_of(&self) -> bool {
         true
     }
+
+    /// See <https://www.postgresql.org/docs/current/sql-select.html#SQL-FROM>
+    fn supports_table_inheritance_markers(&self) -> bool {
+        true
+    }
 }

src/parser/mod.rs

@@ -17437,6 +17437,10 @@ impl<'a> Parser<'a> {
             // Stage reference: @mystage or @namespace.stage (e.g. Snowflake)
             self.parse_snowflake_stage_table_factor()
         } else {
+            let inheritance_markers = self.dialect.supports_table_inheritance_markers();
+            let only = inheritance_markers && self.parse_keyword(Keyword::ONLY);
+            let only_parenthesized = only && self.consume_token(&Token::LParen);
+
             // Unquoted reserved keywords, e.g. `ON`, are not table names. The
             // keyword is consumed first so that `is_table_factor` peeks past it,
             // as it does at the trailing comma check.
@@ -17453,6 +17457,19 @@ impl<'a> Parser<'a> {
 
             let name = self.parse_object_name(true)?;
 
+            let inheritance = if only {
+                if only_parenthesized {
+                    self.expect_token(&Token::RParen)?;
+                }
+                Some(TableInheritance::Only {
+                    parenthesized: only_parenthesized,
+                })
+            } else if inheritance_markers && self.consume_token(&Token::Mul) {
+                Some(TableInheritance::IncludeDescendants)
+            } else {
+                None
+            };
+
             let json_path = match &self.peek_token_ref().token {
                 Token::LBracket if self.dialect.supports_partiql() => Some(self.parse_json_path()?),
                 _ => None,
@@ -17530,6 +17547,7 @@ impl<'a> Parser<'a> {
                 json_path,
                 sample,
                 index_directives,
+                inheritance,
             };
 
             while let Some(kw) = self.parse_one_of_keywords(&[Keyword::PIVOT, Keyword::UNPIVOT]) {
@@ -17580,6 +17598,7 @@ impl<'a> Parser<'a> {
             json_path: None,
             sample: None,
             index_directives: vec![],
+            inheritance: None,
         })
     }
 

src/test_utils.rs

@@ -391,6 +391,7 @@ pub fn table(name: impl Into<String>) -> TableFactor {
         json_path: None,
         sample: None,
         index_directives: vec![],
+        inheritance: None,
     }
 }
 
@@ -406,6 +407,7 @@ pub fn table_from_name(name: ObjectName) -> TableFactor {
         json_path: None,
         sample: None,
         index_directives: vec![],
+        inheritance: None,
     }
 }
 
@@ -425,6 +427,7 @@ pub fn table_with_alias(
         json_path: None,
         sample: None,
         index_directives: vec![],
+        inheritance: None,
     }
 }
 

tests/sqlparser_bigquery.rs

@@ -1768,6 +1768,7 @@ fn parse_table_time_travel() {
                 json_path: None,
                 sample: None,
                 index_directives: vec![],
+                inheritance: None,
             },
             joins: vec![]
         },]
@@ -1883,6 +1884,7 @@ fn parse_merge() {
                     json_path: None,
                     sample: None,
                     index_directives: vec![],
+                    inheritance: None,
                 },
                 table
             );
@@ -1898,6 +1900,7 @@ fn parse_merge() {
                     json_path: None,
                     sample: None,
                     index_directives: vec![],
+                    inheritance: None,
                 },
                 source
             );