sqlglot.parsers.starrocks

  1from __future__ import annotations
  2
  3
  4from sqlglot import exp, parser
  5from sqlglot.dialects.dialect import (
  6    build_date_delta_with_interval,
  7    build_like,
  8    build_timestamp_trunc,
  9)
 10from sqlglot.helper import is_int, seq_get
 11from sqlglot.parsers.mysql import MySQLParser
 12from sqlglot.tokens import TokenType
 13
 14# https://docs.starrocks.io/docs/sql-reference/sql-statements/table_bucket_part_index/CREATE_TABLE/
 15AGGREGATE_COLUMN_CONSTRAINTS = ("SUM", "MAX", "MIN", "REPLACE", "REPLACE_IF_NOT_NULL", "HLL_UNION")
 16
 17
 18def _build_time_slice(args: list[exp.Expr]) -> exp.TimeSlice | None:
 19    # TIME_SLICE(dt, INTERVAL n unit [, FLOOR | CEIL])
 20    # https://docs.starrocks.io/docs/sql-reference/sql-functions/date-time-functions/time_slice/
 21    if len(args) < 2:
 22        return None
 23
 24    if isinstance(seq_get(args, 1), exp.Interval):
 25        time_slice = build_date_delta_with_interval(exp.TimeSlice)(args)
 26        boundary = seq_get(args, 2)
 27    else:
 28        time_slice = exp.TimeSlice.from_arg_list(args)
 29        boundary = seq_get(args, 3)
 30    if not time_slice:
 31        return None
 32
 33    amount = time_slice.expression
 34    if amount.is_string and is_int(amount.name):
 35        time_slice.set("expression", exp.Literal.number(amount.name))
 36
 37    if boundary:
 38        time_slice.set(
 39            "kind", exp.Literal.string("END" if boundary.name.upper() == "CEIL" else "START")
 40        )
 41    return time_slice
 42
 43
 44class StarRocksParser(MySQLParser):
 45    IDENTIFY_PIVOT_STRINGS = True
 46    PIVOT_COLUMN_NAMING = "agg_name_if_multiple"
 47
 48    # Unlike MySQL, dropping a column requires the COLUMN keyword
 49    ALTER_DROP_REQUIRES_COLUMN = True
 50
 51    # StarRocks supports LEFT SEMI JOIN and LEFT ANTI JOIN natively
 52    # https://docs.starrocks.io/docs/sql-reference/sql-statements/table_bucket_part_index/SELECT/SELECT_JOIN/
 53    TABLE_ALIAS_TOKENS = MySQLParser.TABLE_ALIAS_TOKENS - {TokenType.ANTI, TokenType.SEMI}
 54
 55    # The keys of a map{key: value, ...} literal are expressions, not struct field names
 56    MAP_KEYS_ARE_ARBITRARY_EXPRESSIONS = True
 57
 58    FUNCTIONS = {
 59        **MySQLParser.FUNCTIONS,
 60        "ADDDATE": build_date_delta_with_interval(exp.DateAdd, default_unit="DAY"),
 61        "DATE_ADD": build_date_delta_with_interval(exp.DateAdd, default_unit="DAY"),
 62        "DATE_SUB": build_date_delta_with_interval(exp.DateSub, default_unit="DAY"),
 63        "SUBDATE": build_date_delta_with_interval(exp.DateSub, default_unit="DAY"),
 64        "DATE_TRUNC": build_timestamp_trunc,
 65        "DATE_DIFF": lambda args: exp.DateDiff(
 66            this=seq_get(args, 1), expression=seq_get(args, 2), unit=seq_get(args, 0)
 67        ),
 68        "ARRAY_FLATTEN": exp.Flatten.from_arg_list,
 69        "LIKE": build_like(exp.Like),
 70        "TO_DAYS": exp.ToDays.from_arg_list,
 71        "REGEXP": exp.RegexpLike.from_arg_list,
 72        # StarRocks' MAP() is a variadic constructor: MAP(k1, v1, k2, v2, ...)
 73        # https://docs.starrocks.io/docs/sql-reference/sql-functions/map-functions/map/
 74        "MAP": parser.build_var_map,
 75        # TABLE(<tvf>) wraps a table function invocation whose arguments are constants
 76        # https://docs.starrocks.io/docs/sql-reference/sql-functions/table-functions/generate_series/
 77        "TABLE": lambda args: exp.TableFromRows(this=seq_get(args, 0)),
 78        "TIME_SLICE": _build_time_slice,
 79    }
 80
 81    PROPERTY_PARSERS = {
 82        **MySQLParser.PROPERTY_PARSERS,
 83        "PROPERTIES": lambda self: self._parse_wrapped_properties(),
 84        "AGGREGATE": lambda self: self._parse_composite_key_property(exp.AggregateKeyProperty),
 85        "UNIQUE": lambda self: self._parse_composite_key_property(exp.UniqueKeyProperty),
 86        "ROLLUP": lambda self: self._parse_rollup_property(),
 87        "REFRESH": lambda self: self._parse_refresh_property(),
 88    }
 89
 90    CONSTRAINT_PARSERS = {
 91        **MySQLParser.CONSTRAINT_PARSERS,
 92        **dict.fromkeys(
 93            AGGREGATE_COLUMN_CONSTRAINTS, lambda self: exp.var(self._prev.text.upper())
 94        ),
 95    }
 96
 97    def _parse_pivot(self) -> exp.Pivot | None:
 98        pivot = super()._parse_pivot()
 99        if pivot is None or pivot.unpivot:
100            return pivot
101
102        names = [agg.alias for agg in pivot.expressions] if len(pivot.expressions) > 1 else [""]
103        if len(names) > 1 and not all(names):
104            # Unaliased aggregate names depend on StarRocks' column resolution.
105            pivot.set("columns", [])
106            return pivot
107
108        columns: list[exp.Identifier] = []
109        for field in pivot.fields[0].expressions:
110            if isinstance(field, exp.PivotAlias):
111                name = field.alias
112            elif isinstance(field, exp.Tuple):
113                name = (
114                    "{"
115                    + ",".join(
116                        e.sql(dialect=self.dialect, comments=False) for e in field.expressions
117                    )
118                    + "}"
119                )
120            else:
121                name = field.sql(dialect=self.dialect, comments=False)
122            columns.extend(exp.to_identifier(f"{name}_{agg}" if agg else name) for agg in names)
123        pivot.set("columns", columns)
124        return pivot
125
126    def _parse_rollup_property(self) -> exp.RollupProperty:
127        # ROLLUP (rollup_name (col1, col2) [FROM from_index] [PROPERTIES (...)], ...)
128        def parse_rollup_index() -> exp.RollupIndex:
129            return self.expression(
130                exp.RollupIndex(
131                    this=self._parse_id_var(),
132                    expressions=self._parse_wrapped_id_vars(),
133                    from_index=self._parse_id_var() if self._match_text_seq("FROM") else None,
134                    properties=self.expression(
135                        exp.Properties(expressions=self._parse_wrapped_properties())
136                    )
137                    if self._match_text_seq("PROPERTIES")
138                    else None,
139                )
140            )
141
142        return self.expression(
143            exp.RollupProperty(expressions=self._parse_wrapped_csv(parse_rollup_index))
144        )
145
146    def _parse_create(self) -> exp.Create | exp.Command:
147        create = super()._parse_create()
148
149        # Starrocks' primary key is defined outside of the schema, so we need to move it there
150        # https://docs.starrocks.io/docs/table_design/table_types/primary_key_table/#usage
151        if isinstance(create, exp.Create) and isinstance(create.this, exp.Schema):
152            props = create.args.get("properties")
153            if props:
154                primary_key = props.find(exp.PrimaryKey)
155                if primary_key:
156                    create.this.append("expressions", primary_key.pop())
157
158        return create
159
160    def _parse_bracket(self, this: exp.Expr | None = None) -> exp.Expr | None:
161        # map{key: value, ...}
162        # https://docs.starrocks.io/docs/sql-reference/data-types/semi_structured/Map/
163        is_map_literal = (
164            isinstance(this, exp.Column)
165            and not this.table
166            and this.name.upper() == "MAP"
167            and self._match(TokenType.L_BRACE, advance=False)
168        )
169
170        bracket = super()._parse_bracket(this)
171        if not is_map_literal or not bracket:
172            return bracket
173
174        # Wrap the literal itself, not any subscripts that follow it, e.g. map{...}[1]
175        struct = bracket
176        while isinstance(struct, exp.Bracket):
177            struct = struct.this
178
179        to_map = self.expression(exp.ToMap(this=struct.copy()))
180        if struct is bracket:
181            return to_map
182
183        struct.replace(to_map)
184        return bracket
185
186    def _parse_unnest(self, with_alias: bool = True) -> exp.Unnest | None:
187        unnest = super()._parse_unnest(with_alias=with_alias)
188
189        if unnest:
190            alias = unnest.args.get("alias")
191
192            if not alias:
193                # Starrocks defaults to naming the table alias as "unnest"
194                alias = exp.TableAlias(this=exp.to_identifier("unnest"))
195                unnest.set("alias", alias)
196
197            # Starrocks defaults to naming the UNNEST column as "unnest" if it's not otherwise
198            # specified. With several arrays every column is named "unnest", so none is added.
199            if not alias.args.get("columns") and len(unnest.expressions) == 1:
200                alias.set("columns", [exp.to_identifier("unnest")])
201
202        return unnest
203
204    def _parse_partitioned_by(self) -> exp.PartitionedByProperty:
205        return self.expression(
206            exp.PartitionedByProperty(
207                this=exp.Schema(
208                    expressions=self._parse_wrapped_csv(self._parse_assignment, optional=True)
209                )
210            )
211        )
212
213    def _parse_partition_property(
214        self,
215    ) -> exp.Expr | None | list[exp.Expr]:
216        expr = super()._parse_partition_property()
217
218        if not expr:
219            return self._parse_partitioned_by()
220
221        if isinstance(expr, exp.Property):
222            return expr
223
224        self._match_l_paren()
225
226        if self._match_text_seq("START", advance=False):
227            create_expressions = self._parse_csv(self._parse_partitioning_granularity_dynamic)
228        else:
229            create_expressions = None
230
231        self._match_r_paren()
232
233        return self.expression(
234            exp.PartitionByRangeProperty(
235                partition_expressions=expr, create_expressions=create_expressions
236            )
237        )
238
239    def _parse_partitioning_granularity_dynamic(self) -> exp.PartitionByRangePropertyDynamic:
240        self._match_text_seq("START")
241        start = self._parse_wrapped(self._parse_string)
242        self._match_text_seq("END")
243        end = self._parse_wrapped(self._parse_string)
244        self._match_text_seq("EVERY")
245        every = self._parse_wrapped(lambda: self._parse_interval() or self._parse_number())
246        return self.expression(
247            exp.PartitionByRangePropertyDynamic(start=start, end=end, every=every)
248        )
249
250    def _parse_refresh_property(self) -> exp.RefreshTriggerProperty:
251        """
252        REFRESH [DEFERRED | IMMEDIATE]
253                [ASYNC | ASYNC [START (<start_time>)] EVERY (INTERVAL <refresh_interval>) | MANUAL]
254        """
255        method = self._match_texts(("DEFERRED", "IMMEDIATE")) and self._prev.text.upper()
256        kind = self._match_texts(("ASYNC", "MANUAL")) and self._prev.text.upper()
257        start = self._match_text_seq("START") and self._parse_wrapped(self._parse_string)
258
259        if self._match_text_seq("EVERY"):
260            self._match_l_paren()
261            self._match_text_seq("INTERVAL")
262            every = self._parse_number()
263            unit = self._parse_var(any_token=True)
264            self._match_r_paren()
265        else:
266            every = None
267            unit = None
268
269        return self.expression(
270            exp.RefreshTriggerProperty(
271                method=method, kind=kind, starts=start, every=every, unit=unit
272            )
273        )
274
275    def _parse_refresh(self) -> exp.Refresh | exp.Command:
276        if not self._match_text_seq("MATERIALIZED", "VIEW"):
277            return super()._parse_refresh()
278
279        # _parse_table would read PARTITION and FORCE as partition selection and an index hint
280        # https://docs.starrocks.io/docs/sql-reference/sql-statements/materialized_view/REFRESH_MATERIALIZED_VIEW/
281        this = self._parse_table_parts()
282
283        partition_start = None
284        partition_end = None
285        if self._match_text_seq("PARTITION", "START"):
286            partition_start = self._parse_wrapped(self._parse_string)
287            self._match_text_seq("END")
288            partition_end = self._parse_wrapped(self._parse_string)
289
290        force = self._match_text_seq("FORCE")
291        mode = None
292        if self._match_text_seq("WITH", "SYNC", "MODE"):
293            mode = "SYNC"
294        elif self._match_text_seq("WITH", "ASYNC", "MODE"):
295            mode = "ASYNC"
296
297        return self.expression(
298            exp.Refresh(
299                this=this,
300                kind="MATERIALIZED VIEW",
301                partition_start=partition_start,
302                partition_end=partition_end,
303                force=force,
304                mode=mode,
305            )
306        )