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        "and",
108        "anti",
109        "append",
110        "array",
111        "array_range",
112        "as",
113        "asc",
114        "at",
115        "authors",
116        "auto",
117        "auto_increment",
118        "backend",
119        "backends",
120        "backup",
121        "begin",
122        "belong",
123        "between",
124        "bigint",
125        "bin",
126        "binary",
127        "binlog",
128        "bitand",
129        "bitmap",
130        "bitmap_union",
131        "bitor",
132        "bitxor",
133        "blob",
134        "boolean",
135        "brief",
136        "broker",
137        "buckets",
138        "build",
139        "builtin",
140        "bulk",
141        "by",
142        "cached",
143        "call",
144        "cancel",
145        "case",
146        "cast",
147        "catalog",
148        "catalogs",
149        "chain",
150        "char",
151        "character",
152        "charset",
153        "check",
154        "clean",
155        "cluster",
156        "clusters",
157        "collate",
158        "collation",
159        "collect",
160        "column",
161        "columns",
162        "comment",
163        "commit",
164        "committed",
165        "compact",
166        "complete",
167        "config",
168        "connection",
169        "connection_id",
170        "consistent",
171        "constraint",
172        "constraints",
173        "convert",
174        "copy",
175        "count",
176        "create",
177        "creation",
178        "cron",
179        "cross",
180        "cube",
181        "current",
182        "current_catalog",
183        "current_date",
184        "current_time",
185        "current_timestamp",
186        "current_user",
187        "data",
188        "database",
189        "databases",
190        "date",
191        "date_add",
192        "date_ceil",
193        "date_diff",
194        "date_floor",
195        "date_sub",
196        "dateadd",
197        "datediff",
198        "datetime",
199        "datetimev2",
200        "datev2",
201        "datetimev1",
202        "datev1",
203        "day",
204        "days_add",
205        "days_sub",
206        "decimal",
207        "decimalv2",
208        "decimalv3",
209        "decommission",
210        "default",
211        "deferred",
212        "delete",
213        "demand",
214        "desc",
215        "describe",
216        "diagnose",
217        "disk",
218        "distinct",
219        "distinctpc",
220        "distinctpcsa",
221        "distributed",
222        "distribution",
223        "div",
224        "do",
225        "doris_internal_table_id",
226        "double",
227        "drop",
228        "dropp",
229        "dual",
230        "duplicate",
231        "dynamic",
232        "else",
233        "enable",
234        "encryptkey",
235        "encryptkeys",
236        "end",
237        "ends",
238        "engine",
239        "engines",
240        "enter",
241        "errors",
242        "events",
243        "every",
244        "except",
245        "exclude",
246        "execute",
247        "exists",
248        "expired",
249        "explain",
250        "export",
251        "extended",
252        "external",
253        "extract",
254        "failed_login_attempts",
255        "false",
256        "fast",
257        "feature",
258        "fields",
259        "file",
260        "filter",
261        "first",
262        "float",
263        "follower",
264        "following",
265        "for",
266        "foreign",
267        "force",
268        "format",
269        "free",
270        "from",
271        "frontend",
272        "frontends",
273        "full",
274        "function",
275        "functions",
276        "generic",
277        "global",
278        "grant",
279        "grants",
280        "graph",
281        "group",
282        "grouping",
283        "groups",
284        "hash",
285        "having",
286        "hdfs",
287        "help",
288        "histogram",
289        "hll",
290        "hll_union",
291        "hostname",
292        "hour",
293        "hub",
294        "identified",
295        "if",
296        "ignore",
297        "immediate",
298        "in",
299        "incremental",
300        "index",
301        "indexes",
302        "infile",
303        "inner",
304        "insert",
305        "install",
306        "int",
307        "integer",
308        "intermediate",
309        "intersect",
310        "interval",
311        "into",
312        "inverted",
313        "ipv4",
314        "ipv6",
315        "is",
316        "is_not_null_pred",
317        "is_null_pred",
318        "isnull",
319        "isolation",
320        "job",
321        "jobs",
322        "join",
323        "json",
324        "jsonb",
325        "key",
326        "keys",
327        "kill",
328        "label",
329        "largeint",
330        "last",
331        "lateral",
332        "ldap",
333        "ldap_admin_password",
334        "left",
335        "less",
336        "level",
337        "like",
338        "limit",
339        "lines",
340        "link",
341        "list",
342        "load",
343        "local",
344        "localtime",
345        "localtimestamp",
346        "location",
347        "lock",
348        "logical",
349        "low_priority",
350        "manual",
351        "map",
352        "match",
353        "match_all",
354        "match_any",
355        "match_phrase",
356        "match_phrase_edge",
357        "match_phrase_prefix",
358        "match_regexp",
359        "materialized",
360        "max",
361        "maxvalue",
362        "memo",
363        "merge",
364        "migrate",
365        "migrations",
366        "min",
367        "minus",
368        "minute",
369        "modify",
370        "month",
371        "mtmv",
372        "name",
373        "names",
374        "natural",
375        "negative",
376        "never",
377        "next",
378        "ngram_bf",
379        "no",
380        "non_nullable",
381        "not",
382        "null",
383        "nulls",
384        "observer",
385        "of",
386        "offset",
387        "on",
388        "only",
389        "open",
390        "optimized",
391        "or",
392        "order",
393        "outer",
394        "outfile",
395        "over",
396        "overwrite",
397        "parameter",
398        "parsed",
399        "partition",
400        "partitions",
401        "password",
402        "password_expire",
403        "password_history",
404        "password_lock_time",
405        "password_reuse",
406        "path",
407        "pause",
408        "percent",
409        "period",
410        "permissive",
411        "physical",
412        "plan",
413        "process",
414        "plugin",
415        "plugins",
416        "policy",
417        "preceding",
418        "prepare",
419        "primary",
420        "proc",
421        "procedure",
422        "processlist",
423        "profile",
424        "properties",
425        "property",
426        "quantile_state",
427        "quantile_union",
428        "query",
429        "quota",
430        "random",
431        "range",
432        "read",
433        "real",
434        "rebalance",
435        "recover",
436        "recycle",
437        "refresh",
438        "references",
439        "regexp",
440        "release",
441        "rename",
442        "repair",
443        "repeatable",
444        "replace",
445        "replace_if_not_null",
446        "replica",
447        "repositories",
448        "repository",
449        "resource",
450        "resources",
451        "restore",
452        "restrictive",
453        "resume",
454        "returns",
455        "revoke",
456        "rewritten",
457        "right",
458        "rlike",
459        "role",
460        "roles",
461        "rollback",
462        "rollup",
463        "routine",
464        "row",
465        "rows",
466        "s3",
467        "sample",
468        "schedule",
469        "scheduler",
470        "schema",
471        "schemas",
472        "second",
473        "select",
474        "semi",
475        "sequence",
476        "serializable",
477        "session",
478        "set",
479        "sets",
480        "shape",
481        "show",
482        "signed",
483        "skew",
484        "smallint",
485        "snapshot",
486        "soname",
487        "split",
488        "sql_block_rule",
489        "start",
490        "starts",
491        "stats",
492        "status",
493        "stop",
494        "storage",
495        "stream",
496        "streaming",
497        "string",
498        "struct",
499        "subdate",
500        "sum",
501        "superuser",
502        "switch",
503        "sync",
504        "system",
505        "table",
506        "tables",
507        "tablesample",
508        "tablet",
509        "tablets",
510        "task",
511        "tasks",
512        "temporary",
513        "terminated",
514        "text",
515        "than",
516        "then",
517        "time",
518        "timestamp",
519        "timestampadd",
520        "timestampdiff",
521        "tinyint",
522        "to",
523        "transaction",
524        "trash",
525        "tree",
526        "triggers",
527        "trim",
528        "true",
529        "truncate",
530        "type",
531        "type_cast",
532        "types",
533        "unbounded",
534        "uncommitted",
535        "uninstall",
536        "union",
537        "unique",
538        "unlock",
539        "unsigned",
540        "update",
541        "use",
542        "user",
543        "using",
544        "value",
545        "values",
546        "varchar",
547        "variables",
548        "variant",
549        "vault",
550        "verbose",
551        "version",
552        "view",
553        "warnings",
554        "week",
555        "when",
556        "where",
557        "whitelist",
558        "with",
559        "work",
560        "workload",
561        "write",
562        "xor",
563        "year",
564    }
565
566    def uniquekeyproperty_sql(
567        self, expression: exp.UniqueKeyProperty, prefix: str = "UNIQUE KEY"
568    ) -> str:
569        create_stmt = expression.find_ancestor(exp.Create)
570        if create_stmt and create_stmt.args["properties"].find(exp.MaterializedProperty):
571            return super().uniquekeyproperty_sql(expression, prefix="KEY")
572
573        return super().uniquekeyproperty_sql(expression)
574
575    def partitionrange_sql(self, expression: exp.PartitionRange) -> str:
576        name = self.sql(expression, "this")
577        values = expression.expressions
578
579        if len(values) != 1:
580            # Multiple values: use VALUES [ ... )
581            if values and isinstance(values[0], list):
582                values_sql = ", ".join(
583                    f"({', '.join(self.sql(v) for v in inner)})" for inner in values
584                )
585            else:
586                values_sql = ", ".join(f"({self.sql(v)})" for v in values)
587
588            return f"PARTITION {name} VALUES [{values_sql})"
589
590        return f"PARTITION {name} VALUES LESS THAN ({self.sql(values[0])})"
591
592    def partitionbyrangepropertydynamic_sql(
593        self, expression: exp.PartitionByRangePropertyDynamic
594    ) -> str:
595        # Generates: FROM ("start") TO ("end") INTERVAL N UNIT
596        start = self.sql(expression, "start")
597        end = self.sql(expression, "end")
598        every = expression.args.get("every")
599
600        if every:
601            number = self.sql(every, "this")
602            interval = f"INTERVAL {number} {self.sql(every, 'unit')}"
603        else:
604            interval = ""
605
606        return f"FROM ({start}) TO ({end}) {interval}"
607
608    def partitionedbyproperty_sql(self, expression: exp.PartitionedByProperty) -> str:
609        this = expression.this
610        if isinstance(this, exp.Schema):
611            return f"PARTITION BY ({self.expressions(this, flat=True)})"
612        return f"PARTITION BY ({self.sql(this)})"
613
614    def table_sql(self, expression: exp.Table, sep: str = " AS ") -> str:
615        """Override table_sql to avoid AS keyword in UPDATE and DELETE statements."""
616        ancestor = expression.find_ancestor(exp.Update, exp.Delete, exp.Select)
617        if not isinstance(ancestor, exp.Select):
618            sep = " "
619        return super().table_sql(expression, sep=sep)