diff options
Diffstat (limited to 'sqlglot/dialects/drill.py')
| -rw-r--r-- | sqlglot/dialects/drill.py | 33 |
1 files changed, 14 insertions, 19 deletions
diff --git a/sqlglot/dialects/drill.py b/sqlglot/dialects/drill.py index d0a0251..1730eaf 100644 --- a/sqlglot/dialects/drill.py +++ b/sqlglot/dialects/drill.py @@ -1,6 +1,7 @@ from __future__ import annotations import re +import typing as t from sqlglot import exp, generator, parser, tokens from sqlglot.dialects.dialect import ( @@ -16,35 +17,29 @@ from sqlglot.dialects.dialect import ( ) -def _to_timestamp(args): - # TO_TIMESTAMP accepts either a single double argument or (text, text) - if len(args) == 1 and args[0].is_number: - return exp.UnixToTime.from_arg_list(args) - return format_time_lambda(exp.StrToTime, "drill")(args) - - -def _str_to_time_sql(self, expression): +def _str_to_time_sql(self: generator.Generator, expression: exp.TsOrDsToDate) -> str: return f"STRPTIME({self.sql(expression, 'this')}, {self.format_time(expression)})" -def _ts_or_ds_to_date_sql(self, expression): +def _ts_or_ds_to_date_sql(self: generator.Generator, expression: exp.TsOrDsToDate) -> str: time_format = self.format_time(expression) if time_format and time_format not in (Drill.time_format, Drill.date_format): return f"CAST({_str_to_time_sql(self, expression)} AS DATE)" return f"CAST({self.sql(expression, 'this')} AS DATE)" -def _date_add_sql(kind): - def func(self, expression): +def _date_add_sql(kind: str) -> t.Callable[[generator.Generator, exp.DateAdd | exp.DateSub], str]: + def func(self: generator.Generator, expression: exp.DateAdd | exp.DateSub) -> str: this = self.sql(expression, "this") - unit = expression.text("unit").upper() or "DAY" - expression = self.sql(expression, "expression") - return f"DATE_{kind}({this}, INTERVAL '{expression}' {unit})" + unit = exp.Var(this=expression.text("unit").upper() or "DAY") + return ( + f"DATE_{kind}({this}, {self.sql(exp.Interval(this=expression.expression, unit=unit))})" + ) return func -def if_sql(self, expression): +def if_sql(self: generator.Generator, expression: exp.If) -> str: """ Drill requires backticks around certain SQL reserved words, IF being one of them, This function adds the backticks around the keyword IF. @@ -61,7 +56,7 @@ def if_sql(self, expression): return f"`IF`({expressions})" -def _str_to_date(self, expression): +def _str_to_date(self: generator.Generator, expression: exp.StrToDate) -> str: this = self.sql(expression, "this") time_format = self.format_time(expression) if time_format == Drill.date_format: @@ -111,7 +106,7 @@ class Drill(Dialect): class Tokenizer(tokens.Tokenizer): QUOTES = ["'"] IDENTIFIERS = ["`"] - ESCAPES = ["\\"] + STRING_ESCAPES = ["\\"] ENCODE = "utf-8" class Parser(parser.Parser): @@ -168,10 +163,10 @@ class Drill(Dialect): exp.TimeToStr: lambda self, e: f"TO_CHAR({self.sql(e, 'this')}, {self.format_time(e)})", exp.TimeToUnix: rename_func("UNIX_TIMESTAMP"), exp.TryCast: no_trycast_sql, - exp.TsOrDsAdd: lambda self, e: f"DATE_ADD(CAST({self.sql(e, 'this')} AS DATE), INTERVAL '{self.sql(e, 'expression')}' DAY)", + exp.TsOrDsAdd: lambda self, e: f"DATE_ADD(CAST({self.sql(e, 'this')} AS DATE), {self.sql(exp.Interval(this=e.expression, unit=exp.Var(this='DAY')))})", exp.TsOrDsToDate: _ts_or_ds_to_date_sql, exp.TsOrDiToDi: lambda self, e: f"CAST(SUBSTR(REPLACE(CAST({self.sql(e, 'this')} AS VARCHAR), '-', ''), 1, 8) AS INT)", } - def normalize_func(self, name): + def normalize_func(self, name: str) -> str: return name if re.match(exp.SAFE_IDENTIFIER_RE, name) else f"`{name}`" |