sqlglot.dialects.tsql

  1from __future__ import annotations
  2
  3import re
  4import typing as t
  5
  6from sqlglot import exp, generator, parser, tokens
  7from sqlglot.dialects.dialect import (
  8    Dialect,
  9    max_or_greatest,
 10    min_or_least,
 11    parse_date_delta,
 12    rename_func,
 13)
 14from sqlglot.expressions import DataType
 15from sqlglot.helper import seq_get
 16from sqlglot.time import format_time
 17from sqlglot.tokens import TokenType
 18
 19FULL_FORMAT_TIME_MAPPING = {
 20    "weekday": "%A",
 21    "dw": "%A",
 22    "w": "%A",
 23    "month": "%B",
 24    "mm": "%B",
 25    "m": "%B",
 26}
 27
 28DATE_DELTA_INTERVAL = {
 29    "year": "year",
 30    "yyyy": "year",
 31    "yy": "year",
 32    "quarter": "quarter",
 33    "qq": "quarter",
 34    "q": "quarter",
 35    "month": "month",
 36    "mm": "month",
 37    "m": "month",
 38    "week": "week",
 39    "ww": "week",
 40    "wk": "week",
 41    "day": "day",
 42    "dd": "day",
 43    "d": "day",
 44}
 45
 46
 47DATE_FMT_RE = re.compile("([dD]{1,2})|([mM]{1,2})|([yY]{1,4})|([hH]{1,2})|([sS]{1,2})")
 48
 49# N = Numeric, C=Currency
 50TRANSPILE_SAFE_NUMBER_FMT = {"N", "C"}
 51
 52
 53def _format_time_lambda(exp_class, full_format_mapping=None, default=None):
 54    def _format_time(args):
 55        return exp_class(
 56            this=seq_get(args, 1),
 57            format=exp.Literal.string(
 58                format_time(
 59                    seq_get(args, 0).name or (TSQL.time_format if default is True else default),
 60                    {**TSQL.time_mapping, **FULL_FORMAT_TIME_MAPPING}
 61                    if full_format_mapping
 62                    else TSQL.time_mapping,
 63                )
 64            ),
 65        )
 66
 67    return _format_time
 68
 69
 70def _parse_format(args):
 71    fmt = seq_get(args, 1)
 72    number_fmt = fmt.name in TRANSPILE_SAFE_NUMBER_FMT or not DATE_FMT_RE.search(fmt.this)
 73    if number_fmt:
 74        return exp.NumberToStr(this=seq_get(args, 0), format=fmt)
 75    return exp.TimeToStr(
 76        this=seq_get(args, 0),
 77        format=exp.Literal.string(
 78            format_time(fmt.name, TSQL.format_time_mapping)
 79            if len(fmt.name) == 1
 80            else format_time(fmt.name, TSQL.time_mapping)
 81        ),
 82    )
 83
 84
 85def _parse_eomonth(args):
 86    date = seq_get(args, 0)
 87    month_lag = seq_get(args, 1)
 88    unit = DATE_DELTA_INTERVAL.get("month")
 89
 90    if month_lag is None:
 91        return exp.LastDateOfMonth(this=date)
 92
 93    # Remove month lag argument in parser as its compared with the number of arguments of the resulting class
 94    args.remove(month_lag)
 95
 96    return exp.LastDateOfMonth(this=exp.DateAdd(this=date, expression=month_lag, unit=unit))
 97
 98
 99def _parse_hashbytes(args):
