feat(postgres): root operator closes #2940

tobymao · Feb 9, 2024 · 876e075 · 876e075
1 parent e71d489
commit 876e075
Show file tree

Hide file tree

Showing 5 changed files with 13 additions and 0 deletions.
diff --git a/sqlglot/dialects/postgres.py b/sqlglot/dialects/postgres.py
@@ -244,6 +244,8 @@ class Tokenizer(tokens.Tokenizer):
             "@@": TokenType.DAT,
             "@>": TokenType.AT_GT,
             "<@": TokenType.LT_AT,
+            "|/": TokenType.PIPE_SLASH,
+            "||/": TokenType.DPIPE_SLASH,
             "BEGIN": TokenType.COMMAND,
             "BEGIN TRANSACTION": TokenType.BEGIN,
             "BIGSERIAL": TokenType.BIGSERIAL,

diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -4626,6 +4626,11 @@ class CountIf(AggFunc):
     _sql_names = ["COUNT_IF", "COUNTIF"]
 
 
+# cube root
+class Cbrt(Func):
+    pass
+
+
 class CurrentDate(Func):
     arg_types = {"this": False}
 

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -638,6 +638,8 @@ class Parser(metaclass=_Parser):
         TokenType.NOT: lambda self: self.expression(exp.Not, this=self._parse_equality()),
         TokenType.TILDA: lambda self: self.expression(exp.BitwiseNot, this=self._parse_unary()),
         TokenType.DASH: lambda self: self.expression(exp.Neg, this=self._parse_unary()),
+        TokenType.PIPE_SLASH: lambda self: self.expression(exp.Sqrt, this=self._parse_unary()),
+        TokenType.DPIPE_SLASH: lambda self: self.expression(exp.Cbrt, this=self._parse_unary()),
     }
 
     PRIMARY_PARSERS = {

diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
@@ -57,6 +57,8 @@ class TokenType(AutoName):
     AMP = auto()
     DPIPE = auto()
     PIPE = auto()
+    PIPE_SLASH = auto()
+    DPIPE_SLASH = auto()
     CARET = auto()
     TILDA = auto()
     ARROW = auto()

diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
@@ -8,6 +8,8 @@ class TestPostgres(Validator):
     dialect = "postgres"
 
     def test_postgres(self):
+        self.validate_identity("|/ x", "SQRT(x)")
+        self.validate_identity("||/ x", "CBRT(x)")
         expr = parse_one(
             "SELECT * FROM r CROSS JOIN LATERAL UNNEST(ARRAY[1]) AS s(location)", read="postgres"
         )