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    # Unlike MySQL, dropping a column requires the COLUMN keyword
 29    ALTER_DROP_REQUIRES_COLUMN = True
 30
 31    # https://doris.apache.org/docs/sql-manual/basic-element/operators/conditional-operators/full-text-search-operators
 32    MATCH_OPERATORS = {
 33        "MATCH",
 34        "MATCH_ALL",
 35        "MATCH_ANY",
 36        "MATCH_PHRASE",
 37        "MATCH_PHRASE_EDGE",
 38        "MATCH_PHRASE_PREFIX",
 39        "MATCH_REGEXP",
 40    }
 41
 42    FUNCTIONS = {
 43        **MySQLParser.FUNCTIONS,
 44        "ADDDATE": build_date_delta_with_interval(exp.DateAdd, default_unit="DAY"),
 45        "COLLECT_SET": exp.ArrayUniqueAgg.from_arg_list,
 46        "DATE_ADD": build_date_delta_with_interval(exp.DateAdd, default_unit="DAY"),
 47        "DATE_SUB": build_date_delta_with_interval(exp.DateSub, default_unit="DAY"),
 48        "DATE_TRUNC": _build_date_trunc,
 49        "L2_DISTANCE": exp.EuclideanDistance.from_arg_list,
 50        "MONTHS_ADD": exp.AddMonths.from_arg_list,
 51        "REGEXP": exp.RegexpLike.from_arg_list,
 52        "SUBDATE": build_date_delta_with_interval(exp.DateSub, default_unit="DAY"),
 53        "TO_DATE": exp.TsOrDsToDate.from_arg_list,
 54    }
 55
 56    FUNCTION_PARSERS = {
 57        k: v for k, v in MySQLParser.FUNCTION_PARSERS.items() if k != "GROUP_CONCAT"
 58    }
 59
 60    NO_PAREN_FUNCTIONS = {
 61        k: v for k, v in MySQLParser.NO_PAREN_FUNCTIONS.items() if k != TokenType.CURRENT_DATE
 62    }
 63
 64    PROPERTY_PARSERS = {
 65        **MySQLParser.PROPERTY_PARSERS,
 66        "PROPERTIES": lambda self: self._parse_wrapped_properties(),
 67        "UNIQUE": lambda self: self._parse_composite_key_property(exp.UniqueKeyProperty),
 68        # Plain KEY without UNIQUE/DUPLICATE/AGGREGATE prefixes should be treated as UniqueKeyProperty with unique=False
 69        "KEY": lambda self: self._parse_composite_key_property(exp.UniqueKeyProperty),
 70        "BUILD": lambda self: self._parse_build_property(),
 71        "REFRESH": lambda self: self._parse_refresh_property(),
 72    }
 73
 74    def _parse_range(self, this: exp.Expr | None = None, parse_is: bool = False) -> exp.Expr | None:
 75        this = this or self._parse_bitwise()
 76        index = self._index
 77        negate = self._match(TokenType.NOT)
 78
 79        if self._curr.token_type == TokenType.VAR and self._match_texts(self.MATCH_OPERATORS):
 80            kind = self._prev.text.upper().partition("_")[2] or None
 81            expression = self._parse_bitwise()
 82            if not expression:
 83                # Preserve implicit aliases named after MATCH operators.
 84                self._retreat(index)
 85                return super()._parse_range(this, parse_is=parse_is)
 86
 87            analyzer = None
 88            if self._match_text_seq("USING", "ANALYZER"):
 89                analyzer = self._parse_string() or self._parse_id_var(any_token=False)
 90                if not analyzer:
 91                    self.raise_error("Expected an analyzer name after USING ANALYZER")
 92
 93            this = self.expression(
 94                exp.Match(this=this, expression=expression, kind=kind, analyzer=analyzer)
 95            )
 96            if negate:
 97                this = self.expression(exp.Not(this=this))
 98        elif negate:
 99            self._retreat(index)
100
101        return super()._parse_range(this, parse_is=parse_is)
102
103    def _parse_partition_property(
104        self,
105    ) -> exp.Expr | None | list[exp.Expr]:
106        expr = super()._parse_partition_property()
107
108        if not expr:
109            return self._parse_partitioned_by()
110
111        if isinstance(expr, exp.Property):
112            return expr
113
114        self._match_l_paren()
115
116        if self._match_text_seq("FROM", advance=False):
117            create_expressions = self._parse_csv(self._parse_partitioning_granularity_dynamic)
118        else:
119            create_expressions = None
120
121        self._match_r_paren()
122
123        return self.expression(
124            exp.PartitionByRangeProperty(
125                partition_expressions=expr, create_expressions=create_expressions
126            )
127        )
128
129    def _parse_partitioning_granularity_dynamic(self) -> exp.PartitionByRangePropertyDynamic:
130        self._match_text_seq("FROM")
131        start = self._parse_wrapped(self._parse_string)
132        self._match_text_seq("TO")
133        end = self._parse_wrapped(self._parse_string)
134        self._match_text_seq("INTERVAL")
135        number = self._parse_number()
136        unit = self._parse_var(any_token=True)
137        every = self.expression(exp.Interval(this=number, unit=unit))
138        return self.expression(
139            exp.PartitionByRangePropertyDynamic(start=start, end=end, every=every)
140        )
141
142    def _parse_partition_range_value(self) -> exp.Expr | None:
143        expr = super()._parse_partition_range_value()
144
145        if isinstance(expr, exp.Partition):
146            return expr
147
148        self._match_text_seq("VALUES")
149        name = expr
150
151        # Doris-specific bracket syntax: VALUES [(...), (...))
152        self._match(TokenType.L_BRACKET)
153        values = self._parse_csv(lambda: self._parse_wrapped_csv(self._parse_expression))
154
155        self._match(TokenType.R_BRACKET)
156        self._match(TokenType.R_PAREN)
157
158        part_range = self.expression(exp.PartitionRange(this=name, expressions=values))
159        return self.expression(exp.Partition(expressions=[part_range]))
160
161    def _parse_build_property(self) -> exp.BuildProperty:
162        return self.expression(exp.BuildProperty(this=self._parse_var(upper=True)))
163
164    def _parse_refresh_property(self) -> exp.RefreshTriggerProperty:
165        method = self._parse_var(upper=True)
166
167        self._match(TokenType.ON)
168
169        kind = self._match_texts(("MANUAL", "COMMIT", "SCHEDULE")) and self._prev.text.upper()
170        every = self._match_text_seq("EVERY") and self._parse_number()
171        unit = self._parse_var(any_token=True) if every else None
172        starts = self._match_text_seq("STARTS") and self._parse_string()
173
174        return self.expression(
175            exp.RefreshTriggerProperty(
176                method=method, kind=kind, every=every, unit=unit, starts=starts
177            )
178        )