Repository: TobikoData/sqlmesh
Pull request #2901 (2024-07-15 00:03:07)

Feat: Add support for defaults in user audits
I was writing an audit and I figured it would be nice to have defaults for the variables.
I tried it and it almost already works, except it was loading the values as strings instead of expressions.

The goal is to be able to define an audit with defaults like this:
audit (
    name generic_sums_add_up,
    defaults (
        percentage_absolute_difference_accepted=0.01,
    )
);

Since defaults has to be a dictionary of string to Expression, it seems to make sense to parse the string values into Expression.

defaults: t.Dict[str, exp.Expression]

I tried the code and it works, but there are some points I'm not sure about.
1. Is it possible to support not needing a trailing comma when there's only one key in the dictionary?
2. Is sqlglot.parse the right way to do it? How would we get the audit's configured dialect and the project's dialect to pass them in to parse?
3. Is it correct to take out the first element of what parse returns? I tried with a default value that was a list like (0.01, 0.02) and it kept the whole list so that seemed fine.

The formatting of the code is probably not right. I saw there's a make file but not sure how to run it is there a page on that you can point me to? Anything on standards for the codebase would also help.

I could also provide a draft for documentation (I know where the docs file is) and some tests (not sure where I'd put them).

CLAassistant: [![CLA assistant check](https://cla-assistant.io/pull/badge/not_signed)](https://cla-assistant.io/TobikoData/sqlmesh?pullRequest=2901) <br/>Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our [Contributor License Agreement](https://cla-assistant.io/TobikoData/sqlmesh?pullRequest=2901) before we can accept your contribution.<br/><sub>You have signed the CLA already but the status is still pending? Let us [recheck](https://cla-assistant.io/check/TobikoData/sqlmesh?pullRequest=2901) it.</sub>
alexischicoinedatadeveloper: I added the change to parse non expressions in the branch for Tuple and Array, but when I passed in a single element with no trailing comma it came back as a Paren so maybe that's not needed?
tobymao: well do you need to add the case for Paren?
alexischicoinedatadeveloper: > well do you need to add the case for Paren?

Without addressing the Paren case it throws this error when the trailing comma isn't there.
 sqlmesh.utils.errors.AuditConfigError: Defaults must be a tuple of exp.EQ or a dict at '.'
alexischicoinedatadeveloper: I saw another part of the code get the dialect like this does it make sense here as well?
https://github.com/TobikoData/sqlmesh/blob/98d18093dae793bc1c0c90b61cf7351d17cada19/sqlmesh/utils/pydantic.py#L78
tobymao: > I saw another part of the code get the dialect like this does it make sense here as well?
> 
> https://github.com/TobikoData/sqlmesh/blob/98d18093dae793bc1c0c90b61cf7351d17cada19/sqlmesh/utils/pydantic.py#L78

yea i think so, can you use it?
achicoine-coveo: > > I saw another part of the code get the dialect like this does it make sense here as well?
> > https://github.com/TobikoData/sqlmesh/blob/98d18093dae793bc1c0c90b61cf7351d17cada19/sqlmesh/utils/pydantic.py#L78
> 
> yea i think so, can you use it?

To use the function directly?
from sqlmesh.core.model import model

    dialect = (values if isinstance(values, dict) else values.data).get("dialect")
    return model._dialect if dialect is None else dialect

I'm just not sure what this model part does. Would that part also make sense for this usage? Maybe that would be different than the dialect defined in the audit?
tobymao: > > > I saw another part of the code get the dialect like this does it make sense here as well?
> > > https://github.com/TobikoData/sqlmesh/blob/98d18093dae793bc1c0c90b61cf7351d17cada19/sqlmesh/utils/pydantic.py#L78
> > 
> > 
> > yea i think so, can you use it?
> 
> To use the function directly?
> 
> ```python
>     from sqlmesh.core.model import model
> 
>     dialect = (values if isinstance(values, dict) else values.data).get("dialect")
>     return model._dialect if dialect is None else dialect
> ```
> 
> I'm just not sure what this model part does. Would that part also make sense for this usage? Maybe that would be different than the dialect defined in the audit?

yea can we try this?
achicoine-coveo: > > > > I saw another part of the code get the dialect like this does it make sense here as well?
> > > > https://github.com/TobikoData/sqlmesh/blob/98d18093dae793bc1c0c90b61cf7351d17cada19/sqlmesh/utils/pydantic.py#L78
> > > 
> > > 
> > > yea i think so, can you use it?
> > 
> > 
> > To use the function directly?
> > ```python
> >     from sqlmesh.core.model import model
> > 
> >     dialect = (values if isinstance(values, dict) else values.data).get("dialect")
> >     return model._dialect if dialect is None else dialect
> > ```
> > 
> > 
> >     
> >       
> >     
> > 
> >       
> >     
> > 
> >     
> >   
> > I'm just not sure what this model part does. Would that part also make sense for this usage? Maybe that would be different than the dialect defined in the audit?
> 
> yea can we try this?

Ok I'll put it in the code. I'll try to see if I can see the value of dialect from model and what happens if it's different from the one in audit.
achicoine-coveo: There's a test that fails, but I don't see how it's related to the changes.

Patch:
diff --git a/docs/concepts/audits.md b/docs/concepts/audits.md
index 61643803..5e1a3670 100644
--- a/docs/concepts/audits.md
+++ b/docs/concepts/audits.md
@@ -75,6 +75,19 @@ Notice how `column` and `threshold` parameters have been set. These values will
 
 Note that the same audit can be applied more than once to the a model using different sets of parameters.
 
+Generic audits can define default values as follows:
+```sql linenums="1"
+AUDIT (
+  name does_not_exceed_threshold,
+  defaults (
+    threshold = 10,
+    column = id
+  )
+);
+SELECT * FROM @this_model
+WHERE @column >= @threshold;
+```
+
 ### Naming
 We recommended avoiding SQL keywords when naming audit parameters. Quote any audit argument that is also a SQL keyword.
 
diff --git a/sqlmesh/core/audit/definition.py b/sqlmesh/core/audit/definition.py
index aa98f6e7..f4c2e1f8 100644
--- a/sqlmesh/core/audit/definition.py
+++ b/sqlmesh/core/audit/definition.py
@@ -34,6 +34,7 @@ from sqlmesh.utils.metaprogramming import Executable
 from sqlmesh.utils.pydantic import (
     PydanticModel,
     field_validator,
+    get_dialect,
     model_validator,
     model_validator_v1_args,
 )
@@ -148,11 +149,19 @@ def audit_string_validator(cls: t.Type, v: t.Any) -> t.Optional[str]:
 
 
 @field_validator("defaults", mode="before", check_fields=False)
-def audit_map_validator(cls: t.Type, v: t.Any) -> t.Dict[str, t.Any]:
+def audit_map_validator(cls: t.Type, v: t.Any, values: t.Any) -> t.Dict[str, t.Any]:
+    if isinstance(v, exp.Paren):
+        return dict([_maybe_parse_arg_pair(v.unnest())])
     if isinstance(v, (exp.Tuple, exp.Array)):
         return dict(map(_maybe_parse_arg_pair, v.expressions))
     elif isinstance(v, dict):
-        return v
+        dialect = get_dialect(values)
+        return {
+            key: value
+            if isinstance(value, exp.Expression)
+            else d.parse_one(str(value), dialect=dialect)
+            for key, value in v.items()
+        }
     else:
         raise_config_error(
             "Defaults must be a tuple of exp.EQ or a dict", error_type=AuditConfigError


Test patch:
diff --git a/tests/core/test_audit.py b/tests/core/test_audit.py
index 9ab97a43..e07f757d 100644
--- a/tests/core/test_audit.py
+++ b/tests/core/test_audit.py
@@ -303,6 +303,49 @@ def test_load_multiple(assert_exp_eq):
     )
 
 
+def test_load_with_dictionary_defaults():
+    expressions = parse(
+        """
+        AUDIT (
+            name my_audit,
+            dialect spark,
+            defaults (
+                field1 = some_column,
+                field2 = 3
+            ),
+        );
+
+        SELECT 1
+    """
+    )
+
+    audit = load_audit(expressions, dialect="spark")
+    assert audit.defaults.keys() == {"field1", "field2"}
+    for value in audit.defaults.values():
+        assert isinstance(value, exp.Expression)
+
+
+def test_load_with_single_defaults():
+    # testing it also works with a single default with no trailing comma
+    expressions = parse(
+        """
+        AUDIT (
+            name my_audit,
+            defaults (
+                field1 = some_column
+            ),
+        );
+
+        SELECT 1
+    """
+    )
+
+    audit = load_audit(expressions, dialect="duckdb")
+    assert audit.defaults.keys() == {"field1"}
+    for value in audit.defaults.values():
+        assert isinstance(value, exp.Expression)
+
+
 def test_no_audit_statement():
     expressions = parse(
         """