add sqlglot dialect test case

Author: Yibo-Chen13

tests/integration_tests/test_sqlglot_dialect.py

@@ -0,0 +1,46 @@
+
+from sqlglot import exp, parse_one
+from timeplus_connect.tp_superset.sql_parse import TimeplusSqlglotDialect
+
+
+dialect = TimeplusSqlglotDialect()
+
+def validate_identity(sql: str) -> exp.Expression:
+    ast: exp.Expression = parse_one(sql, read=dialect)
+
+    assert ast is not None
+    assert ast.sql(dialect=dialect) == sql
+
+    return ast
+
+def test_simple_sql():
+    validate_identity("CAST(1 AS bool)")
+    validate_identity("SELECT to_string(CHAR(104.1, 101, 108.9, 108.9, 111, 32))")
+    validate_identity("@macro").assert_is(exp.Parameter).this.assert_is(exp.Var)
+    validate_identity("SELECT to_float(like)")
+    validate_identity("SELECT like")
+    validate_identity("SELECT EXTRACT(YEAR FROM to_datetime('2023-02-01'))")
+    validate_identity("extract(haystack, pattern)")
+    validate_identity("SELECT * FROM x LIMIT 1 UNION ALL SELECT * FROM y")
+    validate_identity("SELECT CAST(x AS tuple(string, array(nullable(float64))))")
+    validate_identity("count_if(x)")
+    validate_identity("x = y")
+    validate_identity("x <> y")
+    validate_identity("SELECT * FROM (SELECT a FROM b SAMPLE 0.01)")
+    validate_identity("SELECT * FROM (SELECT a FROM b SAMPLE 1 / 10 OFFSET 1 / 2)")
+    validate_identity("SELECT sum(foo * bar) FROM bla SAMPLE 10000000")
+    validate_identity("CAST(x AS nested(ID uint32, Serial uint32, EventTime DateTime))")
+    validate_identity("CAST(x AS enum('hello' = 1, 'world' = 2))")
+    validate_identity("CAST(x AS enum('hello', 'world'))")
+    validate_identity("CAST(x AS enum('hello' = 1, 'world'))")
+    validate_identity("CAST(x AS enum8('hello' = -123, 'world'))")
+    validate_identity("CAST(x AS fixed_string(1))")
+    validate_identity("CAST(x AS low_cardinality(fixed_string))")
+    validate_identity("SELECT is_nan(1.0)")
+    validate_identity("SELECT start_with('Spider-Man', 'Spi')")
+    validate_identity("SELECT xor(TRUE, FALSE)")
+    validate_identity("CAST(['hello'], 'array(enum8(''hello'' = 1))')")
+    validate_identity("SELECT x, COUNT() FROM y GROUP BY x WITH TOTALS")
+    validate_identity("SELECT INTERVAL t.days DAY")
+    validate_identity("SELECT match('abc', '([a-z]+)')")
+    validate_identity("SELECT window_start, avg(price) AS avg_price FROM tumble(coinbase, 10s) GROUP BY window_start")

tests/test_requirements.txt

@@ -17,3 +17,4 @@ pandas
 zstandard
 lz4
 pyjwt[crypto]==2.10.1
+sqlglot

timeplus_connect/tp_sqlglot/__init__.py

@@ -0,0 +1 @@
+import timeplus_connect.tp_sqlglot.dialect

timeplus_connect/tp_superset/sql_parse.py renamed to timeplus_connect/tp_sqlglot/dialect.py

@@ -226,11 +226,11 @@ class Tokenizer(tokens.Tokenizer):
             "ENUM8": TokenType.ENUM8,
             "ENUM16": TokenType.ENUM16,
             "FINAL": TokenType.FINAL,
-            "FIXEDSTRING": TokenType.FIXEDSTRING,
+            "FIXED_STRING": TokenType.FIXEDSTRING,
             "FLOAT32": TokenType.FLOAT,
             "FLOAT64": TokenType.DOUBLE,
             "GLOBAL": TokenType.GLOBAL,
-            "LOWCARDINALITY": TokenType.LOWCARDINALITY,
+            "LOW_CARDINALITY": TokenType.LOWCARDINALITY,
             "MAP": TokenType.MAP,
             "NESTED": TokenType.NESTED,
             "SAMPLE": TokenType.TABLE_SAMPLE,
