diff --git a/sqlmesh/core/audit/builtin.py b/sqlmesh/core/audit/builtin.py index b4eaab8d50..5838d7bf1d 100644 --- a/sqlmesh/core/audit/builtin.py +++ b/sqlmesh/core/audit/builtin.py @@ -318,10 +318,14 @@ def create_non_blocking_copy(audit: Audit) -> Audit: # valid_http_method(column=column_name) valid_http_method_audit = ModelAudit( name="valid_http_method", + defaults={"condition": exp.true()}, query=""" SELECT * FROM @this_model -WHERE NOT @column IN ('GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'HEAD', 'OPTIONS', 'TRACE', 'CONNECT') +WHERE @AND( + NOT @column IN ('GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'HEAD', 'OPTIONS', 'TRACE', 'CONNECT'), + @condition, +) """, ) @@ -432,7 +436,7 @@ def create_non_blocking_copy(audit: Audit) -> Audit: @column, (@column - mean_@column) / NULLIF(stddev_@column, 0) AS z_score FROM @this_model, stats -WHERE ABS((@column - mean_@column) / NULLIF(stddev_@column, 0)) > @threshold +WHERE @AND(ABS((@column - mean_@column) / NULLIF(stddev_@column, 0)) > @threshold, @condition) """, ) diff --git a/tests/core/test_audit.py b/tests/core/test_audit.py index b5226563b5..ac2e8abe16 100644 --- a/tests/core/test_audit.py +++ b/tests/core/test_audit.py @@ -697,6 +697,89 @@ def test_pattern_audits(model: Model): ) +def test_z_score_audit(model: Model): + rendered_query = model.render_audit_query( + builtin.z_score_audit, + column=exp.to_column("a"), + threshold=exp.convert(3), + ) + assert ( + rendered_query.sql() + == """WITH "stats" AS (SELECT AVG("a") AS "mean_a", STDDEV("a") AS "stddev_a" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_0" WHERE TRUE) SELECT "a" AS "a", ("a" - "mean_a") / NULLIF("stddev_a", 0) AS "z_score" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_1", "stats" AS "stats" WHERE ABS(("a" - "mean_a") / NULLIF("stddev_a", 0)) > 3 AND TRUE""" + ) + + # The condition must filter the rows being checked, not just the rows used to compute the stats + rendered_query = model.render_audit_query( + builtin.z_score_audit, + column=exp.to_column("a"), + threshold=exp.convert(3), + condition=parse_one("b > 0"), + ) + assert ( + rendered_query.sql() + == """WITH "stats" AS (SELECT AVG("a") AS "mean_a", STDDEV("a") AS "stddev_a" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_0" WHERE "b" > 0) SELECT "a" AS "a", ("a" - "mean_a") / NULLIF("stddev_a", 0) AS "z_score" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_1", "stats" AS "stats" WHERE ABS(("a" - "mean_a") / NULLIF("stddev_a", 0)) > 3 AND "b" > 0""" + ) + + +def test_valid_http_method_audit(model: Model): + rendered_query = model.render_audit_query( + builtin.valid_http_method_audit, + column=exp.to_column("a"), + ) + assert ( + rendered_query.sql() + == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_0" WHERE NOT "a" IN ('GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'HEAD', 'OPTIONS', 'TRACE', 'CONNECT') AND TRUE""" + ) + + rendered_query = model.render_audit_query( + builtin.valid_http_method_audit, + column=exp.to_column("a"), + condition=parse_one("b > 0"), + ) + assert ( + rendered_query.sql() + == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_0" WHERE NOT "a" IN ('GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'HEAD', 'OPTIONS', 'TRACE', 'CONNECT') AND "b" > 0""" + ) + + +def test_builtin_audits_condition_filters_audited_rows(): + import duckdb + + model = create_sql_model( + "db.test_model", + parse_one("SELECT region, v, method FROM db.src"), + kind=FullKind(), + dialect="duckdb", + ) + connection = duckdb.connect() + connection.execute( + "CREATE SCHEMA db; CREATE TABLE db.test_model AS " + "SELECT * FROM (VALUES ('us', 10, 'GET'), ('us', 10, 'POST'), ('us', 11, 'PUT'), " + "('eu', 1000, 'get'), ('eu', 1000, 'post')) AS t(region, v, method)" + ) + condition = parse_one("region = 'us'", dialect="duckdb") + + z_score_query = model.render_audit_query( + builtin.z_score_audit, + column=exp.to_column("v"), + threshold=exp.convert(1), + condition=condition, + ) + # Only the rows matching the condition are compared against the stats computed for them + assert connection.execute(z_score_query.sql("duckdb")).fetchall() == [ + (11, pytest.approx(1.1547005)) + ] + + valid_http_method_query = model.render_audit_query( + builtin.valid_http_method_audit, + column=exp.to_column("method"), + condition=condition, + ) + assert connection.execute(valid_http_method_query.sql("duckdb")).fetchall() == [] + + connection.close() + + def test_standalone_audit(model: Model, assert_exp_eq): audit = StandaloneAudit( name="test_audit", query=parse_one(f"SELECT * FROM {model.name} WHERE col IS NULL")