diff --git a/superset/sql/dialects/__init__.py b/superset/sql/dialects/__init__.py index f7b82538115..b0f19d545f2 100644 --- a/superset/sql/dialects/__init__.py +++ b/superset/sql/dialects/__init__.py @@ -20,6 +20,7 @@ from .dremio import Dremio from .firebolt import Firebolt, FireboltOld from .opensearch import OpenSearch from .pinot import Pinot +from .trino import Trino from .vertica import Vertica __all__ = [ @@ -29,5 +30,6 @@ __all__ = [ "FireboltOld", "OpenSearch", "Pinot", + "Trino", "Vertica", ] diff --git a/superset/sql/dialects/trino.py b/superset/sql/dialects/trino.py new file mode 100644 index 00000000000..93a207a13c0 --- /dev/null +++ b/superset/sql/dialects/trino.py @@ -0,0 +1,273 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import typing as t + +from sqlglot import exp +from sqlglot.dialects.trino import Trino as SqlglotTrino +from sqlglot.tokens import Token, TokenType + +# Keywords that open a block terminated by ``END`` in Trino SQL routines +# (https://trino.io/docs/current/udf/sql.html). ``CASE`` is included because +# both the ``CASE`` statement and the ``CASE`` expression are terminated by +# ``END``, so counting them keeps the depth balanced either way. +BLOCK_OPENERS = {"BEGIN", "CASE", "IF", "LOOP", "REPEAT", "WHILE"} + +# Keywords that are also scalar functions in Trino (e.g. ``IF(a, b, c)`` and +# ``REPEAT('a', 3)``). When immediately followed by ``(`` they are function +# calls, not block openers. +AMBIGUOUS_OPENERS = {"IF", "REPEAT"} + +BODY_KEYWORDS = ("RETURN", "BEGIN") + + +class InlineUDF(exp.CTE): + """ + An inline SQL user-defined function declared in a ``WITH`` clause. + + Trino supports declaring UDFs inline as part of a query:: + + WITH FUNCTION meaning_of_life() + RETURNS tinyint + BEGIN + DECLARE a tinyint DEFAULT CAST(6 AS tinyint); + DECLARE b tinyint DEFAULT CAST(7 AS tinyint); + RETURN a * b; + END + SELECT meaning_of_life() + + The function definition is stored verbatim as an opaque string (wrapped + in an ``exp.Var`` so that AST traversal helpers see an expression), since + sqlglot has no representation for SQL routine bodies. Trino does not + allow queries inside SQL UDF bodies, so no table references are hidden + by the opaque representation. + + This subclasses ``exp.CTE`` because ``sqlglot.parser.Parser._parse_with`` + only collects ``exp.CTE`` instances into the ``WITH`` clause. + """ + + arg_types = {"this": True} + + +class Trino(SqlglotTrino): + """ + Custom Trino dialect with support for inline SQL UDFs. + + sqlglot cannot parse Trino SQL routine syntax; see + https://github.com/tobymao/sqlglot/issues/5178. There are two separate + problems: + + 1. The parser splits statements on every semicolon, including the ones + inside a ``BEGIN ... END`` routine body. + 2. The ``FUNCTION`` specification in a ``WITH`` clause is not valid CTE + syntax. + + This dialect keeps routine bodies intact when splitting statements, and + parses inline function specifications into opaque `InlineUDF` nodes that + regenerate verbatim. + + Note that sqlglot's ``Dialect`` metaclass registers subclasses by class + name, so once this module is imported this class also replaces the + built-in dialect for string-based lookups (``dialect="trino"``). This is + intentional, and consistent with how other Superset dialects (e.g. + ``Dremio``) shadow their sqlglot counterparts: the extensions are purely + additive, only activating on syntax that fails to parse upstream. + """ + + class Parser(SqlglotTrino.Parser): + @staticmethod + def _block_depth_delta( + text: str, + prev_text: str, + next_token: Token | None, + ) -> int: + """ + Compute the block nesting change contributed by a routine token. + """ + if text in BLOCK_OPENERS: + if prev_text == "END": + return 0 # block terminator, e.g. `END IF`, `END CASE` + if ( + text in AMBIGUOUS_OPENERS + and next_token + and next_token.token_type == TokenType.L_PAREN + ): + return 0 # scalar function call, e.g. `IF(a, b, c)` + return 1 + if text == "END": + return -1 + return 0 + + def _parse( + self, + parse_method: t.Callable[..., exp.Expression | None], + raw_tokens: list[Token], + sql: str | None = None, + ) -> list[exp.Expression | None]: + """ + Split tokens into statements, keeping routine bodies intact. + + This is a copy of ``sqlglot.parser.Parser._parse`` with one + change: when a statement starts with ``WITH FUNCTION``, ``CREATE + FUNCTION``, or ``CREATE OR REPLACE FUNCTION``, semicolons inside + ``BEGIN ... END`` blocks do not split the statement. + """ + self.reset() + self.sql = sql or "" + + total = len(raw_tokens) + chunks: list[list[Token]] = [[]] + routine_mode = False + depth = 0 + prev_text = "" + + for i, token in enumerate(raw_tokens): + if token.token_type == TokenType.SEMICOLON and depth <= 0: + if token.comments: + chunks.append([token]) + if i < total - 1: + chunks.append([]) + routine_mode = False + depth = 0 + prev_text = "" + continue + + chunk = chunks[-1] + chunk.append(token) + + if token.token_type == TokenType.FUNCTION and not routine_mode: + heads = [tok.token_type for tok in chunk[:-1]] + routine_mode = heads in ( + [TokenType.WITH], + [TokenType.CREATE], + [TokenType.CREATE, TokenType.OR, TokenType.REPLACE], + ) + elif routine_mode: + text = token.text.upper() + next_token = raw_tokens[i + 1] if i < total - 1 else None + depth += self._block_depth_delta(text, prev_text, next_token) + + prev_text = token.text.upper() + + self._chunks = chunks + return self._parse_batch_statements( + parse_method=parse_method, + sep_first_statement=False, + ) + + def _parse_cte(self) -> exp.CTE | None: + """ + Parse a single entry in a ``WITH`` clause. + + An entry starting with the ``FUNCTION`` keyword followed by an + identifier is an inline UDF specification; anything else + (including a CTE named "function") is handled by sqlglot. + """ + if ( + self._curr + and self._curr.token_type == TokenType.FUNCTION + and self._next + and self._next.token_type + not in (TokenType.ALIAS, TokenType.L_PAREN, TokenType.COMMA) + ): + return self._parse_inline_udf() + + return super()._parse_cte() + + def _parse_inline_udf(self) -> InlineUDF: + """ + Consume an inline UDF specification and return it verbatim. + + The specification is ``FUNCTION name(params) RETURNS type`` plus + optional routine characteristics, followed by a body that is + either ``RETURN expression`` or a ``BEGIN ... END`` block. + """ + start = self._curr + self._advance() + + # scan for the start of the function body, skipping over the + # signature, return type, and routine characteristics + paren_depth = 0 + body: str | None = None + while self._curr: + token_type = self._curr.token_type + text = self._curr.text.upper() + if token_type == TokenType.L_PAREN: + paren_depth += 1 + elif token_type == TokenType.R_PAREN: + paren_depth -= 1 + elif paren_depth == 0 and text in BODY_KEYWORDS: + body = text + break + self._advance() + + if body is None: + self.raise_error( + "Expected RETURN or BEGIN in inline function specification" + ) + + if body == "RETURN": + self._advance() + if not self._parse_expression(): + self.raise_error("Expected expression after RETURN") + else: + self._consume_block() + + raw = self.sql[start.start : self._prev.end + 1] + return self.expression(InlineUDF(this=exp.Var(this=raw)), token=start) + + def _consume_block(self) -> None: + """ + Consume a ``BEGIN ... END`` block, tracking nested blocks. + """ + depth = 0 + while self._curr: + text = self._curr.text.upper() + if text in BLOCK_OPENERS: + if ( + text in AMBIGUOUS_OPENERS + and self._next + and self._next.token_type == TokenType.L_PAREN + ): + pass # scalar function call, e.g. `IF(a, b, c)` + else: + depth += 1 + self._advance() + elif text == "END": + depth -= 1 + self._advance() + if ( + depth > 0 + and self._curr + and self._curr.text.upper() in BLOCK_OPENERS + ): + # block terminator, e.g. `END IF`, `END CASE` + self._advance() + if depth == 0: + return + else: + self._advance() + + self.raise_error("Unbalanced BEGIN/END in inline function specification") + + class Generator(SqlglotTrino.Generator): + TRANSFORMS = { + **SqlglotTrino.Generator.TRANSFORMS, + InlineUDF: lambda self, e: e.this.name, + } diff --git a/superset/sql/parse.py b/superset/sql/parse.py index c470e43e2da..1c446fcc71c 100644 --- a/superset/sql/parse.py +++ b/superset/sql/parse.py @@ -48,7 +48,15 @@ from sqlglot.optimizer.scope import ( ) from superset.exceptions import QueryClauseValidationException, SupersetParseError -from superset.sql.dialects import DB2, Dremio, Firebolt, OpenSearch, Pinot, Vertica +from superset.sql.dialects import ( + DB2, + Dremio, + Firebolt, + OpenSearch, + Pinot, + Trino, + Vertica, +) if TYPE_CHECKING: from superset.models.core import Database @@ -153,7 +161,7 @@ SQLGLOT_DIALECTS = { "superset": Dialects.SQLITE, # "taosws": ??? "teradatasql": Dialects.TERADATA, - "trino": Dialects.TRINO, + "trino": Trino, "vertica": Vertica, # "ydb" is a plugin dialect (ydb-sqlglot-plugin) auto-discovered via entry_points, # hence a string name rather than a class reference like the built-in dialects. diff --git a/tests/unit_tests/sql/dialects/trino_tests.py b/tests/unit_tests/sql/dialects/trino_tests.py new file mode 100644 index 00000000000..c2d68472f3f --- /dev/null +++ b/tests/unit_tests/sql/dialects/trino_tests.py @@ -0,0 +1,239 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import pytest +import sqlglot + +from superset.exceptions import SupersetParseError +from superset.sql.dialects.trino import InlineUDF, Trino +from superset.sql.parse import SQLScript, SQLStatement, Table + +# example from https://trino.io/docs/current/udf/sql/begin.html, reported in +# https://github.com/apache/superset/issues/26162 +INLINE_UDF = """ +WITH FUNCTION meaning_of_life() + RETURNS tinyint + BEGIN + DECLARE a tinyint DEFAULT CAST(6 AS tinyint); + DECLARE b tinyint DEFAULT CAST(7 AS tinyint); + RETURN a * b; + END +SELECT meaning_of_life() +""".strip() + + +def test_inline_udf_is_single_statement() -> None: + """ + Semicolons inside the routine body must not split the statement. + """ + statements = sqlglot.parse(INLINE_UDF, dialect=Trino) + assert len(statements) == 1 + assert len(list(statements[0].find_all(InlineUDF))) == 1 + + +def test_inline_udf_generates_verbatim() -> None: + """ + The function specification should be preserved verbatim, and the + generated SQL should be parseable again. + """ + statement = sqlglot.parse_one(INLINE_UDF, dialect=Trino) + generated = statement.sql(dialect=Trino) + assert ( + """ +WITH FUNCTION meaning_of_life() + RETURNS tinyint + BEGIN + DECLARE a tinyint DEFAULT CAST(6 AS tinyint); + DECLARE b tinyint DEFAULT CAST(7 AS tinyint); + RETURN a * b; + END +""".strip() + in generated + ) + assert sqlglot.parse_one(generated, dialect=Trino) + + +def test_inline_udf_return_form() -> None: + """ + Test functions whose body is a single ``RETURN`` expression, including + multiple comma-separated functions in one ``WITH`` clause. + """ + sql = """ +WITH + FUNCTION hello(name varchar) + RETURNS varchar + RETURN format('Hello %s!', name), + FUNCTION bye() + RETURNS varchar + RETURN 'Bye!' +SELECT hello('Finn') || ' and ' || bye() + """.strip() + statement = sqlglot.parse_one(sql, dialect=Trino) + assert len(list(statement.find_all(InlineUDF))) == 2 + + generated = statement.sql(dialect=Trino) + assert "RETURN format('Hello %s!', name)" in generated + assert "RETURN 'Bye!'" in generated + + +@pytest.mark.parametrize( + "sql", + [ + """ +WITH FUNCTION classify(a bigint) + RETURNS varchar + BEGIN + CASE a + WHEN 0 THEN RETURN 'zero'; + WHEN 1 THEN RETURN 'one'; + ELSE RETURN 'more than one or negative'; + END CASE; + RETURN NULL; + END +SELECT classify(x) FROM some_table + """, + """ +WITH FUNCTION classify(a bigint) + RETURNS varchar + BEGIN + IF a > 100 THEN + RETURN 'big'; + ELSEIF a > 0 THEN + RETURN 'small'; + END IF; + RETURN 'negative'; + END +SELECT classify(x) FROM some_table + """, + """ +WITH FUNCTION classify(a bigint) + RETURNS varchar + BEGIN + WHILE a < 100 DO + SET a = a + 1; + END WHILE; + RETURN IF(a = 100, 'hundred', 'other'); + END +SELECT classify(x) FROM some_table + """, + ], +) +def test_inline_udf_nested_blocks(sql: str) -> None: + """ + Test nested blocks: ``CASE ... END CASE``, ``IF ... END IF``, + ``WHILE ... END WHILE``, and scalar ``IF()`` function calls. + """ + statements = sqlglot.parse(sql.strip(), dialect=Trino) + assert len(statements) == 1 + + +def test_cte_named_function_still_works() -> None: + """ + A CTE named "function" must still be parsed as a regular CTE. + """ + sql = "WITH function AS (SELECT 1 AS x) SELECT x FROM function" + statement = sqlglot.parse_one(sql, dialect=Trino) + assert not list(statement.find_all(InlineUDF)) + assert statement.sql(dialect=Trino) == sql + + +def test_unbalanced_body_raises() -> None: + """ + An unterminated routine body should raise a parse error. + """ + sql = "WITH FUNCTION f() RETURNS int BEGIN RETURN 1; SELECT f()" + with pytest.raises(sqlglot.errors.ParseError): + sqlglot.parse(sql, dialect=Trino) + + +def test_missing_body_raises() -> None: + """ + A function specification without a body should raise a parse error. + """ + sql = "WITH FUNCTION f() RETURNS int SELECT f()" + with pytest.raises(sqlglot.errors.ParseError): + sqlglot.parse(sql, dialect=Trino) + + +def test_sqlscript_inline_udf() -> None: + """ + Integration with the Superset parsing API (reproduces #26162). + """ + script = SQLScript(INLINE_UDF, "trino") + assert len(script.statements) == 1 + assert not script.has_mutation() + + statement = script.statements[0] + assert statement.is_select() + assert statement.format() == statement.format() # deterministic + + +def test_sqlscript_inline_udf_multiple_statements() -> None: + """ + Statements after the UDF query should still be split correctly. + """ + script = SQLScript(f"{INLINE_UDF};\nSELECT 42", "trino") + assert len(script.statements) == 2 + + +def test_sqlstatement_extract_tables() -> None: + """ + Tables referenced by the main query should still be extracted. + """ + sql = """ +WITH FUNCTION doubleup(x integer) + RETURNS integer + BEGIN + RETURN x * 2; + END +SELECT doubleup(some_column) FROM some_table + """.strip() + statement = SQLStatement(sql, "trino") + assert statement.tables == {Table("some_table")} + + +def test_sqlstatement_regular_queries_unaffected() -> None: + """ + Regular Trino queries should parse exactly as before. + """ + script = SQLScript( + "WITH t AS (SELECT 1 AS x) SELECT * FROM t; SELECT 2", + "trino", + ) + assert len(script.statements) == 2 + assert script.statements[0].tables == set() + + with pytest.raises(SupersetParseError): + SQLStatement("SELECT * FROM", "trino") + + +def test_create_function_not_split() -> None: + """ + ``CREATE FUNCTION`` bodies should not be split on semicolons either. + """ + sql = """ +CREATE FUNCTION meaning_of_life() + RETURNS tinyint + BEGIN + DECLARE a tinyint DEFAULT CAST(6 AS tinyint); + DECLARE b tinyint DEFAULT CAST(7 AS tinyint); + RETURN a * b; + END; +SELECT 42 + """.strip() + statements = sqlglot.parse(sql, dialect=Trino) + assert len(statements) == 2