sqlglot.parsers.teradata

  1from __future__ import annotations
  2
  3import typing as t
  4
  5from sqlglot import exp, parser
  6from sqlglot.helper import seq_get
  7from sqlglot.tokens import TokenType
  8from sqlglot.trie import new_trie
  9
 10# Ranking and navigation functions don't take a frame and FIRST_VALUE/LAST_VALUE have their own default
 11NON_FRAMED_WINDOW_FUNCTIONS = (
 12    exp.CumeDist,
 13    exp.DenseRank,
 14    exp.FirstValue,
 15    exp.Lag,
 16    exp.LastValue,
 17    exp.Lead,
 18    exp.NthValue,
 19    exp.Ntile,
 20    exp.PercentRank,
 21    exp.Rank,
 22)
 23
 24
 25def is_window_aggregate(window: exp.Window) -> bool:
 26    return isinstance(window.this, exp.AggFunc) and not isinstance(
 27        window.this, NON_FRAMED_WINDOW_FUNCTIONS
 28    )
 29
 30
 31class TeradataParser(parser.Parser):
 32    TABLESAMPLE_CSV = True
 33    VALUES_FOLLOWED_BY_PAREN = False
 34
 35    TABLE_ALIAS_TOKENS = parser.Parser.TABLE_ALIAS_TOKENS | {
 36        TokenType.ANTI,
 37        TokenType.SEMI,
 38    }
 39
 40    CHARSET_TRANSLATORS: t.ClassVar = {
 41        "GRAPHIC_TO_KANJISJIS",
 42        "GRAPHIC_TO_LATIN",
 43        "GRAPHIC_TO_UNICODE",
 44        "GRAPHIC_TO_UNICODE_PadSpace",
 45        "KANJI1_KanjiEBCDIC_TO_UNICODE",
 46        "KANJI1_KanjiEUC_TO_UNICODE",
 47        "KANJI1_KANJISJIS_TO_UNICODE",
 48        "KANJI1_SBC_TO_UNICODE",
 49        "KANJISJIS_TO_GRAPHIC",
 50        "KANJISJIS_TO_LATIN",
 51        "KANJISJIS_TO_UNICODE",
 52        "LATIN_TO_GRAPHIC",
 53        "LATIN_TO_KANJISJIS",
 54        "LATIN_TO_UNICODE",
 55        "LOCALE_TO_UNICODE",
 56        "UNICODE_TO_GRAPHIC",
 57        "UNICODE_TO_GRAPHIC_PadGraphic",
 58        "UNICODE_TO_GRAPHIC_VarGraphic",
 59        "UNICODE_TO_KANJI1_KanjiEBCDIC",
 60        "UNICODE_TO_KANJI1_KanjiEUC",
 61        "UNICODE_TO_KANJI1_KANJISJIS",
 62        "UNICODE_TO_KANJI1_SBC",
 63        "UNICODE_TO_KANJISJIS",
 64        "UNICODE_TO_LATIN",
 65        "UNICODE_TO_LOCALE",
 66        "UNICODE_TO_UNICODE_FoldSpace",
 67        "UNICODE_TO_UNICODE_Fullwidth",
 68        "UNICODE_TO_UNICODE_Halfwidth",
 69        "UNICODE_TO_UNICODE_NFC",
 70        "UNICODE_TO_UNICODE_NFD",
 71        "UNICODE_TO_UNICODE_NFKC",
 72        "UNICODE_TO_UNICODE_NFKD",
 73    }
 74
 75    FUNC_TOKENS = (parser.Parser.FUNC_TOKENS - {TokenType.REPLACE}) | {TokenType.MOD}
 76
 77    STATEMENT_PARSERS = {
 78        **parser.Parser.STATEMENT_PARSERS,
 79        TokenType.DATABASE: lambda self: self.expression(
 80            exp.Use(this=self._parse_table(schema=False))
 81        ),
 82        TokenType.EXECUTE: lambda self: self._parse_execute(),
 83        TokenType.REPLACE: lambda self: self._parse_create(),
 84        TokenType.LOCK: lambda self: self._parse_locking_statement(),
 85    }
 86
 87    def _parse_execute(self) -> exp.Execute | exp.Command:
 88        start = self._prev
 89        this = self._parse_table_parts(schema=True)
 90        expressions = []
 91
 92        if self._match(TokenType.L_PAREN) and not self._match(TokenType.R_PAREN):
 93            while True:
 94                argument = self._parse_assignment()
 95                if argument is None:
 96                    # Do not discard omitted arguments or replace them with NULL.
 97                    return self._parse_as_command(start)
 98                expressions.append(argument)
 99                if not self._match(TokenType.COMMA):