100    kind, data = args
101    kind = kind.name.upper() if kind.is_string else ""
102
103    if kind == "MD5":
104        args.pop(0)
105        return exp.MD5(this=data)
106    if kind in ("SHA", "SHA1"):
107        args.pop(0)
108        return exp.SHA(this=data)
109    if kind == "SHA2_256":
110        return exp.SHA2(this=data, length=exp.Literal.number(256))
111    if kind == "SHA2_512":
112        return exp.SHA2(this=data, length=exp.Literal.number(512))
113    return exp.func("HASHBYTES", *args)
114
115
116def generate_date_delta_with_unit_sql(self, e):
117    func = "DATEADD" if isinstance(e, exp.DateAdd) else "DATEDIFF"
118    return self.func(func, e.text("unit"), e.expression, e.this)
119
120
121def _format_sql(self, e):
122    fmt = (
123        e.args["format"]
124        if isinstance(e, exp.NumberToStr)
125        else exp.Literal.string(format_time(e.text("format"), TSQL.inverse_time_mapping))
126    )
127    return self.func("FORMAT", e.this, fmt)
128
129
130def _string_agg_sql(self, e):
131    e = e.copy()
132
133    this = e.this
134    distinct = e.find(exp.Distinct)
135    if distinct:
136        # exp.Distinct can appear below an exp.Order or an exp.GroupConcat expression
137        self.unsupported("T-SQL STRING_AGG doesn't support DISTINCT.")
138        this = distinct.pop().expressions[0]
139
140    order = ""
141    if isinstance(e.this, exp.Order):
142        if e.this.this:
143            this = e.this.this.pop()
144        order = f" WITHIN GROUP ({self.sql(e.this)[1:]})"  # Order has a leading space
145
146    separator = e.args.get("separator") or exp.Literal.string(",")
147    return f"STRING_AGG({self.format_args(this, separator)}){order}"
148
149
150class TSQL(Dialect):
151    null_ordering = "nulls_are_small"
152    time_format = "'yyyy-mm-dd hh:mm:ss'"
153
154    time_mapping = {
155        "year": "%Y",
156        "qq": "%q",
157        "q": "%q",
158        "quarter": "%q",
159        "dayofyear": "%j",
160        "day": "%d",
161        "dy": "%d",
162        "y": "%Y",
163        "week": "%W",
164        "ww": "%W",
165        "wk": "%W",
166        "hour": "%h",
167        "hh": "%I",
168        "minute": "%M",
169        "mi": "%M",
170        "n": "%M",
171        "second": "%S",
172        "ss": "%S",
173        "s": "%-S",
174        "millisecond": "%f",
175        "ms": "%f",
176        "weekday": "%W",
177        "dw": "%W",
178        "month": "%m",
179        "mm": "%M",
180        "m": "%-M",
181        "Y": "%Y",
182        "YYYY": "%Y",
183        "YY": "%y",
184        "MMMM": "%B",
185        "MMM": "%b",
186        "MM": "%m",
187        "M": "%-m",
188        "dd": "%d",
189        "d": "%-d",
190        "HH": "%H",
191        "H": "%-H",
192        "h": "%-I",
193        "S": "%f",
194        "yyyy": "%Y",
195        "yy": "%y",
196    }
197
198    convert_format_mapping = {
199        "0": "%b %d %Y %-I:%M%p",
200        "1": "%m/%d/%y",
201        "2": "%y.%m.%d",
202        "3": "%d/%m/%y",
203        "4": "%d.%m.%y",
204        "5": "%d-%m-%y",
205        "6": "%d %b %y",
206        "7": "%b %d, %y",
207        "8": "%H:%M:%S",
208        "9": "%b %d %Y %-I:%M:%S:%f%p",
209        "10": "mm-dd-yy",
210        "11": "yy/mm/dd",
211        "12": "yymmdd",
212        "13": "%d %b %Y %H:%M:ss:%f",
213        "14": "%H:%M:%S:%f",
214        "20": "%Y-%m-%d %H:%M:%S",
215        "21": "%Y-%m-%d %H:%M:%S.%f",
216        "22": "%m/%d/%y %-I:%M:%S %p",
217        "23": "%Y-%m-%d",
218        "24": "%H:%M:%S",
219        "25": "%Y-%m-%d %H:%M:%S.%f",
220        "100": "%b %d %Y %-I:%M%p",
221        "101": "%m/%d/%Y",
222        "102": "%Y.%m.%d",
223        "103": "%d/%m/%Y",
224        "104": "%d.%m.%Y",
225        "105": "%d-%m-%Y",
226        "106": "%d %b %Y",
227        "107": "%b %d, %Y",
228        "108": "%H:%M:%S",
229        "109": "%b %d %Y %-I:%M:%S:%f%p",
230        "110": "%m-%d-%Y",
231        "111": "%Y/%m/%d",
232        "112": "%Y%m%d",
233        "113": "%d %b %Y %H:%M:%S:%f",
234        "114": "%H:%M:%S:%f",
235        "120": "%Y-%m-%d %H:%M:%S",
236        "121": "%Y-%m-%d %H:%M:%S.%f",
237    }
238    # not sure if complete
239    format_time_mapping = {
240        "y": "%B %Y",
241        "d": "%m/%d/%Y",
242        "H": "%-H",
243        "h": "%-I",
244        "s": "%Y-%m-%d %H:%M:%S",
245        "D": "%A,%B,%Y",
246        "f": "%A,%B,%Y %-I:%M %p",
247        "F": "%A,%B,%Y %-I:%M:%S %p",
248        "g": "%m/%d/%Y %-I:%M %p",
249        "G": "%m/%d/%Y %-I:%M:%S %p",
250        "M": "%B %-d",
251        "m": "%B %-d",
252        "O": "%Y-%m-%dT%H:%M:%S",
253        "u": "%Y-%M-%D %H:%M:%S%z",
254        "U": "%A, %B %D, %Y %H:%M:%S%z",
255        "T": "%-I:%M:%S %p",
256        "t": "%-I:%M",
257        "Y": "%a %Y",
258    }
259
260    class Tokenizer(tokens.Tokenizer):
261        IDENTIFIERS = ['"', ("[", "]")]
262
263        QUOTES = ["'", '"']
264
265        KEYWORDS = {
266            **tokens.Tokenizer.KEYWORDS,
267            "DATETIME2": TokenType.DATETIME,
268            "DATETIMEOFFSET": TokenType.TIMESTAMPTZ,
269            "DECLARE": TokenType.COMMAND,
270            "IMAGE": TokenType.IMAGE,
271            "MONEY": TokenType.MONEY,
272            "NTEXT": TokenType.TEXT,
273            "NVARCHAR(MAX)": TokenType.TEXT,
274            "PRINT": TokenType.COMMAND,
275            "PROC": TokenType.PROCEDURE,
276            "REAL": TokenType.FLOAT,
277            "ROWVERSION": TokenType.ROWVERSION,
278            "SMALLDATETIME": TokenType.DATETIME,
279            "SMALLMONEY": TokenType.SMALLMONEY,
280            "SQL_VARIANT": TokenType.VARIANT,
281            "TIME": TokenType.TIMESTAMP,
282            "TOP": TokenType.TOP,
283            "UNIQUEIDENTIFIER": TokenType.UNIQUEIDENTIFIER,
284            "VARCHAR(MAX)": TokenType.TEXT,
285            "XML": TokenType.XML,
286            "SYSTEM_USER": TokenType.CURRENT_USER,
287        }
288
289        # TSQL allows @, # to appear as a variable/identifier prefix
290        SINGLE_TOKENS = tokens.Tokenizer.SINGLE_TOKENS.copy()
291        SINGLE_TOKENS.pop("#")
292
293    class Parser(parser.Parser):
294        FUNCTIONS = {
295            **parser.Parser.FUNCTIONS,  # type: ignore
296            "CHARINDEX": lambda args: exp.StrPosition(
297                this=seq_get(args, 1),
298                substr=seq_get(args, 0),
299                position=seq_get(args, 2),
300            ),
301            "DATEADD": parse_date_delta(exp.DateAdd, unit_mapping=DATE_DELTA_INTERVAL),
302            "DATEDIFF": parse_date_delta(exp.DateDiff, unit_mapping=DATE_DELTA_INTERVAL),
303            "DATENAME": _format_time_lambda(exp.TimeToStr, full_format_mapping=True),
304            "DATEPART": _format_time_lambda(exp.TimeToStr),
305            "EOMONTH": _parse_eomonth,
306            "FORMAT": _parse_format,
307            "GETDATE": exp.CurrentTimestamp.from_arg_list,
308            "HASHBYTES": _parse_hashbytes,
309            "IIF": exp.If.from_arg_list,
310            "ISNULL": exp.Coalesce.from_arg_list,
311            "JSON_VALUE": exp.JSONExtractScalar.from_arg_list,
312            "LEN": exp.Length.from_arg_list,
313            "REPLICATE": exp.Repeat.from_arg_list,
314            "SQUARE": lambda args: exp.Pow(this=seq_get(args, 0), expression=exp.Literal.number(2)),
315            "SYSDATETIME": exp.CurrentTimestamp.from_arg_list,
316            "SUSER_NAME": exp.CurrentUser.from_arg_list,
317            "SUSER_SNAME": exp.CurrentUser.from_arg_list,
318            "SYSTEM_USER": exp.CurrentUser.from_arg_list,
319        }
320
321        JOIN_HINTS = {
322            "LOOP",
323            "HASH",
324            "MERGE",
325            "REMOTE",
326        }
327
328        VAR_LENGTH_DATATYPES = {
329            DataType.Type.NVARCHAR,
330            DataType.Type.VARCHAR,
331            DataType.Type.CHAR,
332            DataType.Type.NCHAR,
333        }
334
335        RETURNS_TABLE_TOKENS = parser.Parser.ID_VAR_TOKENS - {  # type: ignore
336            TokenType.TABLE,
337            *parser.Parser.TYPE_TOKENS,  # type: ignore
338        }
339
340        STATEMENT_PARSERS = {
341            **parser.Parser.STATEMENT_PARSERS,  # type: ignore
342            TokenType.END: lambda self: self._parse_command(),
343        }
344
345        LOG_BASE_FIRST = False
346        LOG_DEFAULTS_TO_LN = True
347
348        def _parse_system_time(self) -> t.Optional[exp.Expression]:
349            if not self._match_text_seq("FOR", "SYSTEM_TIME"):
350                return None
351
352            if self._match_text_seq("AS", "OF"):
353                system_time = self.expression(
354                    exp.SystemTime, this=self._parse_bitwise(), kind="AS OF"
355                )
356            elif self._match_set((TokenType.FROM, TokenType.BETWEEN)):
357                kind = self._prev.text
358                this = self._parse_bitwise()
359                self._match_texts(("TO", "AND"))
360                expression = self._parse_bitwise()
361                system_time = self.expression(
362                    exp.SystemTime, this=this, expression=expression, kind=kind
363                )
364            elif self._match_text_seq("CONTAINED", "IN"):
365                args = self._parse_wrapped_csv(self._parse_bitwise)
366                system_time = self.expression(
367                    exp.SystemTime,
368                    this=seq_get(args, 0),
369                    expression=seq_get(args, 1),
370                    kind="CONTAINED IN",
371                )
372            elif self._match(TokenType.ALL):
373                system_time = self.expression(exp.SystemTime, kind="ALL")
374            else:
375                system_time = None
376                self.raise_error("Unable to parse FOR SYSTEM_TIME clause")
377
378            return system_time
379
380        def _parse_table_parts(self, schema: bool = False) -> exp.Expression:
381            table = super()._parse_table_parts(schema=schema)
382            table.set("system_time", self._parse_system_time())
383            return table
384
385        def _parse_returns(self) -> exp.Expression:
386            table = self._parse_id_var(any_token=False, tokens=self.RETURNS_TABLE_TOKENS)
387            returns = super()._parse_returns()
388            returns.set("table", table)
389            return returns
390
391        def _parse_convert(self, strict: bool) -> t.Optional[exp.Expression]:
392            to = self._parse_types()
393            self._match(TokenType.COMMA)
394            this = self._parse_conjunction()
395
396            if not to or not this:
397                return None
398
399            # Retrieve length of datatype and override to default if not specified
400            if seq_get(to.expressions, 0) is None and to.this in self.VAR_LENGTH_DATATYPES:
401                to = exp.DataType.build(to.this, expressions=[exp.Literal.number(30)], nested=False)
402
403            # Check whether a conversion with format is applicable
404            if self._match(TokenType.COMMA):
405                format_val = self._parse_number()
406                format_val_name = format_val.name if format_val else ""
407
408                if format_val_name not in TSQL.convert_format_mapping:
409                    raise ValueError(
410                        f"CONVERT function at T-SQL does not support format style {format_val_name}"
411                    )
412
413                format_norm = exp.Literal.string(TSQL.convert_format_mapping[format_val_name])
414
415                # Check whether the convert entails a string to date format
416                if to.this == DataType.Type.DATE:
417                    return self.expression(exp.StrToDate, this=this, format=format_norm)
418                # Check whether the convert entails a string to datetime format
419                elif to.this == DataType.Type.DATETIME:
420                    return self.expression(exp.StrToTime, this=this, format=format_norm)
421                # Check whether the convert entails a date to string format
422                elif to.this in self.VAR_LENGTH_DATATYPES:
423                    return self.expression(
424                        exp.Cast if strict else exp.TryCast,
425                        to=to,
426                        this=self.expression(exp.TimeToStr, this=this, format=format_norm),
427                    )
428                elif to.this == DataType.Type.TEXT:
429                    return self.expression(exp.TimeToStr, this=this, format=format_norm)
430
431            # Entails a simple cast without any format requirement
432            return self.expression(exp.Cast if strict else exp.TryCast, this=this, to=to)
433
434        def _parse_user_defined_function(
435            self, kind: t.Optional[TokenType] = None
436        ) -> t.Optional[exp.Expression]:
437            this = super()._parse_user_defined_function(kind=kind)
438
439            if (
440                kind == TokenType.FUNCTION
441                or isinstance(this, exp.UserDefinedFunction)
442                or self._match(TokenType.ALIAS, advance=False)
443            ):
444                return this
445
446            expressions = self._parse_csv(self._parse_function_parameter)
447            return self.expression(exp.UserDefinedFunction, this=this, expressions=expressions)
448
449    class Generator(generator.Generator):
450        LOCKING_READS_SUPPORTED = True
451
452        TYPE_MAPPING = {
453            **generator.Generator.TYPE_MAPPING,  # type: ignore
454            exp.DataType.Type.INT: "INTEGER",
455            exp.DataType.Type.DECIMAL: "NUMERIC",
456            exp.DataType.Type.DATETIME: "DATETIME2",
457            exp.DataType.Type.VARIANT: "SQL_VARIANT",
458        }
459
460        TRANSFORMS = {
461            **generator.Generator.TRANSFORMS,  # type: ignore
462            exp.DateAdd: generate_date_delta_with_unit_sql,
463            exp.DateDiff: generate_date_delta_with_unit_sql,
464            exp.CurrentDate: rename_func("GETDATE"),
465            exp.CurrentTimestamp: rename_func("GETDATE"),
466            exp.If: rename_func("IIF"),
467            exp.NumberToStr: _format_sql,
468            exp.TimeToStr: _format_sql,
469            exp.GroupConcat: _string_agg_sql,
470            exp.Max: max_or_greatest,
471            exp.MD5: lambda self, e: self.func("HASHBYTES", exp.Literal.string("MD5"), e.this),
472            exp.Min: min_or_least,
473            exp.SHA: lambda self, e: self.func("HASHBYTES", exp.Literal.string("SHA1"), e.this),
474            exp.SHA2: lambda self, e: self.func(
475                "HASHBYTES", exp.Literal.string(f"SHA2_{e.args.get('length', 256)}"), e.this
476            ),
477        }
478
479        TRANSFORMS.pop(exp.ReturnsProperty)
480
481        PROPERTIES_LOCATION = {
482            **generator.Generator.PROPERTIES_LOCATION,  # type: ignore
483            exp.VolatileProperty: exp.Properties.Location.UNSUPPORTED,
484        }
485
486        LIMIT_FETCH = "FETCH"
487
488        def offset_sql(self, expression: exp.Offset) -> str:
489            return f"{super().offset_sql(expression)} ROWS"
490
491        def systemtime_sql(self, expression: exp.SystemTime) -> str:
492            kind = expression.args["kind"]
493            if kind == "ALL":
494                return "FOR SYSTEM_TIME ALL"
495
496            start = self.sql(expression, "this")
497            if kind == "AS OF":
498                return f"FOR SYSTEM_TIME AS OF {start}"
499
500            end = self.sql(expression, "expression")
501            if kind == "FROM":
502                return f"FOR SYSTEM_TIME FROM {start} TO {end}"
503            if kind == "BETWEEN":
504                return f"FOR SYSTEM_TIME BETWEEN {start} AND {end}"
505
506            return f"FOR SYSTEM_TIME CONTAINED IN ({start}, {end})"
507
508        def returnsproperty_sql(self, expression: exp.ReturnsProperty) -> str:
509            table = expression.args.get("table")
510            table = f"{table} " if table else ""
511            return f"RETURNS {table}{self.sql(expression, 'this')}"