Repository: tobymao/sqlglot

Issue:
Parse error with 'extract' function (Postgresql)
Hello,

This code:
```
import sqlglot
sqlglot.transpile("select extract(month from '2021-01-31'::timestamp without time zone)", read='postgres', pretty=True)
```
raises an error:
ParseError: Expected ). Line 0, Col: 26.
select LPAD(extract(month from '2021-01-31'::timestamp without time zone)::text,

The query works without any issue in Postgresql.

Here is the complete traceback:
[tb_sqlglot.txt](https://github.com/tobymao/sqlglot/files/6883277/tb_sqlglot.txt)

Thanks.

Best,

Discussion:

Patch:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index e3cca63b..a37d1cfa 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -278,7 +278,7 @@ class Dot(Binary):
 
 
 class DPipe(Binary):
-    token_type = TokenType.D_PIPE
+    token_type = TokenType.DPIPE
 
 
 class EQ(Binary):
@@ -385,6 +385,11 @@ class Decimal(Expression):
     arg_types = {'precision': False, 'scale': False}
 
 
+class Extract(Expression):
+    token_type = TokenType.EXTRACT
+    arg_types = {'this': True, 'expression': True}
+
+
 class In(Expression):
     token_type = TokenType.IN
     arg_types = {'this': True, 'expressions': True}
@@ -395,6 +400,11 @@ class Interval(Expression):
     arg_types = {'this': True, 'unit': True}
 
 
+class Timestamp(Expression):
+    token_type = TokenType.TIMESTAMP
+    arg_types = {'tz': False}
+
+
 # Functions
 class Func(Expression):
     token_type = TokenType.FUNC
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 8c1a0728..683101b4 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -344,6 +344,11 @@ class Generator:
         )
         return f"DECIMAL({args})"
 
+    def extract_sql(self, expression):
+        this = self.sql(expression, 'this')
+        expression_sql = self.sql(expression, 'expression')
+        return f"EXTRACT({this} FROM {expression_sql})"
+
     def if_sql(self, expression):
         return self.case_sql(exp.Case(ifs=[expression], default=expression.args['false']))
 
@@ -365,6 +370,10 @@ class Generator:
     def not_sql(self, expression):
         return f"NOT {self.sql(expression, 'this')}"
 
+    def timestamp_sql(self, expression):
+        tz = 'TZ' if expression.args.get('tz') else ''
+        return f"TIMESTAMP{tz}"
+
     def alias_sql(self, expression):
         to_sql = self.sql(expression, 'alias')
         to_sql = f" AS {to_sql}" if to_sql else ''
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 3ffbf51a..bc1b0590 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -60,6 +60,7 @@ class Parser:
         TokenType.BINARY,
         TokenType.JSON,
         TokenType.TIMESTAMP,
+        TokenType.TIMESTAMPTZ,
         TokenType.DATE,
         TokenType.ARRAY,
         TokenType.DECIMAL,
@@ -327,7 +328,7 @@ class Parser:
     def _parse_value(self):
         if not self._match(TokenType.L_PAREN):
             self.raise_error('Expected ( for values')
-        expressions = self._parse_csv(self._parse_primary)
+        expressions = self._parse_csv(self._parse_time)
         if not self._match(TokenType.R_PAREN):
             self.raise_error('Expected ) for values')
         return exp.Tuple(expressions=expressions)
@@ -518,7 +519,7 @@ class Parser:
 
     def _parse_ordered(self):
         return exp.Ordered(
-            this=self._parse_primary(),
+            this=self._parse_time(),
             desc=not self._match(TokenType.ASC) and self._match(TokenType.DESC),
         )
 
@@ -591,16 +592,105 @@ class Parser:
             return exp.BitwiseNot(this=self._parse_unary())
         if self._match(TokenType.DASH):
             return exp.Neg(this=self._parse_unary())
-        return self._parse_special()
+        return self._parse_time()
 
-    def _parse_special(self):
-        if self._match(TokenType.CAST):
-            return self._parse_cast()
-        if self._match(TokenType.CASE):
+    def _parse_time(self):
+        if self._match(TokenType.INTERVAL):
+            return exp.Interval(this=self._match(TokenType.STRING, TokenType.NUMBER), unit=self._match(TokenType.VAR))
+
+        time_type = self._parse_time_type()
+        this = self._parse_primary()
+
+        if time_type:
+            return exp.Cast(this=this, to=time_type)
+
+        if self._match(TokenType.DCOLON):
+            return exp.Cast(this=this, to=self._parse_time_type())
+
+        return this
+
+    def _parse_time_type(self):
+        if self._match(TokenType.DATE):
+            return self._prev
+        if self._match(TokenType.TIMESTAMP, TokenType.TIMESTAMPTZ):
+            tz = self._match(TokenType.WITH)
+            self._match(TokenType.WITHOUT)
+            self._match(TokenType.TIME)
+            self._match(TokenType.ZONE)
+            return exp.Timestamp(tz=tz)
+        return None
+
+    def _parse_primary(self):
+        if self._match(*self.PRIMARY_TOKENS):
+            return self._prev
+
+        if self._match(TokenType.L_PAREN):
+            paren = self._prev
+            this = self._parse_conjunction()
+
+            if not self._match(TokenType.R_PAREN):
+                self.raise_error('Expecting )', paren)
+            return exp.Paren(this=this)
+
+        if self._curr is None:
+            return self.raise_error('Expecting expression')
+
+        return self._parse_column()
+
+    def _parse_column(self):
+        db = None
+        table = None
+
+        if self._curr.token_type in (TokenType.OVER, TokenType.R_PAREN, TokenType.R_BRACKET, TokenType.WHEN):
+            return None
+
+        self._advance()
+        this = self._parse_function(self._prev)
+
+        if self._match(TokenType.DOT):
+            table = this
+            if not self._match(*self.COLUMN_TOKENS):
+                self.raise_error('Expected column name')
+            this = self._prev
+
+            if self._match(TokenType.DOT):
+                db = table
+                table = this
+                if not self._match(*self.COLUMN_TOKENS):
+                    self.raise_error('Expected column name')
+                this = self._prev
+
+        if not isinstance(this, exp.Func) and this.token_type not in (TokenType.ARRAY, TokenType.MAP):
+            this = exp.Column(this=this, db=db, table=table)
+
+        return self._parse_brackets(this)
+
+    def _parse_function(self, this):
+        if this.token_type == TokenType.CASE:
             return self._parse_case()
-        if self._match(TokenType.COUNT):
-            return self._parse_count()
-        return self._parse_primary()
+
+        if not self._match(TokenType.L_PAREN):
+            return this
+
+        if this.token_type == TokenType.CAST:
+            this = self._parse_cast()
+        elif this.token_type == TokenType.COUNT:
+            this = self._parse_count()
+        elif this.token_type == TokenType.EXTRACT:
+            this = self._parse_extract()
+        else:
+            args = self._parse_csv(self._parse_conjunction)
+            function = self.functions.get(this.text.upper())
+
+            if not callable(function):
+                this = exp.Anonymous(this=this, expressions=args)
+            else:
+                this = function(args)
+
+        if not self._match(TokenType.R_PAREN):
+            self.raise_error('Expected )')
+
+        return self._parse_window(this)
 
     def _parse_case(self):
         ifs = []
@@ -623,52 +713,34 @@ class Parser:
         return self._parse_brackets(exp.Case(this=expression, ifs=ifs, default=default))
 
     def _parse_count(self):
-        if not self._match(TokenType.L_PAREN):
-            self.raise_error("Expected ( after COUNT", self._prev)
+        return exp.Count(
+            distinct=self._match(TokenType.DISTINCT),
+            this=self._parse_conjunction(),
+        )
 
-        distinct = self._match(TokenType.DISTINCT)
-        this = self._parse_conjunction()
+    def _parse_extract(self):
+        this = self._match(TokenType.VAR)
 
-        if not self._match(TokenType.R_PAREN):
-            self.raise_error("Expected ) after COUNT")
+        if not self._match(TokenType.FROM):
+            self.raise_error('Expected FROM after EXTRACT', self._prev)
 
-        return self._parse_window(exp.Count(this=this, distinct=distinct))
+        expression = self._parse_time()
 
-    def _parse_cast(self):
-        if not self._match(TokenType.L_PAREN):
-            self.raise_error("Expected ( after CAST", self._prev)
+        return exp.Extract(this=this, expression=expression)
 
+    def _parse_cast(self):
         this = self._parse_conjunction()
 
         if not self._match(TokenType.ALIAS):
-            self.raise_error("Expected AS after CAST")
+            self.raise_error('Expected AS after CAST')
 
         if not self._match(*self.TYPE_TOKENS):
-            self.raise_error("Expected type after CAST")
-
-        to = self._parse_function(self._parse_brackets(self._prev))
+            self.raise_error('Expected TYPE after CAST')
 
-        if not self._match(TokenType.R_PAREN):
-            self.raise_error("Expected ) after CAST")
-
-        return exp.Cast(this=this, to=to)
-
-    def _parse_function(self, this):
-        if not self._match(TokenType.L_PAREN):
-            return this
-
-        args = self._parse_csv(self._parse_conjunction)
-        function = self.functions.get(this.text.upper())
-
-        if not self._match(TokenType.R_PAREN):
-            self.raise_error('Expected )')
-
-        if not callable(function):
-            this = exp.Anonymous(this=this, expressions=args)
-        else:
-            this = function(args)
-
-        return self._parse_window(this)
+        return exp.Cast(
+            this=this,
+            to=self._parse_function(self._parse_brackets(self._prev)),
+        )
 
     def _parse_window(self, this):
         if not self._match(TokenType.OVER):
@@ -680,7 +752,7 @@ class Parser:
         partition = None
 
         if self._match(TokenType.PARTITION):
-            partition = self._parse_csv(self._parse_primary)
+            partition = self._parse_csv(self._parse_time)
 
         order = self._parse_order()
 
@@ -706,60 +778,13 @@ class Parser:
 
         return exp.Window(this=this, partition=partition, order=order, spec=spec)
 
-    def _parse_primary(self):
-        if self._match(*self.PRIMARY_TOKENS):
-            return self._prev
-
-        time_type = self._match(TokenType.TIMESTAMP, TokenType.DATE)
-
-        if time_type:
-            return exp.Cast(this=self._parse_primary(), to=time_type)
-
-        interval = self._parse_interval()
-
-        if interval:
-            return interval
-
-        if self._match(TokenType.L_PAREN):
-            paren = self._prev
-            this = self._parse_conjunction()
-
-            if not self._match(TokenType.R_PAREN):
-                self.raise_error('Expecting )', paren)
-            return exp.Paren(this=this)
-
-        if self._curr is None:
-            return self.raise_error('Expecting expression')
-
-        return self._parse_column()
-
-    def _parse_column(self):
-        db = None
-        table = None
-
-        if self._curr.token_type in (TokenType.OVER, TokenType.R_PAREN, TokenType.R_BRACKET, TokenType.WHEN):
-            return None
-
-        self._advance()
-        this = self._parse_function(self._prev)
-
-        if self._match(TokenType.DOT):
-            table = this
-            if not self._match(*self.COLUMN_TOKENS):
-                self.raise_error('Expected column name')
-            this = self._prev
-
-            if self._match(TokenType.DOT):
-                db = table
-                table = this
-                if not self._match(*self.COLUMN_TOKENS):
-                    self.raise_error('Expected column name')
-                this = self._prev
-
-        if not isinstance(this, exp.Func) and this.token_type not in (TokenType.ARRAY, TokenType.MAP):
-            this = exp.Column(this=this, db=db, table=table)
+    def _parse_window_spec(self):
+        self._match(TokenType.BETWEEN)
 
-        return self._parse_brackets(this)
+        return {
+            'value': self._match(TokenType.UNBOUNDED, TokenType.CURRENT_ROW) or self._parse_time(),
+            'side': self._match(TokenType.PRECEDING, TokenType.FOLLOWING),
+        }
 
     def _parse_brackets(self, this):
         if not self._match(TokenType.L_BRACKET):
@@ -782,14 +807,6 @@ class Parser:
             this = exp.Dot(this=this, expression=self._parse_id_var())
         return this
 
-    def _parse_window_spec(self):
-        self._match(TokenType.BETWEEN)
-
-        return {
-            'value': self._match(TokenType.UNBOUNDED, TokenType.CURRENT_ROW) or self._parse_primary(),
-            'side': self._match(TokenType.PRECEDING, TokenType.FOLLOWING),
-        }
-
     def _parse_alias(self, this):
         self._match(TokenType.ALIAS)
 
@@ -802,13 +819,6 @@ class Parser:
     def _parse_id_var(self):
         return self._match(*self.ID_VAR_TOKENS)
 
-    def _parse_interval(self):
-        if self._match(TokenType.INTERVAL):
-            this = self._match(TokenType.STRING, TokenType.NUMBER)
-            unit = self._match(TokenType.VAR)
-            return exp.Interval(this=this, unit=unit)
-        return None
-
     def _parse_csv(self, parse):
         items = [parse()]
 
diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
index 658dcfd8..7cb7cb41 100644
--- a/sqlglot/tokens.py
+++ b/sqlglot/tokens.py
@@ -18,6 +18,7 @@ class TokenType(AutoName):
     DASH = auto()
     PLUS = auto()
     COLON = auto()
+    DCOLON = auto()
     SEMICOLON = auto()
     STAR = auto()
     SLASH = auto()
@@ -31,7 +32,7 @@ class TokenType(AutoName):
     AND = auto()
     OR = auto()
     AMP = auto()
-    D_PIPE = auto()
+    DPIPE = auto()
     PIPE = auto()
     CARET = auto()
     TILDA = auto()
@@ -63,6 +64,7 @@ class TokenType(AutoName):
     BINARY = auto()
     JSON = auto()
     TIMESTAMP = auto()
+    TIMESTAMPTZ = auto()
     DATE = auto()
 
     # keywords
@@ -88,6 +90,7 @@ class TokenType(AutoName):
     ELSE = auto()
     END = auto()
     EXISTS = auto()
+    EXTRACT = auto()
     FOLLOWING = auto()
     FORMAT = auto()
     FULL = auto()
@@ -128,6 +131,7 @@ class TokenType(AutoName):
     SPEC = auto()
     STORED = auto()
     TEMPORARY = auto()
+    TIME = auto()
     TUPLE = auto()
     THEN = auto()
     UNBOUNDED = auto()
@@ -138,6 +142,8 @@ class TokenType(AutoName):
     WHEN = auto()
     WHERE = auto()
     WITH = auto()
+    WITHOUT = auto()
+    ZONE = auto()
 
 
 class Token:
@@ -201,7 +207,8 @@ class Tokenizer:
         '/*': TokenType.COMMENT_START,
         '*/': TokenType.COMMENT_END,
         '==': TokenType.EQ,
-        '||': TokenType.D_PIPE,
+        '::': TokenType.DCOLON,
+        '||': TokenType.DPIPE,
         '>=': TokenType.GTE,
         '<=': TokenType.LTE,
         '<>': TokenType.NEQ,
@@ -228,6 +235,7 @@ class Tokenizer:
         'ELSE': TokenType.ELSE,
         'END': TokenType.END,
         'EXISTS': TokenType.EXISTS,
+        'EXTRACT': TokenType.EXTRACT,
         'FORMAT': TokenType.FORMAT,
         'FULL': TokenType.FULL,
         'FOLLOWING': TokenType.FOLLOWING,
@@ -267,6 +275,7 @@ class Tokenizer:
         'TABLE': TokenType.TABLE,
         'TEMPORARY': TokenType.TEMPORARY,
         'THEN': TokenType.THEN,
+        'TIME': TokenType.TIME,
         'UNBOUNDED': TokenType.UNBOUNDED,
         'UNION': TokenType.UNION,
         'UNNEST': TokenType.UNNEST,
@@ -275,6 +284,8 @@ class Tokenizer:
         'WHEN': TokenType.WHEN,
         'WHERE': TokenType.WHERE,
         'WITH': TokenType.WITH,
+        'WITHOUT': TokenType.WITHOUT,
+        'ZONE': TokenType.ZONE,
 
         'ARRAY': TokenType.ARRAY,
         'BOOL': TokenType.BOOLEAN,
@@ -304,6 +315,7 @@ class Tokenizer:
         'TEXT': TokenType.TEXT,
         'BINARY': TokenType.BINARY,
         'TIMESTAMP': TokenType.TIMESTAMP,
+        'TIMESTAMPTZ': TokenType.TIMESTAMPTZ,
         'DATE': TokenType.DATE,
     }
 


