sqlglot.generators.starrocks

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp, transforms
  6from sqlglot.dialects.dialect import (
  7    approx_count_distinct_sql,
  8    arrow_json_extract_sql,
  9    rename_func,
 10    weekstart_unit_to_str,
 11    inline_array_sql,
 12    property_sql,
 13    var_map_sql,
 14)
 15from sqlglot.generators.mysql import MySQLGenerator
 16
 17
 18def _eliminate_between_in_delete(expression: exp.Expr) -> exp.Expr:
 19    """
 20    StarRocks doesn't support BETWEEN in DELETE statements, so we convert
 21    BETWEEN expressions to explicit comparisons.
 22
 23    https://docs.starrocks.io/docs/sql-reference/sql-statements/table_bucket_part_index/DELETE/#parameters
 24
 25    Example:
 26        >>> from sqlglot import parse_one
 27        >>> expr = parse_one("DELETE FROM t WHERE x BETWEEN 1 AND 10")
 28        >>> print(_eliminate_between_in_delete(expr).sql(dialect="starrocks"))
 29        DELETE FROM t WHERE x >= 1 AND x <= 10
 30    """
 31    if where := expression.args.get("where"):
 32        for between in where.find_all(exp.Between):
 33            between.replace(
 34                exp.and_(
 35                    exp.GTE(this=between.this.copy(), expression=between.args["low"]),
 36                    exp.LTE(this=between.this.copy(), expression=between.args["high"]),
 37                    copy=False,
 38                )
 39            )
 40    return expression
 41
 42
 43# https://docs.starrocks.io/docs/sql-reference/sql-functions/spatial-functions/st_distance_sphere/
 44def st_distance_sphere(self, expression: exp.StDistance) -> str:
 45    point1 = expression.this
 46    point2 = expression.expression
 47
 48    point1_x = self.func("ST_X", point1)
 49    point1_y = self.func("ST_Y", point1)
 50    point2_x = self.func("ST_X", point2)
 51    point2_y = self.func("ST_Y", point2)
 52
 53    return self.func("ST_Distance_Sphere", point1_x, point1_y, point2_x, point2_y)
 54
 55
 56class StarRocksGenerator(MySQLGenerator):
 57    EXCEPT_INTERSECT_SUPPORT_ALL_CLAUSE = False
 58    JSON_TYPE_REQUIRED_FOR_EXTRACTION = False
 59    VARCHAR_REQUIRES_SIZE = False
 60    PARSE_JSON_NAME: str | None = "PARSE_JSON"
 61    WITH_PROPERTIES_PREFIX = "PROPERTIES"
 62    UPDATE_STATEMENT_SUPPORTS_FROM = True
 63    INSERT_OVERWRITE = " OVERWRITE"
 64
 65    # StarRocks doesn't support "IS TRUE/FALSE" syntax.
 66    IS_BOOL_ALLOWED = False
 67    # StarRocks doesn't support renaming a table with a database.
 68    RENAME_TABLE_WITH_DB = False
 69
 70    CAST_MAPPING: t.ClassVar[dict[exp.DType, str]] = {}
 71
 72    TYPE_MAPPING = {
 73        **MySQLGenerator.TYPE_MAPPING,
 74        exp.DType.INT128: "LARGEINT",
 75        exp.DType.TEXT: "STRING",
 76        exp.DType.TIMESTAMP: "DATETIME",
 77        exp.DType.TIMESTAMPTZ: "DATETIME",
 78    }
 79
 80    SQL_SECURITY_VIEW_LOCATION: t.ClassVar = exp.Properties.Location.POST_SCHEMA
 81
 82    PROPERTIES_LOCATION = {
 83        **MySQLGenerator.PROPERTIES_LOCATION,
 84        exp.PrimaryKey: exp.Properties.Location.POST_SCHEMA,
 85        exp.UniqueKeyProperty: exp.Properties.Location.POST_SCHEMA,
 86        exp.RollupProperty: exp.Properties.Location.POST_SCHEMA,
 87        exp.PartitionedByProperty: exp.Properties.Location.POST_SCHEMA,
 88    }
 89
 90    TRANSFORMS = {
 91        # StarRocks uses the native TRIM(str, chars)/LTRIM/RTRIM function form, not
 92        # MySQL's TRIM(chars FROM str) syntax, so it falls back to the base generator.
 93        **{
 94            k: v for k, v in MySQLGenerator.TRANSFORMS.items() if k not in (exp.DateTrunc, exp.Trim)
 95        },
 96        exp.ArgMax: rename_func("MAX_BY"),
 97        exp.ArgMin: rename_func("MIN_BY"),
 98        exp.Array: inline_array_sql,
 99        exp.ArrayAgg: rename_func("ARRAY_AGG"),
100        # a <@ b (ArrayContainedBy) is equivalent to ARRAY_CONTAINS_ALL(b, a)
101        exp.ArrayContainedBy: lambda self, e: self.func("ARRAY_CONTAINS_ALL", e.expression, e.this),
102        exp.ArrayContainsAll: rename_func("ARRAY_CONTAINS_ALL"),
103        exp.ArrayFilter: rename_func("ARRAY_FILTER"),
104        exp.ArrayToString: rename_func("ARRAY_JOIN"),
105        exp.ApproxDistinct: approx_count_distinct_sql,
106        exp.CurrentVersion: lambda *_: "CURRENT_VERSION()",
107        exp.DateDiff: lambda self, e: self.func(
108            "DATE_DIFF", weekstart_unit_to_str(self, e), e.this, e.expression
109        ),
110        exp.Delete: transforms.preprocess([_eliminate_between_in_delete]),
111        exp.Flatten: rename_func("ARRAY_FLATTEN"),
112        exp.JSONExtractScalar: arrow_json_extract_sql,
113        exp.JSONExtract: arrow_json_extract_sql,
114        # Both MAP forms (two-array MAP([keys], [values]) and variadic MAP(k1, v1, ...))
115        # generate StarRocks' variadic MAP(k1, v1, k2, v2, ...) constructor
116        exp.Map: lambda self, e: var_map_sql(self, e, "MAP"),
117        exp.Property: property_sql,
118        exp.RegexpLike: rename_func("REGEXP"),
119        # Inherited from MySQL, minus operations StarRocks supports natively
120        # (QUALIFY, FULL OUTER JOIN, SEMI/ANTI JOIN)
121        exp.Select: transforms.preprocess(
122            [
123                transforms.eliminate_distinct_on,
124                transforms.unnest_generate_date_array_using_recursive_cte,
125            ]
126        ),
127        exp.SchemaCommentProperty: lambda self, e: self.naked_property(e),
128        exp.SqlSecurityProperty: lambda self, e: f"SECURITY {self.sql(e.this)}",
129        exp.StDistance: st_distance_sphere,
130        exp.StrToUnix: lambda self, e: self.func("UNIX_TIMESTAMP", e.this, self.format_time(e)),
131        exp.TimestampTrunc: lambda self, e: self.func(
132            "DATE_TRUNC", weekstart_unit_to_str(self, e), e.this
133        ),
134        exp.TimeStrToDate: rename_func("TO_DATE"),
135        exp.UnixToStr: lambda self, e: self.func("FROM_UNIXTIME", e.this, self.format_time(e)),
136        exp.UnixToTime: rename_func("FROM_UNIXTIME"),
137        exp.VarMap: lambda self, e: var_map_sql(self, e, "MAP"),
138    }
139
140    # https://docs.starrocks.io/docs/sql-reference/sql-statements/keywords/#reserved-keywords
141    RESERVED_KEYWORDS = {
142        "add",
143        "all",
144        "alter",
145        "analyze",
146        "and",
147        "array",
148        "as",
149        "asc",
150        "between",
151        "bigint",
152        "bitmap",
153        "both",
154        "by",
155        "case",
156        "char",
157        "character",
158        "check",
159        "collate",
160        "column",
161        "compaction",
162        "convert",
163        "create",
164        "cross",
165        "cube",
166        "current_date",
167        "current_role",
168        "current_time",
169        "current_timestamp",
170        "current_user",
171        "database",
172        "databases",
173        "decimal",
174        "decimalv2",
175        "decimal32",
176        "decimal64",
177        "decimal128",
178        "default",
179        "deferred",
180        "delete",
181        "dense_rank",
182        "desc",
183        "describe",
184        "distinct",
185        "double",
186        "drop",
187        "dual",
188        "else",
189        "except",
190        "exists",
191        "explain",
192        "false",
193        "first_value",
194        "float",
195        "for",
196        "force",
197        "from",
198        "full",
199        "function",
200        "grant",
201        "group",
202        "grouping",
203        "grouping_id",
204        "groups",
205        "having",
206        "hll",
207        "host",
208        "if",
209        "ignore",
210        "immediate",
211        "in",
212        "index",
213        "infile",
214        "inner",
215        "insert",
216        "int",
217        "integer",
218        "intersect",
219        "into",
220        "is",
221        "join",
222        "json",
223        "key",
224        "keys",
225        "kill",
226        "lag",
227        "largeint",
228        "last_value",
229        "lateral",
230        "lead",
231        "left",
232        "like",
233        "limit",
234        "load",
235        "localtime",
236        "localtimestamp",
237        "maxvalue",
238        "minus",
239        "mod",
240        "not",
241        "ntile",
242        "null",
243        "on",
244        "or",
245        "order",
246        "outer",
247        "outfile",
248        "over",
249        "partition",
250        "percentile",
251        "primary",
252        "procedure",
253        "qualify",
254        "range",
255        "rank",
256        "read",
257        "regexp",
258        "release",
259        "rename",
260        "replace",
261        "revoke",
262        "right",
263        "rlike",
264        "row",
265        "row_number",
266        "rows",
267        "schema",
268        "schemas",
269        "select",
270        "set",
271        "set_var",
272        "show",
273        "smallint",
274        "system",
275        "table",
276        "terminated",
277        "text",
278        "then",
279        "tinyint",
280        "to",
281        "true",
282        "union",
283        "unique",
284        "unsigned",
285        "update",
286        "use",
287        "using",
288        "values",
289        "varchar",
290        "when",
291        "where",
292        "with",
293    }
294
295    def create_sql(self, expression: exp.Create) -> str:
296        # Starrocks' primary key is defined outside of the schema, so we need to move it there
297        schema = expression.this
298        if isinstance(schema, exp.Schema):
299            primary_key = schema.find(exp.PrimaryKey)
300
301            if primary_key:
302                props = expression.args.get("properties")
303
304                if not props:
305                    props = exp.Properties(expressions=[])
306                    expression.set("properties", props)
307
308                # Verify if the first one is an engine property. Is true then insert it after the engine,
309                # otherwise insert it at the beginning
310                engine = props.find(exp.EngineProperty)
311                engine_index = (engine.index or 0) if engine else -1
312                props.set("expressions", primary_key.pop(), engine_index + 1, overwrite=False)
313
314        return super().create_sql(expression)
315
316    def partitionedbyproperty_sql(self, expression: exp.PartitionedByProperty) -> str:
317        this = expression.this
318        if isinstance(this, exp.Schema):
319            # For MVs, StarRocks needs outer parentheses.
320            create = expression.find_ancestor(exp.Create)
321
322            sql = self.expressions(this, flat=True)
323            if (create and create.kind == "VIEW") or all(
324                isinstance(col, (exp.Column, exp.Identifier)) for col in this.expressions
325            ):
326                sql = f"({sql})"
327
328            return f"PARTITION BY {sql}"
329
330        return f"PARTITION BY {self.sql(this)}"
331
332    def clusterproperty_sql(self, expression: exp.ClusterProperty) -> str:
333        """Generate StarRocks ORDER BY clause for clustering."""
334        if expression.this:
335            self.unsupported(f"Unsupported CLUSTER BY {self.sql(expression, 'this')}")
336            return ""
337        expressions = self.expressions(expression, flat=True)
338        return f"ORDER BY ({expressions})"
339
340    def refreshtriggerproperty_sql(self, expression: exp.RefreshTriggerProperty) -> str:
341        """Generate StarRocks REFRESH clause for materialized views.
342        There is a little difference of the syntax between StarRocks and Doris.
343        """
344        method = self.sql(expression, "method")
345        method = f" {method}" if method else ""
346        kind = self.sql(expression, "kind")
347        kind = f" {kind}" if kind else ""
348        starts = self.sql(expression, "starts")
349        starts = f" START ({starts})" if starts else ""
350        every = self.sql(expression, "every")
351        unit = self.sql(expression, "unit")
352        every = f" EVERY (INTERVAL {every} {unit})" if every and unit else ""
353
354        return f"REFRESH{method}{kind}{starts}{every}"