sqlglot.parsers.dremio

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import expressions as exp
  6from sqlglot import parser
  7from sqlglot.dialects.dialect import (
  8    build_date_delta,
  9    build_formatted_time,
 10    build_timetostr_or_tochar,
 11)
 12from sqlglot.helper import seq_get
 13from sqlglot.tokens import TokenType
 14
 15if t.TYPE_CHECKING:
 16    from sqlglot.dialects.dialect import DialectType
 17
 18DATE_DELTA = t.Union[exp.DateAdd, exp.DateSub]
 19
 20
 21def to_char_is_numeric_handler(args: list, dialect: DialectType) -> exp.TimeToStr | exp.ToChar:
 22    expression = build_timetostr_or_tochar(args, dialect)
 23    fmt = seq_get(args, 1)
 24
 25    if fmt and isinstance(expression, exp.ToChar) and fmt.is_string and "#" in fmt.name:
 26        # Only mark as numeric if format is a literal containing #
 27        expression.set("is_numeric", True)
 28
 29    return expression
 30
 31
 32def build_date_delta_with_cast_interval(
 33    expression_class: type[DATE_DELTA],
 34) -> t.Callable[[list[exp.Expr]], exp.Expr]:
 35    fallback_builder = build_date_delta(expression_class)
 36
 37    def _builder(args):
 38        if len(args) == 2:
 39            date_arg, interval_arg = args
 40
 41            if (
 42                isinstance(interval_arg, exp.Cast)
 43                and isinstance(interval_arg.to, exp.DataType)
 44                and isinstance(interval_arg.to.this, exp.Interval)
 45            ):
 46                return expression_class(
 47                    this=date_arg,
 48                    expression=interval_arg.this,
 49                    unit=interval_arg.to.this.unit,
 50                )
 51
 52            return expression_class(this=date_arg, expression=interval_arg)
 53
 54        return fallback_builder(args)
 55
 56    return _builder
 57
 58
 59def datetype_handler(args: list[exp.Expr], dialect: DialectType) -> exp.Expr:
 60    from sqlglot.dialects.dialect import Dialect
 61
 62    year, month, day = args
 63
 64    if all(isinstance(arg, exp.Literal) and arg.is_int for arg in (year, month, day)):
 65        date_str = f"{int(year.this):04d}-{int(month.this):02d}-{int(day.this):02d}"
 66        return exp.Date(this=exp.Literal.string(date_str))
 67
 68    dialect = Dialect.get_or_raise(dialect)
 69
 70    return exp.Cast(
 71        this=exp.Concat(
 72            expressions=[
 73                year,
 74                exp.Literal.string("-"),
 75                month,
 76                exp.Literal.string("-"),
 77                day,
 78            ],
 79            coalesce=dialect.CONCAT_COALESCE,
 80        ),
 81        to=exp.DType.DATE.into_expr(),
 82    )
 83
 84
 85class DremioParser(parser.Parser):
 86    LOG_DEFAULTS_TO_LN = True
 87
 88    TABLE_ALIAS_TOKENS = parser.Parser.TABLE_ALIAS_TOKENS | {
 89        TokenType.ANTI,
 90        TokenType.SEMI,
 91    }
 92
 93    NO_PAREN_FUNCTION_PARSERS = {
 94        **parser.Parser.NO_PAREN_FUNCTION_PARSERS,
 95        "CURRENT_DATE_UTC": lambda self: self._parse_current_date_utc(),
 96    }
 97
 98    FUNCTIONS = {
 99        **parser.Parser.FUNCTIONS,
100        "ARRAY_GENERATE_RANGE": exp.GenerateSeries.from_arg_list,
101        "BIT_AND": exp.BitwiseAndAgg.from_arg_list,
102        "BIT_OR": exp.BitwiseOrAgg.from_arg_list,
103        "DATE_ADD": build_date_delta_with_cast_interval(exp.DateAdd),
104        "DATE_FORMAT": build_formatted_time(exp.TimeToStr, "dremio"),
105        "DATE_SUB": build_date_delta_with_cast_interval(exp.DateSub),
106        "REGEXP_MATCHES": exp.RegexpLike.from_arg_list,
107        "REPEATSTR": exp.Repeat.from_arg_list,
108        "TO_CHAR": to_char_is_numeric_handler,
109        "TO_DATE": build_formatted_time(exp.TsOrDsToDate, "dremio"),
110        "DATE_PART": exp.Extract.from_arg_list,
111        "DATETYPE": datetype_handler,
112    }
113
114    def _parse_current_date_utc(self) -> exp.Cast:
115        if self._match(TokenType.L_PAREN):
116            self._match_r_paren()
117
118        return exp.Cast(
119            this=exp.AtTimeZone(
120                this=exp.CurrentTimestamp(),
121                zone=exp.Literal.string("UTC"),
122            ),
123            to=exp.DType.DATE.into_expr(),
124        )