125 lines
4.4 KiB
Python
125 lines
4.4 KiB
Python
from sqlglot import exp
|
|
from sqlglot.dialects.dialect import (
|
|
Dialect,
|
|
inline_array_sql,
|
|
no_ilike_sql,
|
|
rename_func,
|
|
)
|
|
from sqlglot.generator import Generator
|
|
from sqlglot.helper import list_get
|
|
from sqlglot.parser import Parser
|
|
from sqlglot.tokens import Tokenizer, TokenType
|
|
|
|
|
|
def _date_add(expression_class):
|
|
def func(args):
|
|
interval = list_get(args, 1)
|
|
return expression_class(
|
|
this=list_get(args, 0),
|
|
expression=interval.this,
|
|
unit=interval.args.get("unit"),
|
|
)
|
|
|
|
return func
|
|
|
|
|
|
def _date_add_sql(data_type, kind):
|
|
def func(self, expression):
|
|
this = self.sql(expression, "this")
|
|
unit = self.sql(expression, "unit") or "'day'"
|
|
expression = self.sql(expression, "expression")
|
|
return f"{data_type}_{kind}({this}, INTERVAL {expression} {unit})"
|
|
|
|
return func
|
|
|
|
|
|
class BigQuery(Dialect):
|
|
unnest_column_only = True
|
|
|
|
class Tokenizer(Tokenizer):
|
|
QUOTES = [
|
|
(prefix + quote, quote) if prefix else quote
|
|
for quote in ["'", '"', '"""', "'''"]
|
|
for prefix in ["", "r", "R"]
|
|
]
|
|
IDENTIFIERS = ["`"]
|
|
ESCAPE = "\\"
|
|
HEX_STRINGS = [("0x", ""), ("0X", "")]
|
|
|
|
KEYWORDS = {
|
|
**Tokenizer.KEYWORDS,
|
|
"CURRENT_DATETIME": TokenType.CURRENT_DATETIME,
|
|
"CURRENT_TIME": TokenType.CURRENT_TIME,
|
|
"GEOGRAPHY": TokenType.GEOGRAPHY,
|
|
"INT64": TokenType.BIGINT,
|
|
"FLOAT64": TokenType.DOUBLE,
|
|
"QUALIFY": TokenType.QUALIFY,
|
|
"UNKNOWN": TokenType.NULL,
|
|
"WINDOW": TokenType.WINDOW,
|
|
}
|
|
|
|
class Parser(Parser):
|
|
FUNCTIONS = {
|
|
**Parser.FUNCTIONS,
|
|
"DATE_ADD": _date_add(exp.DateAdd),
|
|
"DATETIME_ADD": _date_add(exp.DatetimeAdd),
|
|
"TIME_ADD": _date_add(exp.TimeAdd),
|
|
"TIMESTAMP_ADD": _date_add(exp.TimestampAdd),
|
|
"DATE_SUB": _date_add(exp.DateSub),
|
|
"DATETIME_SUB": _date_add(exp.DatetimeSub),
|
|
"TIME_SUB": _date_add(exp.TimeSub),
|
|
"TIMESTAMP_SUB": _date_add(exp.TimestampSub),
|
|
}
|
|
|
|
NO_PAREN_FUNCTIONS = {
|
|
**Parser.NO_PAREN_FUNCTIONS,
|
|
TokenType.CURRENT_DATETIME: exp.CurrentDatetime,
|
|
TokenType.CURRENT_TIME: exp.CurrentTime,
|
|
}
|
|
|
|
class Generator(Generator):
|
|
TRANSFORMS = {
|
|
exp.Array: inline_array_sql,
|
|
exp.ArraySize: rename_func("ARRAY_LENGTH"),
|
|
exp.DateAdd: _date_add_sql("DATE", "ADD"),
|
|
exp.DateSub: _date_add_sql("DATE", "SUB"),
|
|
exp.DatetimeAdd: _date_add_sql("DATETIME", "ADD"),
|
|
exp.DatetimeSub: _date_add_sql("DATETIME", "SUB"),
|
|
exp.ILike: no_ilike_sql,
|
|
exp.TimeAdd: _date_add_sql("TIME", "ADD"),
|
|
exp.TimeSub: _date_add_sql("TIME", "SUB"),
|
|
exp.TimestampAdd: _date_add_sql("TIMESTAMP", "ADD"),
|
|
exp.TimestampSub: _date_add_sql("TIMESTAMP", "SUB"),
|
|
exp.VariancePop: rename_func("VAR_POP"),
|
|
}
|
|
|
|
TYPE_MAPPING = {
|
|
**Generator.TYPE_MAPPING,
|
|
exp.DataType.Type.TINYINT: "INT64",
|
|
exp.DataType.Type.SMALLINT: "INT64",
|
|
exp.DataType.Type.INT: "INT64",
|
|
exp.DataType.Type.BIGINT: "INT64",
|
|
exp.DataType.Type.DECIMAL: "NUMERIC",
|
|
exp.DataType.Type.FLOAT: "FLOAT64",
|
|
exp.DataType.Type.DOUBLE: "FLOAT64",
|
|
exp.DataType.Type.BOOLEAN: "BOOL",
|
|
exp.DataType.Type.TEXT: "STRING",
|
|
exp.DataType.Type.VARCHAR: "STRING",
|
|
exp.DataType.Type.NVARCHAR: "STRING",
|
|
}
|
|
|
|
def in_unnest_op(self, unnest):
|
|
return self.sql(unnest)
|
|
|
|
def union_op(self, expression):
|
|
return f"UNION{' DISTINCT' if expression.args.get('distinct') else ' ALL'}"
|
|
|
|
def except_op(self, expression):
|
|
if not expression.args.get("distinct", False):
|
|
self.unsupported("EXCEPT without DISTINCT is not supported in BigQuery")
|
|
return f"EXCEPT{' DISTINCT' if expression.args.get('distinct') else ' ALL'}"
|
|
|
|
def intersect_op(self, expression):
|
|
if not expression.args.get("distinct", False):
|
|
self.unsupported("INTERSECT without DISTINCT is not supported in BigQuery")
|
|
return f"INTERSECT{' DISTINCT' if expression.args.get('distinct') else ' ALL'}"
|