sqlglot.generators.doris

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp
  6from sqlglot.dialects.dialect import (
  7    approx_count_distinct_sql,
  8    property_sql,
  9    rename_func,
 10    time_format,
 11    weekstart_unit_to_str,
 12)
 13from sqlglot.generators.mysql import MySQLGenerator
 14
 15
 16def _lag_lead_sql(self, expression: exp.Lag | exp.Lead) -> str:
 17    return self.func(
 18        "LAG" if isinstance(expression, exp.Lag) else "LEAD",
 19        expression.this,
 20        expression.args.get("offset") or exp.Literal.number(1),
 21        expression.args.get("default") or exp.null(),
 22    )
 23
 24
 25class DorisGenerator(MySQLGenerator):
 26    LAST_DAY_SUPPORTS_DATE_PART = False
 27    SUPPORTS_ALTER_COLUMN_NULLABILITY = False
 28    VARCHAR_REQUIRES_SIZE = False
 29    WITH_PROPERTIES_PREFIX = "PROPERTIES"
 30    RENAME_TABLE_WITH_DB = False
 31    UPDATE_STATEMENT_SUPPORTS_FROM = True
 32
 33    TYPE_MAPPING = {
 34        **MySQLGenerator.TYPE_MAPPING,
 35        exp.DType.TEXT: "STRING",
 36        exp.DType.TIMESTAMP: "DATETIME",
 37        exp.DType.TIMESTAMPTZ: "DATETIME",
 38    }
 39
 40    PROPERTIES_LOCATION = {
 41        **MySQLGenerator.PROPERTIES_LOCATION,
 42        exp.UniqueKeyProperty: exp.Properties.Location.POST_SCHEMA,
 43        exp.PartitionedByProperty: exp.Properties.Location.POST_SCHEMA,
 44        exp.BuildProperty: exp.Properties.Location.POST_SCHEMA,
 45    }
 46
 47    CAST_MAPPING: t.ClassVar[dict[exp.DType, str]] = {}
 48    TIMESTAMP_FUNC_TYPES: t.ClassVar[set[exp.DType]] = set()
 49
 50    TRANSFORMS = {
 51        **MySQLGenerator.TRANSFORMS,
 52        exp.AddMonths: rename_func("MONTHS_ADD"),
 53        exp.ApproxDistinct: approx_count_distinct_sql,
 54        exp.ArgMax: rename_func("MAX_BY"),
 55        exp.ArgMin: rename_func("MIN_BY"),
 56        exp.ArrayAgg: rename_func("COLLECT_LIST"),
 57        exp.ArrayToString: rename_func("ARRAY_JOIN"),
 58        exp.ArrayUniqueAgg: rename_func("COLLECT_SET"),
 59        exp.CurrentDate: lambda self, _: self.func("CURRENT_DATE"),
 60        exp.CurrentTimestamp: lambda self, _: self.func("NOW"),
 61        exp.DateTrunc: lambda self, e: self.func(
 62            "DATE_TRUNC", e.this, weekstart_unit_to_str(self, e)
 63        ),
 64        exp.EuclideanDistance: rename_func("L2_DISTANCE"),
 65        exp.GroupConcat: lambda self, e: self.func(
 66            "GROUP_CONCAT", e.this, e.args.get("separator") or exp.Literal.string(",")
 67        ),
 68        exp.JSONExtractScalar: lambda self, e: self.func("JSON_EXTRACT", e.this, e.expression),
 69        exp.Lag: _lag_lead_sql,
 70        exp.Lead: _lag_lead_sql,
 71        exp.Map: rename_func("ARRAY_MAP"),
 72        exp.Property: property_sql,
 73        exp.RegexpLike: rename_func("REGEXP"),
 74        exp.RegexpSplit: rename_func("SPLIT_BY_STRING"),
 75        exp.SchemaCommentProperty: lambda self, e: self.naked_property(e),
 76        exp.Split: rename_func("SPLIT_BY_STRING"),
 77        exp.StringToArray: rename_func("SPLIT_BY_STRING"),
 78        exp.StrToUnix: lambda self, e: self.func("UNIX_TIMESTAMP", e.this, self.format_time(e)),
 79        exp.TimeStrToDate: rename_func("TO_DATE"),
 80        exp.TsOrDsAdd: lambda self, e: self.func("DATE_ADD", e.this, e.expression),
 81        exp.TsOrDsToDate: lambda self, e: self.func("TO_DATE", e.this),
 82        exp.TimeToUnix: rename_func("UNIX_TIMESTAMP"),
 83        exp.TimestampTrunc: lambda self, e: self.func(
 84            "DATE_TRUNC", e.this, weekstart_unit_to_str(self, e)
 85        ),
 86        exp.UnixToStr: lambda self, e: self.func(
 87            "FROM_UNIXTIME", e.this, time_format("doris")(self, e)
 88        ),
 89        exp.UnixToTime: rename_func("FROM_UNIXTIME"),
 90    }
 91
 92    # https://github.com/apache/doris/blob/e4f41dbf1ec03f5937fdeba2ee1454a20254015b/fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisLexer.g4#L93
 93    RESERVED_KEYWORDS = {
 94        "account_lock",
 95        "account_unlock",
 96        "add",
 97        "adddate",
 98        "admin",
 99        "after",
100        "agg_state",
101        "aggregate",
102        "alias",
103        "all",
104        "alter",
105        "analyze",
106        "analyzed",
107        "analyzer",
108        "and",
109        "anti",
110        "append",
111        "array",
112        "array_range",
113        "as",
114        "asc",
115        "asof",
116        "at",
117        "authors",
118        "auto",
119        "auto_increment",
120        "backend",
121        "backends",
122        "backup",
123        "begin",
124        "belong",
125        "between",
126        "bigint",
127        "bin",
128        "binary",
129        "binlog",
130        "bitand",
131        "bitmap",
132        "bitmap_union",
133        "bitor",
134        "bitxor",
135        "blob",
136        "boolean",
137        "both",
138        "brief",
139        "broker",
140        "buckets",
141        "build",
142        "builtin",
143        "bulk",
144        "by",
145        "cached",
146        "call",
147        "cancel",
148        "case",
149        "cast",
150        "catalog",
151        "catalogs",
152        "chain",
153        "char",
154        "char_filter",
155        "character",
156        "charset",
157        "check",
158        "clean",
159        "cluster",
160        "clusters",
161        "collate",
162        "collation",
163        "collect",
164        "column",
165        "columns",
166        "comment",
167        "commit",
168        "committed",
169        "compact",
170        "complete",
171        "config",
172        "connection",
173        "connection_id",
174        "consistent",
175        "constraint",
176        "constraints",
177        "convert",
178        "copy",
179        "count",
180        "create",
181        "creation",
182        "cron",
183        "cross",
184        "cube",
185        "current",
186        "current_catalog",
187        "current_date",
188        "current_time",
189        "current_timestamp",
190        "current_user",
191        "data",
192        "database",
193        "databases",
194        "date",
195        "date_add",
196        "date_ceil",
197        "date_diff",
198        "date_floor",
199        "date_sub",
200        "dateadd",
201        "datediff",
202        "datetime",
203        "datetimev2",
204        "datev2",
205        "datetimev1",
206        "datev1",
207        "day",
208        "days_add",
209        "days_sub",
210        "decimal",
211        "decimalv2",
212        "decimalv3",
213        "decommission",
214        "default",
215        "deferred",
216        "delete",
217        "demand",
218        "desc",
219        "describe",
220        "diagnose",
221        "disk",
222        "distinct",
223        "distinctpc",
224        "distinctpcsa",
225        "distributed",
226        "distribution",
227        "div",
228        "do",
229        "doris_internal_table_id",
230        "double",
231        "drop",
232        "dropp",
233        "dual",
234        "dump",
235        "duplicate",
236        "dynamic",
237        "else",
238        "enable",
239        "encryptkey",
240        "encryptkeys",
241        "end",
242        "ends",
243        "engine",
244        "engines",
245        "enter",
246        "errors",
247        "events",
248        "every",
249        "except",
250        "exclude",
251        "execute",
252        "exists",
253        "expired",
254        "explain",
255        "export",
256        "extended",
257        "external",
258        "extract",
259        "failed_login_attempts",
260        "false",
261        "fast",
262        "feature",
263        "fields",
264        "file",
265        "filter",
266        "first",
267        "float",
268        "follower",
269        "following",
270        "for",
271        "foreign",
272        "force",
273        "format",
274        "free",
275        "from",
276        "frontend",
277        "frontends",
278        "full",
279        "function",
280        "functions",
281        "generic",
282        "global",
283        "grant",
284        "grants",
285        "graph",
286        "group",
287        "grouping",
288        "groups",
289        "hash",
290        "having",
291        "hdfs",
292        "help",
293        "histogram",
294        "hll",
295        "hll_union",
296        "hostname",
297        "hour",
298        "hub",
299        "identified",
300        "if",
301        "ignore",
302        "immediate",
303        "in",
304        "incremental",
305        "index",
306        "indexes",
307        "infile",
308        "inner",
309        "insert",
310        "install",
311        "int",
312        "integer",
313        "intermediate",
314        "intersect",
315        "interval",
316        "into",
317        "inverted",
318        "ipv4",
319        "ipv6",
320        "is",
321        "is_not_null_pred",
322        "is_null_pred",
323        "isnull",
324        "isolation",
325        "job",
326        "jobs",
327        "join",
328        "json",
329        "jsonb",
330        "key",
331        "keys",
332        "kill",
333        "label",
334        "largeint",
335        "last",
336        "lateral",
337        "layout",
338        "ldap",
339        "ldap_admin_password",
340        "leading",
341        "left",
342        "less",
343        "level",
344        "like",
345        "limit",
346        "lines",
347        "link",
348        "list",
349        "load",
350        "local",
351        "localtime",
352        "localtimestamp",
353        "location",
354        "lock",
355        "logical",
356        "low_priority",
357        "manual",
358        "map",
359        "match",
360        "match_all",
361        "match_any",
362        "match_condition",
363        "match_phrase",
364        "match_phrase_edge",
365        "match_phrase_prefix",
366        "match_regexp",
367        "materialized",
368        "max",
369        "maxvalue",
370        "memo",
371        "merge",
372        "migrate",
373        "migrations",
374        "min",
375        "minus",
376        "minute",
377        "modify",
378        "month",
379        "mtmv",
380        "name",
381        "names",
382        "natural",
383        "negative",
384        "never",
385        "next",
386        "ngram_bf",
387        "no",
388        "no_use_mv",
389        "non_nullable",
390        "not",
391        "null",
392        "nulls",
393        "observer",
394        "of",
395        "offset",
396        "on",
397        "only",
398        "open",
399        "optimized",
400        "or",
401        "order",
402        "outer",
403        "outfile",
404        "over",
405        "overwrite",
406        "parameter",
407        "parsed",
408        "partition",
409        "partitions",
410        "password",
411        "password_expire",
412        "password_history",
413        "password_lock_time",
414        "password_reuse",
415        "path",
416        "pause",
417        "percent",
418        "period",
419        "permissive",
420        "physical",
421        "plan",
422        "play",
423        "process",
424        "plugin",
425        "plugins",
426        "policy",
427        "preceding",
428        "prepare",
429        "primary",
430        "proc",
431        "procedure",
432        "processlist",
433        "profile",
434        "properties",
435        "property",
436        "quantile_state",
437        "quantile_union",
438        "query",
439        "quota",
440        "random",
441        "range",
442        "read",
443        "real",
444        "rebalance",
445        "recover",
446        "recycle",
447        "refresh",
448        "references",
449        "regexp",
450        "release",
451        "rename",
452        "repair",
453        "repeatable",
454        "replace",
455        "replace_if_not_null",
456        "replica",
457        "repositories",
458        "repository",
459        "resource",
460        "resources",
461        "restore",
462        "restrictive",
463        "resume",
464        "returns",
465        "revoke",
466        "rewritten",
467        "right",
468        "rlike",
469        "role",
470        "roles",
471        "rollback",
472        "rollup",
473        "routine",
474        "row",
475        "rows",
476        "s3",
477        "sample",
478        "schedule",
479        "scheduler",
480        "schema",
481        "schemas",
482        "second",
483        "select",
484        "semi",
485        "sequence",
486        "serializable",
487        "session",
488        "set",
489        "sets",
490        "shape",
491        "show",
492        "signed",
493        "skew",
494        "smallint",
495        "snapshot",
496        "soname",
497        "split",
498        "sql_block_rule",
499        "start",
500        "starts",
501        "stats",
502        "status",
503        "stop",
504        "storage",
505        "stream",
506        "streaming",
507        "string",
508        "struct",
509        "subdate",
510        "sum",
511        "superuser",
512        "switch",
513        "sync",
514        "system",
515        "table",
516        "tables",
517        "tablesample",
518        "tablet",
519        "tablets",
520        "task",
521        "tasks",
522        "temporary",
523        "terminated",
524        "text",
525        "than",
526        "then",
527        "time",
528        "timestamp",
529        "timestampadd",
530        "timestampdiff",
531        "tinyint",
532        "to",
533        "token_filter",
534        "tokenizer",
535        "trailing",
536        "transaction",
537        "trash",
538        "tree",
539        "triggers",
540        "trim",
541        "true",
542        "truncate",
543        "try_cast",
544        "type",
545        "type_cast",
546        "types",
547        "unbounded",
548        "uncommitted",
549        "uninstall",
550        "union",
551        "unique",
552        "unlock",
553        "unsigned",
554        "update",
555        "use",
556        "use_mv",
557        "user",
558        "using",
559        "value",
560        "values",
561        "varchar",
562        "variables",
563        "variant",
564        "vault",
565        "verbose",
566        "version",
567        "view",
568        "warnings",
569        "week",
570        "when",
571        "where",
572        "whitelist",
573        "with",
574        "work",
575        "workload",
576        "write",
577        "xor",
578        "year",
579    }
580
581    def uniquekeyproperty_sql(
582        self, expression: exp.UniqueKeyProperty, prefix: str = "UNIQUE KEY"
583    ) -> str:
584        create_stmt = expression.find_ancestor(exp.Create)
585        if create_stmt and create_stmt.args["properties"].find(exp.MaterializedProperty):
586            return super().uniquekeyproperty_sql(expression, prefix="KEY")
587
588        return super().uniquekeyproperty_sql(expression)
589
590    def partitionrange_sql(self, expression: exp.PartitionRange) -> str:
591        name = self.sql(expression, "this")
592        values = expression.expressions
593
594        if len(values) != 1:
595            # Multiple values: use VALUES [ ... )
596            if values and isinstance(values[0], list):
597                values_sql = ", ".join(
598                    f"({', '.join(self.sql(v) for v in inner)})" for inner in values
599                )
600            else:
601                values_sql = ", ".join(f"({self.sql(v)})" for v in values)
602
603            return f"PARTITION {name} VALUES [{values_sql})"
604
605        return f"PARTITION {name} VALUES LESS THAN ({self.sql(values[0])})"
606
607    def partitionbyrangepropertydynamic_sql(
608        self, expression: exp.PartitionByRangePropertyDynamic
609    ) -> str:
610        # Generates: FROM ("start") TO ("end") INTERVAL N UNIT
611        start = self.sql(expression, "start")
612        end = self.sql(expression, "end")
613        every = expression.args.get("every")
614
615        if every:
616            number = self.sql(every, "this")
617            interval = f"INTERVAL {number} {self.sql(every, 'unit')}"
618        else:
619            interval = ""
620
621        return f"FROM ({start}) TO ({end}) {interval}"
622
623    def partitionedbyproperty_sql(self, expression: exp.PartitionedByProperty) -> str:
624        this = expression.this
625        if isinstance(this, exp.Schema):
626            return f"PARTITION BY ({self.expressions(this, flat=True)})"
627        return f"PARTITION BY ({self.sql(this)})"
628
629    def table_sql(self, expression: exp.Table, sep: str = " AS ") -> str:
630        """Override table_sql to avoid AS keyword in UPDATE and DELETE statements."""
631        ancestor = expression.find_ancestor(exp.Update, exp.Delete, exp.Select)
632        if not isinstance(ancestor, exp.Select):
633            sep = " "
634        return super().table_sql(expression, sep=sep)