Repository: tobymao/sqlglot

Issue:
Snowflake SELECT ... AT/BEFORE syntax (time travel) does not parse in 20.1.0
In 20.1.0, the SELECT AT/BEFORE (time travel) examples from the documentation raise a ParserError. The CREATE TABLE/SCHEMA/DATABASE examples transpile successfully.

Documentation/references:
* https://docs.snowflake.com/en/user-guide/data-time-travel#querying-historical-data
* https://docs.snowflake.com/en/sql-reference/constructs/at-before 
* Issue for time travel for other databases: https://github.com/tobymao/sqlglot/issues/2128

```
examples = [
    "SELECT * FROM my_table AT(TIMESTAMP => 'Fri, 01 May 2015 16:20:00 -0700'::timestamp);",
    "SELECT * FROM my_table AT(TIMESTAMP => TO_TIMESTAMP(1432669154242, 3));",
    "SELECT * FROM my_table AT(OFFSET => -60*5) AS T WHERE T.flag = 'valid';",
    "SELECT * FROM my_table AT (STATEMENT=>$query_id_var);",
    "SELECT * FROM my_table AT(TIMESTAMP => 'Fri, 01 May 2015 16:20:00 -0700'::timestamp_tz);",
    "SELECT * FROM my_table AT (OFFSET => -60*5);",
    "SELECT * FROM my_table AT(STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726');",
    "SELECT * FROM my_table BEFORE(STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726');",
    "SELECT * FROM my_table BEFORE(STATEMENT=>$query_id_var);",
    "SELECT * FROM my_table BEFORE (TIMESTAMP => 'Fri, 01 May 2015 16:20:00 -0700'::timestamp_tz);",
    "SELECT * FROM my_table BEFORE(OFFSET => -60*5);",
    "SELECT * FROM my_table BEFORE (STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726');",
    """
    SELECT oldt.* , newt.*
    FROM my_table BEFORE(STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726') AS oldt
    FULL OUTER JOIN my_table AT(STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726') AS newt
    ON oldt.id = newt.id
    WHERE oldt.id IS NULL OR newt.id IS NULL;
    """,
    # these work today:
    "CREATE TABLE restored_table CLONE my_table AT(TIMESTAMP => 'Sat, 09 May 2015 01:01:00 +0300'::timestamp_tz);",
    "CREATE SCHEMA restored_schema CLONE my_schema AT(OFFSET => -3600);",
    "CREATE DATABASE restored_db CLONE my_db BEFORE(STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726');",
]

import sqlglot
import sqlglot.errors

print(sqlglot.__version__)
for i, s in enumerate(examples):
    try:
        t = sqlglot.transpile(
            s.strip(), read="snowflake", write="snowflake", pretty=True
        )
        print(i, "okay", t)
    except sqlglot.errors.ParseError as e:
        print(i, "error", e)
```

