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