From 8bd2483531a6c0311e8fe93a22f0eaa44a1ab976 Mon Sep 17 00:00:00 2001 From: Rodrigo-Palma Date: Thu, 17 Sep 2026 12:00:21 -0300 Subject: [PATCH] fix(macros): keep SAFE_ADD, SAFE_SUB and SAFE_DIV grouped inside surrounding operators The optimizer collapses the CASE built by SAFE_ADD/SAFE_SUB into its ELSE branch without parentheses, so an enclosing operator swallows an operand. Emit the ELSE arithmetic as an explicit Paren. SAFE_DIV now returns its quotient wrapped in a Paren so x / @SAFE_DIV(a, b) keeps its grouping. Closes #5649 Signed-off-by: Rodrigo-Palma --- docs/concepts/macros/sqlmesh_macros.md | 6 ++-- sqlmesh/core/macros.py | 33 +++++++++++------ tests/core/test_macros.py | 49 ++++++++++++++++++++++++++ 3 files changed, 75 insertions(+), 13 deletions(-) diff --git a/docs/concepts/macros/sqlmesh_macros.md b/docs/concepts/macros/sqlmesh_macros.md index 5459d79ca8..871ff43e80 100644 --- a/docs/concepts/macros/sqlmesh_macros.md +++ b/docs/concepts/macros/sqlmesh_macros.md @@ -887,7 +887,7 @@ would be rendered as: ```sql linenums="1" SELECT - CASE WHEN a IS NULL AND b IS NULL AND c IS NULL THEN NULL ELSE COALESCE(a, 0) + COALESCE(b, 0) + COALESCE(c, 0) END + CASE WHEN a IS NULL AND b IS NULL AND c IS NULL THEN NULL ELSE (COALESCE(a, 0) + COALESCE(b, 0) + COALESCE(c, 0)) END FROM foo ``` @@ -906,7 +906,7 @@ would be rendered as: ```sql linenums="1" SELECT - CASE WHEN a IS NULL AND b IS NULL AND c IS NULL THEN NULL ELSE COALESCE(a, 0) - COALESCE(b, 0) - COALESCE(c, 0) END + CASE WHEN a IS NULL AND b IS NULL AND c IS NULL THEN NULL ELSE (COALESCE(a, 0) - COALESCE(b, 0) - COALESCE(c, 0)) END FROM foo ``` @@ -925,7 +925,7 @@ would be rendered as: ```sql linenums="1" SELECT - a / NULLIF(b, 0) + (a / NULLIF(b, 0)) FROM foo ``` diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py index 2e995003bf..4751bd85db 100644 --- a/sqlmesh/core/macros.py +++ b/sqlmesh/core/macros.py @@ -1052,12 +1052,11 @@ def safe_add(_: MacroEvaluator, *fields: exp.Expr) -> exp.Case: >>> from sqlmesh.core.macros import MacroEvaluator >>> sql = "SELECT @SAFE_ADD(a, b) FROM foo" >>> MacroEvaluator().transform(parse_one(sql)).sql() - 'SELECT CASE WHEN a IS NULL AND b IS NULL THEN NULL ELSE COALESCE(a, 0) + COALESCE(b, 0) END FROM foo' + 'SELECT CASE WHEN a IS NULL AND b IS NULL THEN NULL ELSE (COALESCE(a, 0) + COALESCE(b, 0)) END FROM foo' """ - return ( - exp.Case() - .when(exp.and_(*(field.is_(exp.null()) for field in fields)), exp.null()) - .else_(reduce(lambda a, b: a + b, [exp.func("COALESCE", field, 0) for field in fields])) # type: ignore + return _null_if_all_null( + fields, + reduce(lambda a, b: a + b, [exp.func("COALESCE", field, 0) for field in fields]), # type: ignore ) @@ -1070,17 +1069,30 @@ def safe_sub(_: MacroEvaluator, *fields: exp.Expr) -> exp.Case: >>> from sqlmesh.core.macros import MacroEvaluator >>> sql = "SELECT @SAFE_SUB(a, b) FROM foo" >>> MacroEvaluator().transform(parse_one(sql)).sql() - 'SELECT CASE WHEN a IS NULL AND b IS NULL THEN NULL ELSE COALESCE(a, 0) - COALESCE(b, 0) END FROM foo' + 'SELECT CASE WHEN a IS NULL AND b IS NULL THEN NULL ELSE (COALESCE(a, 0) - COALESCE(b, 0)) END FROM foo' + """ + return _null_if_all_null( + fields, + reduce(lambda a, b: a - b, [exp.func("COALESCE", field, 0) for field in fields]), # type: ignore + ) + + +def _null_if_all_null(fields: t.Sequence[exp.Expr], arithmetic: exp.Expr) -> exp.Case: + """Returns NULL when every field is NULL, otherwise the result of the arithmetic. + + The arithmetic is parenthesized because the optimizer replaces the CASE with this branch + when the condition is statically false (e.g. `1 IS NULL`), and without the parentheses + the operation would bind to the operators around the macro call. """ return ( exp.Case() .when(exp.and_(*(field.is_(exp.null()) for field in fields)), exp.null()) - .else_(reduce(lambda a, b: a - b, [exp.func("COALESCE", field, 0) for field in fields])) # type: ignore + .else_(exp.paren(arithmetic, copy=False)) ) @macro() -def safe_div(_: MacroEvaluator, numerator: exp.Expr, denominator: exp.Expr) -> exp.Div: +def safe_div(_: MacroEvaluator, numerator: exp.Expr, denominator: exp.Expr) -> exp.Paren: """Divides numbers, returns null if the denominator is 0. Example: @@ -1088,9 +1100,10 @@ def safe_div(_: MacroEvaluator, numerator: exp.Expr, denominator: exp.Expr) -> e >>> from sqlmesh.core.macros import MacroEvaluator >>> sql = "SELECT @SAFE_DIV(a, b) FROM foo" >>> MacroEvaluator().transform(parse_one(sql)).sql() - 'SELECT a / NULLIF(b, 0) FROM foo' + 'SELECT (a / NULLIF(b, 0)) FROM foo' """ - return numerator / exp.func("NULLIF", denominator, 0) + # The quotient must stay a single operand of whatever operator surrounds the macro call + return exp.paren(numerator / exp.func("NULLIF", denominator, 0), copy=False) @macro() diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py index 0b3bcf70ee..fd5a6daf7e 100644 --- a/tests/core/test_macros.py +++ b/tests/core/test_macros.py @@ -11,6 +11,7 @@ from sqlmesh.utils.errors import SQLMeshError from sqlmesh.utils.metaprogramming import Executable from sqlmesh.core.macros import RuntimeStage +from sqlmesh.core.model import load_sql_based_model @pytest.fixture @@ -1313,3 +1314,51 @@ def render(dialect: str, hash_function: str) -> str: render("snowflake", "SHA256") == "SELECT SHA256(CONCAT(COALESCE(CAST(a AS VARCHAR), '_sqlmesh_surrogate_key_null_'))) FROM foo" ) + + +@pytest.mark.parametrize( + "projection, expected", + [ + # GitHub issue #5649: once the optimizer resolves `1 IS NULL`, the CASE collapses into its + # ELSE branch, and the arithmetic must stay grouped inside the surrounding multiplication. + ( + "(@SAFE_SUB(price, amount_off)) * (@SAFE_SUB(1, percent_off / 100))", + 'CASE WHEN "s"."amount_off" IS NULL AND "s"."price" IS NULL THEN NULL ELSE (COALESCE("s"."price", 0) - COALESCE("s"."amount_off", 0)) END * (COALESCE(1, 0) - COALESCE("s"."percent_off" / 100, 0))', + ), + ("x * @SAFE_SUB(1, y)", '"s"."x" * (COALESCE(1, 0) - COALESCE("s"."y", 0))'), + ("-@SAFE_SUB(1, y)", '-(COALESCE(1, 0) - COALESCE("s"."y", 0))'), + ("x - @SAFE_ADD(1, y)", '"s"."x" - (COALESCE(1, 0) + COALESCE("s"."y", 0))'), + ("x * @SAFE_ADD(1, y)", '"s"."x" * (COALESCE(1, 0) + COALESCE("s"."y", 0))'), + ( + "@SAFE_DIV(@SAFE_SUB(1, y), x)", + '(COALESCE(1, 0) - COALESCE("s"."y", 0)) / NULLIF("s"."x", 0)', + ), + # The quotient is a single operand of the enclosing operator. + ("x / @SAFE_DIV(price, y)", '"s"."x" / ("s"."price" / NULLIF("s"."y", 0))'), + ("x * @SAFE_DIV(price, y)", '"s"."x" * ("s"."price" / NULLIF("s"."y", 0))'), + # Standalone usages: the optimizer drops the redundant parentheses around the quotient, + # while the ELSE branch keeps its grouping. + ("@SAFE_DIV(price, y)", '"s"."price" / NULLIF("s"."y", 0)'), + ( + "@SAFE_SUB(price, amount_off) + 1", + 'CASE WHEN "s"."amount_off" IS NULL AND "s"."price" IS NULL THEN NULL ELSE (COALESCE("s"."price", 0) - COALESCE("s"."amount_off", 0)) END + 1', + ), + ], +) +def test_safe_arithmetic_macros_keep_precedence_after_optimization( + projection: str, expected: str +) -> None: + model = load_sql_based_model( + d.parse( + f""" + MODEL (name db.safe_arithmetic); + + SELECT {projection} AS result + FROM (SELECT 1 AS x, 2 AS y, 100 AS price, 25 AS amount_off, 20 AS percent_off) AS s + """ + ) + ) + + rendered_projection = model.render_query_or_raise().selects[0] + + assert rendered_projection.sql() == f'{expected} AS "result"'