Feat(postgres): add support for all range/multirange types (#1718)

tobymao · Jun 3, 2023 · 5867fc4 · 5867fc4
1 parent 4decebb
commit 5867fc4
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 0 deletions.
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -3222,6 +3222,18 @@ class Type(AutoName):
         DATE = auto()
         DATETIME = auto()
         DATETIME64 = auto()
+        INT4RANGE = auto()
+        INT4MULTIRANGE = auto()
+        INT8RANGE = auto()
+        INT8MULTIRANGE = auto()
+        NUMRANGE = auto()
+        NUMMULTIRANGE = auto()
+        TSRANGE = auto()
+        TSMULTIRANGE = auto()
+        TSTZRANGE = auto()
+        TSTZMULTIRANGE = auto()
+        DATERANGE = auto()
+        DATEMULTIRANGE = auto()
         DECIMAL = auto()
         DOUBLE = auto()
         FLOAT = auto()

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -155,6 +155,18 @@ class Parser(metaclass=_Parser):
         TokenType.DATETIME,
         TokenType.DATETIME64,
         TokenType.DATE,
+        TokenType.INT4RANGE,
+        TokenType.INT4MULTIRANGE,
+        TokenType.INT8RANGE,
+        TokenType.INT8MULTIRANGE,
+        TokenType.NUMRANGE,
+        TokenType.NUMMULTIRANGE,
+        TokenType.TSRANGE,
+        TokenType.TSMULTIRANGE,
+        TokenType.TSTZRANGE,
+        TokenType.TSTZMULTIRANGE,
+        TokenType.DATERANGE,
+        TokenType.DATEMULTIRANGE,
         TokenType.DECIMAL,
         TokenType.BIGDECIMAL,
         TokenType.UUID,

diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
@@ -113,6 +113,18 @@ class TokenType(AutoName):
     DATETIME = auto()
     DATETIME64 = auto()
     DATE = auto()
+    INT4RANGE = auto()
+    INT4MULTIRANGE = auto()
+    INT8RANGE = auto()
+    INT8MULTIRANGE = auto()
+    NUMRANGE = auto()
+    NUMMULTIRANGE = auto()
+    TSRANGE = auto()
+    TSMULTIRANGE = auto()
+    TSTZRANGE = auto()
+    TSTZMULTIRANGE = auto()
+    DATERANGE = auto()
+    DATEMULTIRANGE = auto()
     UUID = auto()
     GEOGRAPHY = auto()
     NULLABLE = auto()
@@ -669,6 +681,18 @@ class Tokenizer(metaclass=_Tokenizer):
         "TIMESTAMPLTZ": TokenType.TIMESTAMPLTZ,
         "DATE": TokenType.DATE,
         "DATETIME": TokenType.DATETIME,
+        "INT4RANGE": TokenType.INT4RANGE,
+        "INT4MULTIRANGE": TokenType.INT4MULTIRANGE,
+        "INT8RANGE": TokenType.INT8RANGE,
+        "INT8MULTIRANGE": TokenType.INT8MULTIRANGE,
+        "NUMRANGE": TokenType.NUMRANGE,
+        "NUMMULTIRANGE": TokenType.NUMMULTIRANGE,
+        "TSRANGE": TokenType.TSRANGE,
+        "TSMULTIRANGE": TokenType.TSMULTIRANGE,
+        "TSTZRANGE": TokenType.TSTZRANGE,
+        "TSTZMULTIRANGE": TokenType.TSTZMULTIRANGE,
+        "DATERANGE": TokenType.DATERANGE,
+        "DATEMULTIRANGE": TokenType.DATEMULTIRANGE,
         "UNIQUE": TokenType.UNIQUE,
         "STRUCT": TokenType.STRUCT,
         "VARIANT": TokenType.VARIANT,

diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
@@ -7,6 +7,7 @@ class TestPostgres(Validator):
     dialect = "postgres"
 
     def test_ddl(self):
+        self.validate_identity("CREATE TABLE public.y (x TSTZRANGE NOT NULL)")
         self.validate_identity("CREATE TABLE test (foo HSTORE)")
         self.validate_identity("CREATE TABLE test (foo JSONB)")
         self.validate_identity("CREATE TABLE test (foo VARCHAR(64)[])")
@@ -85,6 +86,18 @@ def test_ddl(self):
             )
 
     def test_postgres(self):
+        self.validate_identity("CAST(x AS INT4RANGE)")
+        self.validate_identity("CAST(x AS INT4MULTIRANGE)")
+        self.validate_identity("CAST(x AS INT8RANGE)")
+        self.validate_identity("CAST(x AS INT8MULTIRANGE)")
+        self.validate_identity("CAST(x AS NUMRANGE)")
+        self.validate_identity("CAST(x AS NUMMULTIRANGE)")
+        self.validate_identity("CAST(x AS TSRANGE)")
+        self.validate_identity("CAST(x AS TSMULTIRANGE)")
+        self.validate_identity("CAST(x AS TSTZRANGE)")
+        self.validate_identity("CAST(x AS TSTZMULTIRANGE)")
+        self.validate_identity("CAST(x AS DATERANGE)")
+        self.validate_identity("CAST(x AS DATEMULTIRANGE)")
         self.validate_identity(
             """LAST_VALUE("col1") OVER (ORDER BY "col2" RANGE BETWEEN INTERVAL '1 day' PRECEDING AND '1 month' FOLLOWING)"""
         )