Output (expected is all are okay):
```
20.1.0
0 error Expecting ). Line 1, Col: 38.
  SELECT * FROM my_table AT(TIMESTAMP => 'Fri, 01 May 2015 16:20:00 -0700'::timestamp);
1 error Expecting ). Line 1, Col: 38.
  SELECT * FROM my_table AT(TIMESTAMP => TO_TIMESTAMP(1432669154242, 3));
2 error Expecting ). Line 1, Col: 35.
  SELECT * FROM my_table AT(OFFSET => -60*5) AS T WHERE T.flag = 'valid';
3 error Expecting ). Line 1, Col: 38.
  SELECT * FROM my_table AT (STATEMENT=>$query_id_var);
4 error Expecting ). Line 1, Col: 38.
  SELECT * FROM my_table AT(TIMESTAMP => 'Fri, 01 May 2015 16:20:00 -0700'::timestamp_tz);
5 error Expecting ). Line 1, Col: 36.
  SELECT * FROM my_table AT (OFFSET => -60*5);
6 error Expecting ). Line 1, Col: 38.
  SELECT * FROM my_table AT(STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726');
7 error Expecting ). Line 1, Col: 42.
  SELECT * FROM my_table BEFORE(STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726');
8 error Expecting ). Line 1, Col: 41.
  SELECT * FROM my_table BEFORE(STATEMENT=>$query_id_var);
9 error Expecting ). Line 1, Col: 43.
  SELECT * FROM my_table BEFORE (TIMESTAMP => 'Fri, 01 May 2015 16:20:00 -0700'::timestamp_tz);
10 error Expecting ). Line 1, Col: 39.
  SELECT * FROM my_table BEFORE(OFFSET => -60*5);
11 error Expecting ). Line 1, Col: 43.
  SELECT * FROM my_table BEFORE (STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726');
12 error Expecting ). Line 2, Col: 34.
  SELECT oldt.* , newt.*
    FROM my_table BEFORE(STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726') AS oldt
    FULL OUTER JOIN my_table AT(STATEMENT => '8e5d0
13 okay ["CREATE TABLE restored_table CLONE my_table AT (TIMESTAMP => CAST('Sat, 09 May 2015 01:01:00 +0300' AS TIMESTAMPTZ))"]
14 okay ['CREATE SCHEMA restored_schema CLONE my_schema AT (OFFSET => -3600)']
15 okay ["CREATE DATABASE restored_db CLONE my_db BEFORE (STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726')"]
```

Comments:

Patch:
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index 42e8c661..fca42d48 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -369,12 +369,35 @@ class Snowflake(Dialect):
 
             return lateral
 
+        def _parse_at_before(self, table: exp.Table) -> exp.Table:
+            # https://docs.snowflake.com/en/sql-reference/constructs/at-before
+            index = self._index
+            if self._match_texts(("AT", "BEFORE")):
+                this = self._prev.text.upper()
+                kind = (
+                    self._match(TokenType.L_PAREN)
+                    and self._match_texts(self.HISTORICAL_DATA_KIND)
+                    and self._prev.text.upper()
+                )
+                expression = self._match(TokenType.FARROW) and self._parse_bitwise()
+
+                if expression:
+                    self._match_r_paren()
+                    when = self.expression(
+                        exp.HistoricalData, this=this, kind=kind, expression=expression
+                    )
+                    table.set("when", when)
+                else:
+                    self._retreat(index)
+
+            return table
+
         def _parse_table_parts(self, schema: bool = False) -> exp.Table:
             # https://docs.snowflake.com/en/user-guide/querying-stage
-            if self._match_text_seq("@", advance=False):
-                table: t.Optional[exp.Expression] = self._parse_location_path()
-            elif self._match(TokenType.STRING, advance=False):
+            if self._match(TokenType.STRING, advance=False):
                 table = self._parse_string()
+            elif self._match_text_seq("@", advance=False):
+                table = self._parse_location_path()
             else:
                 table = None
 
@@ -393,9 +416,11 @@ class Snowflake(Dialect):
 
                     self._match(TokenType.COMMA)
 
-                return self.expression(exp.Table, this=table, format=file_format, pattern=pattern)
+                table = self.expression(exp.Table, this=table, format=file_format, pattern=pattern)
+            else:
+                table = super()._parse_table_parts(schema=schema)
 
