Fix: generate_surrogate_key returns hex strings for SHA256/SHA512 on Presto and Trino (#5888)

Signed-off-by: Pawan Singh Kapkoti <42340841+Pawansingh3889@users.noreply.github.com>
Signed-off-by: Pawansingh3889 <pawansinghkapkoti@gmail.com>
Co-authored-by: Pawansingh3889 <pawansinghkapkoti@gmail.com>

Author: Pawansingh3889

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.

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"
+    )