sqlglot.dialects.tsql

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