sqlglot.generators.starrocks

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