sqlglot.parsers.starrocks

  1from __future__ import annotations
  2
  3
  4from sqlglot import exp
  5from sqlglot.dialects.dialect import build_date_delta_with_interval, build_timestamp_trunc
  6from sqlglot.helper import seq_get
  7from sqlglot.parsers.mysql import MySQLParser
  8from sqlglot.tokens import TokenType
  9
 10
 11class StarRocksParser(MySQLParser):
 12    # StarRocks supports LEFT SEMI JOIN and LEFT ANTI JOIN natively
 13    # https://docs.starrocks.io/docs/sql-reference/sql-statements/table_bucket_part_index/SELECT/SELECT_JOIN/
 14    TABLE_ALIAS_TOKENS = MySQLParser.TABLE_ALIAS_TOKENS - {TokenType.ANTI, TokenType.SEMI}
 15
 16    FUNCTIONS = {
 17        **MySQLParser.FUNCTIONS,
 18        "ADDDATE": build_date_delta_with_interval(exp.DateAdd, default_unit="DAY"),
 19        "DATE_ADD": build_date_delta_with_interval(exp.DateAdd, default_unit="DAY"),
 20        "DATE_SUB": build_date_delta_with_interval(exp.DateSub, default_unit="DAY"),
 21        "SUBDATE": build_date_delta_with_interval(exp.DateSub, default_unit="DAY"),
 22        "DATE_TRUNC": build_timestamp_trunc,
 23        "DATEDIFF": lambda args: exp.DateDiff(
 24            this=seq_get(args, 0), expression=seq_get(args, 1), unit=exp.Literal.string("DAY")
 25        ),
 26        "DATE_DIFF": lambda args: exp.DateDiff(
 27            this=seq_get(args, 1), expression=seq_get(args, 2), unit=seq_get(args, 0)
 28        ),
 29        "ARRAY_FLATTEN": exp.Flatten.from_arg_list,
 30        "REGEXP": exp.RegexpLike.from_arg_list,
 31    }
 32
 33    PROPERTY_PARSERS = {
 34        **MySQLParser.PROPERTY_PARSERS,
 35        "PROPERTIES": lambda self: self._parse_wrapped_properties(),
 36        "UNIQUE": lambda self: self._parse_composite_key_property(exp.UniqueKeyProperty),
 37        "ROLLUP": lambda self: self._parse_rollup_property(),
 38        "REFRESH": lambda self: self._parse_refresh_property(),
 39    }
 40
 41    def _parse_rollup_property(self) -> exp.RollupProperty:
 42        # ROLLUP (rollup_name (col1, col2) [FROM from_index] [PROPERTIES (...)], ...)
 43        def parse_rollup_index() -> exp.RollupIndex:
 44            return self.expression(
 45                exp.RollupIndex(
 46                    this=self._parse_id_var(),
 47                    expressions=self._parse_wrapped_id_vars(),
 48                    from_index=self._parse_id_var() if self._match_text_seq("FROM") else None,
 49                    properties=self.expression(
 50                        exp.Properties(expressions=self._parse_wrapped_properties())
 51                    )
 52                    if self._match_text_seq("PROPERTIES")
 53                    else None,
 54                )
 55            )
 56
 57        return self.expression(
 58            exp.RollupProperty(expressions=self._parse_wrapped_csv(parse_rollup_index))
 59        )
 60
 61    def _parse_create(self) -> exp.Create | exp.Command:
 62        create = super()._parse_create()
 63
 64        # Starrocks' primary key is defined outside of the schema, so we need to move it there
 65        # https://docs.starrocks.io/docs/table_design/table_types/primary_key_table/#usage
 66        if isinstance(create, exp.Create) and isinstance(create.this, exp.Schema):
 67            props = create.args.get("properties")
 68            if props:
 69                primary_key = props.find(exp.PrimaryKey)
 70                if primary_key:
 71                    create.this.append("expressions", primary_key.pop())
 72
 73        return create
 74
 75    def _parse_unnest(self, with_alias: bool = True) -> exp.Unnest | None:
 76        unnest = super()._parse_unnest(with_alias=with_alias)
 77
 78        if unnest:
 79            alias = unnest.args.get("alias")
 80
 81            if not alias:
 82                # Starrocks defaults to naming the table alias as "unnest"
 83                alias = exp.TableAlias(
 84                    this=exp.to_identifier("unnest"), columns=[exp.to_identifier("unnest")]
 85                )
 86                unnest.set("alias", alias)
 87            elif not alias.args.get("columns"):
 88                # Starrocks defaults to naming the UNNEST column as "unnest"
 89                # if it's not otherwise specified
 90                alias.set("columns", [exp.to_identifier("unnest")])
 91
 92        return unnest
 93
 94    def _parse_partitioned_by(self) -> exp.PartitionedByProperty:
 95        return self.expression(
 96            exp.PartitionedByProperty(
 97                this=exp.Schema(
 98                    expressions=self._parse_wrapped_csv(self._parse_assignment, optional=True)
 99                )
100            )
101        )
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("START", 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("START")
131        start = self._parse_wrapped(self._parse_string)
132        self._match_text_seq("END")
133        end = self._parse_wrapped(self._parse_string)
134        self._match_text_seq("EVERY")
135        every = self._parse_wrapped(lambda: self._parse_interval() or self._parse_number())
136        return self.expression(
137            exp.PartitionByRangePropertyDynamic(start=start, end=end, every=every)
138        )
139
140    def _parse_refresh_property(self) -> exp.RefreshTriggerProperty:
141        """
142        REFRESH [DEFERRED | IMMEDIATE]
143                [ASYNC | ASYNC [START (<start_time>)] EVERY (INTERVAL <refresh_interval>) | MANUAL]
144        """
145        method = self._match_texts(("DEFERRED", "IMMEDIATE")) and self._prev.text.upper()
146        kind = self._match_texts(("ASYNC", "MANUAL")) and self._prev.text.upper()
147        start = self._match_text_seq("START") and self._parse_wrapped(self._parse_string)
148
149        if self._match_text_seq("EVERY"):
150            self._match_l_paren()
151            self._match_text_seq("INTERVAL")
152            every = self._parse_number()
153            unit = self._parse_var(any_token=True)
154            self._match_r_paren()
155        else:
156            every = None
157            unit = None
158
159        return self.expression(
160            exp.RefreshTriggerProperty(
161                method=method, kind=kind, starts=start, every=every, unit=unit
162            )
163        )