Test patch:
diff --git a/tests/test_transpile.py b/tests/test_transpile.py
index 20bf3dc3..63a2e1d4 100644
--- a/tests/test_transpile.py
+++ b/tests/test_transpile.py
@@ -1,3 +1,4 @@
+# pylint: disable=too-many-statements
 import os
 import unittest
 
@@ -54,13 +55,31 @@ class TestTranspile(unittest.TestCase):
         self.validate('a NOT BETWEEN b AND c', 'NOT a BETWEEN b AND c')
         self.validate('a NOT IN (1, 2)', 'NOT a IN (1, 2)')
 
+    def test_extract(self):
+        self.validate("EXTRACT(day FROM '2020-01-01'::TIMESTAMP)", "EXTRACT(day FROM CAST('2020-01-01' AS TIMESTAMP))")
+        self.validate("EXTRACT(timezone FROM '2020-01-01'::TIMESTAMP)", "EXTRACT(timezone FROM CAST('2020-01-01' AS TIMESTAMP))")
+        self.validate(
+            "EXTRACT(year FROM '2020-01-01'::TIMESTAMP WITH TIME ZONE)",
+            "EXTRACT(year FROM CAST('2020-01-01' AS TIMESTAMPTZ))",
+        )
+        self.validate(
+            "extract(month from '2021-01-31'::timestamp without time zone)",
+            "EXTRACT(month FROM CAST('2021-01-31' AS TIMESTAMP))",
+        )
+
     def test_if(self):
         self.validate('SELECT IF(a > 1, 1, 0) FROM foo', 'SELECT CASE WHEN a > 1 THEN 1 ELSE 0 END FROM foo')
         self.validate('SELECT IF(a > 1, 1) FROM foo', 'SELECT CASE WHEN a > 1 THEN 1 END FROM foo')
 
     def test_time(self):
         self.validate("TIMESTAMP '2020-01-01'", "CAST('2020-01-01' AS TIMESTAMP)")