100                    break
101            self._match_r_paren()
102
103        if self._curr:
104            # Preserve other EXECUTE forms, such as EXECUTE IMMEDIATE.
105            return self._parse_as_command(start)
106
107        return self.expression(exp.Execute(this=this, expressions=expressions))
108
109    def _parse_locking_statement(self) -> exp.LockingStatement:
110        # Reuse exp.LockingProperty parsing for the lock kind, type etc
111        locking_property = self._parse_locking()
112        wrapped_query = self._parse_select()
113
114        if not wrapped_query:
115            self.raise_error("Expected SELECT statement after LOCKING clause")
116
117        return self.expression(
118            exp.LockingStatement(this=locking_property, expression=wrapped_query)
119        )
120
121    SET_PARSERS = {
122        **parser.Parser.SET_PARSERS,
123        "QUERY_BAND": lambda self: self._parse_query_band(),
124    }
125
126    SET_TRIE = new_trie(key.split(" ") for key in SET_PARSERS)
127
128    FUNCTION_PARSERS = {
129        **parser.Parser.FUNCTION_PARSERS,
130        # https://docs.teradata.com/r/SQL-Functions-Operators-Exprs-and-Predicates/June-2017/Data-Type-Conversions/TRYCAST
131        "TRYCAST": parser.Parser.FUNCTION_PARSERS["TRY_CAST"],
132        "RANGE_N": lambda self: self._parse_rangen(),
133        "TRANSLATE": lambda self: self._parse_translate(),
134    }
135
136    FUNCTIONS = {
137        **parser.Parser.FUNCTIONS,
138        "CARDINALITY": exp.ArraySize.from_arg_list,
139        "RANDOM": lambda args: exp.Rand(lower=seq_get(args, 0), upper=seq_get(args, 1)),
140    }
141
142    EXPONENT = {
143        TokenType.DSTAR: exp.Pow,
144    }
145
146    def _parse_translate(self) -> exp.TranslateCharacters:
147        this = self._parse_assignment()
148        self._match(TokenType.USING)
149        self._match_texts(self.CHARSET_TRANSLATORS)
150
151        return self.expression(
152            exp.TranslateCharacters(
153                this=this,
154                expression=self._prev.text.upper(),
155                with_error=self._match_text_seq("WITH", "ERROR"),
156            )
157        )
158
159    # FROM before SET in Teradata UPDATE syntax
160    # https://docs.teradata.com/r/Enterprise_IntelliFlex_VMware/Teradata-VantageTM-SQL-Data-Manipulation-Language-17.20/Statement-Syntax/UPDATE/UPDATE-Syntax-Basic-Form-FROM-Clause
161    def _parse_update(self) -> exp.Update:
162        return self.expression(
163            exp.Update(
164                this=self._parse_table(alias_tokens=self.UPDATE_ALIAS_TOKENS),
165                from_=self._parse_from(joins=True),
166                expressions=self._match(TokenType.SET)
167                and self._parse_csv(self._parse_update_assignment),
168                where=self._parse_where(),
169            )
170        )
171
172    def _parse_table_sample(self, as_modifier: bool = False) -> exp.TableSample | None:
173        sample = super()._parse_table_sample(as_modifier=as_modifier)
174
175        # A single integer sample is a row count, while a fraction is a proportion of rows
176        # https://docs.teradata.com/r/Enterprise_IntelliFlex_VMware/SQL-Data-Manipulation-Language/SELECT-Statements/SAMPLE-Clause/SAMPLE-Clause-Syntax
177        if sample and len(sample.expressions) == 1 and sample.expressions[0].is_number:
178            num = sample.expressions[0].pop()
179            if num.is_int:
180                sample.set("size", num)
181            else:
182                percent = format(num.to_py().scaleb(2).normalize(), "f")
183                sample.set("percent", exp.Literal.number(percent))
184
185        return sample
186
187    def _parse_rangen(self):
188        this = self._parse_id_var()
189        self._match(TokenType.BETWEEN)
190
191        expressions = self._parse_csv(self._parse_assignment)
192        each = self._match_text_seq("EACH") and self._parse_assignment()
193
194        return self.expression(exp.RangeN(this=this, expressions=expressions, each=each))
195
196    def _parse_query_band(self) -> exp.QueryBand:
197        # Parse: SET QUERY_BAND = 'key=value;key2=value2;' FOR SESSION|TRANSACTION
198        # Also supports: SET QUERY_BAND = 'key=value;' UPDATE FOR SESSION|TRANSACTION
199        # Also supports: SET QUERY_BAND = NONE FOR SESSION|TRANSACTION
200        self._match(TokenType.EQ)
201
202        # Handle both string literals and NONE keyword
203        if self._match_text_seq("NONE"):
204            query_band_string: exp.Expr | None = exp.Var(this="NONE")
205        else:
206            query_band_string = self._parse_string()
207
208        update = self._match_text_seq("UPDATE")
209        self._match_text_seq("FOR")
210
211        # Handle scope - can be SESSION, TRANSACTION, VOLATILE, or SESSION VOLATILE
212        if self._match_text_seq("SESSION", "VOLATILE"):
213            scope = "SESSION VOLATILE"
214        elif self._match_texts(("SESSION", "TRANSACTION")):
215            scope = self._prev.text.upper()
216        else:
217            scope = None
218
219        return self.expression(exp.QueryBand(this=query_band_string, scope=scope, update=update))
220
221    def _parse_index_params(self) -> exp.IndexParameters:
222        this = super()._parse_index_params()
223
224        if this.args.get("on"):
225            this.set("on", None)
226            self._retreat(self._index - 2)
227        return this
228
229    def _parse_function(
230        self,
231        functions: dict[str, t.Callable] | None = None,
232        anonymous: bool = False,
233        optional_parens: bool = True,
234        any_token: bool = False,
235    ) -> exp.Expr | None:
236        # Column data attributes are not function arguments.
237        if (
238            self._next
239            and self._next.token_type == TokenType.L_PAREN
240            and self._index + 3 < len(self._tokens)
241            and self._tokens[self._index + 2].token_type in {TokenType.FORMAT, TokenType.VAR}
242            and self._tokens[self._index + 2].text.upper() in {"FORMAT", "TITLE"}
243            and self._tokens[self._index + 3].token_type != TokenType.L_PAREN
244        ):
245            return None
246
247        return super()._parse_function(
248            functions=functions,
249            anonymous=anonymous,
250            optional_parens=optional_parens,
251            any_token=any_token,
252        )
253
254    def _parse_column_ops(self, this: exp.Expr | None) -> exp.Expr | None:
255        this = super()._parse_column_ops(this)
256
257        while (
258            this
259            and self._curr.token_type == TokenType.L_PAREN
260            and self._next.token_type in {TokenType.FORMAT, TokenType.VAR}
261            and self._next.text.upper() in {"FORMAT", "TITLE"}
262        ):
263            self._advance()
264            attributes: dict[str, exp.Expr | None] = {}
265            while True:
266                attribute = self._curr.text.upper()
267                if not self._match_texts(("FORMAT", "TITLE")):
268                    self.raise_error("Expected FORMAT or TITLE")
269                    break
270                if attribute in attributes:
271                    self.raise_error(f"Duplicate {attribute} attribute")
272                value = self._parse_string()
273                if not value:
274                    self.raise_error(f"Expected string after {attribute}")
275                elif attribute == "TITLE" and not value.is_string:
276                    self.raise_error("Expected string literal after TITLE")
277                attributes[attribute] = value
278                if not self._match(TokenType.COMMA):
279                    break
280            self._match_r_paren()
281            if "FORMAT" in attributes:
282                this = self.expression(exp.FormatPhrase(this=this, format=attributes["FORMAT"]))
283            if "TITLE" in attributes:
284                this = self.expression(exp.TeradataTitle(this=this, title=attributes["TITLE"]))
285            this = super()._parse_column_ops(this)
286
287        return this
288
289    def _parse_window(self, this: exp.Expr | None, alias: bool = False) -> exp.Expr | None:
290        window = super()._parse_window(this, alias)
291
292        # Unlike ANSI, Teradata window aggregates with ORDER BY and no ROWS clause span the whole partition
293        # https://docs.teradata.com/r/Enterprise_IntelliFlex_VMware/SQL-Functions-Expressions-and-Predicates/Ordered-Analytical/Window-Aggregate-Functions/The-Window-Specification/ANSI-Compliance
294        if (
295            isinstance(window, exp.Window)
296            and window.args.get("order")
297            and not window.args.get("spec")
298            and is_window_aggregate(window)
299        ):
300            window.set(
301                "spec",
302                exp.WindowSpec(
303                    kind="ROWS",
304                    start="UNBOUNDED",
305                    start_side="PRECEDING",
306                    end="UNBOUNDED",
307                    end_side="FOLLOWING",
308                ),
309            )
310
311        return window