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    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("DATE_DIFF", unit_to_str(e), e.this, e.expression),
107        exp.Delete: transforms.preprocess([_eliminate_between_in_delete]),
108        exp.Flatten: rename_func("ARRAY_FLATTEN"),
109        exp.JSONExtractScalar: arrow_json_extract_sql,
110        exp.JSONExtract: arrow_json_extract_sql,
111        # Both MAP forms (two-array MAP([keys], [values]) and variadic MAP(k1, v1, ...))
112        # generate StarRocks' variadic MAP(k1, v1, k2, v2, ...) constructor
113        exp.Map: lambda self, e: var_map_sql(self, e, "MAP"),
114        exp.Property: property_sql,
115        exp.RegexpLike: rename_func("REGEXP"),
116        # Inherited from MySQL, minus operations StarRocks supports natively
117        # (QUALIFY, FULL OUTER JOIN, SEMI/ANTI JOIN)
118        exp.Select: transforms.preprocess(
119            [
120                transforms.eliminate_distinct_on,
121                transforms.unnest_generate_date_array_using_recursive_cte,
122            ]
123        ),
124        exp.SchemaCommentProperty: lambda self, e: self.naked_property(e),
125        exp.SqlSecurityProperty: lambda self, e: f"SECURITY {self.sql(e.this)}",
126        exp.StDistance: st_distance_sphere,
127        exp.StrToUnix: lambda self, e: self.func("UNIX_TIMESTAMP", e.this, self.format_time(e)),
128        exp.TimestampTrunc: lambda self, e: self.func("DATE_TRUNC", unit_to_str(e), e.this),
129        exp.TimeStrToDate: rename_func("TO_DATE"),
130        exp.UnixToStr: lambda self, e: self.func("FROM_UNIXTIME", e.this, self.format_time(e)),
131        exp.UnixToTime: rename_func("FROM_UNIXTIME"),
132        exp.VarMap: lambda self, e: var_map_sql(self, e, "MAP"),
133    }
134
135    # https://docs.starrocks.io/docs/sql-reference/sql-statements/keywords/#reserved-keywords
136    RESERVED_KEYWORDS = {
137        "add",
138        "all",
139        "alter",
140        "analyze",
141        "and",
142        "array",
143        "as",
144        "asc",
145        "between",
146        "bigint",
147        "bitmap",
148        "both",
149        "by",
150        "case",
151        "char",
152        "character",
153        "check",
154        "collate",
155        "column",
156        "compaction",
157        "convert",
158        "create",
159        "cross",
160        "cube",
161        "current_date",
162        "current_role",
163        "current_time",
164        "current_timestamp",
165        "current_user",
166        "database",
167        "databases",
168        "decimal",
169        "decimalv2",
170        "decimal32",
171        "decimal64",
172        "decimal128",
173        "default",
174        "deferred",
175        "delete",
176        "dense_rank",
177        "desc",
178        "describe",
179        "distinct",
180        "double",
181        "drop",
182        "dual",
183        "else",
184        "except",
185        "exists",
186        "explain",
187        "false",
188        "first_value",
189        "float",
190        "for",
191        "force",
192        "from",
193        "full",
194        "function",
195        "grant",
196        "group",
197        "grouping",
198        "grouping_id",
199        "groups",
200        "having",
201        "hll",
202        "host",
203        "if",
204        "ignore",
205        "immediate",
206        "in",
207        "index",
208        "infile",
209        "inner",
210        "insert",
211        "int",
212        "integer",
213        "intersect",
214        "into",
215        "is",
216        "join",
217        "json",
218        "key",
219        "keys",
220        "kill",
221        "lag",
222        "largeint",
223        "last_value",
224        "lateral",
225        "lead",
226        "left",
227        "like",
228        "limit",
229        "load",
230        "localtime",
231        "localtimestamp",
232        "maxvalue",
233        "minus",
234        "mod",
235        "not",
236        "ntile",
237        "null",
238        "on",
239        "or",
240        "order",
241        "outer",
242        "outfile",
243        "over",
244        "partition",
245        "percentile",
246        "primary",
247        "procedure",
248        "qualify",
249        "range",
250        "rank",
251        "read",
252        "regexp",
253        "release",
254        "rename",
255        "replace",
256        "revoke",
257        "right",
258        "rlike",
259        "row",
260        "row_number",
261        "rows",
262        "schema",
263        "schemas",
264        "select",
265        "set",
266        "set_var",
267        "show",
268        "smallint",
269        "system",
270        "table",
271        "terminated",
272        "text",
273        "then",
274        "tinyint",
275        "to",
276        "true",
277        "union",
278        "unique",
279        "unsigned",
280        "update",
281        "use",
282        "using",
283        "values",
284        "varchar",
285        "when",
286        "where",
287        "with",
288    }
289
290    def create_sql(self, expression: exp.Create) -> str:
291        # Starrocks' primary key is defined outside of the schema, so we need to move it there
292        schema = expression.this
293        if isinstance(schema, exp.Schema):
294            primary_key = schema.find(exp.PrimaryKey)
295
296            if primary_key:
297                props = expression.args.get("properties")
298
299                if not props:
300                    props = exp.Properties(expressions=[])
301                    expression.set("properties", props)
302
303                # Verify if the first one is an engine property. Is true then insert it after the engine,
304                # otherwise insert it at the beginning
305                engine = props.find(exp.EngineProperty)
306                engine_index = (engine.index or 0) if engine else -1
307                props.set("expressions", primary_key.pop(), engine_index + 1, overwrite=False)
308
309        return super().create_sql(expression)
310
311    def partitionedbyproperty_sql(self, expression: exp.PartitionedByProperty) -> str:
312        this = expression.this
313        if isinstance(this, exp.Schema):
314            # For MVs, StarRocks needs outer parentheses.
315            create = expression.find_ancestor(exp.Create)
316
317            sql = self.expressions(this, flat=True)
318            if (create and create.kind == "VIEW") or all(
319                isinstance(col, (exp.Column, exp.Identifier)) for col in this.expressions
320            ):
321                sql = f"({sql})"
322
323            return f"PARTITION BY {sql}"
324
325        return f"PARTITION BY {self.sql(this)}"
326
327    def clusterproperty_sql(self, expression: exp.ClusterProperty) -> str:
328        """Generate StarRocks ORDER BY clause for clustering."""
329        if expression.this:
330            self.unsupported(f"Unsupported CLUSTER BY {self.sql(expression, 'this')}")
331            return ""
332        expressions = self.expressions(expression, flat=True)
333        return f"ORDER BY ({expressions})"
334
335    def refreshtriggerproperty_sql(self, expression: exp.RefreshTriggerProperty) -> str:
336        """Generate StarRocks REFRESH clause for materialized views.
337        There is a little difference of the syntax between StarRocks and Doris.
338        """
339        method = self.sql(expression, "method")
340        method = f" {method}" if method else ""
341        kind = self.sql(expression, "kind")
342        kind = f" {kind}" if kind else ""
343        starts = self.sql(expression, "starts")
344        starts = f" START ({starts})" if starts else ""
345        every = self.sql(expression, "every")
346        unit = self.sql(expression, "unit")
347        every = f" EVERY (INTERVAL {every} {unit})" if every and unit else ""
348
349        return f"REFRESH{method}{kind}{starts}{every}"