+        self.validate("TIMESTAMP WITH TIME ZONE '2020-01-01'", "CAST('2020-01-01' AS TIMESTAMPTZ)")
+        self.validate("TIMESTAMP WITHOUT TIME ZONE '2020-01-01'", "CAST('2020-01-01' AS TIMESTAMP)")
+        self.validate("'2020-01-01'::TIMESTAMP", "CAST('2020-01-01' AS TIMESTAMP)")
+        self.validate("'2020-01-01'::TIMESTAMP WITHOUT TIME ZONE", "CAST('2020-01-01' AS TIMESTAMP)")
+        self.validate("'2020-01-01'::TIMESTAMP WITH TIME ZONE", "CAST('2020-01-01' AS TIMESTAMPTZ)")
         self.validate("DATE '2020-01-01'", "CAST('2020-01-01' AS DATE)")
+        self.validate("'2020-01-01'::DATE", "CAST('2020-01-01' AS DATE)")
         self.validate("STR_TO_TIME('x', 'y')", "STRPTIME('x', 'y')", write='duckdb')
         self.validate("STR_TO_UNIX('x', 'y')", "EPOCH(STRPTIME('x', 'y'))", write='duckdb')
         self.validate("TIME_TO_STR(x, 'y')", "STRFTIME(x, 'y')", write='duckdb')