fix(trino): correctly detect parenthesized IF blocks in routine bodies

Add type annotations to the module-level constants and local counters
in the Trino dialect, and fix the `IF (a > b) THEN` vs. scalar `IF(...)`
ambiguity: an IF immediately followed by `(` is now classified as a
procedural block only when the matching closing paren is followed by
THEN, otherwise as a scalar function call. Adds a regression test for
the parenthesized condition case.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Evan
2026-07-20 18:02:48 -07:00
parent 43813d0681
commit ac63ff56cd
2 changed files with 73 additions and 22 deletions

View File

@@ -128,6 +128,19 @@ WITH FUNCTION classify(a bigint)
END WHILE;
RETURN IF(a = 100, 'hundred', 'other');
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
""",
],
@@ -135,7 +148,8 @@ 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.
``WHILE ... END WHILE``, scalar ``IF()`` function calls, and a
parenthesized ``IF (...) THEN`` condition.
"""
statements = sqlglot.parse(sql.strip(), dialect=Trino)
assert len(statements) == 1