Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 55 additions & 7 deletions sqlmesh/core/macros.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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


Expand All @@ -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}

Expand All @@ -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.
Expand All @@ -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.
Expand Down
60 changes: 56 additions & 4 deletions tests/core/test_macros.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)