sqlglot.executor.env

  1import datetime
  2import inspect
  3import re
  4import statistics
  5from functools import wraps
  6
  7from sqlglot import exp
  8from sqlglot.generator import Generator
  9from sqlglot.helper import PYTHON_VERSION
 10
 11
 12class reverse_key:
 13    def __init__(self, obj):
 14        self.obj = obj
 15
 16    def __eq__(self, other):
 17        return other.obj == self.obj
 18
 19    def __lt__(self, other):
 20        return other.obj < self.obj
 21
 22
 23def filter_nulls(func, empty_null=True):
 24    @wraps(func)
 25    def _func(values):
 26        filtered = tuple(v for v in values if v is not None)
 27        if not filtered and empty_null:
 28            return None
 29        return func(filtered)
 30
 31    return _func
 32
 33
 34def null_if_any(*required):
 35    """
 36    Decorator that makes a function return `None` if any of the `required` arguments are `None`.
 37
 38    This also supports decoration with no arguments, e.g.:
 39
 40        @null_if_any
 41        def foo(a, b): ...
 42
 43    In which case all arguments are required.
 44    """
 45    f = None
 46    if len(required) == 1 and callable(required[0]):
 47        f = required[0]
 48        required = ()
 49
 50    def decorator(func):
 51        if required:
 52            required_indices = [
 53                i for i, param in enumerate(inspect.signature(func).parameters) if param in required
 54            ]
 55
 56            def predicate(*args):
 57                return any(args[i] is None for i in required_indices)
 58
 59        else:
 60
 61            def predicate(*args):
 62                return any(a is None for a in args)
 63
 64        @wraps(func)
 65        def _func(*args):
 66            if predicate(*args):
 67                return None
 68            return func(*args)
 69
 70        return _func
 71
 72    if f:
 73        return decorator(f)
 74
 75    return decorator
 76
 77
 78@null_if_any("substr", "this")
 79def str_position(substr, this, position=None):
 80    position = position - 1 if position is not None else position
 81    return this.find(substr, position) + 1
 82
 83
 84@null_if_any("this")
 85def substring(this, start=None, length=None):
 86    if start is None:
 87        return this
 88    elif start == 0:
 89        return ""
 90    elif start < 0:
 91        start = len(this) + start
 92    else:
 93        start -= 1
 94
 95    end = None if length is None else start + length
 96
 97    return this[start:end]
 98
 99
100@null_if_any
101def cast(this, to):
102    if to == exp.DataType.Type.DATE:
103        return datetime.date.fromisoformat(this)
104    if to == exp.DataType.Type.DATETIME:
105        return datetime.datetime.fromisoformat(this)
106    if to in exp.DataType.TEXT_TYPES:
107        return str(this)
108    if to in {exp.DataType.Type.FLOAT, exp.DataType.Type.DOUBLE}:
109        return float(this)
110    if to in exp.DataType.NUMERIC_TYPES:
111        return int(this)
112    raise NotImplementedError(f"Casting to '{to}' not implemented.")
113
114
115def ordered(this, desc, nulls_first):
116    if desc:
117        return reverse_key(this)
118    return this
119
120
121@null_if_any
122def interval(this, unit):
123    unit = unit.lower()
124    plural = unit + "s"
125    if plural in Generator.TIME_PART_SINGULARS:
126        unit = plural
127    return datetime.timedelta(**{unit: float(this)})
128
129
130ENV = {
131    "exp": exp,
132    # aggs
133    "ARRAYAGG": list,
134    "AVG": filter_nulls(statistics.fmean if PYTHON_VERSION >= (3, 8) else statistics.mean),  # type: ignore
135    "COUNT": filter_nulls(lambda acc: sum(1 for _ in acc), False),
136    "MAX": filter_nulls(max),
137    "MIN": filter_nulls(min),
138    "SUM": filter_nulls(sum),
139    # scalar functions
140    "ABS": null_if_any(lambda this: abs(this)),
141    "ADD": null_if_any(lambda e, this: e + this),
142    "ARRAYANY": null_if_any(lambda arr, func: any(func(e) for e in arr)),
143    "BETWEEN": null_if_any(lambda this, low, high: low <= this and this <= high),
144    "BITWISEAND": null_if_any(lambda this, e: this & e),
145    "BITWISELEFTSHIFT": null_if_any(lambda this, e: this << e),
146    "BITWISEOR": null_if_any(lambda this, e: this | e),
147    "BITWISERIGHTSHIFT": null_if_any(lambda this, e: this >> e),
148    "BITWISEXOR": null_if_any(lambda this, e: this ^ e),
149    "CAST": cast,
150    "COALESCE": lambda *args: next((a for a in args if a is not None), None),
151    "CONCAT": null_if_any(lambda *args: "".join(args)),
152    "CONCATWS": null_if_any(lambda this, *args: this.join(args)),
153    "DIV": null_if_any(lambda e, this: e / this),
154    "EQ": null_if_any(lambda this, e: this == e),
155    "EXTRACT": null_if_any(lambda this, e: getattr(e, this)),
156    "GT": null_if_any(lambda this, e: this > e),
157    "GTE": null_if_any(lambda this, e: this >= e),
158    "IFNULL": lambda e, alt: alt if e is None else e,
159    "IF": lambda predicate, true, false: true if predicate else false,
160    "INTDIV": null_if_any(lambda e, this: e // this),
161    "INTERVAL": interval,
162    "LIKE": null_if_any(
163        lambda this, e: bool(re.match(e.replace("_", ".").replace("%", ".*"), this))
164    ),
165    "LOWER": null_if_any(lambda arg: arg.lower()),
166    "LT": null_if_any(lambda this, e: this < e),
167    "LTE": null_if_any(lambda this, e: this <= e),
168    "MOD": null_if_any(lambda e, this: e % this),
169    "MUL": null_if_any(lambda e, this: e * this),
170    "NEQ": null_if_any(lambda this, e: this != e),
171    "ORD": null_if_any(ord),
172    "ORDERED": ordered,
173    "POW": pow,
174    "STRPOSITION": str_position,
175    "SUB": null_if_any(lambda e, this: e - this),
176    "SUBSTRING": substring,
177    "TIMESTRTOTIME": null_if_any(lambda arg: datetime.datetime.fromisoformat(arg)),
178    "UPPER": null_if_any(lambda arg: arg.upper()),
179    "YEAR": null_if_any(lambda arg: arg.year),
180    "MONTH": null_if_any(lambda arg: arg.month),
181    "DAY": null_if_any(lambda arg: arg.day),
182    "CURRENTDATETIME": datetime.datetime.now,
183    "CURRENTTIMESTAMP": datetime.datetime.now,
184    "CURRENTTIME": datetime.datetime.now,
185    "CURRENTDATE": datetime.date.today,
186}