1"""sqlglot expressions functions."""
  2
  3from __future__ import annotations
  4
  5import typing as t
  6
  7from sqlglot.expressions.core import (
  8    Expression,
  9    Func,
 10    Binary,
 11    SubqueryPredicate,
 12    ExpOrStr,
 13    maybe_parse,
 14    maybe_copy,
 15)
 16
 17# Re-export from focused submodules (backward compatibility)
 18from sqlglot.expressions.math import *  # noqa: F401,F403
 19from sqlglot.expressions.string import *  # noqa: F401,F403
 20from sqlglot.expressions.temporal import *  # noqa: F401,F403
 21from sqlglot.expressions.aggregate import *  # noqa: F401,F403
 22from sqlglot.expressions.array import *  # noqa: F401,F403
 23from sqlglot.expressions.json import *  # noqa: F401,F403
 24
 25if t.TYPE_CHECKING:
 26    from sqlglot.expressions.datatypes import DataType, DATA_TYPE
 27    from typing_extensions import Unpack
 28    from sqlglot._typing import ParserArgs
 29
 30
 31# Cast / type conversion
 32
 33
 34class Cast(Expression, Func):
 35    is_cast: t.ClassVar[bool] = True
 36    arg_types = {
 37        "this": True,
 38        "to": True,
 39        "format": False,
 40        "safe": False,
 41        "action": False,
 42        "default": False,
 43    }
 44
 45    @property
 46    def name(self) -> str:
 47        return self.this.name
 48
 49    @property
 50    def to(self) -> DataType:
 51        return self.args["to"]
 52
 53    @property
 54    def output_name(self) -> str:
 55        return self.name
 56
 57    def is_type(self, *dtypes: DATA_TYPE) -> bool:
 58        """
 59        Checks whether this Cast's DataType matches one of the provided data types. Nested types
 60        like arrays or structs will be compared using "structural equivalence" semantics, so e.g.
 61        array<int> != array<float>.
 62
 63        Args:
 64            dtypes: the data types to compare this Cast's DataType to.
 65
 66        Returns:
 67            True, if and only if there is a type in `dtypes` which is equal to this Cast's DataType.
 68        """
 69        return self.to.is_type(*dtypes)
 70
 71
 72class TryCast(Cast):
 73    arg_types = {
 74        **Cast.arg_types,
 75        "requires_string": False,
 76        "null_on_text_overflow": False,
 77        "probe_date_format": False,
 78    }
 79
 80
 81class JSONCast(Cast):
 82    pass
 83
 84
 85class CastToStrType(Expression, Func):
 86    arg_types = {"this": True, "to": True}
 87
 88
 89class Convert(Expression, Func):
 90    arg_types = {"this": True, "expression": True, "style": False, "safe": False}
 91
 92
 93# Conditional
 94
 95
 96class If(Expression, Func):
 97    arg_types = {"this": True, "true": True, "false": False}
 98    _sql_names = ["IF", "IIF"]
 99
