sqlglot.typing.duckdb

  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
 12DATETIME_WITHOUT_TZ = {
 13    exp.DType.DATE,
 14    exp.DType.TIMESTAMP,
 15    exp.DType.TIMESTAMPNTZ,
 16    exp.DType.TIMESTAMP_S,
 17    exp.DType.TIMESTAMP_MS,
 18    exp.DType.TIMESTAMP_NS,
 19}
 20
 21
 22def _annotate_at_time_zone(self: TypeAnnotator, expression: exp.AtTimeZone) -> exp.Expr:
 23    this = expression.this
 24
 25    if this.is_type(*DATETIME_WITHOUT_TZ):
 26        return self._set_type(expression, exp.DType.TIMESTAMPTZ)
 27    if this.is_type(exp.DType.TIMESTAMPTZ):
 28        return self._set_type(expression, exp.DType.TIMESTAMP)
 29    if this.is_type(exp.DType.TIMETZ):
 30        return self._set_type(expression, exp.DType.TIMETZ)
 31
 32    return self._set_type(expression, exp.DType.UNKNOWN)
 33
 34
 35def _annotate_date_trunc(
 36    self: TypeAnnotator, expression: exp.DateTrunc | exp.TimestampTrunc
 37) -> exp.Expr:
 38    if expression.this.is_type(
 39        exp.DType.DATE,
 40        exp.DType.TIMESTAMP_S,
 41        exp.DType.TIMESTAMP_MS,
 42        exp.DType.TIMESTAMP_NS,
 43    ):
 44        return self._set_type(expression, exp.DType.TIMESTAMP)
 45
 46    return self._set_type(expression, expression.this.type)
 47
 48
 49EXPRESSION_METADATA = {
 50    **EXPRESSION_METADATA,
 51    **{
 52        expr_type: {"returns": exp.DType.BIGINT}
 53        for expr_type in {
 54            exp.BitLength,
 55            exp.DateDiff,
 56            exp.Day,
 57            exp.DayOfMonth,
 58            exp.DayOfWeek,
 59            exp.DayOfWeekIso,
 60            exp.DayOfYear,
 61            exp.Extract,
 62            exp.Hour,
 63            exp.Length,
 64            exp.Minute,
 65            exp.Month,
 66            exp.Quarter,
 67            exp.Second,
 68            exp.Week,
 69            exp.Year,
 70        }
 71    },
 72    **{
 73        expr_type: {"returns": exp.DType.INT128}
 74        for expr_type in {
 75            exp.CountIf,
 76            exp.Factorial,
 77        }
 78    },
 79    **{
 80        expr_type: {"returns": exp.DType.DOUBLE}
 81        for expr_type in {
 82            exp.Atan2,
 83            exp.JarowinklerSimilarity,
 84            exp.TimeToUnix,
 85        }
 86    },
 87    **{
 88        expr_type: {"returns": exp.DType.VARCHAR}
 89        for expr_type in {
 90            exp.Format,
 91            exp.Reverse,
 92            exp.Decode,
 93        }
 94    },
 95    **{
 96        expr_type: {"returns": exp.DType.VARBINARY}
 97        for expr_type in {
 98            exp.Encode,
 99            exp.Unhex,
100        }
101    },
102    **{
103        expr_type: {"annotator": lambda self, e: self._annotate_by_args(e, "this")}
104        for expr_type in {
105            exp.BitwiseAndAgg,
106            exp.BitwiseOrAgg,
107            exp.BitwiseXorAgg,
108            exp.PercentileDisc,
109        }
110    },
111    exp.AtTimeZone: {"annotator": _annotate_at_time_zone},
112    exp.DateBin: {"annotator": lambda self, e: self._annotate_by_args(e, "expression")},
113    exp.DateTrunc: {"annotator": _annotate_date_trunc},
114    exp.TimestampTrunc: {"annotator": _annotate_date_trunc},
115    exp.Localtimestamp: {"returns": exp.DType.TIMESTAMP},
116    exp.ToDays: {"returns": exp.DType.INTERVAL},
117    exp.TimeFromParts: {"returns": exp.DType.TIME},
118    exp.TimestampFromParts: {"returns": exp.DType.TIMESTAMP},
119}