72 lines
2.8 KiB
Python
72 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
from sqlglot import exp
|
|
from sqlglot.dialects.dialect import (
|
|
approx_count_distinct_sql,
|
|
arrow_json_extract_sql,
|
|
parse_timestamp_trunc,
|
|
rename_func,
|
|
time_format,
|
|
)
|
|
from sqlglot.dialects.mysql import MySQL
|
|
|
|
|
|
class Doris(MySQL):
|
|
DATE_FORMAT = "'yyyy-MM-dd'"
|
|
DATEINT_FORMAT = "'yyyyMMdd'"
|
|
TIME_FORMAT = "'yyyy-MM-dd HH:mm:ss'"
|
|
|
|
class Parser(MySQL.Parser):
|
|
FUNCTIONS = {
|
|
**MySQL.Parser.FUNCTIONS,
|
|
"COLLECT_SET": exp.ArrayUniqueAgg.from_arg_list,
|
|
"DATE_TRUNC": parse_timestamp_trunc,
|
|
"REGEXP": exp.RegexpLike.from_arg_list,
|
|
"TO_DATE": exp.TsOrDsToDate.from_arg_list,
|
|
}
|
|
|
|
class Generator(MySQL.Generator):
|
|
CAST_MAPPING = {}
|
|
|
|
TYPE_MAPPING = {
|
|
**MySQL.Generator.TYPE_MAPPING,
|
|
exp.DataType.Type.TEXT: "STRING",
|
|
exp.DataType.Type.TIMESTAMP: "DATETIME",
|
|
exp.DataType.Type.TIMESTAMPTZ: "DATETIME",
|
|
}
|
|
|
|
LAST_DAY_SUPPORTS_DATE_PART = False
|
|
|
|
TIMESTAMP_FUNC_TYPES = set()
|
|
|
|
TRANSFORMS = {
|
|
**MySQL.Generator.TRANSFORMS,
|
|
exp.ApproxDistinct: approx_count_distinct_sql,
|
|
exp.ArgMax: rename_func("MAX_BY"),
|
|
exp.ArgMin: rename_func("MIN_BY"),
|
|
exp.ArrayAgg: rename_func("COLLECT_LIST"),
|
|
exp.ArrayUniqueAgg: rename_func("COLLECT_SET"),
|
|
exp.CurrentTimestamp: lambda *_: "NOW()",
|
|
exp.DateTrunc: lambda self, e: self.func(
|
|
"DATE_TRUNC", e.this, "'" + e.text("unit") + "'"
|
|
),
|
|
exp.JSONExtractScalar: arrow_json_extract_sql,
|
|
exp.JSONExtract: arrow_json_extract_sql,
|
|
exp.Map: rename_func("ARRAY_MAP"),
|
|
exp.RegexpLike: rename_func("REGEXP"),
|
|
exp.RegexpSplit: rename_func("SPLIT_BY_STRING"),
|
|
exp.StrToUnix: lambda self, e: f"UNIX_TIMESTAMP({self.sql(e, 'this')}, {self.format_time(e)})",
|
|
exp.Split: rename_func("SPLIT_BY_STRING"),
|
|
exp.TimeStrToDate: rename_func("TO_DATE"),
|
|
exp.ToChar: lambda self, e: f"DATE_FORMAT({self.sql(e, 'this')}, {self.format_time(e)})",
|
|
exp.TsOrDsAdd: lambda self, e: f"DATE_ADD({self.sql(e, 'this')}, {self.sql(e, 'expression')})", # Only for day level
|
|
exp.TsOrDsToDate: lambda self, e: self.func("TO_DATE", e.this),
|
|
exp.TimeToUnix: rename_func("UNIX_TIMESTAMP"),
|
|
exp.TimestampTrunc: lambda self, e: self.func(
|
|
"DATE_TRUNC", e.this, "'" + e.text("unit") + "'"
|
|
),
|
|
exp.UnixToStr: lambda self, e: self.func(
|
|
"FROM_UNIXTIME", e.this, time_format("doris")(self, e)
|
|
),
|
|
exp.UnixToTime: rename_func("FROM_UNIXTIME"),
|
|
}
|