Repository: tobymao/sqlglot

Issue:
Unpivot with a table alias leads to error
Hey there, and thanks for a great project!

I noticed there's a syntax error when using a PIVOT / UNPIVOT keyword (this example uses BigQuery syntax): 

```sql
   select 1 from a.b.table1 as t
    unpivot (
        (c3)
        for c4
        in (a,b)
    )
```

Leads to the following parsing error:
```bash
$ python3 -m sqlglot --read bigquery ./unpivot.sql
Traceback (most recent call last):
  File "/usr/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/home/e/.local/lib/python3.8/site-packages/sqlglot/__main__.py", line 69, in <module>
    sqls = sqlglot.transpile(
  File "/home/e/.local/lib/python3.8/site-packages/sqlglot/__init__.py", line 177, in transpile
    for expression in parse(sql, read, error_level=error_level)
  File "/home/e/.local/lib/python3.8/site-packages/sqlglot/__init__.py", line 65, in parse
    return dialect.parse(sql, **opts)
  File "/home/e/.local/lib/python3.8/site-packages/sqlglot/dialects/dialect.py", line 163, in parse
    return self.parser(**opts).parse(self.tokenizer.tokenize(sql), sql)
  File "/home/e/.local/lib/python3.8/site-packages/sqlglot/parser.py", line 745, in parse
    return self._parse(
  File "/home/e/.local/lib/python3.8/site-packages/sqlglot/parser.py", line 811, in _parse
    self.raise_error("Invalid expression / Unexpected token")
  File "/home/e/.local/lib/python3.8/site-packages/sqlglot/parser.py", line 854, in raise_error
    raise error
sqlglot.errors.ParseError: Invalid expression / Unexpected token. Line 1, Col: 1.
  ./unpivot.sql
```


I would love to hear your thoughts about this, or possibly help implementing it, if you could point me in the right direction.

Thanks!


Discussion:

Resolving patch:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -2504,6 +2504,7 @@ class Tag(Expression):
 class Pivot(Expression):
     arg_types = {
         "this": False,
+        "alias": False,
         "expressions": True,
         "field": True,
         "unpivot": True,
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -1003,10 +1003,6 @@ def table_sql(self, expression: exp.Table, sep: str = " AS ") -> str:
         system_time = expression.args.get("system_time")
         system_time = f" {self.sql(expression, 'system_time')}" if system_time else ""
 
-        if alias and pivots:
-            pivots = f"{pivots}{alias}"
-            alias = ""
-
         return f"{table}{system_time}{alias}{hints}{laterals}{joins}{pivots}"
 
     def tablesample_sql(self, expression: exp.TableSample) -> str:
@@ -1034,11 +1030,13 @@ def tablesample_sql(self, expression: exp.TableSample) -> str:
 
     def pivot_sql(self, expression: exp.Pivot) -> str:
         this = self.sql(expression, "this")
+        alias = self.sql(expression, "alias")
+        alias = f" AS {alias}" if alias else ""
         unpivot = expression.args.get("unpivot")
         direction = "UNPIVOT" if unpivot else "PIVOT"
         expressions = self.expressions(expression, key="expressions")
         field = self.sql(expression, "field")
-        return f"{this} {direction}({expressions} FOR {field})"
+        return f"{this} {direction}({expressions} FOR {field}){alias}"
 
     def tuple_sql(self, expression: exp.Tuple) -> str:
         return f"({self.expressions(expression, flat=True)})"
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -2042,6 +2042,9 @@ def _parse_table(
         if alias:
             this.set("alias", alias)
 
+        if not this.args.get("pivots"):
+            this.set("pivots", self._parse_pivots())
+
         if self._match_pair(TokenType.WITH, TokenType.L_PAREN):
             this.set(
                 "hints",
@@ -2182,7 +2185,12 @@ def _parse_pivot(self) -> t.Optional[exp.Expression]:
 
         self._match_r_paren()
 
-        return self.expression(exp.Pivot, expressions=expressions, field=field, unpivot=unpivot)
+        pivot = self.expression(exp.Pivot, expressions=expressions, field=field, unpivot=unpivot)
+
+        if not self._match_set((TokenType.PIVOT, TokenType.UNPIVOT), advance=False):
+            pivot.set("alias", self._parse_table_alias())
+
+        return pivot
 
     def _parse_where(self, skip_where_token: bool = False) -> t.Optional[exp.Expression]:
         if not skip_where_token and not self._match(TokenType.WHERE):
@@ -3783,12 +3791,13 @@ def _match(self, token_type, advance=True):
 
         return None
 
-    def _match_set(self, types):
+    def _match_set(self, types, advance=True):
         if not self._curr:
             return None
 
         if self._curr.token_type in types:
-            self._advance()
+            if advance:
+                self._advance()
             return True
 
         return None
@@ -3816,9 +3825,10 @@ def _match_r_paren(self, expression=None):
         if expression and self._prev_comments:
             expression.comments = self._prev_comments
 
-    def _match_texts(self, texts):
+    def _match_texts(self, texts, advance=True):
         if self._curr and self._curr.text.upper() in texts:
-            self._advance()
+            if advance:
+                self._advance()
             return True
         return False
 


Test patch:
diff --git a/tests/fixtures/identity.sql b/tests/fixtures/identity.sql
--- a/tests/fixtures/identity.sql
+++ b/tests/fixtures/identity.sql
@@ -305,6 +305,7 @@ SELECT a FROM test TABLESAMPLE(100 ROWS)
 SELECT a FROM test TABLESAMPLE BERNOULLI (50)
 SELECT a FROM test TABLESAMPLE SYSTEM (75)
 SELECT a FROM test PIVOT(SUM(x) FOR y IN ('z', 'q'))
+SELECT 1 FROM a.b.table1 AS t UNPIVOT((c3) FOR c4 IN (a, b))
 SELECT a FROM test PIVOT(SOMEAGG(x, y, z) FOR q IN (1))
 SELECT a FROM test PIVOT(SUM(x) FOR y IN ('z', 'q')) PIVOT(MAX(b) FOR c IN ('d'))
 SELECT a FROM (SELECT a, b FROM test) PIVOT(SUM(x) FOR y IN ('z', 'q'))
diff --git a/tests/test_transpile.py b/tests/test_transpile.py
--- a/tests/test_transpile.py
+++ b/tests/test_transpile.py
@@ -28,6 +28,11 @@ def test_alias(self):
         self.assertEqual(transpile("SELECT 1 current_datetime")[0], "SELECT 1 AS current_datetime")
         self.assertEqual(transpile("SELECT 1 row")[0], "SELECT 1 AS row")
 
+        self.assertEqual(
+            transpile("SELECT 1 FROM a.b.table1 t UNPIVOT((c3) FOR c4 IN (a, b))")[0],
+            "SELECT 1 FROM a.b.table1 AS t UNPIVOT((c3) FOR c4 IN (a, b))",
+        )
+
         for key in ("union", "over", "from", "join"):
             with self.subTest(f"alias {key}"):
                 self.validate(f"SELECT x AS {key}", f"SELECT x AS {key}")