100
101class Case(Expression, Func):
102    arg_types = {"this": False, "ifs": True, "default": False}
103
104    def when(
105        self,
106        condition: ExpOrStr,
107        then: ExpOrStr,
108        copy: bool = True,
109        **opts: Unpack[ParserArgs],
110    ) -> Case:
111        instance = maybe_copy(self, copy)
112        instance.append(
113            "ifs",
114            If(
115                this=maybe_parse(condition, copy=copy, **opts),
116                true=maybe_parse(then, copy=copy, **opts),
117            ),
118        )
119        return instance
120
121    def else_(self, condition: ExpOrStr, copy: bool = True, **opts: Unpack[ParserArgs]) -> Case:
122        instance = maybe_copy(self, copy)
123        instance.set("default", maybe_parse(condition, copy=copy, **opts))
124        return instance
125
126
127class Coalesce(Expression, Func):
128    arg_types = {"this": True, "expressions": False, "is_nvl": False, "is_null": False}
129    is_var_len_args = True
130    _sql_names = ["COALESCE", "IFNULL", "NVL"]
131
132
133class DecodeCase(Expression, Func):
134    arg_types = {"expressions": True}
135    is_var_len_args = True
136
137
138class EqualNull(Expression, Func):
139    arg_types = {"this": True, "expression": True}
140
141
142class Greatest(Expression, Func):
143    arg_types = {"this": True, "expressions": False, "ignore_nulls": True}
144    is_var_len_args = True
145
146
147class Least(Expression, Func):
148    arg_types = {"this": True, "expressions": False, "ignore_nulls": True}
149    is_var_len_args = True
150
151
152class Nullif(Expression, Func):
153    arg_types = {"this": True, "expression": True}
154
155
156class ObjectTransform(Expression, Func):
157    arg_types = {"this": True, "keep": False, "set_": False}
158
159
160class Nvl2(Expression, Func):
161    arg_types = {"this": True, "true": True, "false": False}
162
163
164class Try(Expression, Func):
165    pass
166
167
168# Predicates / misc functions
169
170
171class Collate(Expression, Binary, Func):
172    pass
173
174
175class Collation(Expression, Func):
176    pass
177
178
179class ConnectByRoot(Expression, Func):
180    pass
181
182
183class CheckXml(Expression, Func):
184    arg_types = {"this": True, "disable_auto_convert": False}
185
186
187class Exists(Expression, Func, SubqueryPredicate):
188    arg_types = {"this": True, "expression": False}
189
190
191# Type coercions / lax types
192
193
194class Float64(Expression, Func):
195    arg_types = {"this": True, "expression": False}
196
197
198class Int64(Expression, Func):
199    pass
200
201
202class IsArray(Expression, Func):
203    pass
204
205
206class IsNullValue(Expression, Func):
207    pass
208
209
210class LaxBool(Expression, Func):
211    pass
212
213
214class LaxFloat64(Expression, Func):
215    pass
216
217
218class LaxInt64(Expression, Func):
219    pass
220
221
222class LaxString(Expression, Func):
223    pass
224
225
226class ToBoolean(Expression, Func):
227    arg_types = {"this": True, "safe": False}
228
229
230class ToVariant(Expression, Func):
231    pass
232
233
234# Session / context functions
235
236
237class CurrentAccount(Expression, Func):
238    arg_types = {}
239
240
241class CurrentAccountName(Expression, Func):
242    arg_types = {}
243
244
245class CurrentAvailableRoles(Expression, Func):
246    arg_types = {}
247
248
249class CurrentCatalog(Expression, Func):
250    arg_types = {}
251
252
253class CurrentClient(Expression, Func):
254    arg_types = {}
255
256
257class CurrentDatabase(Expression, Func):
258    arg_types = {}
259
260
261class CurrentIpAddress(Expression, Func):
262    arg_types = {}
263
264
265class CurrentOrganizationName(Expression, Func):
266    arg_types = {}
267
268
269class CurrentOrganizationUser(Expression, Func):
270    arg_types = {}
271
272
273class CurrentRegion(Expression, Func):
274    arg_types = {}
275
276
277class CurrentRole(Expression, Func):
278    arg_types = {}
279
280
281class CurrentRoleType(Expression, Func):
282    arg_types = {}
283
284
285class CurrentSchema(Expression, Func):
286    arg_types = {"this": False}
287
288
289class CurrentSchemas(Expression, Func):
290    arg_types = {"this": False}
291
292
293class CurrentSecondaryRoles(Expression, Func):
294    arg_types = {}
295
296
297class CurrentSession(Expression, Func):
298    arg_types = {}
299
300
301class CurrentStatement(Expression, Func):
302    arg_types = {}
303
304
305class CurrentTransaction(Expression, Func):
306    arg_types = {}
307
308
309class CurrentUser(Expression, Func):
310    arg_types = {"this": False}
311
312
313class CurrentUserId(Expression, Func):
314    arg_types = {}
315
316
317class CurrentVersion(Expression, Func):
318    arg_types = {}
319
320
321class CurrentWarehouse(Expression, Func):
322    arg_types = {}
323
324
325class SessionUser(Expression, Func):
326    arg_types = {}
327
328
329# ML / AI
330
331
332class AIClassify(Expression, Func):
333    arg_types = {"this": True, "categories": True, "config": False}
334    _sql_names = ["AI_CLASSIFY"]
335
336
337class AIEmbed(Expression, Func):
338    arg_types = {"expressions": True}
339    is_var_len_args = True
340    _sql_names = ["AI_EMBED"]
341
342
343class AISimilarity(Expression, Func):
344    arg_types = {"expressions": True}
345    is_var_len_args = True
346    _sql_names = ["AI_SIMILARITY"]
347
348
349class AIGenerate(Expression, Func):
350    arg_types = {"expressions": True}
351    is_var_len_args = True
352    _sql_names = ["AI_GENERATE"]
353
354
355class FeaturesAtTime(Expression, Func):
356    arg_types = {"this": True, "time": False, "num_rows": False, "ignore_feature_nulls": False}
357
358
359class GenerateEmbedding(Expression, Func):
360    arg_types = {"this": True, "expression": True, "params_struct": False, "is_text": False}
361
362
363class GenerateText(Expression, Func):
364    arg_types = {"this": True, "expression": False, "params_struct": False}
365
366
367class GenerateTable(Expression, Func):
368    arg_types = {"this": True, "expression": False, "params_struct": False}
369
370
371class GenerateBool(Expression, Func):
372    arg_types = {"this": True, "expression": False, "params_struct": False}
373
374
375class GenerateInt(Expression, Func):
376    arg_types = {"this": True, "expression": False, "params_struct": False}
377
378
379class GenerateDouble(Expression, Func):
380    arg_types = {"this": True, "expression": False, "params_struct": False}
381
382
383class MLForecast(Expression, Func):
384    arg_types = {"this": True, "expression": False, "params_struct": False}
385
386
387class AIForecast(Expression, Func):
388    arg_types = {
389        "this": True,
390        "data_col": False,
391        "timestamp_col": False,
392        "model": False,
393        "id_cols": False,
394        "horizon": False,
395        "forecast_end_timestamp": False,
396        "confidence_level": False,
397        "output_historical_time_series": False,
398        "context_window": False,
399    }
400
401
402class MLTranslate(Expression, Func):
403    arg_types = {"this": True, "expression": True, "params_struct": True}
404
405
406class Predict(Expression, Func):
407    arg_types = {"this": True, "expression": True, "params_struct": False}
408
409
410class VectorSearch(Expression, Func):
411    arg_types = {
412        "this": True,
413        "column_to_search": True,
414        "query_table": True,
415        "query_column_to_search": False,
416        "top_k": False,
417        "distance_type": False,
418        "options": False,
419    }
420
421
422# Data reading
423
424
425class ReadCSV(Expression, Func):
426    _sql_names = ["READ_CSV"]
427    is_var_len_args = True
428    arg_types = {"this": True, "expressions": False}
429
430
431class ReadParquet(Expression, Func):
432    is_var_len_args = True
433    arg_types = {"expressions": True}
434
435
436# XML
437
438
439class XMLElement(Expression, Func):
440    _sql_names = ["XMLELEMENT"]
441    arg_types = {"this": True, "expressions": False, "evalname": False}
442
443
444class GetIgnoreCase(Expression, Func):
445    arg_types = {"this": True, "expression": True}
446
447
448class XMLGet(Expression, Func):
449    _sql_names = ["XMLGET"]
450    arg_types = {"this": True, "expression": True, "instance": False}
451
452
453class XMLTable(Expression, Func):
454    arg_types = {
455        "this": True,
456        "namespaces": False,
457        "passing": False,
458        "columns": False,
459        "by_ref": False,
460    }
461
462
463# Network / domain
464
465
466class Host(Expression, Func):
467    pass
468
469
470class NetFunc(Expression, Func):
471    pass
472
473
474class ParseIp(Expression, Func):
475    arg_types = {"this": True, "type": True, "permissive": False}
476
477
478class RegDomain(Expression, Func):
479    pass
480
481
482# Misc utility
483
484
485class Columns(Expression, Func):
486    arg_types = {"this": True, "unpack": False}
487
488
489class Normal(Expression, Func):
490    arg_types = {"this": True, "stddev": True, "gen": True}
491
492
493class Rand(Expression, Func):
494    _sql_names = ["RAND", "RANDOM"]
495    arg_types = {"this": False, "lower": False, "upper": False}
496
497
498class Randn(Expression, Func):
499    arg_types = {"this": False}
500
501
502class Randstr(Expression, Func):
503    arg_types = {"this": True, "generator": False}
504
505
506class RangeBucket(Expression, Func):
507    arg_types = {"this": True, "expression": True}
508
509
510class RangeN(Expression, Func):
511    arg_types = {"this": True, "expressions": True, "each": False}
512
513
514class Seq1(Expression, Func):
515    arg_types = {"this": False}
516
517
518class Seq2(Expression, Func):
519    arg_types = {"this": False}
520
521
522class Seq4(Expression, Func):
523    arg_types = {"this": False}
524
525
526class Seq8(Expression, Func):
527    arg_types = {"this": False}
528
529
530class Uniform(Expression, Func):
531    arg_types = {"this": True, "expression": True, "gen": False, "seed": False}
532
533
534class Uuid(Expression, Func):
535    _sql_names = ["UUID", "GEN_RANDOM_UUID", "GENERATE_UUID", "UUID_STRING"]
536
537    arg_types = {"this": False, "name": False, "is_string": False}
538
539
540class WeekStart(Expression, Func):
541    pass
542
543
544class WidthBucket(Expression, Func):
545    arg_types = {
546        "this": True,
547        "min_value": False,
548        "max_value": False,
549        "num_buckets": False,
550        "threshold": False,
551    }
552
553
554class Zipf(Expression, Func):
555    arg_types = {"this": True, "elementcount": True, "gen": True}