diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py index 3b52d4caa2..2e995003bf 100644 --- a/sqlmesh/core/macros.py +++ b/sqlmesh/core/macros.py @@ -965,17 +965,84 @@ 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")) + elif isinstance(func, exp.Anonymous) and _is_presto_family(evaluator.dialect): + # Athena runs the Trino engine, so sha256() takes varbinary there too, + # but its parser has no SHA256/SHA512 entry: exp.func returns an + # Anonymous node, so neither branch above fires and the surrogate key + # keeps the bare SHA256(varchar) form reported in #5871. Unlike the + # probe below, this is not a pin-era workaround — Athena still parses + # to Anonymous on sqlglot versions that carry tobymao/sqlglot#7824. + # + # Anonymous is the catch-all for every unrecognised function name, and + # hash_function is caller-supplied, so the name is checked rather than + # assumed: an unknown hash must pass through untouched. + length = _SHA2_DIGEST_LENGTHS.get(func.name.upper()) + if length is not None: + func = exp.SHA2(this=concat, length=exp.Literal.number(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 +# Dialects that model string and binary hashes separately, so a bare +# SHA256(varchar) is a type error rather than a hex-string surrogate key. +# Athena is on the list because it runs the Trino engine. +_PRESTO_FAMILY = frozenset({"presto", "trino", "athena"}) + +# The SHA-2 digest widths a surrogate key may ask for, by function name. +_SHA2_DIGEST_LENGTHS = {"SHA256": 256, "SHA512": 512} + + +def _is_presto_family(dialect: DialectType) -> bool: + """Whether this dialect is Presto, Trino or Athena.""" + return (str(dialect) if dialect else "").split(",")[0].strip().lower() in _PRESTO_FAMILY + + +@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. + """ + if not _is_presto_family(dialect): + 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 28216d6b9e..0b3bcf70ee 100644 --- a/tests/core/test_macros.py +++ b/tests/core/test_macros.py @@ -1233,3 +1233,83 @@ 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 + + # 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 = generate_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" + rendered = MacroEvaluator(dialect=dialect).transform(parse_one(sql, dialect=dialect)) + assert isinstance(rendered, exp.Expr) + return rendered.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. + # Athena is included: it runs the Trino engine and hits the same + # sha256(varbinary) failure, but its parser has no SHA256/SHA512 entry, so + # exp.func hands back exp.Anonymous rather than exp.SHA2/exp.SHA2Digest. + # That is true on every sqlglot version tested, before and after #7824, so + # the Anonymous path is not a pin-era workaround the way the probe is. + for _dialect in ("trino", "presto", "athena"): + 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" + ) + + # Anonymous is sqlglot's catch-all for an unrecognised function name, so + # the conversion is keyed on the name: an unknown hash_function must pass + # through untouched rather than be reinterpreted as a SHA-2 digest. + assert ( + render("athena", "MYHASH") + == "SELECT MYHASH(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" + )