sqlglot.parsers.doris

  1from __future__ import annotations
  2
  3
  4from sqlglot import exp
  5from sqlglot.dialects.dialect import build_date_delta_with_interval
  6from sqlglot.helper import seq_get
  7from sqlglot.parsers.mysql import MySQLParser
  8from sqlglot.tokens import TokenType
  9
 10
 11# Accept both DATE_TRUNC(datetime, unit) and DATE_TRUNC(unit, datetime)
 12def _build_date_trunc(args: list[exp.Expr]) -> exp.Expr:
 13    a0, a1 = seq_get(args, 0), seq_get(args, 1)
 14
 15    def _is_unit_like(e: exp.Expr | None) -> bool:
 16        if not (isinstance(e, exp.Literal) and e.is_string):
 17            return False
 18        text = e.this
 19        return not any(ch.isdigit() for ch in text)
 20
 21    # Determine which argument is the unit
 22    unit, this = (a0, a1) if _is_unit_like(a0) else (a1, a0)
 23
 24    return exp.TimestampTrunc(this=this, unit=unit)
 25
 26
 27class DorisParser(MySQLParser):
 28    FUNCTIONS = {
 29        **MySQLParser.FUNCTIONS,
 30        "ADDDATE": build_date_delta_with_interval(exp.DateAdd, default_unit="DAY"),
 31        "COLLECT_SET": exp.ArrayUniqueAgg.from_arg_list,
 32        "DATE_ADD": build_date_delta_with_interval(exp.DateAdd, default_unit="DAY"),
 33        "DATE_SUB": build_date_delta_with_interval(exp.DateSub, default_unit="DAY"),
 34        "DATE_TRUNC": _build_date_trunc,
 35        "L2_DISTANCE": exp.EuclideanDistance.from_arg_list,
 36        "MONTHS_ADD": exp.AddMonths.from_arg_list,
 37        "REGEXP": exp.RegexpLike.from_arg_list,
 38        "SUBDATE": build_date_delta_with_interval(exp.DateSub, default_unit="DAY"),
 39        "TO_DATE": exp.TsOrDsToDate.from_arg_list,
 40    }
 41
 42    FUNCTION_PARSERS = {
 43        k: v for k, v in MySQLParser.FUNCTION_PARSERS.items() if k != "GROUP_CONCAT"
 44    }
 45
 46    NO_PAREN_FUNCTIONS = {
 47        k: v for k, v in MySQLParser.NO_PAREN_FUNCTIONS.items() if k != TokenType.CURRENT_DATE
 48    }
 49
 50    PROPERTY_PARSERS = {
 51        **MySQLParser.PROPERTY_PARSERS,
 52        "PROPERTIES": lambda self: self._parse_wrapped_properties(),
 53        "UNIQUE": lambda self: self._parse_composite_key_property(exp.UniqueKeyProperty),
 54        # Plain KEY without UNIQUE/DUPLICATE/AGGREGATE prefixes should be treated as UniqueKeyProperty with unique=False
 55        "KEY": lambda self: self._parse_composite_key_property(exp.UniqueKeyProperty),
 56        "BUILD": lambda self: self._parse_build_property(),
 57        "REFRESH": lambda self: self._parse_refresh_property(),
 58    }
 59
 60    def _parse_partition_property(
 61        self,
 62    ) -> exp.Expr | None | list[exp.Expr]:
 63        expr = super()._parse_partition_property()
 64
 65        if not expr:
 66            return self._parse_partitioned_by()
 67
 68        if isinstance(expr, exp.Property):
 69            return expr
 70
 71        self._match_l_paren()
 72
 73        if self._match_text_seq("FROM", advance=False):
 74            create_expressions = self._parse_csv(self._parse_partitioning_granularity_dynamic)
 75        else:
 76            create_expressions = None
 77
 78        self._match_r_paren()
 79
 80        return self.expression(
 81            exp.PartitionByRangeProperty(
 82                partition_expressions=expr, create_expressions=create_expressions
 83            )
 84        )
 85
 86    def _parse_partitioning_granularity_dynamic(self) -> exp.PartitionByRangePropertyDynamic:
 87        self._match_text_seq("FROM")
 88        start = self._parse_wrapped(self._parse_string)
 89        self._match_text_seq("TO")
 90        end = self._parse_wrapped(self._parse_string)
 91        self._match_text_seq("INTERVAL")
 92        number = self._parse_number()
 93        unit = self._parse_var(any_token=True)
 94        every = self.expression(exp.Interval(this=number, unit=unit))
 95        return self.expression(
 96            exp.PartitionByRangePropertyDynamic(start=start, end=end, every=every)
 97        )
 98
 99    def _parse_partition_range_value(self) -> exp.Expr | None:
100        expr = super()._parse_partition_range_value()
101
102        if isinstance(expr, exp.Partition):
103            return expr
104
105        self._match_text_seq("VALUES")
106        name = expr
107
108        # Doris-specific bracket syntax: VALUES [(...), (...))
109        self._match(TokenType.L_BRACKET)
110        values = self._parse_csv(lambda: self._parse_wrapped_csv(self._parse_expression))
111
112        self._match(TokenType.R_BRACKET)
113        self._match(TokenType.R_PAREN)
114
115        part_range = self.expression(exp.PartitionRange(this=name, expressions=values))
116        return self.expression(exp.Partition(expressions=[part_range]))
117
118    def _parse_build_property(self) -> exp.BuildProperty:
119        return self.expression(exp.BuildProperty(this=self._parse_var(upper=True)))
120
121    def _parse_refresh_property(self) -> exp.RefreshTriggerProperty:
122        method = self._parse_var(upper=True)
123
124        self._match(TokenType.ON)
125
126        kind = self._match_texts(("MANUAL", "COMMIT", "SCHEDULE")) and self._prev.text.upper()
127        every = self._match_text_seq("EVERY") and self._parse_number()
128        unit = self._parse_var(any_token=True) if every else None
129        starts = self._match_text_seq("STARTS") and self._parse_string()
130
131        return self.expression(
132            exp.RefreshTriggerProperty(
133                method=method, kind=kind, every=every, unit=unit, starts=starts
134            )
135        )