Databricks: Support Timetravel With "TIMESTAMP AS OF" (#2134)

Co-authored-by: Ifeanyi Ubah <[email protected]>

Author: JamesVorder

src/ast/query.rs

@@ -2241,6 +2241,10 @@ 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 `TIMESTAMP AS OF`.
+    /// 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 a function.
     /// For example: `SELECT * FROM tbl AT(TIMESTAMP => '2020-08-14 09:30:00')`
     Function(Expr),
@@ -2250,6 +2254,7 @@ 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::TimestampAsOf(e) => write!(f, "TIMESTAMP AS OF {e}")?,
             TableVersion::Function(func) => write!(f, "{func}")?,
         }
         Ok(())

src/dialect/databricks.rs

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

src/parser/mod.rs

@@ -15493,6 +15493,9 @@ impl<'a> Parser<'a> {
                 let func_name = self.parse_object_name(true)?;
                 let func = self.parse_function(func_name)?;
                 return Ok(Some(TableVersion::Function(func)));
+            } else if self.parse_keywords(&[Keyword::TIMESTAMP, Keyword::AS, Keyword::OF]) {
+                let expr = self.parse_expr()?;
+                return Ok(Some(TableVersion::TimestampAsOf(expr)));
             }
         }
         Ok(None)

tests/sqlparser_bigquery.rs

@@ -1739,7 +1739,7 @@ fn parse_table_time_travel() {
                 args: None,
                 with_hints: vec![],
                 version: Some(TableVersion::ForSystemTimeAsOf(Expr::Value(
-                    Value::SingleQuotedString(version).with_empty_span()
+                    Value::SingleQuotedString(version.clone()).with_empty_span()
                 ))),
                 partitions: vec![],
                 with_ordinality: false,

tests/sqlparser_databricks.rs

@@ -366,3 +366,17 @@ fn data_type_timestamp_ntz() {
         s => panic!("Unexpected statement: {s:?}"),
     }
 }
+
+#[test]
+fn parse_table_time_travel() {
+    all_dialects_where(|d| d.supports_timestamp_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(
+        "SELECT 1 FROM t1 TIMESTAMP AS OF CURRENT_TIMESTAMP() - INTERVAL 12 HOURS",
+    );
+
+    assert!(databricks()
+        .parse_sql_statements("SELECT 1 FROM t1 FOR TIMESTAMP AS OF 'some_timestamp'")
+        .is_err());
+}