@altimateai/altimate-code 0.5.1 → 0.5.3

package/dbt-tools/dist/altimate_python_packages/altimate_packages/sqlglot/dialects/doris.py

@@ -0,0 +1,561 @@
+from __future__ import annotations
+
+from sqlglot import exp
+from sqlglot.dialects.dialect import (
+    approx_count_distinct_sql,
+    build_timestamp_trunc,
+    rename_func,
+    time_format,
+    unit_to_str,
+)
+from sqlglot.dialects.mysql import MySQL
+
+
+def _lag_lead_sql(self, expression: exp.Lag | exp.Lead) -> str:
+    return self.func(
+        "LAG" if isinstance(expression, exp.Lag) else "LEAD",
+        expression.this,
+        expression.args.get("offset") or exp.Literal.number(1),
+        expression.args.get("default") or exp.null(),
+    )
+
+
+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": build_timestamp_trunc,
+            "MONTHS_ADD": exp.AddMonths.from_arg_list,
+            "REGEXP": exp.RegexpLike.from_arg_list,
+            "TO_DATE": exp.TsOrDsToDate.from_arg_list,
+        }
+
+        FUNCTION_PARSERS = MySQL.Parser.FUNCTION_PARSERS.copy()
+        FUNCTION_PARSERS.pop("GROUP_CONCAT")
+
+    class Generator(MySQL.Generator):
+        LAST_DAY_SUPPORTS_DATE_PART = False
+        VARCHAR_REQUIRES_SIZE = False
+
+        TYPE_MAPPING = {
+            **MySQL.Generator.TYPE_MAPPING,
+            exp.DataType.Type.TEXT: "STRING",
+            exp.DataType.Type.TIMESTAMP: "DATETIME",
+            exp.DataType.Type.TIMESTAMPTZ: "DATETIME",
+        }
+
+        CAST_MAPPING = {}
+        TIMESTAMP_FUNC_TYPES = set()
+
+        TRANSFORMS = {
+            **MySQL.Generator.TRANSFORMS,
+            exp.AddMonths: rename_func("MONTHS_ADD"),
+            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.ArrayToString: rename_func("ARRAY_JOIN"),
+            exp.ArrayUniqueAgg: rename_func("COLLECT_SET"),
+            exp.CurrentTimestamp: lambda self, _: self.func("NOW"),
+            exp.DateTrunc: lambda self, e: self.func("DATE_TRUNC", e.this, unit_to_str(e)),
+            exp.GroupConcat: lambda self, e: self.func(
+                "GROUP_CONCAT", e.this, e.args.get("separator") or exp.Literal.string(",")
+            ),
+            exp.JSONExtractScalar: lambda self, e: self.func("JSON_EXTRACT", e.this, e.expression),
+            exp.Lag: _lag_lead_sql,
+            exp.Lead: _lag_lead_sql,
+            exp.Map: rename_func("ARRAY_MAP"),
+            exp.RegexpLike: rename_func("REGEXP"),
+            exp.RegexpSplit: rename_func("SPLIT_BY_STRING"),
+            exp.Split: rename_func("SPLIT_BY_STRING"),
+            exp.StringToArray: rename_func("SPLIT_BY_STRING"),
+            exp.StrToUnix: lambda self, e: self.func("UNIX_TIMESTAMP", e.this, self.format_time(e)),
+            exp.TimeStrToDate: rename_func("TO_DATE"),
+            exp.TsOrDsAdd: lambda self, e: self.func("DATE_ADD", e.this, e.expression),
+            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, unit_to_str(e)),
+            exp.UnixToStr: lambda self, e: self.func(
+                "FROM_UNIXTIME", e.this, time_format("doris")(self, e)
+            ),
+            exp.UnixToTime: rename_func("FROM_UNIXTIME"),
+        }
+
+        # https://github.com/apache/doris/blob/e4f41dbf1ec03f5937fdeba2ee1454a20254015b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4#L93
+        RESERVED_KEYWORDS = {
+            "account_lock",
+            "account_unlock",
+            "add",
+            "adddate",
+            "admin",
+            "after",
+            "agg_state",
+            "aggregate",
+            "alias",
+            "all",
+            "alter",
+            "analyze",
+            "analyzed",
+            "and",
+            "anti",
+            "append",
+            "array",
+            "array_range",
+            "as",
+            "asc",
+            "at",
+            "authors",
+            "auto",
+            "auto_increment",
+            "backend",
+            "backends",
+            "backup",
+            "begin",
+            "belong",
+            "between",
+            "bigint",
+            "bin",
+            "binary",
+            "binlog",
+            "bitand",
+            "bitmap",
+            "bitmap_union",
+            "bitor",
+            "bitxor",
+            "blob",
+            "boolean",
+            "brief",
+            "broker",
+            "buckets",
+            "build",
+            "builtin",
+            "bulk",
+            "by",
+            "cached",
+            "call",
+            "cancel",
+            "case",
+            "cast",
+            "catalog",
+            "catalogs",
+            "chain",
+            "char",
+            "character",
+            "charset",
+            "check",
+            "clean",
+            "cluster",
+            "clusters",
+            "collate",
+            "collation",
+            "collect",
+            "column",
+            "columns",
+            "comment",
+            "commit",
+            "committed",
+            "compact",
+            "complete",
+            "config",
+            "connection",
+            "connection_id",
+            "consistent",
+            "constraint",
+            "constraints",
+            "convert",
+            "copy",
+            "count",
+            "create",
+            "creation",
+            "cron",
+            "cross",
+            "cube",
+            "current",
+            "current_catalog",
+            "current_date",
+            "current_time",
+            "current_timestamp",
+            "current_user",
+            "data",
+            "database",
+            "databases",
+            "date",
+            "date_add",
+            "date_ceil",
+            "date_diff",
+            "date_floor",
+            "date_sub",
+            "dateadd",
+            "datediff",
+            "datetime",
+            "datetimev2",
+            "datev2",
+            "datetimev1",
+            "datev1",
+            "day",
+            "days_add",
+            "days_sub",
+            "decimal",
+            "decimalv2",
+            "decimalv3",
+            "decommission",
+            "default",
+            "deferred",
+            "delete",
+            "demand",
+            "desc",
+            "describe",
+            "diagnose",
+            "disk",
+            "distinct",
+            "distinctpc",
+            "distinctpcsa",
+            "distributed",
+            "distribution",
+            "div",
+            "do",
+            "doris_internal_table_id",
+            "double",
+            "drop",
+            "dropp",
+            "dual",
+            "duplicate",
+            "dynamic",
+            "else",
+            "enable",
+            "encryptkey",
+            "encryptkeys",
+            "end",
+            "ends",
+            "engine",
+            "engines",
+            "enter",
+            "errors",
+            "events",
+            "every",
+            "except",
+            "exclude",
+            "execute",
+            "exists",
+            "expired",
+            "explain",
+            "export",
+            "extended",
+            "external",
+            "extract",
+            "failed_login_attempts",
+            "false",
+            "fast",
+            "feature",
+            "fields",
+            "file",
+            "filter",
+            "first",
+            "float",
+            "follower",
+            "following",
+            "for",
+            "foreign",
+            "force",
+            "format",
+            "free",
+            "from",
+            "frontend",
+            "frontends",
+            "full",
+            "function",
+            "functions",
+            "generic",
+            "global",
+            "grant",
+            "grants",
+            "graph",
+            "group",
+            "grouping",
+            "groups",
+            "hash",
+            "having",
+            "hdfs",
+            "help",
+            "histogram",
+            "hll",
+            "hll_union",
+            "hostname",
+            "hour",
+            "hub",
+            "identified",
+            "if",
+            "ignore",
+            "immediate",
+            "in",
+            "incremental",
+            "index",
+            "indexes",
+            "infile",
+            "inner",
+            "insert",
+            "install",
+            "int",
+            "integer",
+            "intermediate",
+            "intersect",
+            "interval",
+            "into",
+            "inverted",
+            "ipv4",
+            "ipv6",
+            "is",
+            "is_not_null_pred",
+            "is_null_pred",
+            "isnull",
+            "isolation",
+            "job",
+            "jobs",
+            "join",
+            "json",
+            "jsonb",
+            "key",
+            "keys",
+            "kill",
+            "label",
+            "largeint",
+            "last",
+            "lateral",
+            "ldap",
+            "ldap_admin_password",
+            "left",
+            "less",
+            "level",
+            "like",
+            "limit",
+            "lines",
+            "link",
+            "list",
+            "load",
+            "local",
+            "localtime",
+            "localtimestamp",
+            "location",
+            "lock",
+            "logical",
+            "low_priority",
+            "manual",
+            "map",
+            "match",
+            "match_all",
+            "match_any",
+            "match_phrase",
+            "match_phrase_edge",
+            "match_phrase_prefix",
+            "match_regexp",
+            "materialized",
+            "max",
+            "maxvalue",
+            "memo",
+            "merge",
+            "migrate",
+            "migrations",
+            "min",
+            "minus",
+            "minute",
+            "modify",
+            "month",
+            "mtmv",
+            "name",
+            "names",
+            "natural",
+            "negative",
+            "never",
+            "next",
+            "ngram_bf",
+            "no",
+            "non_nullable",
+            "not",
+            "null",
+            "nulls",
+            "observer",
+            "of",
+            "offset",
+            "on",
+            "only",
+            "open",
+            "optimized",
+            "or",
+            "order",
+            "outer",
+            "outfile",
+            "over",
+            "overwrite",
+            "parameter",
+            "parsed",
+            "partition",
+            "partitions",
+            "password",
+            "password_expire",
+            "password_history",
+            "password_lock_time",
+            "password_reuse",
+            "path",
+            "pause",
+            "percent",
+            "period",
+            "permissive",
+            "physical",
+            "plan",
+            "process",
+            "plugin",
+            "plugins",
+            "policy",
+            "preceding",
+            "prepare",
+            "primary",
+            "proc",
+            "procedure",
+            "processlist",
+            "profile",
+            "properties",
+            "property",
+            "quantile_state",
+            "quantile_union",
+            "query",
+            "quota",
+            "random",
+            "range",
+            "read",
+            "real",
+            "rebalance",
+            "recover",
+            "recycle",
+            "refresh",
+            "references",
+            "regexp",
+            "release",
+            "rename",
+            "repair",
+            "repeatable",
+            "replace",
+            "replace_if_not_null",
+            "replica",
+            "repositories",
+            "repository",
+            "resource",
+            "resources",
+            "restore",
+            "restrictive",
+            "resume",
+            "returns",
+            "revoke",
+            "rewritten",
+            "right",
+            "rlike",
+            "role",
+            "roles",
+            "rollback",
+            "rollup",
+            "routine",
+            "row",
+            "rows",
+            "s3",
+            "sample",
+            "schedule",
+            "scheduler",
+            "schema",
+            "schemas",
+            "second",
+            "select",
+            "semi",
+            "sequence",
+            "serializable",
+            "session",
+            "set",
+            "sets",
+            "shape",
+            "show",
+            "signed",
+            "skew",
+            "smallint",
+            "snapshot",
+            "soname",
+            "split",
+            "sql_block_rule",
+            "start",
+            "starts",
+            "stats",
+            "status",
+            "stop",
+            "storage",
+            "stream",
+            "streaming",
+            "string",
+            "struct",
+            "subdate",
+            "sum",
+            "superuser",
+            "switch",
+            "sync",
+            "system",
+            "table",
+            "tables",
+            "tablesample",
+            "tablet",
+            "tablets",
+            "task",
+            "tasks",
+            "temporary",
+            "terminated",
+            "text",
+            "than",
+            "then",
+            "time",
+            "timestamp",
+            "timestampadd",
+            "timestampdiff",
+            "tinyint",
+            "to",
+            "transaction",
+            "trash",
+            "tree",
+            "triggers",
+            "trim",
+            "true",
+            "truncate",
+            "type",
+            "type_cast",
+            "types",
+            "unbounded",
+            "uncommitted",
+            "uninstall",
+            "union",
+            "unique",
+            "unlock",
+            "unsigned",
+            "update",
+            "use",
+            "user",
+            "using",
+            "value",
+            "values",
+            "varchar",
+            "variables",
+            "variant",
+            "vault",
+            "verbose",
+            "version",
+            "view",
+            "warnings",
+            "week",
+            "when",
+            "where",
+            "whitelist",
+            "with",
+            "work",
+            "workload",
+            "write",
+            "xor",
+            "year",
+        }