diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py index ad8b922a07..7d920abb54 100644 --- a/sqlmesh/core/macros.py +++ b/sqlmesh/core/macros.py @@ -965,13 +965,22 @@ 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")) return func diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py index 0135cd8ca5..cd5a7fdb61 100644 --- a/tests/core/test_macros.py +++ b/tests/core/test_macros.py @@ -1178,3 +1178,54 @@ 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" + ) + # Trino/Presto render SHA256/SHA512 surrogate keys as + # LOWER(TO_HEX(SHA256(TO_UTF8(...)))) once sqlglot ships + # tobymao/sqlglot#7824; the string assertions for that belong with the + # sqlglot version bump.