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