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