sqlglot.dialects.bigquery

  1from __future__ import annotations
  2
  3import logging
  4import re
  5import typing as t
  6
  7from sqlglot import exp, generator, parser, tokens, transforms
  8from sqlglot._typing import E
  9from sqlglot.dialects.dialect import (
 10    Dialect,
 11    binary_from_function,
 12    datestrtodate_sql,
 13    format_time_lambda,
 14    inline_array_sql,
 15    max_or_greatest,
 16    min_or_least,
 17    no_ilike_sql,
 18    parse_date_delta_with_interval,
 19    regexp_replace_sql,
 20    rename_func,
 21    timestrtotime_sql,
 22    ts_or_ds_to_date_sql,
 23)
 24from sqlglot.helper import seq_get, split_num_words
 25from sqlglot.tokens import TokenType
 26
 27logger = logging.getLogger("sqlglot")
 28
 29
 30def _date_add_sql(
 31    data_type: str, kind: str
 32) -> t.Callable[[generator.Generator, exp.Expression], str]:
 33    def func(self, expression):
 34        this = self.sql(expression, "this")
 35        unit = expression.args.get("unit")
 36        unit = exp.var(unit.name.upper() if unit else "DAY")
 37        interval = exp.Interval(this=expression.expression, unit=unit)
 38        return f"{data_type}_{kind}({this}, {self.sql(interval)})"
 39
 40    return func
 41
 42
 43def _derived_table_values_to_unnest(self: generator.Generator, expression: exp.Values) -> str:
 44    if not expression.find_ancestor(exp.From, exp.Join):
 45        return self.values_sql(expression)
 46
 47    alias = expression.args.get("alias")
 48
 49    structs = [
 50        exp.Struct(
 51            expressions=[
 52                exp.alias_(value, column_name)
 53                for value, column_name in zip(
 54                    t.expressions,
 55                    alias.columns
 56                    if alias and alias.columns
 57                    else (f"_c{i}" for i in range(len(t.expressions))),
 58                )
 59            ]
 60        )
 61        for t in expression.find_all(exp.Tuple)
 62    ]
 63
 64    return self.unnest_sql(exp.Unnest(expressions=[exp.Array(expressions=structs)]))
 65
 66
 67def _returnsproperty_sql(self: generator.Generator, expression: exp.ReturnsProperty) -> str:
 68    this = expression.this
 69    if isinstance(this, exp.Schema):
 70        this = f"{this.this} <{self.expressions(this)}>"
 71    else:
 72        this = self.sql(this)
 73    return f"RETURNS {this}"
 74
 75
 76def _create_sql(self: generator.Generator, expression: exp.Create) -> str:
 77    kind = expression.args["kind"]
 78    returns = expression.find(exp.ReturnsProperty)
 79    if kind.upper() == "FUNCTION" and returns and returns.args.get("is_table"):
 80        expression = expression.copy()
 81        expression.set("kind", "TABLE FUNCTION")
 82        if isinstance(
 83            expression.expression,
 84            (
 85                exp.Subquery,
 86                exp.Literal,
 87            ),
 88        ):
 89            expression.set("expression", expression.expression.this)
 90
 91        return self.create_sql(expression)
 92
 93    return self.create_sql(expression)
 94
 95
 96def _unqualify_unnest(expression: exp.Expression) -> exp.Expression:
 97    """Remove references to unnest table aliases since bigquery doesn't allow them.
 98
 99    These are added by the optimizer's qualify_column step.
100    """
101    from sqlglot.optimizer.scope import Scope
102
103    if isinstance(expression, exp.Select):
104        for unnest in expression.find_all(exp.Unnest):
105            if isinstance(unnest.parent, (exp.From, exp.Join)) and unnest.alias:
106                for column in Scope(expression).find_all(exp.Column):
107                    if column.table == unnest.alias:
108                        column.set("table", None)
109
110    return expression
111
112
113# https://issuetracker.google.com/issues/162294746
114# workaround for bigquery bug when grouping by an expression and then ordering
115# WITH x AS (SELECT 1 y)
116# SELECT y + 1 z
117# FROM x
118# GROUP BY x + 1
119# ORDER by z
120def _alias_ordered_group(expression: exp.Expression) -> exp.Expression:
121    if isinstance(expression, exp.Select):
122        group = expression.args.get("group")
123        order = expression.args.get("order")
124
125        if group and order:
126            aliases = {
127                select.this: select.args["alias"]
128                for select in expression.selects
129                if isinstance(select, exp.Alias)
130            }
131
132            for e in group.expressions:
133                alias = aliases.get(e)
134
135                if alias:
136                    e.replace(exp.column(alias))
137
138    return expression
139
140
141def _pushdown_cte_column_names(expression: exp.Expression) -> exp.Expression:
142    """BigQuery doesn't allow column names when defining a CTE, so we try to push them down."""
143    if isinstance(expression, exp.CTE) and expression.alias_column_names:
144        cte_query = expression.this
145
146        if cte_query.is_star:
147            logger.warning(
148                "Can't push down CTE column names for star queries. Run the query through"
149                " the optimizer or use 'qualify' to expand the star projections first."
150            )
151            return expression
152
153        column_names = expression.alias_column_names
154        expression.args["alias"].set("columns", None)
155
156        for name, select in zip(column_names, cte_query.selects):
157            to_replace = select
158
159            if isinstance(select, exp.Alias):
160                select = select.this
161
162            # Inner aliases are shadowed by the CTE column names
163            to_replace.replace(exp.alias_(select, name))
164
165    return expression
166
167
168def _parse_timestamp(args: t.List) -> exp.StrToTime:
169    this = format_time_lambda(exp.StrToTime, "bigquery")([seq_get(args, 1), seq_get(args, 0)])
170    this.set("zone", seq_get(args, 2))
171    return this
172
173
174def _parse_date(args: t.List) -> exp.Date | exp.DateFromParts:
175    expr_type = exp.DateFromParts if len(args) == 3 else exp.Date
176    return expr_type.from_arg_list(args)
177
178
179def _parse_to_hex(args: t.List) -> exp.Hex | exp.MD5:
180    # TO_HEX(MD5(..)) is common in BigQuery, so it's parsed into MD5 to simplify its transpilation
181    arg = seq_get(args, 0)
182    return exp.MD5(this=arg.this) if isinstance(arg, exp.MD5Digest) else exp.Hex(this=arg)
183
184
185class BigQuery(Dialect):
186    UNNEST_COLUMN_ONLY = True
187
188    # https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#case_sensitivity
189    RESOLVES_IDENTIFIERS_AS_UPPERCASE = None
190
191    # bigquery udfs are case sensitive
192    NORMALIZE_FUNCTIONS = False
193
194    TIME_MAPPING = {
195        "%D": "%m/%d/%y",
196    }
197
198    FORMAT_MAPPING = {
199        "DD": "%d",
200        "MM": "%m",
201        "MON": "%b",
202        "MONTH": "%B",
203        "YYYY": "%Y",
204        "YY": "%y",
205        "HH": "%I",
206        "HH12": "%I",
207        "HH24": "%H",
208        "MI": "%M",
209        "SS": "%S",
210        "SSSSS": "%f",
211        "TZH": "%z",
212    }
213
214    @classmethod
215    def normalize_identifier(cls, expression: E) -> E:
216        # In BigQuery, CTEs aren't case-sensitive, but table names are (by default, at least).
217        # The following check is essentially a heuristic to detect tables based on whether or
218        # not they're qualified.
219        if isinstance(expression, exp.Identifier):
220            parent = expression.parent
221
222            while isinstance(parent, exp.Dot):
223                parent = parent.parent
224
225            if (
226                not isinstance(parent, exp.UserDefinedFunction)
227                and not (isinstance(parent, exp.Table) and parent.db)
228                and not expression.meta.get("is_table")
229            ):
230                expression.set("this", expression.this.lower())
231
232        return expression
233
234    class Tokenizer(tokens.Tokenizer):
235        QUOTES = ["'", '"', '"""', "'''"]
236        COMMENTS = ["--", "#", ("/*", "*/")]
237        IDENTIFIERS = ["`"]
238        STRING_ESCAPES = ["\\"]
239
240        HEX_STRINGS = [("0x", ""), ("0X", "")]
241
242        BYTE_STRINGS = [
243            (prefix + q, q) for q in t.cast(t.List[str], QUOTES) for prefix in ("b", "B")
244        ]
245
246        RAW_STRINGS = [
247            (prefix + q, q) for q in t.cast(t.List[str], QUOTES) for prefix in ("r", "R")
248        ]
249
250        KEYWORDS = {
251            **tokens.Tokenizer.KEYWORDS,
252            "ANY TYPE": TokenType.VARIANT,
253            "BEGIN": TokenType.COMMAND,
254            "BEGIN TRANSACTION": TokenType.BEGIN,
255            "CURRENT_DATETIME": TokenType.CURRENT_DATETIME,
256            "BYTES": TokenType.BINARY,
257            "DECLARE": TokenType.COMMAND,
258            "FLOAT64": TokenType.DOUBLE,
259            "INT64": TokenType.BIGINT,
260            "RECORD": TokenType.STRUCT,
261            "TIMESTAMP": TokenType.TIMESTAMPTZ,
262            "NOT DETERMINISTIC": TokenType.VOLATILE,
263            "UNKNOWN": TokenType.NULL,
264        }
265        KEYWORDS.pop("DIV")
266
267    class Parser(parser.Parser):
268        PREFIXED_PIVOT_COLUMNS = True
269
270        LOG_BASE_FIRST = False
271        LOG_DEFAULTS_TO_LN = True
272
273        FUNCTIONS = {
274            **parser.Parser.FUNCTIONS,
275            "DATE": _parse_date,
276            "DATE_ADD": parse_date_delta_with_interval(exp.DateAdd),
277            "DATE_SUB": parse_date_delta_with_interval(exp.DateSub),
278            "DATE_TRUNC": lambda args: exp.DateTrunc(
279                unit=exp.Literal.string(str(seq_get(args, 1))),
280                this=seq_get(args, 0),
281            ),
282            "DATETIME_ADD": parse_date_delta_with_interval(exp.DatetimeAdd),
283            "DATETIME_SUB": parse_date_delta_with_interval(exp.DatetimeSub),
284            "DIV": binary_from_function(exp.IntDiv),
285            "GENERATE_ARRAY": exp.GenerateSeries.from_arg_list,
286            "MD5": exp.MD5Digest.from_arg_list,
287            "TO_HEX": _parse_to_hex,
288            "PARSE_DATE": lambda args: format_time_lambda(exp.StrToDate, "bigquery")(
289                [seq_get(args, 1), seq_get(args, 0)]
290            ),
291            "PARSE_TIMESTAMP": _parse_timestamp,
292            "REGEXP_CONTAINS": exp.RegexpLike.from_arg_list,
293            "REGEXP_EXTRACT": lambda args: exp.RegexpExtract(
294                this=seq_get(args, 0),
295                expression=seq_get(args, 1),
296                position=seq_get(args, 2),
297                occurrence=seq_get(args, 3),
298                group=exp.Literal.number(1)
299                if re.compile(str(seq_get(args, 1))).groups == 1
300                else None,
301            ),
302            "SPLIT": lambda args: exp.Split(
303                # https://cloud.google.com/bigquery/docs/reference/standard-sql/string_functions#split
304                this=seq_get(args, 0),
305                expression=seq_get(args, 1) or exp.Literal.string(","),
306            ),
307            "TIME_ADD": parse_date_delta_with_interval(exp.TimeAdd),
308            "TIME_SUB": parse_date_delta_with_interval(exp.TimeSub),
309            "TIMESTAMP_ADD": parse_date_delta_with_interval(exp.TimestampAdd),
310            "TIMESTAMP_SUB": parse_date_delta_with_interval(exp.TimestampSub),
311            "TO_JSON_STRING": exp.JSONFormat.from_arg_list,
312        }
313
314        FUNCTION_PARSERS = {
315            **parser.Parser.FUNCTION_PARSERS,
316            "ARRAY": lambda self: self.expression(exp.Array, expressions=[self._parse_statement()]),
317        }
318        FUNCTION_PARSERS.pop("TRIM")
319
320        NO_PAREN_FUNCTIONS = {
321            **parser.Parser.NO_PAREN_FUNCTIONS,
322            TokenType.CURRENT_DATETIME: exp.CurrentDatetime,
323        }
324
325        NESTED_TYPE_TOKENS = {
326            *parser.Parser.NESTED_TYPE_TOKENS,
327            TokenType.TABLE,
328        }
329
330        ID_VAR_TOKENS = {
331            *parser.Parser.ID_VAR_TOKENS,
332            TokenType.VALUES,
333        }
334
335        PROPERTY_PARSERS = {
336            **parser.Parser.PROPERTY_PARSERS,
337            "NOT DETERMINISTIC": lambda self: self.expression(
338                exp.StabilityProperty, this=exp.Literal.string("VOLATILE")
339            ),
340            "OPTIONS": lambda self: self._parse_with_property(),
341        }
342
343        CONSTRAINT_PARSERS = {
344            **parser.Parser.CONSTRAINT_PARSERS,
345            "OPTIONS": lambda self: exp.Properties(expressions=self._parse_with_property()),
346        }
347
348        def _parse_table_part(self, schema: bool = False) -> t.Optional[exp.Expression]:
349            this = super()._parse_table_part(schema=schema)
350
351            # https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#table_names
352            if isinstance(this, exp.Identifier):
353                table_name = this.name
354                while self._match(TokenType.DASH, advance=False) and self._next:
355                    self._advance(2)
356                    table_name += f"-{self._prev.text}"
357
358                this = exp.Identifier(this=table_name, quoted=this.args.get("quoted"))
359
360            return this
361
362        def _parse_table_parts(self, schema: bool = False) -> exp.Table:
363            table = super()._parse_table_parts(schema=schema)
364            if isinstance(table.this, exp.Identifier) and "." in table.name:
365                catalog, db, this, *rest = (
366                    t.cast(t.Optional[exp.Expression], exp.to_identifier(x))
367                    for x in split_num_words(table.name, ".", 3)
368                )
369
370                if rest and this:
371                    this = exp.Dot.build(t.cast(t.List[exp.Expression], [this, *rest]))
372
373                table = exp.Table(this=this, db=db, catalog=catalog)
374
375            return table
376
377    class Generator(generator.Generator):
378        EXPLICIT_UNION = True
379        INTERVAL_ALLOWS_PLURAL_FORM = False
380        JOIN_HINTS = False
381        QUERY_HINTS = False
382        TABLE_HINTS = False
383        LIMIT_FETCH = "LIMIT"
384        RENAME_TABLE_WITH_DB = False
385        ESCAPE_LINE_BREAK = True
386
387        TRANSFORMS = {
388            **generator.Generator.TRANSFORMS,
389            exp.ApproxDistinct: rename_func("APPROX_COUNT_DISTINCT"),
390            exp.ArraySize: rename_func("ARRAY_LENGTH"),
391            exp.Cast: transforms.preprocess([transforms.remove_precision_parameterized_types]),
392            exp.Create: _create_sql,
393            exp.CTE: transforms.preprocess([_pushdown_cte_column_names]),
394            exp.DateAdd: _date_add_sql("DATE", "ADD"),
395            exp.DateDiff: lambda self, e: f"DATE_DIFF({self.sql(e, 'this')}, {self.sql(e, 'expression')}, {self.sql(e.args.get('unit', 'DAY'))})",
396            exp.DateFromParts: rename_func("DATE"),
397            exp.DateStrToDate: datestrtodate_sql,
398            exp.DateSub: _date_add_sql("DATE", "SUB"),
399            exp.DatetimeAdd: _date_add_sql("DATETIME", "ADD"),
400            exp.DatetimeSub: _date_add_sql("DATETIME", "SUB"),
401            exp.DateTrunc: lambda self, e: self.func("DATE_TRUNC", e.this, e.text("unit")),
402            exp.GenerateSeries: rename_func("GENERATE_ARRAY"),
403            exp.GroupConcat: rename_func("STRING_AGG"),
404            exp.Hex: rename_func("TO_HEX"),
405            exp.ILike: no_ilike_sql,
406            exp.IntDiv: rename_func("DIV"),
407            exp.JSONFormat: rename_func("TO_JSON_STRING"),
408            exp.Max: max_or_greatest,
409            exp.MD5: lambda self, e: self.func("TO_HEX", self.func("MD5", e.this)),
410            exp.MD5Digest: rename_func("MD5"),
411            exp.Min: min_or_least,
412            exp.PartitionedByProperty: lambda self, e: f"PARTITION BY {self.sql(e, 'this')}",
413            exp.RegexpExtract: lambda self, e: self.func(
414                "REGEXP_EXTRACT",
415                e.this,
416                e.expression,
417                e.args.get("position"),
418                e.args.get("occurrence"),
419            ),
420            exp.RegexpReplace: regexp_replace_sql,
421            exp.RegexpLike: rename_func("REGEXP_CONTAINS"),
422            exp.ReturnsProperty: _returnsproperty_sql,
423            exp.Select: transforms.preprocess(
424                [
425                    transforms.explode_to_unnest,
426                    _unqualify_unnest,
427                    transforms.eliminate_distinct_on,
428                    _alias_ordered_group,
429                ]
430            ),
431            exp.StabilityProperty: lambda self, e: f"DETERMINISTIC"
432            if e.name == "IMMUTABLE"
433            else "NOT DETERMINISTIC",
434            exp.StrToDate: lambda self, e: f"PARSE_DATE({self.format_time(e)}, {self.sql(e, 'this')})",
435            exp.StrToTime: lambda self, e: self.func(
436                "PARSE_TIMESTAMP", self.format_time(e), e.this, e.args.get("zone")
437            ),
438            exp.TimeAdd: _date_add_sql("TIME", "ADD"),
439            exp.TimeSub: _date_add_sql("TIME", "SUB"),
440            exp.TimestampAdd: _date_add_sql("TIMESTAMP", "ADD"),
441            exp.TimestampSub: _date_add_sql("TIMESTAMP", "SUB"),
442            exp.TimeStrToTime: timestrtotime_sql,
443            exp.Trim: lambda self, e: self.func(f"TRIM", e.this, e.expression),
444            exp.TsOrDsAdd: _date_add_sql("DATE", "ADD"),
445            exp.TsOrDsToDate: ts_or_ds_to_date_sql("bigquery"),
446            exp.Unhex: rename_func("FROM_HEX"),
447            exp.Values: _derived_table_values_to_unnest,
448            exp.VariancePop: rename_func("VAR_POP"),
449        }
450
451        TYPE_MAPPING = {
452            **generator.Generator.TYPE_MAPPING,
453            exp.DataType.Type.BIGDECIMAL: "BIGNUMERIC",
454            exp.DataType.Type.BIGINT: "INT64",
455            exp.DataType.Type.BINARY: "BYTES",
456            exp.DataType.Type.BOOLEAN: "BOOL",
457            exp.DataType.Type.CHAR: "STRING",
458            exp.DataType.Type.DECIMAL: "NUMERIC",
459            exp.DataType.Type.DOUBLE: "FLOAT64",
460            exp.DataType.Type.FLOAT: "FLOAT64",
461            exp.DataType.Type.INT: "INT64",
462            exp.DataType.Type.NCHAR: "STRING",
463            exp.DataType.Type.NVARCHAR: "STRING",
464            exp.DataType.Type.SMALLINT: "INT64",
465            exp.DataType.Type.TEXT: "STRING",
466            exp.DataType.Type.TIMESTAMP: "DATETIME",
467            exp.DataType.Type.TIMESTAMPTZ: "TIMESTAMP",
468            exp.DataType.Type.TIMESTAMPLTZ: "TIMESTAMP",
469            exp.DataType.Type.TINYINT: "INT64",
470            exp.DataType.Type.VARBINARY: "BYTES",
471            exp.DataType.Type.VARCHAR: "STRING",
472            exp.DataType.Type.VARIANT: "ANY TYPE",
473        }
474
475        PROPERTIES_LOCATION = {
476            **generator.Generator.PROPERTIES_LOCATION,
477            exp.PartitionedByProperty: exp.Properties.Location.POST_SCHEMA,
478            exp.VolatileProperty: exp.Properties.Location.UNSUPPORTED,
479        }
480
481        # from: https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#reserved_keywords
482        RESERVED_KEYWORDS = {
483            *generator.Generator.RESERVED_KEYWORDS,
484            "all",
485            "and",
486            "any",
487            "array",
488            "as",
489            "asc",
490            "assert_rows_modified",
491            "at",
492            "between",
493            "by",
494            "case",
495            "cast",
496            "collate",
497            "contains",
498            "create",
499            "cross",
500            "cube",
501            "current",
502            "default",
503            "define",
504            "desc",
505            "distinct",
506            "else",
507            "end",
508            "enum",
509            "escape",
510            "except",
511            "exclude",
512            "exists",
513            "extract",
514            "false",
515            "fetch",
516            "following",
517            "for",
518            "from",
519            "full",
520            "group",
521            "grouping",
522            "groups",
523            "hash",
524            "having",
525            "if",
526            "ignore",
527            "in",
528            "inner",
529            "intersect",
530            "interval",
531            "into",
532            "is",
533            "join",
534            "lateral",
535            "left",
536            "like",
537            "limit",
538            "lookup",
539            "merge",
540            "natural",
541            "new",
542            "no",
543            "not",
544            "null",
545            "nulls",
546            "of",
547            "on",
548            "or",
549            "order",
550            "outer",
551            "over",
552            "partition",
553            "preceding",
554            "proto",
555            "qualify",
556            "range",
557            "recursive",
558            "respect",
559            "right",
560            "rollup",
561            "rows",
562            "select",
563            "set",
564            "some",
565            "struct",
566            "tablesample",
567            "then",
568            "to",
569            "treat",
570            "true",
571            "unbounded",
572            "union",
573            "unnest",
574            "using",
575            "when",
576            "where",
577            "window",
578            "with",
579            "within",
580        }
581
582        def attimezone_sql(self, expression: exp.AtTimeZone) -> str:
583            parent = expression.parent
584
585            # BigQuery allows CAST(.. AS {STRING|TIMESTAMP} [FORMAT <fmt> [AT TIME ZONE <tz>]]).
586            # Only the TIMESTAMP one should use the below conversion, when AT TIME ZONE is included.
587            if not isinstance(parent, exp.Cast) or not parent.to.is_type("text"):
588                return self.func(
589                    "TIMESTAMP", self.func("DATETIME", expression.this, expression.args.get("zone"))
590                )
591
592            return super().attimezone_sql(expression)
593
594        def trycast_sql(self, expression: exp.TryCast) -> str:
595            return self.cast_sql(expression, safe_prefix="SAFE_")
596
597        def cte_sql(self, expression: exp.CTE) -> str:
598            if expression.alias_column_names:
599                self.unsupported("Column names in CTE definition are not supported.")
600            return super().cte_sql(expression)
601
602        def array_sql(self, expression: exp.Array) -> str:
603            first_arg = seq_get(expression.expressions, 0)
604            if isinstance(first_arg, exp.Subqueryable):
605                return f"ARRAY{self.wrap(self.sql(first_arg))}"
606
607            return inline_array_sql(self, expression)
608
609        def transaction_sql(self, *_) -> str:
610            return "BEGIN TRANSACTION"
611
612        def commit_sql(self, *_) -> str:
613            return "COMMIT TRANSACTION"
614
615        def rollback_sql(self, *_) -> str:
616            return "ROLLBACK TRANSACTION"
617
618        def in_unnest_op(self, expression: exp.Unnest) -> str:
619            return self.sql(expression)
620
621        def except_op(self, expression: exp.Except) -> str:
622            if not expression.args.get("distinct", False):
623                self.unsupported("EXCEPT without DISTINCT is not supported in BigQuery")
624            return f"EXCEPT{' DISTINCT' if expression.args.get('distinct') else ' ALL'}"
625
626        def intersect_op(self, expression: exp.Intersect) -> str:
627            if not expression.args.get("distinct", False):
628                self.unsupported("INTERSECT without DISTINCT is not supported in BigQuery")
629            return f"INTERSECT{' DISTINCT' if expression.args.get('distinct') else ' ALL'}"
630
631        def with_properties(self, properties: exp.Properties) -> str:
632            return self.properties(properties, prefix=self.seg("OPTIONS"))