Databricks: Support Timetravel With "VERSION AS OF" (apache#2155)

Author: JamesVorder

src/ast/query.rs

@@ -2398,6 +2398,10 @@ pub enum TableVersion {
     /// Databricks supports this syntax.
     /// For example: `SELECT * FROM tbl TIMESTAMP AS OF CURRENT_TIMESTAMP() - INTERVAL 1 HOUR`
     TimestampAsOf(Expr),
+    /// When the table version is defined using `VERSION AS OF`.
+    /// Databricks supports this syntax.
+    /// For example: `SELECT * FROM tbl VERSION AS OF 2`
+    VersionAsOf(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),
@@ -2408,6 +2412,7 @@ impl Display for TableVersion {
         match self {
             TableVersion::ForSystemTimeAsOf(e) => write!(f, "FOR SYSTEM_TIME AS OF {e}")?,
             TableVersion::TimestampAsOf(e) => write!(f, "TIMESTAMP AS OF {e}")?,
+            TableVersion::VersionAsOf(e) => write!(f, "VERSION AS OF {e}")?,
             TableVersion::Function(func) => write!(f, "{func}")?,
         }
         Ok(())

src/dialect/bigquery.rs

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

src/dialect/databricks.rs

@@ -48,7 +48,7 @@ impl Dialect for DatabricksDialect {
     }
 
     /// <https://docs.databricks.com/gcp/en/delta/history#delta-time-travel-syntax>
-    fn supports_timestamp_versioning(&self) -> bool {
+    fn supports_table_versioning(&self) -> bool {
         true
     }
 

src/dialect/mod.rs

@@ -1072,7 +1072,7 @@ pub trait Dialect: Debug + Any {
 
     /// 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 {
+    fn supports_table_versioning(&self) -> bool {
         false
     }
 

src/dialect/mssql.rs

@@ -107,7 +107,7 @@ impl Dialect for MsSqlDialect {
     }
 
     /// 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 {
+    fn supports_table_versioning(&self) -> bool {
         true
     }
 

src/dialect/snowflake.rs

@@ -540,7 +540,7 @@ impl Dialect for SnowflakeDialect {
     }
 
     /// See: <https://docs.snowflake.com/en/sql-reference/constructs/at-before>
-    fn supports_timestamp_versioning(&self) -> bool {
+    fn supports_table_versioning(&self) -> bool {
         true
     }
 

src/parser/mod.rs

@@ -15731,7 +15731,7 @@ impl<'a> Parser<'a> {
 
     /// 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.dialect.supports_table_versioning() {
             if self.parse_keywords(&[Keyword::FOR, Keyword::SYSTEM_TIME, Keyword::AS, Keyword::OF])
             {
                 let expr = self.parse_expr()?;
@@ -15743,6 +15743,9 @@ impl<'a> Parser<'a> {
             } else if self.parse_keywords(&[Keyword::TIMESTAMP, Keyword::AS, Keyword::OF]) {
                 let expr = self.parse_expr()?;
                 return Ok(Some(TableVersion::TimestampAsOf(expr)));
+            } else if self.parse_keywords(&[Keyword::VERSION, Keyword::AS, Keyword::OF]) {
+                let expr = Expr::Value(self.parse_number_value()?);
+                return Ok(Some(TableVersion::VersionAsOf(expr)));
             }
         }
         Ok(None)

tests/sqlparser_databricks.rs

@@ -371,14 +371,21 @@ fn data_type_timestamp_ntz() {
 
 #[test]
 fn parse_table_time_travel() {
-    all_dialects_where(|d| d.supports_timestamp_versioning())
+    all_dialects_where(|d| d.supports_table_versioning())
         .verified_only_select("SELECT 1 FROM t1 TIMESTAMP AS OF '2018-10-18T22:15:12.013Z'");
 
-    all_dialects_where(|d| d.supports_timestamp_versioning()).verified_only_select(
+    all_dialects_where(|d| d.supports_table_versioning()).verified_only_select(
         "SELECT 1 FROM t1 TIMESTAMP AS OF CURRENT_TIMESTAMP() - INTERVAL 12 HOURS",
     );
 
+    all_dialects_where(|d| d.supports_table_versioning())
+        .verified_only_select("SELECT 1 FROM t1 VERSION AS OF 1");
+
     assert!(databricks()
         .parse_sql_statements("SELECT 1 FROM t1 FOR TIMESTAMP AS OF 'some_timestamp'")
         .is_err());
+
+    assert!(all_dialects_where(|d| d.supports_table_versioning())
+        .parse_sql_statements("SELECT 1 FROM t1 VERSION AS OF 1 - 2",)
+        .is_err())
 }