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