fix(macros): generate_surrogate_key renders valid hex strings on mysql, tsql, starrocks, snowflake

Signed-off-by: ptimizeroracle <contact@binblok.com>

Author: ptimizeroracle

sqlmesh/core/macros.py

@@ -981,13 +981,16 @@ def generate_surrogate_key(
         # 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.
+    elif isinstance(func, exp.Anonymous):
+        # Some dialects' parsers have no entry for the SHA-2 functions, so
+        # exp.func returns an Anonymous node and neither branch above fires:
+        # on MySQL and StarRocks the bare SHA256(varchar) is a runtime error
+        # because those engines only spell the function SHA2(expr, length),
+        # and Athena runs the Trino engine where sha256() takes varbinary
+        # (#5871). Mapping the untyped name to exp.SHA2 with the canonical
+        # digest length renders the valid call on every one of them. Unlike
+        # the probe below, this is not a pin-era workaround: these parsers
+        # hand back Anonymous on every sqlglot version.
         #
         # Anonymous is the catch-all for every unrecognised function name, and
         # hash_function is caller-supplied, so the name is checked rather than
@@ -1013,6 +1016,26 @@ def generate_surrogate_key(
             )
         )
 
+    if isinstance(func, (exp.MD5, exp.SHA, exp.SHA2)) and _renders_varbinary(evaluator.dialect):
+        # T-SQL renders every one of these as HASHBYTES, which returns
+        # VARBINARY rather than the hex string the surrogate key promises.
+        # CONVERT style 2 strips the 0x prefix and LOWER() restores the
+        # lowercase hex the other dialects return, so keys hash identically
+        # everywhere. The probe keeps this branch inert if the tsql generator
+        # ever emits the conversion itself.
+        digest_bits = 128
+        if isinstance(func, exp.SHA):
+            digest_bits = 160
+        elif isinstance(func, exp.SHA2) and func.args.get("length") is not None:
+            digest_bits = int(str(func.args["length"].name))  # type: ignore[union-attr]
+        func = exp.Lower(
+            this=exp.Convert(
+                this=exp.DataType.build(f"VARCHAR({digest_bits // 4})"),
+                expression=func,
+                style=exp.Literal.number(2),
+            )
+        )
+
     return func
 
 
