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