From 9982dbaf83c6dbc76e46effa49b9c8bf07119c25 Mon Sep 17 00:00:00 2001 From: Evan Date: Mon, 6 Jul 2026 11:22:45 -0700 Subject: [PATCH] fix(trino): support inline SQL UDFs (WITH FUNCTION ... BEGIN ... END) sqlglot cannot parse Trino SQL routine syntax, so queries declaring inline UDFs failed to parse in SQL Lab: the parser splits statements on every semicolon (including the ones inside BEGIN ... END routine bodies) and has no grammar for FUNCTION specifications in a WITH clause. The upstream issue (tobymao/sqlglot#5178) was closed as low priority, so this extends the Trino dialect on the Superset side. The custom dialect keeps routine bodies intact when splitting statements and parses inline function specifications into opaque InlineUDF nodes that regenerate verbatim. Trino does not allow queries inside SQL UDF bodies, so the opaque representation hides no table references from Superset's security checks. The extensions only activate on syntax that fails to parse today, so existing queries are unaffected. Fixes #26162 Co-Authored-By: Claude Opus 4.8 --- superset/sql/dialects/__init__.py | 2 + superset/sql/dialects/trino.py | 273 +++++++++++++++++++ superset/sql/parse.py | 12 +- tests/unit_tests/sql/dialects/trino_tests.py | 239 ++++++++++++++++ 4 files changed, 524 insertions(+), 2 deletions(-) create mode 100644 superset/sql/dialects/trino.py create mode 100644 tests/unit_tests/sql/dialects/trino_tests.py 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