39 lines
1.5 KiB
Python
39 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
from sqlglot import exp
|
|
from sqlglot.dialects.dialect import arrow_json_extract_sql, rename_func
|
|
from sqlglot.dialects.mysql import MySQL
|
|
from sqlglot.helper import seq_get
|
|
|
|
|
|
class StarRocks(MySQL):
|
|
class Parser(MySQL.Parser): # type: ignore
|
|
FUNCTIONS = {
|
|
**MySQL.Parser.FUNCTIONS,
|
|
"DATE_TRUNC": lambda args: exp.TimestampTrunc(
|
|
this=seq_get(args, 1), unit=seq_get(args, 0)
|
|
),
|
|
}
|
|
|
|
class Generator(MySQL.Generator): # type: ignore
|
|
TYPE_MAPPING = {
|
|
**MySQL.Generator.TYPE_MAPPING, # type: ignore
|
|
exp.DataType.Type.TEXT: "STRING",
|
|
exp.DataType.Type.TIMESTAMP: "DATETIME",
|
|
exp.DataType.Type.TIMESTAMPTZ: "DATETIME",
|
|
}
|
|
|
|
TRANSFORMS = {
|
|
**MySQL.Generator.TRANSFORMS, # type: ignore
|
|
exp.JSONExtractScalar: arrow_json_extract_sql,
|
|
exp.JSONExtract: arrow_json_extract_sql,
|
|
exp.DateDiff: rename_func("DATEDIFF"),
|
|
exp.StrToUnix: lambda self, e: f"UNIX_TIMESTAMP({self.sql(e, 'this')}, {self.format_time(e)})",
|
|
exp.TimestampTrunc: lambda self, e: self.func(
|
|
"DATE_TRUNC", exp.Literal.string(e.text("unit")), e.this
|
|
),
|
|
exp.TimeStrToDate: rename_func("TO_DATE"),
|
|
exp.UnixToStr: lambda self, e: f"FROM_UNIXTIME({self.sql(e, 'this')}, {self.format_time(e)})",
|
|
exp.UnixToTime: rename_func("FROM_UNIXTIME"),
|
|
}
|
|
TRANSFORMS.pop(exp.DateTrunc)
|