@@ -965,17 +965,84 @@ def generate_surrogate_key(
965965 )
966966 )
967967
968+ concat = exp .func ("CONCAT" , * string_fields )
969+ # The argument is always a string; annotating it here lets generators that
970+ # split string/binary hash semantics (Presto, Trino) wrap the encode.
971+ concat .type = exp .DataType .build ("text" )
972+
968973 func = exp .func (
969974 hash_function .name ,
970- exp . func ( "CONCAT" , * string_fields ) ,
975+ concat ,
971976 dialect = evaluator .dialect ,
972977 )
973978 if isinstance (func , exp .MD5Digest ):
974979 func = exp .MD5 (this = func .this )
980+ elif isinstance (func , exp .SHA2Digest ):
981+ # Same split as MD5/MD5Digest: the surrogate key must be a hex string,
982+ # not a binary digest, on every dialect.
983+ func = exp .SHA2 (this = func .this , length = func .args .get ("length" ))
984+ elif isinstance (func , exp .Anonymous ) and _is_presto_family (evaluator .dialect ):
985+ # Athena runs the Trino engine, so sha256() takes varbinary there too,
986+ # but its parser has no SHA256/SHA512 entry: exp.func returns an
987+ # Anonymous node, so neither branch above fires and the surrogate key
988+ # keeps the bare SHA256(varchar) form reported in #5871. Unlike the
989+ # probe below, this is not a pin-era workaround — Athena still parses
990+ # to Anonymous on sqlglot versions that carry tobymao/sqlglot#7824.
991+ #
992+ # Anonymous is the catch-all for every unrecognised function name, and
993+ # hash_function is caller-supplied, so the name is checked rather than
994+ # assumed: an unknown hash must pass through untouched.
995+ length = _SHA2_DIGEST_LENGTHS .get (func .name .upper ())
996+ if length is not None :
997+ func = exp .SHA2 (this = concat , length = exp .Literal .number (length ))
998+
999+ if isinstance (func , exp .SHA2 ) and _sha2_renders_binary (evaluator .dialect ):
1000+ # Presto/Trino render a bare SHA256(varchar) for exp.SHA2 on sqlglot
1001+ # versions without tobymao/sqlglot#7824: a type error on Trino, and
1002+ # binary rather than string semantics where it runs. Build the
1003+ # hex-string form explicitly, mirroring what those generators do for
1004+ # MD5: LOWER(TO_HEX(SHA256(TO_UTF8(...)))). The probe keeps this
1005+ # branch inert once sqlglot renders the hex form natively, so the
1006+ # expression is never wrapped twice.
1007+ return exp .Lower (
1008+ this = exp .Hex (
1009+ this = exp .SHA2 (
1010+ this = exp .Encode (this = func .this , charset = exp .Literal .string ("utf-8" )),
1011+ length = func .args .get ("length" ),
1012+ )
1013+ )
1014+ )
9751015
9761016 return func
9771017
9781018
1019+ # Dialects that model string and binary hashes separately, so a bare
1020+ # SHA256(varchar) is a type error rather than a hex-string surrogate key.
1021+ # Athena is on the list because it runs the Trino engine.
1022+ _PRESTO_FAMILY = frozenset ({"presto" , "trino" , "athena" })
1023+
1024+ # The SHA-2 digest widths a surrogate key may ask for, by function name.
1025+ _SHA2_DIGEST_LENGTHS = {"SHA256" : 256 , "SHA512" : 512 }
1026+
1027+
1028+ def _is_presto_family (dialect : DialectType ) -> bool :
1029+ """Whether this dialect is Presto, Trino or Athena."""
1030+ return (str (dialect ) if dialect else "" ).split ("," )[0 ].strip ().lower () in _PRESTO_FAMILY
1031+
1032+
1033+ @lru_cache (maxsize = None )
1034+ def _sha2_renders_binary (dialect : DialectType ) -> bool :
1035+ """Whether this dialect renders exp.SHA2 as a bare binary-semantics call.
1036+
1037+ Only the Presto family models string and binary hashes separately; other
1038+ dialects' SHA256(varchar) already returns a hex string.
1039+ """
1040+ if not _is_presto_family (dialect ):
1041+ return False
1042+ probe = exp .SHA2 (this = exp .column ("_sqlmesh_probe" ), length = exp .Literal .number (256 ))
1043+ return "TO_HEX" not in probe .sql (dialect = dialect )
1044+
1045+
9791046@macro ()
9801047def safe_add (_ : MacroEvaluator , * fields : exp .Expr ) -> exp .Case :
9811048 """Adds numbers together, substitutes nulls for 0s and only returns null if all fields are null.
@@ -1379,15 +1446,17 @@ def resolve_template(
13791446 """
13801447 Generates either a String literal or an exp.Table representing a physical table location, based on rendering the provided template String literal.
13811448
1382- Note: It relies on the @this_model variable being available in the evaluation context (@this_model resolves to an exp.Table object
1383- representing the current physical table).
1449+ Note: It relies on the @this_model variable being available in the evaluation context. @this_model usually resolves to an
1450+ exp.Table object representing the current physical table, but in an audit on a model with a time column it resolves to a
1451+ subquery that selects from that table and filters it down to the audited time range. In that case the placeholders below
1452+ are resolved against the physical table the subquery selects from.
13841453 Therefore, the @resolve_template macro must be used at creation or evaluation time and not at load time.
13851454
13861455 Args:
13871456 template: Template string literal. Can contain the following placeholders:
1388- @{catalog_name} -> replaced with the catalog of the exp.Table returned from @this_model
1389- @{schema_name} -> replaced with the schema of the exp.Table returned from @this_model
1390- @{table_name} -> replaced with the name of the exp.Table returned from @this_model
1457+ @{catalog_name} -> replaced with the catalog of the physical table @this_model refers to
1458+ @{schema_name} -> replaced with the schema of the physical table @this_model refers to
1459+ @{table_name} -> replaced with the name of the physical table @this_model refers to
13911460 mode: What to return.
13921461 'literal' -> return an exp.Literal string
13931462 'table' -> return an exp.Table
@@ -1400,9 +1469,26 @@ def resolve_template(
14001469 >>> evaluator.locals.update({"this_model": exp.to_table("test_catalog.sqlmesh__test.test__test_model__2517971505")})
14011470 >>> evaluator.transform(parse_one(sql)).sql()
14021471 "'s3://data-bucket/prod/test_catalog/sqlmesh__test/test__test_model__2517971505'"
1472+
1473+ The same template resolves to the same location when @this_model is the time-filtered
1474+ subquery that audits on models with a time column receive:
1475+
1476+ >>> table = exp.to_table("test_catalog.sqlmesh__test.test__test_model__2517971505")
1477+ >>> subquery = exp.select("*").from_(table).where(exp.column("ds").eq("2020-01-01")).subquery()
1478+ >>> evaluator.locals.update({"this_model": subquery})
1479+ >>> evaluator.transform(parse_one(sql)).sql()
1480+ "'s3://data-bucket/prod/test_catalog/sqlmesh__test/test__test_model__2517971505'"
14031481 """
14041482 if "this_model" in evaluator .locals :
1405- this_model = exp .to_table (evaluator .locals ["this_model" ], dialect = evaluator .dialect )
1483+ this_model_expr = evaluator .locals ["this_model" ]
1484+ if isinstance (this_model_expr , exp .Subquery ):
1485+ # Audits on models with a time column render @this_model as a subquery that filters the
1486+ # physical table on the audited time range, so resolve against the table it selects from
1487+ from_ = this_model_expr .unnest ().args .get ("from_" )
1488+ if from_ is not None and isinstance (from_ .this , exp .Table ):
1489+ this_model_expr = from_ .this
1490+
1491+ this_model = exp .to_table (this_model_expr , dialect = evaluator .dialect )
14061492 template_str : str = template .this
14071493 result = (
14081494 template_str .replace ("@{catalog_name}" , this_model .catalog )
0 commit comments