@@ -393,7 +393,7 @@ class Parser(parser.Parser):
 
         AGG_FUNCTIONS_SUFFIXES = [
             "if",
-            "aary",
+            "array",
             "array_if",
             "map",
             "simple_state",
@@ -524,9 +524,9 @@ def _parse_types(
             if isinstance(dtype, exp.DataType) and dtype.args.get("nullable") is not True:
                 # Mark every type as non-nullable which is Timeplus's default, unless it's
                 # already marked as nullable. This marker helps us transpile types from other
-                # dialects to Timeplus, so that we can e.g. produce `CAST(x AS nullabe(String))`
+                # dialects to Timeplus, so that we can e.g. produce `CAST(x AS nullable(String))`
                 # from `CAST(x AS TEXT)`. If there is a `NULL` value in `x`, the former would
-                # fail in Timeplus without the `nullabe` type constructor.
+                # fail in Timeplus without the `nullable` type constructor.
                 dtype.set("nullable", False)
 
             return dtype
@@ -563,7 +563,7 @@ def _parse_assignment(self) -> t.Optional[exp.Expression]:
 
         def _parse_query_parameter(self) -> t.Optional[exp.Expression]:
             """
-            Parse a placeholder expression like SELECT {abc: UInt32} or FROM {table: Identifier}
+            Parse a placeholder expression like SELECT {abc: uint32} or FROM {table: Identifier}
             """
             index = self._index
 
@@ -958,7 +958,7 @@ class Generator(generator.Generator):
             exp.ArgMax: arg_max_or_min_no_count("arg_max"),
             exp.ArgMin: arg_max_or_min_no_count("arg_min"),
             exp.Array: inline_array_sql,
-            exp.CastToStrType: rename_func("cast"),
+            exp.CastToStrType: rename_func("CAST"),
             exp.CountIf: rename_func("count_if"),
             exp.CompressColumnConstraint: lambda self,
             e: f"CODEC({self.expressions(e, key='this', flat=True)})",
@@ -1062,7 +1062,7 @@ def cast_sql(self, expression: exp.Cast, safe_prefix: t.Optional[str] = None) ->
         def trycast_sql(self, expression: exp.TryCast) -> str:
             dtype = expression.to
             if not dtype.is_type(*self.NON_NULLABLE_TYPES, check_nullable=True):
-                # Casting x into nullabe(T) appears to behave similarly to TRY_CAST(x AS T)
+                # Casting x into nullable(T) appears to behave similarly to TRY_CAST(x AS T)
                 dtype.set("nullable", True)
 
             return super().cast_sql(expression)
@@ -1110,13 +1110,13 @@ def datatype_sql(self, expression: exp.DataType) -> str:
             else:
                 dtype = super().datatype_sql(expression)
 
-            # This section changes the type to `nullabe(...)` if the following conditions hold:
-            # - It's marked as nullable - this ensures we won't wrap Timeplus types with `nullabe`
+            # This section changes the type to `nullable(...)` if the following conditions hold:
+            # - It's marked as nullable - this ensures we won't wrap Timeplus types with `nullable`
             #   and change their semantics
             # - It's not the key type of a `Map`. This is because Timeplus enforces the following
             #   constraint: "Type of Map key must be a type, that can be represented by integer or
             #   String or FixedString (possibly LowCardinality) or UUID or IPv6"
-            # - It's not a composite type, e.g. `nullabe(array(...))` is not a valid type
+            # - It's not a composite type, e.g. `nullable(array(...))` is not a valid type
             parent = expression.parent
             nullable = expression.args.get("nullable")
             if nullable is True or (
@@ -1128,7 +1128,7 @@ def datatype_sql(self, expression: exp.DataType) -> str:
                 )
                 and not expression.is_type(*self.NON_NULLABLE_TYPES, check_nullable=True)
             ):
-                dtype = f"nullabe({dtype})"
+                dtype = f"nullable({dtype})"
 
             return dtype
 
@@ -1247,3 +1247,7 @@ def is_sql(self, expression: exp.Is) -> str:
                 is_sql = self.wrap(is_sql)
 
             return is_sql
+
+    tokenizer_class = Tokenizer
+    parser_class = Parser
+    generator_class = Generator

timeplus_connect/tp_superset/__init__.py

@@ -1,2 +1 @@
 import timeplus_connect.tp_superset.db_engine_spec
-import timeplus_connect.tp_superset.sql_parse

timeplus_connect/tp_superset/db_engine_spec.py

@@ -2,7 +2,7 @@
 from datetime import datetime
 from typing import Any
 
-# pylint:disable=E0401
+# pylint:disable=E0401, E0611
 from marshmallow import fields, Schema
 from marshmallow.validate import Range
 from superset.db_engine_specs.base import BaseEngineSpec
@@ -13,7 +13,7 @@
 from sqlalchemy.engine.url import URL
 
 
-from timeplus_connect.tp_superset.sql_parse import TimeplusSqlglotDialect
+from timeplus_connect.tp_sqlglot.dialect import TimeplusSqlglotDialect
 
 
 logger = logging.getLogger(__name__)