@@ -1021,6 +1044,11 @@ def generate_surrogate_key(
 # Athena is on the list because it runs the Trino engine.
 _PRESTO_FAMILY = frozenset({"presto", "trino", "athena"})
 
+# Dialects that render every string hash as HASHBYTES, which returns
+# VARBINARY rather than a hex string. Fabric is on the list because it runs
+# the T-SQL engine.
+_TSQL_FAMILY = frozenset({"tsql", "fabric"})
+
 # The SHA-2 digest widths a surrogate key may ask for, by function name.
 _SHA2_DIGEST_LENGTHS = {"SHA256": 256, "SHA512": 512}
 
@@ -1030,6 +1058,11 @@ def _is_presto_family(dialect: DialectType) -> bool:
     return (str(dialect) if dialect else "").split(",")[0].strip().lower() in _PRESTO_FAMILY
 
 
+def _is_tsql_family(dialect: DialectType) -> bool:
+    """Whether this dialect is T-SQL (MSSQL or Fabric)."""
+    return (str(dialect) if dialect else "").split(",")[0].strip().lower() in _TSQL_FAMILY
+
+
 @lru_cache(maxsize=None)
 def _sha2_renders_binary(dialect: DialectType) -> bool:
     """Whether this dialect renders exp.SHA2 as a bare binary-semantics call.
@@ -1043,6 +1076,21 @@ def _sha2_renders_binary(dialect: DialectType) -> bool:
     return "TO_HEX" not in probe.sql(dialect=dialect)
 
 
+@lru_cache(maxsize=None)
+def _renders_varbinary(dialect: DialectType) -> bool:
+    """Whether this dialect renders the string hashes as HASHBYTES (VARBINARY).
+
+    The T-SQL family (MSSQL, Fabric) has no MD5/SHA2 functions: every string
+    hash renders as HASHBYTES, which returns VARBINARY instead of the hex
+    string the surrogate key promises.
+    """
+    if not _is_tsql_family(dialect):
+        return False
+    probe = exp.MD5(this=exp.column("_sqlmesh_probe"))
+    rendered = probe.sql(dialect=dialect)
+    return "HASHBYTES" in rendered and "CONVERT" not in rendered
+
+
 @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

@@ -1302,14 +1302,66 @@ def render(dialect: str, hash_function: str) -> str:
         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")
+    # Snowflake's parser also hands back Anonymous for SHA256, and SHA2 is its
+    # only spelling (digest size optional there, defaulting to 256), so the
+    # mapped form is what the engine accepts.
     assert (
         render("snowflake", "SHA256")
-        == "SELECT SHA256(CONCAT(COALESCE(CAST(a AS VARCHAR), '_sqlmesh_surrogate_key_null_'))) FROM foo"
+        == "SELECT SHA2(CONCAT(COALESCE(CAST(a AS VARCHAR), '_sqlmesh_surrogate_key_null_')), 256) FROM foo"
+    )
+
+
+def test_generate_surrogate_key_hex_string_on_mysql_tsql_starrocks() -> None:
+    """The hex-string invariant on dialects whose parsers hand back Anonymous
+    for the SHA-2 functions and whose engines spell them differently.
+
+    MySQL and StarRocks only accept SHA2(expr, digest_length): a bare
+    SHA256(...) is ERROR 1305 (function does not exist) on MySQL 8.4. T-SQL
+    (MSSQL, Fabric) renders every hash as HASHBYTES, which returns VARBINARY,
+    so the key must be converted to the lowercase hex string the other
+    dialects return (CONVERT style 2, verified byte-identical to DuckDB).
+    """
+
+    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)
+
+    # MySQL and StarRocks: the untyped SHA256/SHA512 names must render as
+    # SHA2(expr, digest_length), the only spelling those engines accept.
+    assert (
+        render("mysql", "SHA256")
+        == "SELECT SHA2(CONCAT(COALESCE(CAST(a AS CHAR), '_sqlmesh_surrogate_key_null_')), 256) FROM foo"
+    )
+    assert (
+        render("mysql", "SHA512")
+        == "SELECT SHA2(CONCAT(COALESCE(CAST(a AS CHAR), '_sqlmesh_surrogate_key_null_')), 512) FROM foo"
+    )
+    assert (
+        render("starrocks", "SHA256")
+        == "SELECT SHA2(CONCAT(COALESCE(CAST(a AS STRING), '_sqlmesh_surrogate_key_null_')), 256) FROM foo"
+    )
+
+    # T-SQL family (MSSQL, Fabric): HASHBYTES returns VARBINARY, so the key is
+    # converted to lowercase hex, with the VARCHAR sized to the digest width.
+    assert (
+        render("tsql", "MD5")
+        == "SELECT LOWER(CONVERT(VARCHAR(32), HASHBYTES('MD5', COALESCE(CAST(a AS VARCHAR(MAX)), '_sqlmesh_surrogate_key_null_')), 2)) FROM foo"
+    )
+    assert (
+        render("tsql", "SHA1")
+        == "SELECT LOWER(CONVERT(VARCHAR(40), HASHBYTES('SHA1', COALESCE(CAST(a AS VARCHAR(MAX)), '_sqlmesh_surrogate_key_null_')), 2)) FROM foo"
+    )
+    assert (
+        render("tsql", "SHA256")
+        == "SELECT LOWER(CONVERT(VARCHAR(64), HASHBYTES('SHA2_256', COALESCE(CAST(a AS VARCHAR(MAX)), '_sqlmesh_surrogate_key_null_')), 2)) FROM foo"
+    )
+    assert (
+        render("fabric", "SHA512")
+        == "SELECT LOWER(CONVERT(VARCHAR(128), HASHBYTES('SHA2_512', COALESCE(CAST(a AS VARCHAR(MAX)), '_sqlmesh_surrogate_key_null_')), 2)) FROM foo"
     )