Edit on GitHub

sqlglot.executor.env

View Source

  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, is_int, seq_get
 10
 11
 12def sql_not(value):
 13    return None if value is None else not bool(value)
 14
 15
 16def sql_and(left, right):
 17    left = left()
 18    left = None if left is None else bool(left)
 19
 20    if left is False:
 21        return False
 22
 23    right = right()
 24    right = None if right is None else bool(right)
 25    if right is False:
 26        return False
 27
 28    return None if left is None or right is None else True
 29
 30
 31def sql_or(left, right):
 32    left = left()
 33    left = None if left is None else bool(left)
 34
 35    if left is True:
 36        return True
 37
 38    right = right()
 39    right = None if right is None else bool(right)
 40    if right is True:
 41        return True
 42
 43    return None if left is None or right is None else False
 44
 45
 46def sql_in(value, *candidates):
 47    if value is None:
 48        return None
 49
 50    has_null = False
 51    for candidate in candidates:
 52        if candidate is None:
 53            has_null = True
 54        elif value == candidate:
 55            return True
 56
 57    return None if has_null else False
 58
 59
 60class reverse_key:
 61    def __init__(self, obj):
 62        self.obj = obj
 63
 64    def __eq__(self, other):
 65        return other.obj == self.obj
 66
 67    def __lt__(self, other):
 68        return other.obj < self.obj
 69
 70
 71def filter_nulls(func, empty_null=True):
 72    @wraps(func)
 73    def _func(values):
 74        filtered = tuple(v for v in values if v is not None)
 75        if not filtered and empty_null:
 76            return None
 77        return func(filtered)
 78
 79    return _func
 80
 81
 82def null_if_any(*required):
 83    """
 84    Decorator that makes a function return `None` if any of the `required` arguments are `None`.
 85
 86    This also supports decoration with no arguments, e.g.:
 87
 88        @null_if_any
 89        def foo(a, b): ...
 90
 91    In which case all arguments are required.
 92    """
 93    f = None
 94    if len(required) == 1 and callable(required[0]):
 95        f = required[0]
 96        required = ()
 97
 98    def decorator(func):
 99        if required:
