feat(Clickhouse): Support large data types (#1568)

* feat(Clickhouse): Support large data types * Update sqlglot/dialects/clickhouse.py --------- Co-authored-by: Toby Mao <[email protected]>
tobymao · May 8, 2023 · c9103fe · c9103fe
1 parent 9f13b6c
commit c9103fe
Show file tree

Hide file tree

Showing 5 changed files with 56 additions and 1 deletion.
diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
@@ -33,10 +33,18 @@ class Tokenizer(tokens.Tokenizer):
             "FINAL": TokenType.FINAL,
             "FLOAT32": TokenType.FLOAT,
             "FLOAT64": TokenType.DOUBLE,
+            "INT8": TokenType.TINYINT,
+            "UINT8": TokenType.UTINYINT,
             "INT16": TokenType.SMALLINT,
+            "UINT16": TokenType.USMALLINT,
             "INT32": TokenType.INT,
+            "UINT32": TokenType.UINT,
             "INT64": TokenType.BIGINT,
-            "INT8": TokenType.TINYINT,
+            "UINT64": TokenType.UBIGINT,
+            "INT128": TokenType.INT128,
+            "UINT128": TokenType.UINT128,
+            "INT256": TokenType.INT256,
+            "UINT256": TokenType.UINT256,
             "TUPLE": TokenType.STRUCT,
         }
 
@@ -123,9 +131,17 @@ class Generator(generator.Generator):
             exp.DataType.Type.ARRAY: "Array",
             exp.DataType.Type.STRUCT: "Tuple",
             exp.DataType.Type.TINYINT: "Int8",
+            exp.DataType.Type.UTINYINT: "UInt8",
             exp.DataType.Type.SMALLINT: "Int16",
+            exp.DataType.Type.USMALLINT: "UInt16",
             exp.DataType.Type.INT: "Int32",
+            exp.DataType.Type.UINT: "UInt32",
             exp.DataType.Type.BIGINT: "Int64",
+            exp.DataType.Type.UBIGINT: "UInt64",
+            exp.DataType.Type.INT128: "Int128",
+            exp.DataType.Type.UINT128: "UInt128",
+            exp.DataType.Type.INT256: "Int256",
+            exp.DataType.Type.UINT256: "UInt256",
             exp.DataType.Type.FLOAT: "Float32",
             exp.DataType.Type.DOUBLE: "Float64",
         }

diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -3006,6 +3006,10 @@ class Type(AutoName):
         USMALLINT = auto()
         BIGINT = auto()
         UBIGINT = auto()
+        INT128 = auto()
+        UINT128 = auto()
+        INT256 = auto()
+        UINT256 = auto()
         FLOAT = auto()
         DOUBLE = auto()
         DECIMAL = auto()
@@ -3059,6 +3063,8 @@ class Type(AutoName):
         Type.TINYINT,
         Type.SMALLINT,
         Type.BIGINT,
+        Type.INT128,
+        Type.INT256,
     }
 
     FLOAT_TYPES = {

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -126,9 +126,17 @@ class Parser(metaclass=_Parser):
         TokenType.BIT,
         TokenType.BOOLEAN,
         TokenType.TINYINT,
+        TokenType.UTINYINT,
         TokenType.SMALLINT,
+        TokenType.USMALLINT,
         TokenType.INT,
+        TokenType.UINT,
         TokenType.BIGINT,
+        TokenType.UBIGINT,
+        TokenType.INT128,
+        TokenType.UINT128,
+        TokenType.INT256,
+        TokenType.UINT256,
         TokenType.FLOAT,
         TokenType.DOUBLE,
         TokenType.CHAR,

diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
@@ -84,6 +84,10 @@ class TokenType(AutoName):
     UINT = auto()
     BIGINT = auto()
     UBIGINT = auto()
+    INT128 = auto()
+    UINT128 = auto()
+    INT256 = auto()
+    UINT256 = auto()
     FLOAT = auto()
     DOUBLE = auto()
     DECIMAL = auto()

diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
@@ -65,3 +65,24 @@ def test_cte(self):
         self.validate_identity("WITH SUM(bytes) AS foo SELECT foo FROM system.parts")
         self.validate_identity("WITH (SELECT foo) AS bar SELECT bar + 5")
         self.validate_identity("WITH test1 AS (SELECT i + 1, j + 1 FROM test1) SELECT * FROM test1")
+
+    def test_signed_and_unsigned_types(self):
+        data_types = [
+            "UInt8",
+            "UInt16",
+            "UInt32",
+            "UInt64",
+            "UInt128",
+            "UInt256",
+            "Int8",
+            "Int16",
+            "Int32",
+            "Int64",
+            "Int128",
+            "Int256",
+        ]
+        for data_type in data_types:
+            self.validate_all(
+                f"pow(2, 32)::{data_type}",
+                write={"clickhouse": f"CAST(POWER(2, 32) AS {data_type})"},
+            )