-            return super()._parse_table_parts(schema=schema)
+            return self._parse_at_before(table)
 
         def _parse_id_var(
             self,
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 19a96df2..6990344e 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -1105,14 +1105,7 @@ class Create(DDL):
 # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#create_table_clone_statement
 # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#create_table_copy
 class Clone(Expression):
-    arg_types = {
-        "this": True,
-        "when": False,
-        "kind": False,
-        "shallow": False,
-        "expression": False,
-        "copy": False,
-    }
+    arg_types = {"this": True, "shallow": False, "copy": False}
 
 
 class Describe(Expression):
@@ -2522,6 +2515,11 @@ class IndexTableHint(Expression):
     arg_types = {"this": True, "expressions": False, "target": False}
 
 
+# https://docs.snowflake.com/en/sql-reference/constructs/at-before
+class HistoricalData(Expression):
+    arg_types = {"this": True, "kind": True, "expression": True}
+
+
 class Table(Expression):
     arg_types = {
         "this": True,
@@ -2538,6 +2536,7 @@ class Table(Expression):
         "pattern": False,
         "index": False,
         "ordinality": False,
+        "when": False,
     }
 
     @property
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index f3f90601..e03462d1 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -862,15 +862,7 @@ class Generator:
         this = self.sql(expression, "this")
         shallow = "SHALLOW " if expression.args.get("shallow") else ""
         keyword = "COPY" if expression.args.get("copy") and self.SUPPORTS_TABLE_COPY else "CLONE"
-        this = f"{shallow}{keyword} {this}"
-        when = self.sql(expression, "when")
-
-        if when:
-            kind = self.sql(expression, "kind")
-            expr = self.sql(expression, "expression")
-            return f"{this} {when} ({kind} => {expr})"
-
-        return this
+        return f"{shallow}{keyword} {this}"
 
     def describe_sql(self, expression: exp.Describe) -> str:
         return f"DESCRIBE {self.sql(expression, 'this')}"
@@ -1400,6 +1392,12 @@ class Generator:
         target = f" FOR {target}" if target else ""
         return f"{this}{target} ({self.expressions(expression, flat=True)})"
 
+    def historicaldata_sql(self, expression: exp.HistoricalData) -> str:
+        this = self.sql(expression, "this")
+        kind = self.sql(expression, "kind")
+        expr = self.sql(expression, "expression")
+        return f"{this} ({kind} => {expr})"
+
     def table_sql(self, expression: exp.Table, sep: str = " AS ") -> str:
         table = ".".join(
             self.sql(part)
@@ -1436,6 +1434,10 @@ class Generator:
             ordinality = f" WITH ORDINALITY{alias}"
             alias = ""
 
+        when = self.sql(expression, "when")
+        if when:
+            table = f"{table} {when}"
+
         return f"{table}{version}{file_format}{alias}{index}{hints}{pivots}{joins}{laterals}{ordinality}"
 
     def tablesample_sql(
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 1fae5959..5399b293 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -907,7 +907,7 @@ class Parser(metaclass=_Parser):
     INSERT_ALTERNATIVES = {"ABORT", "FAIL", "IGNORE", "REPLACE", "ROLLBACK"}
 
     CLONE_KEYWORDS = {"CLONE", "COPY"}
-    CLONE_KINDS = {"TIMESTAMP", "OFFSET", "STATEMENT"}
+    HISTORICAL_DATA_KIND = {"TIMESTAMP", "OFFSET", "STATEMENT", "STREAM"}
 
     OPCLASS_FOLLOW_KEYWORDS = {"ASC", "DESC", "NULLS"}
     OPTYPE_FOLLOW_TOKENS = {TokenType.COMMA, TokenType.R_PAREN}
@@ -1411,23 +1411,8 @@ class Parser(metaclass=_Parser):
 
             if self._match_texts(self.CLONE_KEYWORDS):
                 copy = self._prev.text.lower() == "copy"
-                clone = self._parse_table(schema=True)
-                when = self._match_texts(("AT", "BEFORE")) and self._prev.text.upper()
-                clone_kind = (
-                    self._match(TokenType.L_PAREN)
-                    and self._match_texts(self.CLONE_KINDS)
-                    and self._prev.text.upper()
-                )
-                clone_expression = self._match(TokenType.FARROW) and self._parse_bitwise()
-                self._match(TokenType.R_PAREN)
                 clone = self.expression(
-                    exp.Clone,
-                    this=clone,
-                    when=when,
-                    kind=clone_kind,
-                    shallow=shallow,
-                    expression=clone_expression,
-                    copy=copy,
+                    exp.Clone, this=self._parse_table(schema=True), shallow=shallow, copy=copy
                 )
 
         return self.expression(


Test patch:
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index 29931323..13f32c13 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -72,6 +72,18 @@ WHERE
         self.validate_identity(
             'DESCRIBE TABLE "SNOWFLAKE_SAMPLE_DATA"."TPCDS_SF100TCL"."WEB_SITE" type=stage'
         )
+        self.validate_identity(
+            "SELECT * FROM foo at",
+            "SELECT * FROM foo AS at",
+        )
+        self.validate_identity(
+            "SELECT * FROM foo before",
+            "SELECT * FROM foo AS before",
+        )
+        self.validate_identity(
+            "SELECT * FROM foo at (col)",
+            "SELECT * FROM foo AS at(col)",
+        )
         self.validate_identity(
             "SELECT * FROM unnest(x) with ordinality",
             "SELECT * FROM TABLE(FLATTEN(INPUT => x)) AS _u(seq, key, path, index, value, this)",
@@ -779,6 +791,53 @@ WHERE
             },
         )
 
+    def test_historical_data(self):
+        self.validate_identity("SELECT * FROM my_table AT (STATEMENT => $query_id_var)")
+        self.validate_identity("SELECT * FROM my_table AT (OFFSET => -60 * 5)")
+        self.validate_identity("SELECT * FROM my_table BEFORE (STATEMENT => $query_id_var)")
+        self.validate_identity("SELECT * FROM my_table BEFORE (OFFSET => -60 * 5)")
+        self.validate_identity("CREATE SCHEMA restored_schema CLONE my_schema AT (OFFSET => -3600)")
+        self.validate_identity(
+            "CREATE TABLE restored_table CLONE my_table AT (TIMESTAMP => CAST('Sat, 09 May 2015 01:01:00 +0300' AS TIMESTAMPTZ))",
+        )
+        self.validate_identity(
+            "CREATE DATABASE restored_db CLONE my_db BEFORE (STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726')"
+        )
+        self.validate_identity(
+            "SELECT * FROM my_table AT (TIMESTAMP => TO_TIMESTAMP(1432669154242, 3))"
+        )
+        self.validate_identity(
+            "SELECT * FROM my_table AT (OFFSET => -60 * 5) AS T WHERE T.flag = 'valid'"
+        )
+        self.validate_identity(
+            "SELECT * FROM my_table AT (STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726')"
+        )
+        self.validate_identity(
+            "SELECT * FROM my_table BEFORE (STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726')"
+        )
+        self.validate_identity(
+            "SELECT * FROM my_table AT (TIMESTAMP => 'Fri, 01 May 2015 16:20:00 -0700'::timestamp)",
+            "SELECT * FROM my_table AT (TIMESTAMP => CAST('Fri, 01 May 2015 16:20:00 -0700' AS TIMESTAMPNTZ))",
+        )
+        self.validate_identity(
+            "SELECT * FROM my_table AT(TIMESTAMP => 'Fri, 01 May 2015 16:20:00 -0700'::timestamp_tz)",
+            "SELECT * FROM my_table AT (TIMESTAMP => CAST('Fri, 01 May 2015 16:20:00 -0700' AS TIMESTAMPTZ))",
+        )
+        self.validate_identity(
+            "SELECT * FROM my_table BEFORE (TIMESTAMP => 'Fri, 01 May 2015 16:20:00 -0700'::timestamp_tz);",
+            "SELECT * FROM my_table BEFORE (TIMESTAMP => CAST('Fri, 01 May 2015 16:20:00 -0700' AS TIMESTAMPTZ))",
+        )
+        self.validate_identity(
+            """
+            SELECT oldt.* , newt.*
+            FROM my_table BEFORE(STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726') AS oldt
+            FULL OUTER JOIN my_table AT(STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726') AS newt
+            ON oldt.id = newt.id
+            WHERE oldt.id IS NULL OR newt.id IS NULL;
+            """,
+            "SELECT oldt.*, newt.* FROM my_table BEFORE (STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726') AS oldt FULL OUTER JOIN my_table AT (STATEMENT => '8e5d0ca9-005e-44e6-b858-a8f5b37c5726') AS newt ON oldt.id = newt.id WHERE oldt.id IS NULL OR newt.id IS NULL",
+        )
+
     def test_ddl(self):
         self.validate_identity(
             """create external table et2(