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    # The _PARTITIONTIME and _PARTITIONDATE pseudo-columns are not returned by a SELECT * statement
215    # https://cloud.google.com/bigquery/docs/querying-partitioned-tables#query_an_ingestion-time_partitioned_table
216    PSEUDOCOLUMNS = {"_PARTITIONTIME", "_PARTITIONDATE"}
217
218    @classmethod
219    def normalize_identifier(cls, expression: E) -> E:
220        # In BigQuery, CTEs aren't case-sensitive, but table names are (by default, at least).
221        # The following check is essentially a heuristic to detect tables based on whether or
222        # not they're qualified.
223        if isinstance(expression, exp.Identifier):
224            parent = expression.parent
225
226            while isinstance(parent, exp.Dot):
227                parent = parent.parent
228
229            if (
230                not isinstance(parent, exp.UserDefinedFunction)
231                and not (isinstance(parent, exp.Table) and parent.db)
232                and not expression.meta.get("is_table")
233            ):
234                expression.set("this", expression.this.lower())
235
236        return expression
237
238    class Tokenizer(tokens.Tokenizer):
239        QUOTES = ["'", '"', '"""', "'''"]
240        COMMENTS = ["--", "#", ("/*", "*/")]
241        IDENTIFIERS = ["`"]
242        STRING_ESCAPES = ["\\"]
243
244        HEX_STRINGS = [("0x", ""), ("0X", "")]
245
246        BYTE_STRINGS = [
247            (prefix + q, q) for q in t.cast(t.List[str], QUOTES) for prefix in ("b", "B")
248        ]
249
250        RAW_STRINGS = [
251            (prefix + q, q) for q in t.cast(t.List[str], QUOTES) for prefix in ("r", "R")
252        ]
253
254        KEYWORDS = {
255            **tokens.Tokenizer.KEYWORDS,
256            "ANY TYPE": TokenType.VARIANT,
257            "BEGIN": TokenType.COMMAND,
258            "BEGIN TRANSACTION": TokenType.BEGIN,
259            "CURRENT_DATETIME": TokenType.CURRENT_DATETIME,
260            "BYTES": TokenType.BINARY,
261            "DECLARE": TokenType.COMMAND,
262            "FLOAT64": TokenType.DOUBLE,
263            "INT64": TokenType.BIGINT,
264            "RECORD": TokenType.STRUCT,
265            "TIMESTAMP": TokenType.TIMESTAMPTZ,
266            "NOT DETERMINISTIC": TokenType.VOLATILE,
267            "UNKNOWN": TokenType.NULL,
268        }
269        KEYWORDS.pop("DIV")
270
271    class Parser(parser.Parser):
272        PREFIXED_PIVOT_COLUMNS = True
273
274        LOG_BASE_FIRST = False
275        LOG_DEFAULTS_TO_LN = True
276
277        FUNCTIONS = {
278            **parser.Parser.FUNCTIONS,
279            "DATE": _parse_date,
280            "DATE_ADD": parse_date_delta_with_interval(exp.DateAdd),
281            "DATE_SUB": parse_date_delta_with_interval(exp.DateSub),
282            "DATE_TRUNC": lambda args: exp.DateTrunc(
283                unit=exp.Literal.string(str(seq_get(args, 1))),
284                this=seq_get(args, 0),
285            ),
286            "DATETIME_ADD": parse_date_delta_with_interval(exp.DatetimeAdd),
287            "DATETIME_SUB": parse_date_delta_with_interval(exp.DatetimeSub),
288            "DIV": binary_from_function(exp.IntDiv),
289            "GENERATE_ARRAY": exp.GenerateSeries.from_arg_list,
290            "MD5": exp.MD5Digest.from_arg_list,
291            "TO_HEX": _parse_to_hex,
292            "PARSE_DATE": lambda args: format_time_lambda(exp.StrToDate, "bigquery")(
293                [seq_get(args, 1), seq_get(args, 0)]
294            ),
295            "PARSE_TIMESTAMP": _parse_timestamp,
296            "REGEXP_CONTAINS": exp.RegexpLike.from_arg_list,
297            "REGEXP_EXTRACT": lambda args: exp.RegexpExtract(
298                this=seq_get(args, 0),
299                expression=seq_get(args, 1),
300                position=seq_get(args, 2),
301                occurrence=seq_get(args, 3),
302                group=exp.Literal.number(1)
303                if re.compile(str(seq_get(args, 1))).groups == 1
304                else None,
305            ),
306            "SPLIT": lambda args: exp.Split(
307                # https://cloud.google.com/bigquery/docs/reference/standard-sql/string_functions#split
308                this=seq_get(args, 0),
309                expression=seq_get(args, 1) or exp.Literal.string(","),
310            ),
311            "TIME_ADD": parse_date_delta_with_interval(exp.TimeAdd),
312            "TIME_SUB": parse_date_delta_with_interval(exp.TimeSub),
313            "TIMESTAMP_ADD": parse_date_delta_with_interval(exp.TimestampAdd),
314            "TIMESTAMP_SUB": parse_date_delta_with_interval(exp.TimestampSub),
315            "TO_JSON_STRING": exp.JSONFormat.from_arg_list,
316        }
317
318        FUNCTION_PARSERS = {
319            **parser.Parser.FUNCTION_PARSERS,
320            "ARRAY": lambda self: self.expression(exp.Array, expressions=[self._parse_statement()]),
321        }
322        FUNCTION_PARSERS.pop("TRIM")
323
324        NO_PAREN_FUNCTIONS = {
325            **parser.Parser.NO_PAREN_FUNCTIONS,
326            TokenType.CURRENT_DATETIME: exp.CurrentDatetime,
327        }
328
329        NESTED_TYPE_TOKENS = {
330            *parser.Parser.NESTED_TYPE_TOKENS,
331            TokenType.TABLE,
332        }
333
334        ID_VAR_TOKENS = {
335            *parser.Parser.ID_VAR_TOKENS,
336            TokenType.VALUES,
337        }
338
339        PROPERTY_PARSERS = {
340            **parser.Parser.PROPERTY_PARSERS,
341            "NOT DETERMINISTIC": lambda self: self.expression(
342                exp.StabilityProperty, this=exp.Literal.string("VOLATILE")
343            ),
344            "OPTIONS": lambda self: self._parse_with_property(),
345        }
346
347        CONSTRAINT_PARSERS = {
348            **parser.Parser.CONSTRAINT_PARSERS,
349            "OPTIONS": lambda self: exp.Properties(expressions=self._parse_with_property()),
350        }
351
352        def _parse_table_part(self, schema: bool = False) -> t.Optional[exp.Expression]:
353            this = super()._parse_table_part(schema=schema)
354
355            # https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#table_names
356            if isinstance(this, exp.Identifier):
357                table_name = this.name
358                while self._match(TokenType.DASH, advance=False) and self._next:
359                    self._advance(2)
360                    table_name += f"-{self._prev.text}"
361
362                this = exp.Identifier(this=table_name, quoted=this.args.get("quoted"))
363
364            return this
365
366        def _parse_table_parts(self, schema: bool = False) -> exp.Table:
367            table = super()._parse_table_parts(schema=schema)
368            if isinstance(table.this, exp.Identifier) and "." in table.name:
369                catalog, db, this, *rest = (
370                    t.cast(t.Optional[exp.Expression], exp.to_identifier(x))
371                    for x in split_num_words(table.name, ".", 3)
372                )
373
374                if rest and this:
375                    this = exp.Dot.build(t.cast(t.List[exp.Expression], [this, *rest]))
376
377                table = exp.Table(this=this, db=db, catalog=catalog)
378
379            return table
380
381    class Generator(generator.Generator):
382        EXPLICIT_UNION = True
383        INTERVAL_ALLOWS_PLURAL_FORM = False
384        JOIN_HINTS = False
385        QUERY_HINTS = False
386        TABLE_HINTS = False
387        LIMIT_FETCH = "LIMIT"
388        RENAME_TABLE_WITH_DB = False
389        ESCAPE_LINE_BREAK = True
390
391        TRANSFORMS = {
392            **generator.Generator.TRANSFORMS,
393            exp.ApproxDistinct: rename_func("APPROX_COUNT_DISTINCT"),
394            exp.ArraySize: rename_func("ARRAY_LENGTH"),
395            exp.Cast: transforms.preprocess([transforms.remove_precision_parameterized_types]),
396            exp.Create: _create_sql,
397            exp.CTE: transforms.preprocess([_pushdown_cte_column_names]),
398            exp.DateAdd: _date_add_sql("DATE", "ADD"),
399            exp.DateDiff: lambda self, e: f"DATE_DIFF({self.sql(e, 'this')}, {self.sql(e, 'expression')}, {self.sql(e.args.get('unit', 'DAY'))})",
400            exp.DateFromParts: rename_func("DATE"),
401            exp.DateStrToDate: datestrtodate_sql,
402            exp.DateSub: _date_add_sql("DATE", "SUB"),
403            exp.DatetimeAdd: _date_add_sql("DATETIME", "ADD"),
404            exp.DatetimeSub: _date_add_sql("DATETIME", "SUB"),
405            exp.DateTrunc: lambda self, e: self.func("DATE_TRUNC", e.this, e.text("unit")),
406            exp.GenerateSeries: rename_func("GENERATE_ARRAY"),
407            exp.GroupConcat: rename_func("STRING_AGG"),
408            exp.Hex: rename_func("TO_HEX"),
409            exp.ILike: no_ilike_sql,
410            exp.IntDiv: rename_func("DIV"),
411            exp.JSONFormat: rename_func("TO_JSON_STRING"),
412            exp.Max: max_or_greatest,
413            exp.MD5: lambda self, e: self.func("TO_HEX", self.func("MD5", e.this)),
414            exp.MD5Digest: rename_func("MD5"),
415            exp.Min: min_or_least,
416            exp.PartitionedByProperty: lambda self, e: f"PARTITION BY {self.sql(e, 'this')}",
417            exp.RegexpExtract: lambda self, e: self.func(
418                "REGEXP_EXTRACT",
419                e.this,
420                e.expression,
421                e.args.get("position"),
422                e.args.get("occurrence"),
423            ),
424            exp.RegexpReplace: regexp_replace_sql,
425            exp.RegexpLike: rename_func("REGEXP_CONTAINS"),
426            exp.ReturnsProperty: _returnsproperty_sql,
427            exp.Select: transforms.preprocess(
428                [
429                    transforms.explode_to_unnest,
430                    _unqualify_unnest,
431                    transforms.eliminate_distinct_on,
432                    _alias_ordered_group,
433                ]
434            ),
435            exp.StabilityProperty: lambda self, e: f"DETERMINISTIC"
436            if e.name == "IMMUTABLE"
437            else "NOT DETERMINISTIC",
438            exp.StrToDate: lambda self, e: f"PARSE_DATE({self.format_time(e)}, {self.sql(e, 'this')})",
439            exp.StrToTime: lambda self, e: self.func(
440                "PARSE_TIMESTAMP", self.format_time(e), e.this, e.args.get("zone")
441            ),
442            exp.TimeAdd: _date_add_sql("TIME", "ADD"),
443            exp.TimeSub: _date_add_sql("TIME", "SUB"),
444            exp.TimestampAdd: _date_add_sql("TIMESTAMP", "ADD"),
445            exp.TimestampSub: _date_add_sql("TIMESTAMP", "SUB"),
446            exp.TimeStrToTime: timestrtotime_sql,
447            exp.Trim: lambda self, e: self.func(f"TRIM", e.this, e.expression),
448            exp.TsOrDsAdd: _date_add_sql("DATE", "ADD"),
449            exp.TsOrDsToDate: ts_or_ds_to_date_sql("bigquery"),
450            exp.Unhex: rename_func("FROM_HEX"),
451            exp.Values: _derived_table_values_to_unnest,
452            exp.VariancePop: rename_func("VAR_POP"),
453        }
454
455        TYPE_MAPPING = {
456            **generator.Generator.TYPE_MAPPING,
457            exp.DataType.Type.BIGDECIMAL: "BIGNUMERIC",
458            exp.DataType.Type.BIGINT: "INT64",
459            exp.DataType.Type.BINARY: "BYTES",
460            exp.DataType.Type.BOOLEAN: "BOOL",
461            exp.DataType.Type.CHAR: "STRING",
462            exp.DataType.Type.DECIMAL: "NUMERIC",
463            exp.DataType.Type.DOUBLE: "FLOAT64",
464            exp.DataType.Type.FLOAT: "FLOAT64",
465            exp.DataType.Type.INT: "INT64",
466            exp.DataType.Type.NCHAR: "STRING",
467            exp.DataType.Type.NVARCHAR: "STRING",
468            exp.DataType.Type.SMALLINT: "INT64",
469            exp.DataType.Type.TEXT: "STRING",
470            exp.DataType.Type.TIMESTAMP: "DATETIME",
471            exp.DataType.Type.TIMESTAMPTZ: "TIMESTAMP",
472            exp.DataType.Type.TIMESTAMPLTZ: "TIMESTAMP",
473            exp.DataType.Type.TINYINT: "INT64",
474            exp.DataType.Type.VARBINARY: "BYTES",
475            exp.DataType.Type.VARCHAR: "STRING",
476            exp.DataType.Type.VARIANT: "ANY TYPE",
477        }
478
479        PROPERTIES_LOCATION = {
480            **generator.Generator.PROPERTIES_LOCATION,
481            exp.PartitionedByProperty: exp.Properties.Location.POST_SCHEMA,
482            exp.VolatileProperty: exp.Properties.Location.UNSUPPORTED,
483        }
484
485        # from: https://cloud.google.com/bigquery/docs/reference/standard-sql/lexical#reserved_keywords
486        RESERVED_KEYWORDS = {
487            *generator.Generator.RESERVED_KEYWORDS,
488            "all",
489            "and",
490            "any",
491            "array",
492            "as",
493            "asc",
494            "assert_rows_modified",
495            "at",
496            "between",
497            "by",
498            "case",
499            "cast",
500            "collate",
501            "contains",
502            "create",
503            "cross",
504            "cube",
505            "current",
506            "default",
507            "define",
508            "desc",
509            "distinct",
510            "else",
511            "end",
512            "enum",
513            "escape",
514            "except",
515            "exclude",
516            "exists",
517            "extract",
518            "false",
519            "fetch",
520            "following",
521            "for",
522            "from",
523            "full",
524            "group",
525            "grouping",
526            "groups",
527            "hash",
528            "having",
529            "if",
530            "ignore",
531            "in",
532            "inner",
533            "intersect",
534            "interval",
535            "into",
536            "is",
537            "join",
538            "lateral",
539            "left",
540            "like",
541            "limit",
542            "lookup",
543            "merge",
544            "natural",
545            "new",
546            "no",
547            "not",
548            "null",
549            "nulls",
550            "of",
551            "on",
552            "or",
553            "order",
554            "outer",
555            "over",
556            "partition",
557            "preceding",
558            "proto",
559            "qualify",
560            "range",
561            "recursive",
562            "respect",
563            "right",
564            "rollup",
565            "rows",
566            "select",
567            "set",
568            "some",
569            "struct",
570            "tablesample",
571            "then",
572            "to",
573            "treat",
574            "true",
575            "unbounded",
576            "union",
577            "unnest",
578            "using",
579            "when",
580            "where",
581            "window",
582            "with",
583            "within",
584        }
585
586        def attimezone_sql(self, expression: exp.AtTimeZone) -> str:
587            parent = expression.parent
588
589            # BigQuery allows CAST(.. AS {STRING|TIMESTAMP} [FORMAT <fmt> [AT TIME ZONE <tz>]]).
590            # Only the TIMESTAMP one should use the below conversion, when AT TIME ZONE is included.
591            if not isinstance(parent, exp.Cast) or not parent.to.is_type("text"):
592                return self.func(
593                    "TIMESTAMP", self.func("DATETIME", expression.this, expression.args.get("zone"))
594                )
595
596            return super().attimezone_sql(expression)
597
598        def trycast_sql(self, expression: exp.TryCast) -> str:
599            return self.cast_sql(expression, safe_prefix="SAFE_")
600
601        def cte_sql(self, expression: exp.CTE) -> str:
602            if expression.alias_column_names:
603                self.unsupported("Column names in CTE definition are not supported.")
604            return super().cte_sql(expression)
605
606        def array_sql(self, expression: exp.Array) -> str:
607            first_arg = seq_get(expression.expressions, 0)
608            if isinstance(first_arg, exp.Subqueryable):
609                return f"ARRAY{self.wrap(self.sql(first_arg))}"
610
611            return inline_array_sql(self, expression)
612
613        def transaction_sql(self, *_) -> str:
614            return "BEGIN TRANSACTION"
615
616        def commit_sql(self, *_) -> str:
617            return "COMMIT TRANSACTION"
618
619        def rollback_sql(self, *_) -> str:
620            return "ROLLBACK TRANSACTION"
621
622        def in_unnest_op(self, expression: exp.Unnest) -> str:
623            return self.sql(expression)
624
625        def except_op(self, expression: exp.Except) -> str:
626            if not expression.args.get("distinct", False):
627                self.unsupported("EXCEPT without DISTINCT is not supported in BigQuery")
628            return f"EXCEPT{' DISTINCT' if expression.args.get('distinct') else ' ALL'}"
629
630        def intersect_op(self, expression: exp.Intersect) -> str:
631            if not expression.args.get("distinct", False):
632                self.unsupported("INTERSECT without DISTINCT is not supported in BigQuery")
633            return f"INTERSECT{' DISTINCT' if expression.args.get('distinct') else ' ALL'}"
634
635        def with_properties(self, properties: exp.Properties) -> str:
636            return self.properties(properties, prefix=self.seg("OPTIONS"))