100            required_indices = [
101                i for i, param in enumerate(inspect.signature(func).parameters) if param in required
102            ]
103
104            def predicate(*args):
105                return any(args[i] is None for i in required_indices)
106
107        else:
108
109            def predicate(*args):
110                return any(a is None for a in args)
111
112        @wraps(func)
113        def _func(*args):
114            if predicate(*args):
115                return None
116            return func(*args)
117
118        return _func
119
120    if f:
121        return decorator(f)
122
123    return decorator
124
125
126@null_if_any("this", "substr")
127def str_position(this, substr, position=None):
128    position = position - 1 if position is not None else position
129    return this.find(substr, position) + 1
130
131
132@null_if_any("this")
133def substring(this, start=None, length=None):
134    if start is None:
135        return this
136    elif start == 0:
137        return ""
138    elif start < 0:
139        start = len(this) + start
140    else:
141        start -= 1
142
143    end = None if length is None else start + length
144
145    return this[start:end]
146
147
148@null_if_any
149def cast(this, to):
150    if to == exp.DType.DATE:
151        if isinstance(this, datetime.datetime):
152            return this.date()
153        if isinstance(this, datetime.date):
154            return this
155        if isinstance(this, str):
156            return datetime.date.fromisoformat(this)
157    if to == exp.DType.TIME:
158        if isinstance(this, datetime.datetime):
159            return this.time()
160        if isinstance(this, datetime.time):
161            return this
162        if isinstance(this, str):
163            return datetime.time.fromisoformat(this)
164    if to in (exp.DType.DATETIME, exp.DType.TIMESTAMP):
165        if isinstance(this, datetime.datetime):
166            return this
167        if isinstance(this, datetime.date):
168            return datetime.datetime(this.year, this.month, this.day)
169        if isinstance(this, str):
170            return datetime.datetime.fromisoformat(this)
171    if to == exp.DType.BOOLEAN:
172        return bool(this)
173    if to in exp.DataType.TEXT_TYPES:
174        return str(this)
175    if to in {exp.DType.FLOAT, exp.DType.DOUBLE}:
176        return float(this)
177    if to in exp.DataType.NUMERIC_TYPES:
178        return int(this)
179    raise NotImplementedError(f"Casting {this} to '{to}' not implemented.")
180
181
182def ordered(this, desc, nulls_first):
183    if desc:
184        return reverse_key(this)
185    return this
186
187
188def _like(this, e, flags=0):
189    return bool(
190        re.fullmatch(re.escape(e).replace("_", ".").replace("%", ".*"), this, re.DOTALL | flags)
191    )
192
193
194@null_if_any
195def interval(this, unit):
196    plural = unit + "S"
197    if plural in Generator.TIME_PART_SINGULARS:
198        unit = plural
199    return datetime.timedelta(**{unit.lower(): float(this)})
200
201
202@null_if_any("this", "expression")
203def arraytostring(this, expression, null=None):
204    return expression.join(x for x in (x if x is not None else null for x in this) if x is not None)
205
206
207@null_if_any("this", "expression")
208def jsonextract(this, expression):
209    for path_segment in expression:
210        if isinstance(this, dict):
211            this = this.get(path_segment)
212        elif isinstance(this, list) and is_int(path_segment):
213            this = seq_get(this, int(path_segment))
214        else:
215            raise NotImplementedError(f"Unable to extract value for {this} at {path_segment}.")
216
217        if this is None:
218            break
219
220    return this
221
222
223ENV = {
224    "exp": exp,
225    "AND": sql_and,
226    # aggs
227    "ARRAYAGG": list,
228    "ARRAYUNIQUEAGG": filter_nulls(lambda acc: list(set(acc))),
229    "AVG": filter_nulls(statistics.fmean if PYTHON_VERSION >= (3, 8) else statistics.mean),  # type: ignore
230    "COUNT": filter_nulls(lambda acc: sum(1 for _ in acc), False),
231    "MAX": filter_nulls(max),
232    "MIN": filter_nulls(min),
233    "SUM": filter_nulls(sum),
234    # scalar functions
235    "ABS": null_if_any(lambda this: abs(this)),
236    "ADD": null_if_any(lambda e, this: e + this),
237    "ARRAYANY": null_if_any(lambda arr, func: any(func(e) for e in arr)),
238    "ARRAYTOSTRING": arraytostring,
239    "BETWEEN": null_if_any(lambda this, low, high: low <= this and this <= high),
240    "BITWISEAND": null_if_any(lambda this, e: this & e),
241    "BITWISELEFTSHIFT": null_if_any(lambda this, e: this << e),
242    "BITWISEOR": null_if_any(lambda this, e: this | e),
243    "BITWISERIGHTSHIFT": null_if_any(lambda this, e: this >> e),
244    "BITWISEXOR": null_if_any(lambda this, e: this ^ e),
245    "CAST": cast,
246    "COALESCE": lambda *args: next((a for a in args if a is not None), None),
247    "CONCAT": null_if_any(lambda *args: "".join(args)),
248    "SAFECONCAT": null_if_any(lambda *args: "".join(str(arg) for arg in args)),
249    "CONCATWS": null_if_any(lambda this, *args: this.join(args)),
250    "DATEDIFF": null_if_any(lambda this, expression, *_: (this - expression).days),
251    "DATESTRTODATE": null_if_any(lambda arg: datetime.date.fromisoformat(arg)),
252    "DIV": null_if_any(lambda e, this: e / this),
253    "DOT": null_if_any(lambda e, this: e[this]),
254    "EQ": null_if_any(lambda this, e: this == e),
255    "EXTRACT": null_if_any(lambda this, e: getattr(e, this)),
256    "GT": null_if_any(lambda this, e: this > e),
257    "GTE": null_if_any(lambda this, e: this >= e),
258    "IF": lambda predicate, true, false: true if predicate else false,
259    "IN": sql_in,
260    "INTDIV": null_if_any(lambda e, this: e // this),
261    "INTERVAL": interval,
262    "JSONEXTRACT": jsonextract,
263    "LEFT": null_if_any(lambda this, e: this[:e]),
264    "LENGTH": null_if_any(len),
265    "LIKE": null_if_any(lambda this, e: _like(this, e)),
266    "ILIKE": null_if_any(lambda this, e: _like(this, e, re.IGNORECASE)),
267    "LOWER": null_if_any(lambda arg: arg.lower()),
268    "LT": null_if_any(lambda this, e: this < e),
269    "LTE": null_if_any(lambda this, e: this <= e),
270    "MAP": null_if_any(lambda *args: dict(zip(*args))),  # type: ignore
271    "MOD": null_if_any(lambda e, this: e % this),
272    "MUL": null_if_any(lambda e, this: e * this),
273    "NEQ": null_if_any(lambda this, e: this != e),
274    "ORD": null_if_any(ord),
275    "NOT": sql_not,
276    "OR": sql_or,
277    "ORDERED": ordered,
278    "POW": pow,
279    "RIGHT": null_if_any(lambda this, e: this[-e:]),
280    "ROUND": null_if_any(lambda this, decimals=None, truncate=None: round(this, ndigits=decimals)),
281    "STRPOSITION": str_position,
282    "SUB": null_if_any(lambda e, this: e - this),
283    "SUBSTRING": substring,
284    "TIMESTRTOTIME": null_if_any(lambda arg: datetime.datetime.fromisoformat(arg)),
285    "UPPER": null_if_any(lambda arg: arg.upper()),
286    "YEAR": null_if_any(lambda arg: arg.year),
287    "MONTH": null_if_any(lambda arg: arg.month),
288    "DAY": null_if_any(lambda arg: arg.day),
289    "CURRENTDATETIME": datetime.datetime.now,
290    "CURRENTTIMESTAMP": datetime.datetime.now,
291    "CURRENTTIME": datetime.datetime.now,
292    "CURRENTDATE": datetime.date.today,
293    "STRFTIME": null_if_any(lambda fmt, arg: datetime.datetime.fromisoformat(arg).strftime(fmt)),
294    "STRTOTIME": null_if_any(lambda arg, format: datetime.datetime.strptime(arg, format)),
295    "TRIM": null_if_any(lambda this, e=None: this.strip(e)),
296    "STRUCT": lambda *args: {
297        args[x]: args[x + 1]
298        for x in range(0, len(args), 2)
299        if (args[x + 1] is not None and args[x] is not None)
300    },
301    "UNIXTOTIME": null_if_any(
302        lambda arg: datetime.datetime.fromtimestamp(arg, datetime.timezone.utc)
303    ),
304}

def sql_not(value): View Source

13def sql_not(value):
14    return None if value is None else not bool(value)

def sql_and(left, right): View Source

17def sql_and(left, right):
18    left = left()
19    left = None if left is None else bool(left)
20
21    if left is False:
22        return False
23
24    right = right()
25    right = None if right is None else bool(right)
26    if right is False:
27        return False
28
29    return None if left is None or right is None else True

def sql_or(left, right): View Source

32def sql_or(left, right):
33    left = left()
34    left = None if left is None else bool(left)
35
36    if left is True:
37        return True
38
39    right = right()
40    right = None if right is None else bool(right)
41    if right is True:
42        return True
43
44    return None if left is None or right is None else False

def sql_in(value, *candidates): View Source

47def sql_in(value, *candidates):
48    if value is None:
49        return None
50
51    has_null = False
52    for candidate in candidates:
53        if candidate is None:
54            has_null = True
55        elif value == candidate:
56            return True
57
58    return None if has_null else False

class reverse_key: View Source

61class reverse_key:
62    def __init__(self, obj):
63        self.obj = obj
64
65    def __eq__(self, other):
66        return other.obj == self.obj
67
68    def __lt__(self, other):
69        return other.obj < self.obj

reverse_key(obj) View Source

62    def __init__(self, obj):
63        self.obj = obj

obj

def filter_nulls(func, empty_null=True): View Source

72def filter_nulls(func, empty_null=True):
73    @wraps(func)
74    def _func(values):
75        filtered = tuple(v for v in values if v is not None)
76        if not filtered and empty_null:
77            return None
78        return func(filtered)
79
80    return _func

def null_if_any(*required): View Source

 83def null_if_any(*required):
 84    """
 85    Decorator that makes a function return `None` if any of the `required` arguments are `None`.
 86
 87    This also supports decoration with no arguments, e.g.:
 88
 89        @null_if_any
 90        def foo(a, b): ...
 91
 92    In which case all arguments are required.
 93    """
 94    f = None
 95    if len(required) == 1 and callable(required[0]):
 96        f = required[0]
 97        required = ()
 98
 99    def decorator(func):
100        if required:
101            required_indices = [
102                i for i, param in enumerate(inspect.signature(func).parameters) if param in required
103            ]
104
105            def predicate(*args):
106                return any(args[i] is None for i in required_indices)
107
108        else:
109
110            def predicate(*args):
111                return any(a is None for a in args)
112
113        @wraps(func)
114        def _func(*args):
115            if predicate(*args):
116                return None
117            return func(*args)
118
119        return _func
120
121    if f:
122        return decorator(f)
123
124    return decorator

Decorator that makes a function return None if any of the required arguments are None.

This also supports decoration with no arguments, e.g.:

@null_if_any
def foo(a, b): ...

In which case all arguments are required.

@null_if_any('this', 'substr')

def str_position(this, substr, position=None): View Source

127@null_if_any("this", "substr")
128def str_position(this, substr, position=None):
129    position = position - 1 if position is not None else position
130    return this.find(substr, position) + 1

@null_if_any('this')

def substring(this, start=None, length=None): View Source

133@null_if_any("this")
134def substring(this, start=None, length=None):
135    if start is None:
136        return this
137    elif start == 0:
138        return ""
139    elif start < 0:
140        start = len(this) + start
141    else:
142        start -= 1
143
144    end = None if length is None else start + length
145
146    return this[start:end]

@null_if_any

def cast(this, to): View Source

149@null_if_any
150def cast(this, to):
151    if to == exp.DType.DATE:
152        if isinstance(this, datetime.datetime):
153            return this.date()
154        if isinstance(this, datetime.date):
155            return this
156        if isinstance(this, str):
157            return datetime.date.fromisoformat(this)
158    if to == exp.DType.TIME:
159        if isinstance(this, datetime.datetime):
160            return this.time()
161        if isinstance(this, datetime.time):
162            return this
163        if isinstance(this, str):
164            return datetime.time.fromisoformat(this)
165    if to in (exp.DType.DATETIME, exp.DType.TIMESTAMP):
166        if isinstance(this, datetime.datetime):
167            return this
168        if isinstance(this, datetime.date):
169            return datetime.datetime(this.year, this.month, this.day)
170        if isinstance(this, str):
171            return datetime.datetime.fromisoformat(this)
172    if to == exp.DType.BOOLEAN:
173        return bool(this)
174    if to in exp.DataType.TEXT_TYPES:
175        return str(this)
176    if to in {exp.DType.FLOAT, exp.DType.DOUBLE}:
177        return float(this)
178    if to in exp.DataType.NUMERIC_TYPES:
179        return int(this)
180    raise NotImplementedError(f"Casting {this} to '{to}' not implemented.")

def ordered(this, desc, nulls_first): View Source

183def ordered(this, desc, nulls_first):
184    if desc:
185        return reverse_key(this)
186    return this

@null_if_any

def interval(this, unit): View Source

195@null_if_any
196def interval(this, unit):
197    plural = unit + "S"
198    if plural in Generator.TIME_PART_SINGULARS:
199        unit = plural
200    return datetime.timedelta(**{unit.lower(): float(this)})

@null_if_any('this', 'expression')

def arraytostring(this, expression, null=None): View Source

203@null_if_any("this", "expression")
204def arraytostring(this, expression, null=None):
205    return expression.join(x for x in (x if x is not None else null for x in this) if x is not None)

@null_if_any('this', 'expression')

def jsonextract(this, expression): View Source

208@null_if_any("this", "expression")
209def jsonextract(this, expression):
210    for path_segment in expression:
211        if isinstance(this, dict):
212            this = this.get(path_segment)
213        elif isinstance(this, list) and is_int(path_segment):
214            this = seq_get(this, int(path_segment))
215        else:
216            raise NotImplementedError(f"Unable to extract value for {this} at {path_segment}.")
217
218        if this is None:
219            break
220
221    return this

ENV = {'exp': <module 'sqlglot.expressions' from '/home/runner/work/sqlglot/sqlglot/sqlglot/expressions/__init__.py'>, 'AND': <function sql_and>, 'ARRAYAGG': <class 'list'>, 'ARRAYUNIQUEAGG': <function <lambda>>, 'AVG': <function fmean>, 'COUNT': <function <lambda>>, 'MAX': <function max>, 'MIN': <function min>, 'SUM': <function sum>, 'ABS': <function <lambda>>, 'ADD': <function <lambda>>, 'ARRAYANY': <function <lambda>>, 'ARRAYTOSTRING': <function arraytostring>, 'BETWEEN': <function <lambda>>, 'BITWISEAND': <function <lambda>>, 'BITWISELEFTSHIFT': <function <lambda>>, 'BITWISEOR': <function <lambda>>, 'BITWISERIGHTSHIFT': <function <lambda>>, 'BITWISEXOR': <function <lambda>>, 'CAST': <function cast>, 'COALESCE': <function <lambda>>, 'CONCAT': <function <lambda>>, 'SAFECONCAT': <function <lambda>>, 'CONCATWS': <function <lambda>>, 'DATEDIFF': <function <lambda>>, 'DATESTRTODATE': <function <lambda>>, 'DIV': <function <lambda>>, 'DOT': <function <lambda>>, 'EQ': <function <lambda>>, 'EXTRACT': <function <lambda>>, 'GT': <function <lambda>>, 'GTE': <function <lambda>>, 'IF': <function <lambda>>, 'IN': <function sql_in>, 'INTDIV': <function <lambda>>, 'INTERVAL': <function interval>, 'JSONEXTRACT': <function jsonextract>, 'LEFT': <function <lambda>>, 'LENGTH': <function len>, 'LIKE': <function <lambda>>, 'ILIKE': <function <lambda>>, 'LOWER': <function <lambda>>, 'LT': <function <lambda>>, 'LTE': <function <lambda>>, 'MAP': <function <lambda>>, 'MOD': <function <lambda>>, 'MUL': <function <lambda>>, 'NEQ': <function <lambda>>, 'ORD': <function ord>, 'NOT': <function sql_not>, 'OR': <function sql_or>, 'ORDERED': <function ordered>, 'POW': <built-in function pow>, 'RIGHT': <function <lambda>>, 'ROUND': <function <lambda>>, 'STRPOSITION': <function str_position>, 'SUB': <function <lambda>>, 'SUBSTRING': <function substring>, 'TIMESTRTOTIME': <function <lambda>>, 'UPPER': <function <lambda>>, 'YEAR': <function <lambda>>, 'MONTH': <function <lambda>>, 'DAY': <function <lambda>>, 'CURRENTDATETIME': <built-in method now of type object>, 'CURRENTTIMESTAMP': <built-in method now of type object>, 'CURRENTTIME': <built-in method now of type object>, 'CURRENTDATE': <built-in method today of type object>, 'STRFTIME': <function <lambda>>, 'STRTOTIME': <function <lambda>>, 'TRIM': <function <lambda>>, 'STRUCT': <function <lambda>>, 'UNIXTOTIME': <function <lambda>>}