diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py index ad8b922a07..9f8f3fa0ad 100644 --- a/sqlmesh/core/macros.py +++ b/sqlmesh/core/macros.py @@ -965,17 +965,57 @@ def generate_surrogate_key( ) ) + concat = exp.func("CONCAT", *string_fields) + # The argument is always a string; annotating it here lets generators that + # split string/binary hash semantics (Presto, Trino) wrap the encode. + concat.type = exp.DataType.build("text") + func = exp.func( hash_function.name, - exp.func("CONCAT", *string_fields), + concat, dialect=evaluator.dialect, ) if isinstance(func, exp.MD5Digest): func = exp.MD5(this=func.this) + elif isinstance(func, exp.SHA2Digest): + # Same split as MD5/MD5Digest: the surrogate key must be a hex string, + # not a binary digest, on every dialect. + func = exp.SHA2(this=func.this, length=func.args.get("length")) + + if isinstance(func, exp.SHA2) and _sha2_renders_binary(evaluator.dialect): + # Presto/Trino render a bare SHA256(varchar) for exp.SHA2 on sqlglot + # versions without tobymao/sqlglot#7824: a type error on Trino, and + # binary rather than string semantics where it runs. Build the + # hex-string form explicitly, mirroring what those generators do for + # MD5: LOWER(TO_HEX(SHA256(TO_UTF8(...)))). The probe keeps this + # branch inert once sqlglot renders the hex form natively, so the + # expression is never wrapped twice. + return exp.Lower( + this=exp.Hex( + this=exp.SHA2( + this=exp.Encode(this=func.this, charset=exp.Literal.string("utf-8")), + length=func.args.get("length"), + ) + ) + ) return func +@lru_cache(maxsize=None) +def _sha2_renders_binary(dialect: DialectType) -> bool: + """Whether this dialect renders exp.SHA2 as a bare binary-semantics call. + + Only the Presto family models string and binary hashes separately; other + dialects' SHA256(varchar) already returns a hex string. + """ + dialect_name = (str(dialect) if dialect else "").split(",")[0].strip().lower() + if dialect_name not in ("presto", "trino", "athena"): + return False + probe = exp.SHA2(this=exp.column("_sqlmesh_probe"), length=exp.Literal.number(256)) + return "TO_HEX" not in probe.sql(dialect=dialect) + + @macro() def safe_add(_: MacroEvaluator, *fields: exp.Expr) -> exp.Case: """Adds numbers together, substitutes nulls for 0s and only returns null if all fields are null. diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py index 0135cd8ca5..089dad0034 100644 --- a/tests/core/test_macros.py +++ b/tests/core/test_macros.py @@ -1178,3 +1178,76 @@ def test_macro_coerce_literal_type(macro_evaluator): expression = d.parse_one("@TEST_LITERAL_TYPE(1.0)") with pytest.raises(MacroEvalError, match=".*Coercion failed"): macro_evaluator.transform(expression) + + +def test_generate_surrogate_key_hash_semantics() -> None: + from sqlmesh.core.macros import generate_surrogate_key + + surrogate_key = ( + generate_surrogate_key.func + if hasattr(generate_surrogate_key, "func") + else generate_surrogate_key + ) + + # The macro must always build the string-semantics hash expression, never + # a binary digest, so dialects that model the two separately (Presto and + # Trino after tobymao/sqlglot#7824) can render the hex-string form. + # BigQuery's parser maps SHA256 to SHA2Digest, which exercises the + # conversion on every supported sqlglot version. + func = surrogate_key( + MacroEvaluator(dialect="bigquery"), + exp.column("a"), + hash_function=exp.Literal.string("SHA256"), + ) + assert isinstance(func, exp.SHA2) + + # The hash argument is annotated as text so generators that wrap an + # encode around string inputs (TO_UTF8 on Presto/Trino) can do so without + # a separate annotation pass. + assert func.this.is_type("text") + + def render(dialect: str, hash_function: str) -> str: + sql = f"SELECT @GENERATE_SURROGATE_KEY(a, hash_function := '{hash_function}') FROM foo" + return ( + MacroEvaluator(dialect=dialect).transform(parse_one(sql, dialect=dialect)).sql(dialect) + ) + + # Rendered SQL, stable across supported sqlglot versions. + assert ( + render("bigquery", "SHA256") + == "SELECT SHA256(CONCAT(COALESCE(CAST(a AS STRING), '_sqlmesh_surrogate_key_null_'))) FROM foo" + ) + assert ( + render("duckdb", "SHA256") + == "SELECT SHA256(COALESCE(CAST(a AS TEXT), '_sqlmesh_surrogate_key_null_')) FROM foo" + ) + assert ( + render("trino", "MD5") + == "SELECT LOWER(TO_HEX(MD5(TO_UTF8(CAST(COALESCE(CAST(a AS VARCHAR), '_sqlmesh_surrogate_key_null_') AS VARCHAR))))) FROM foo" + ) + + # The reported bug (#5871): Trino/Presto SHA256/SHA512 surrogate keys must + # be the hex-string form, not a bare SHA256(varchar). The macro-side + # fallback produces it under the current sqlglot pin; once sqlglot renders + # exp.SHA2 this way natively (tobymao/sqlglot#7824), the probe disables + # the fallback and these assertions hold unchanged. + for _dialect in ("trino", "presto"): + assert ( + render(_dialect, "SHA256") + == "SELECT LOWER(TO_HEX(SHA256(TO_UTF8(CAST(COALESCE(CAST(a AS VARCHAR), '_sqlmesh_surrogate_key_null_') AS VARCHAR))))) FROM foo" + ) + assert ( + render(_dialect, "SHA512") + == "SELECT LOWER(TO_HEX(SHA512(TO_UTF8(CAST(COALESCE(CAST(a AS VARCHAR), '_sqlmesh_surrogate_key_null_') AS VARCHAR))))) FROM foo" + ) + + # The fallback is scoped to the Presto family: dialects whose bare + # SHA256(varchar) already returns a hex string are left to sqlglot. + from sqlmesh.core.macros import _sha2_renders_binary + + assert not _sha2_renders_binary("duckdb") + assert not _sha2_renders_binary("bigquery") + assert ( + render("snowflake", "SHA256") + == "SELECT SHA256(CONCAT(COALESCE(CAST(a AS VARCHAR), '_sqlmesh_surrogate_key_null_'))) FROM foo" + )