diff options
Diffstat (limited to 'sqlglot/dialects/bigquery.py')
| -rw-r--r-- | sqlglot/dialects/bigquery.py | 30 |
1 files changed, 21 insertions, 9 deletions
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py index 6a19b46..7fd9e35 100644 --- a/sqlglot/dialects/bigquery.py +++ b/sqlglot/dialects/bigquery.py @@ -2,6 +2,7 @@ from __future__ import annotations +import re import typing as t from sqlglot import exp, generator, parser, tokens, transforms @@ -31,13 +32,6 @@ def _date_add(expression_class: t.Type[E]) -> t.Callable[[t.Sequence], E]: return func -def _date_trunc(args: t.Sequence) -> exp.Expression: - unit = seq_get(args, 1) - if isinstance(unit, exp.Column): - unit = exp.Var(this=unit.name) - return exp.DateTrunc(this=seq_get(args, 0), expression=unit) - - def _date_add_sql( data_type: str, kind: str ) -> t.Callable[[generator.Generator, exp.Expression], str]: @@ -158,11 +152,23 @@ class BigQuery(Dialect): class Parser(parser.Parser): FUNCTIONS = { **parser.Parser.FUNCTIONS, # type: ignore - "DATE_TRUNC": _date_trunc, + "DATE_TRUNC": lambda args: exp.DateTrunc( + unit=exp.Literal.string(seq_get(args, 1).name), # type: ignore + this=seq_get(args, 0), + ), "DATE_ADD": _date_add(exp.DateAdd), "DATETIME_ADD": _date_add(exp.DatetimeAdd), "DIV": lambda args: exp.IntDiv(this=seq_get(args, 0), expression=seq_get(args, 1)), "REGEXP_CONTAINS": exp.RegexpLike.from_arg_list, + "REGEXP_EXTRACT": lambda args: exp.RegexpExtract( + this=seq_get(args, 0), + expression=seq_get(args, 1), + position=seq_get(args, 2), + occurrence=seq_get(args, 3), + group=exp.Literal.number(1) + if re.compile(str(seq_get(args, 1))).groups == 1 + else None, + ), "TIME_ADD": _date_add(exp.TimeAdd), "TIMESTAMP_ADD": _date_add(exp.TimestampAdd), "DATE_SUB": _date_add(exp.DateSub), @@ -214,6 +220,7 @@ class BigQuery(Dialect): exp.DatetimeSub: _date_add_sql("DATETIME", "SUB"), exp.DateDiff: lambda self, e: f"DATE_DIFF({self.sql(e, 'this')}, {self.sql(e, 'expression')}, {self.sql(e.args.get('unit', 'DAY'))})", exp.DateStrToDate: datestrtodate_sql, + exp.DateTrunc: lambda self, e: self.func("DATE_TRUNC", e.this, e.text("unit")), exp.GroupConcat: rename_func("STRING_AGG"), exp.ILike: no_ilike_sql, exp.IntDiv: rename_func("DIV"), @@ -226,11 +233,12 @@ class BigQuery(Dialect): exp.TimestampAdd: _date_add_sql("TIMESTAMP", "ADD"), exp.TimestampSub: _date_add_sql("TIMESTAMP", "SUB"), exp.TimeStrToTime: timestrtotime_sql, + exp.PartitionedByProperty: lambda self, e: f"PARTITION BY {self.sql(e, 'this')}", exp.VariancePop: rename_func("VAR_POP"), exp.Values: _derived_table_values_to_unnest, exp.ReturnsProperty: _returnsproperty_sql, exp.Create: _create_sql, - exp.Trim: lambda self, e: f"TRIM({self.format_args(e.this, e.expression)})", + exp.Trim: lambda self, e: self.func(f"TRIM", e.this, e.expression), exp.VolatilityProperty: lambda self, e: f"DETERMINISTIC" if e.name == "IMMUTABLE" else "NOT DETERMINISTIC", @@ -251,6 +259,10 @@ class BigQuery(Dialect): exp.DataType.Type.VARCHAR: "STRING", exp.DataType.Type.NVARCHAR: "STRING", } + PROPERTIES_LOCATION = { + **generator.Generator.PROPERTIES_LOCATION, # type: ignore + exp.PartitionedByProperty: exp.Properties.Location.POST_SCHEMA, + } EXPLICIT_UNION = True |