sqlglot.generators.doris

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