sqlglot.typing.mysql

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp
  6from sqlglot.typing import EXPRESSION_METADATA
  7
  8if t.TYPE_CHECKING:
  9    from sqlglot.optimizer.annotate_types import TypeAnnotator
 10
 11
 12COMPRESS_LONGBLOB_TYPES = {
 13    exp.DType.TEXT,
 14    exp.DType.MEDIUMTEXT,
 15    exp.DType.LONGTEXT,
 16    exp.DType.BLOB,
 17    exp.DType.MEDIUMBLOB,
 18    exp.DType.LONGBLOB,
 19    exp.DType.JSON,
 20}
 21
 22
 23COMPRESS_VARBINARY_TYPES = {
 24    exp.DType.CHAR,
 25    exp.DType.VARCHAR,
 26    exp.DType.BINARY,
 27    exp.DType.VARBINARY,
 28    exp.DType.TINYBLOB,
 29    exp.DType.ENUM,
 30    exp.DType.INT,
 31    exp.DType.BIGINT,
 32    exp.DType.DECIMAL,
 33    exp.DType.DOUBLE,
 34    exp.DType.DATE,
 35    exp.DType.DATETIME,
 36}
 37
 38
 39def _annotate_reverse(self: TypeAnnotator, expression: exp.Reverse) -> exp.Reverse:
 40    if expression.this.is_type(exp.DType.BINARY, exp.DType.VARBINARY, exp.DType.UNKNOWN):
 41        self._annotate_by_args(expression, "this")
 42    else:
 43        self._set_type(expression, exp.DType.VARCHAR)
 44
 45    return expression
 46
 47
 48def _annotate_truncate(self: TypeAnnotator, expression: exp.Trunc) -> exp.Expr:
 49    if expression.this.is_type(*exp.DataType.TEXT_TYPES):
 50        return self._set_type(expression, exp.DType.DOUBLE)
 51
 52    return self._annotate_by_args(expression, "this")
 53
 54
 55def _annotate_regexp_replace(self: TypeAnnotator, expression: exp.RegexpReplace) -> exp.Expr:
 56    args = (expression.this, expression.expression, expression.args.get("replacement"))
 57
 58    has_binary = False
 59    for arg in args:
 60        if arg is not None:
 61            if arg.is_type(exp.DType.UNKNOWN):
 62                return self._set_type(expression, exp.DType.UNKNOWN)
 63            has_binary = has_binary or arg.is_type(*exp.DataType.BINARY_TYPES)
 64
 65    return self._set_type(expression, exp.DType.LONGBLOB if has_binary else exp.DType.LONGTEXT)
 66
 67
 68def _annotate_compress(self: TypeAnnotator, expression: exp.Compress) -> exp.Expr:
 69    this = expression.this
 70
 71    if this.is_type(*COMPRESS_VARBINARY_TYPES):
 72        return self._set_type(expression, exp.DType.VARBINARY)
 73
 74    if this.is_type(*COMPRESS_LONGBLOB_TYPES):
 75        return self._set_type(expression, exp.DType.LONGBLOB)
 76
 77    if this.is_type(exp.DType.TINYTEXT):
 78        return self._set_type(expression, exp.DType.BLOB)
 79
 80    return self._set_type(expression, exp.DType.UNKNOWN)
 81
 82
 83def _annotate_bit_and(self: TypeAnnotator, expression: exp.BitwiseAndAgg) -> exp.Expr:
 84    this = expression.this
 85
 86    if this.is_type(*exp.DataType.BINARY_TYPES):
 87        return self._set_type(expression, exp.DType.VARBINARY)
 88
 89    return self._set_type(expression, exp.DType.BIGINT)
 90
 91
 92EXPRESSION_METADATA = {
 93    **EXPRESSION_METADATA,
 94    **{
 95        expr_type: {"returns": exp.DType.DOUBLE}
 96        for expr_type in {
 97            exp.Atan2,
 98            exp.MatchAgainst,
 99            exp.StDistance,
100        }
101    },
102    **{
103        expr_type: {"returns": exp.DType.DATETIME}
104        for expr_type in {
105            exp.CurrentTimestamp,
106            exp.ConvertTimezone,
107            exp.Localtime,
108            exp.Localtimestamp,
109            exp.UtcTimestamp,
110        }
111    },
112    **{
113        expr_type: {"returns": exp.DType.DATE}
114        for expr_type in {
115            exp.UtcDate,
116        }
117    },
118    **{
119        expr_type: {"returns": exp.DType.VARCHAR}
120        for expr_type in {
121            exp.Elt,
122            exp.Hex,
123            exp.NumberToStr,  # format()
124            exp.Replace,
125            exp.Stuff,  # insert function
126            exp.SubstringIndex,
127            exp.RegexpSubstr,
128            exp.Collation,
129        }
130    },
131    **{
132        expr_type: {"returns": exp.DType.INT}
133        for expr_type in {
134            exp.Month,
135            exp.Second,
136            exp.Week,
137            exp.Minute,
138        }
139    },
140    **{
141        expr_type: {"returns": exp.DType.BIGINT}
142        for expr_type in {
143            exp.RegexpInstr,
144        }
145    },
146    **{
147        expr_type: {"returns": exp.DType.TIME}
148        for expr_type in {
149            exp.TimeFromParts,
150            exp.UtcTime,
151        }
152    },
153    **{
154        expr_type: {"returns": exp.DType.VARBINARY}
155        for expr_type in {
156            exp.Unhex,
157        }
158    },
159    **{
160        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this")}
161        for expr_type in {
162            exp.Pad,
163            exp.Left,
164            exp.Right,
165        }
166    },
167    exp.Reverse: {"annotator": _annotate_reverse},
168    exp.Trunc: {"annotator": _annotate_truncate},
169    exp.RegexpReplace: {"annotator": _annotate_regexp_replace},
170    exp.Compress: {"annotator": _annotate_compress},
171    exp.BitwiseAndAgg: {"annotator": _annotate_bit_and},
172}