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
182FIRST = 0
183LAST = 1
184NULL_PLACEHOLDER = 0
185
186
187def ordered(this, desc, nulls_first):
188    if this is None:
189        return (FIRST if nulls_first else LAST, NULL_PLACEHOLDER)
190    return (LAST if nulls_first else FIRST, reverse_key(this) if desc else this)
191
192
193def _like(this, e, flags=0):
194    return bool(
195        re.fullmatch(re.escape(e).replace("_", ".").replace("%", ".*"), this, re.DOTALL | flags)
196    )
197
198
199@null_if_any
200def interval(this, unit):
201    plural = unit + "S"
202    if plural in Generator.TIME_PART_SINGULARS:
203        unit = plural
204    return datetime.timedelta(**{unit.lower(): float(this)})
205
206
207@null_if_any
208def arrayconcat(*args):
209    result = []
210    for arg in args:
211        result.extend(arg if isinstance(arg, list) else [arg])
212    return result
213
214
215@null_if_any("this", "expression")
216def arraytostring(this, expression, null=None):
217    return expression.join(x for x in (x if x is not None else null for x in this) if x is not None)
218
219
220@null_if_any("this", "expression")
221def jsonextract(this, expression):
222    for path_segment in expression:
223        if isinstance(this, dict):
224            this = this.get(path_segment)
225        elif isinstance(this, list) and is_int(path_segment):
226            this = seq_get(this, int(path_segment))
227        else:
228            raise NotImplementedError(f"Unable to extract value for {this} at {path_segment}.")
229
230        if this is None:
231            break
232
233    return this
234
235
236ENV = {
237    "exp": exp,
238    "AND": sql_and,
239    # aggs
240    "ARRAYAGG": list,
241    "ARRAYUNIQUEAGG": filter_nulls(lambda acc: list(set(acc))),
242    "AVG": filter_nulls(statistics.fmean if PYTHON_VERSION >= (3, 8) else statistics.mean),  # type: ignore
243    "COUNT": filter_nulls(lambda acc: sum(1 for _ in acc), False),
244    "MAX": filter_nulls(max),
245    "MIN": filter_nulls(min),
246    "SUM": filter_nulls(sum),
247    # scalar functions
248    "ABS": null_if_any(lambda this: abs(this)),
249    "ADD": null_if_any(lambda e, this: e + this),
250    "ARRAYANY": null_if_any(lambda arr, func: any(func(e) for e in arr)),
251    "ARRAYCONCAT": arrayconcat,
252    "ARRAYTOSTRING": arraytostring,
253    "BETWEEN": null_if_any(lambda this, low, high: low <= this and this <= high),
254    "BITWISEAND": null_if_any(lambda this, e: this & e),
255    "BITWISELEFTSHIFT": null_if_any(lambda this, e: this << e),
256    "BITWISEOR": null_if_any(lambda this, e: this | e),
257    "BITWISERIGHTSHIFT": null_if_any(lambda this, e: this >> e),
258    "BITWISEXOR": null_if_any(lambda this, e: this ^ e),
259    "CAST": cast,
260    "COALESCE": lambda *args: next((a for a in args if a is not None), None),
261    "CONCAT": null_if_any(lambda *args: "".join(args)),
262    "SAFECONCAT": null_if_any(lambda *args: "".join(str(arg) for arg in args)),
263    "CONCATWS": null_if_any(lambda this, *args: this.join(args)),
264    "DATEDIFF": null_if_any(lambda this, expression, *_: (this - expression).days),
265    "DATESTRTODATE": null_if_any(lambda arg: datetime.date.fromisoformat(arg)),
266    "DIV": null_if_any(lambda e, this: e / this),
267    "DOT": null_if_any(lambda e, this: e[this]),
268    "EQ": null_if_any(lambda this, e: this == e),
269    "EXTRACT": null_if_any(lambda this, e: getattr(e, this)),
270    "GT": null_if_any(lambda this, e: this > e),
271    "GTE": null_if_any(lambda this, e: this >= e),
272    "IF": lambda predicate, true, false: true if predicate else false,
273    "IN": sql_in,
274    "INT": null_if_any(int),
275    "INTDIV": null_if_any(lambda e, this: e // this),
276    "INTERVAL": interval,
277    "JSONEXTRACT": jsonextract,
278    "LEFT": null_if_any(lambda this, e: this[:e]),
279    "LENGTH": null_if_any(len),
280    "LIKE": null_if_any(lambda this, e: _like(this, e)),
281    "ILIKE": null_if_any(lambda this, e: _like(this, e, re.IGNORECASE)),
282    "LOWER": null_if_any(lambda arg: arg.lower()),
283    "LT": null_if_any(lambda this, e: this < e),
284    "LTE": null_if_any(lambda this, e: this <= e),
285    "MAP": null_if_any(lambda *args: dict(zip(*args))),  # type: ignore
286    "MOD": null_if_any(lambda e, this: e % this),
287    "MUL": null_if_any(lambda e, this: e * this),
288    "NEQ": null_if_any(lambda this, e: this != e),
289    "ORD": null_if_any(ord),
290    "NOT": sql_not,
291    "OR": sql_or,
292    "ORDERED": ordered,
293    "POW": null_if_any(pow),
294    "REVERSE": null_if_any(lambda this: this[::-1]),
295    "RIGHT": null_if_any(lambda this, e: this[-e:]),
296    "ROUND": null_if_any(lambda this, decimals=None, truncate=None: round(this, ndigits=decimals)),
297    "STRPOSITION": str_position,
298    "SUB": null_if_any(lambda e, this: e - this),
299    "SUBSTRING": substring,
300    "TIMESTRTOTIME": null_if_any(lambda arg: datetime.datetime.fromisoformat(arg)),
301    "UPPER": null_if_any(lambda arg: arg.upper()),
302    "YEAR": null_if_any(lambda arg: arg.year),
303    "MONTH": null_if_any(lambda arg: arg.month),
304    "DAY": null_if_any(lambda arg: arg.day),
305    "CURRENTDATETIME": datetime.datetime.now,
306    "CURRENTTIMESTAMP": datetime.datetime.now,
307    "CURRENTTIME": datetime.datetime.now,
308    "CURRENTDATE": datetime.date.today,
309    "STRFTIME": null_if_any(lambda fmt, arg: datetime.datetime.fromisoformat(arg).strftime(fmt)),
310    "STRTOTIME": null_if_any(lambda arg, format: datetime.datetime.strptime(arg, format)),
311    "TRIM": null_if_any(lambda this, e=None: this.strip(e)),
312    "STRUCT": lambda *args: {
313        args[x]: args[x + 1]
314        for x in range(0, len(args), 2)
315        if (args[x + 1] is not None and args[x] is not None)
316    },
317    "UNIXTOTIME": null_if_any(
318        lambda arg: datetime.datetime.fromtimestamp(arg, datetime.timezone.utc)
319    ),
320}

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.")

FIRST = 0

LAST = 1

NULL_PLACEHOLDER = 0

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

188def ordered(this, desc, nulls_first):
189    if this is None:
190        return (FIRST if nulls_first else LAST, NULL_PLACEHOLDER)
191    return (LAST if nulls_first else FIRST, reverse_key(this) if desc else this)

@null_if_any

def interval(this, unit): View Source

200@null_if_any
201def interval(this, unit):
202    plural = unit + "S"
203    if plural in Generator.TIME_PART_SINGULARS:
204        unit = plural
205    return datetime.timedelta(**{unit.lower(): float(this)})

@null_if_any

def arrayconcat(*args): View Source

208@null_if_any
209def arrayconcat(*args):
210    result = []
211    for arg in args:
212        result.extend(arg if isinstance(arg, list) else [arg])
213    return result

@null_if_any('this', 'expression')

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

216@null_if_any("this", "expression")
217def arraytostring(this, expression, null=None):
218    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

221@null_if_any("this", "expression")
222def jsonextract(this, expression):
223    for path_segment in expression:
224        if isinstance(this, dict):
225            this = this.get(path_segment)
226        elif isinstance(this, list) and is_int(path_segment):
227            this = seq_get(this, int(path_segment))
228        else:
229            raise NotImplementedError(f"Unable to extract value for {this} at {path_segment}.")
230
231        if this is None:
232            break
233
234    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>>, 'ARRAYCONCAT': <function arrayconcat>, '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>, 'INT': <function int>, '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': <function pow>, 'REVERSE': <function <lambda>>, '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>>}