Add support for Snowflake AT/BEFORE (apache#1667)

Author: yoavcloud

src/ast/query.rs

@@ -1873,13 +1873,19 @@ impl fmt::Display for TableAliasColumnDef {
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
 pub enum TableVersion {
+    /// When the table version is defined using `FOR SYSTEM_TIME AS OF`.
+    /// For example: `SELECT * FROM tbl FOR SYSTEM_TIME AS OF TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL 1 HOUR)`
     ForSystemTimeAsOf(Expr),
+    /// When the table version is defined using a function.
+    /// For example: `SELECT * FROM tbl AT(TIMESTAMP => '2020-08-14 09:30:00')`
+    Function(Expr),
 }
 
 impl Display for TableVersion {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
             TableVersion::ForSystemTimeAsOf(e) => write!(f, " FOR SYSTEM_TIME AS OF {e}")?,
+            TableVersion::Function(func) => write!(f, " {func}")?,
         }
         Ok(())
     }

src/dialect/bigquery.rs

@@ -77,4 +77,9 @@ impl Dialect for BigQueryDialect {
     fn supports_struct_literal(&self) -> bool {
         true
     }
+
+    // See <https://cloud.google.com/bigquery/docs/access-historical-data>
+    fn supports_timestamp_versioning(&self) -> bool {
+        true
+    }
 }

src/dialect/mod.rs

@@ -834,6 +834,12 @@ pub trait Dialect: Debug + Any {
     fn is_table_factor_alias(&self, explicit: bool, kw: &Keyword, _parser: &mut Parser) -> bool {
         explicit || !keywords::RESERVED_FOR_TABLE_ALIAS.contains(kw)
     }
+
+    /// Returns true if this dialect supports querying historical table data
+    /// by specifying which version of the data to query.
+    fn supports_timestamp_versioning(&self) -> bool {
+        false
+    }
 }
 
 /// This represents the operators for which precedence must be defined

src/dialect/mssql.rs

@@ -90,4 +90,9 @@ impl Dialect for MsSqlDialect {
     fn supports_set_stmt_without_operator(&self) -> bool {
         true
     }
+
+    /// See: <https://learn.microsoft.com/en-us/sql/relational-databases/tables/querying-data-in-a-system-versioned-temporal-table>
+    fn supports_timestamp_versioning(&self) -> bool {
+        true
+    }
 }

src/dialect/snowflake.rs

@@ -307,6 +307,11 @@ impl Dialect for SnowflakeDialect {
             _ => true,
         }
     }
+
+    /// See: <https://docs.snowflake.com/en/sql-reference/constructs/at-before>
+    fn supports_timestamp_versioning(&self) -> bool {
+        true
+    }
 }
 
 fn parse_file_staging_command(kw: Keyword, parser: &mut Parser) -> Result<Statement, ParserError> {

src/parser/mod.rs

@@ -11208,7 +11208,7 @@ impl<'a> Parser<'a> {
             };
 
             // Parse potential version qualifier
-            let version = self.parse_table_version()?;
+            let version = self.maybe_parse_table_version()?;
 
             // Postgres, MSSQL, ClickHouse: table-valued functions:
             let args = if self.consume_token(&Token::LParen) {
@@ -11639,18 +11639,20 @@ impl<'a> Parser<'a> {
         }
     }
 
-    /// Parse a given table version specifier.
-    ///
-    /// For now it only supports timestamp versioning for BigQuery and MSSQL dialects.
-    pub fn parse_table_version(&mut self) -> Result<Option<TableVersion>, ParserError> {
-        if dialect_of!(self is BigQueryDialect | MsSqlDialect)
-            && self.parse_keywords(&[Keyword::FOR, Keyword::SYSTEM_TIME, Keyword::AS, Keyword::OF])
-        {
-            let expr = self.parse_expr()?;
-            Ok(Some(TableVersion::ForSystemTimeAsOf(expr)))
-        } else {
-            Ok(None)
+    /// Parses a the timestamp version specifier (i.e. query historical data)
+    pub fn maybe_parse_table_version(&mut self) -> Result<Option<TableVersion>, ParserError> {
+        if self.dialect.supports_timestamp_versioning() {
+            if self.parse_keywords(&[Keyword::FOR, Keyword::SYSTEM_TIME, Keyword::AS, Keyword::OF])
+            {
+                let expr = self.parse_expr()?;
+                return Ok(Some(TableVersion::ForSystemTimeAsOf(expr)));
+            } else if self.peek_keyword(Keyword::AT) || self.peek_keyword(Keyword::BEFORE) {
+                let func_name = self.parse_object_name(true)?;
+                let func = self.parse_function(func_name)?;
+                return Ok(Some(TableVersion::Function(func)));
+            }
         }
+        Ok(None)
     }
 
     /// Parses MySQL's JSON_TABLE column definition.

tests/sqlparser_snowflake.rs

@@ -3144,3 +3144,10 @@ fn test_sql_keywords_as_select_item_aliases() {
             .is_err());
     }
 }
+
+#[test]
+fn test_timetravel_at_before() {
+    snowflake().verified_only_select("SELECT * FROM tbl AT(TIMESTAMP => '2024-12-15 00:00:00')");
+    snowflake()
+        .verified_only_select("SELECT * FROM tbl BEFORE(TIMESTAMP => '2024-12-15 00:00:00')");
+}