Fix: make macro function argument resolution more robust (#2604)

* Fix: refactor macro function argument resolution so it's more robust

* Fix mypy issue

* Throw MacroEvalError instead of TypeError, add tests

* Update comment

* Fix typo in macro docs (except -> except_)

* Incorporate PR feedback

Author: georgesittas

docs/concepts/macros/sqlmesh_macros.md

@@ -600,12 +600,12 @@ If the column data types are known, the resulting query `CAST`s columns to their
 
 - `relation`: The relation/table whose columns are being selected
 - `alias` (optional): The alias of the relation (if it has one)
-- `except` (optional): A list of columns to exclude
+- `except_` (optional): A list of columns to exclude
 - `prefix` (optional): A string to use as a prefix for all selected column names
 - `suffix` (optional): A string to use as a suffix for all selected column names
 - `quote_identifiers` (optional): Whether to quote the resulting identifiers, defaults to true
 
-Like all SQLMesh macro functions, omitting an argument when calling `@STAR` requires passing all subsequent arguments with their name and the special `:=` keyword operator. For example, we might omit the `alias` argument with `@STAR(foo, except := [c])`. Learn more about macro function arguments [below](#positional-and-keyword-arguments).
+Like all SQLMesh macro functions, omitting an argument when calling `@STAR` requires passing all subsequent arguments with their name and the special `:=` keyword operator. For example, we might omit the `alias` argument with `@STAR(foo, except_ := [c])`. Learn more about macro function arguments [below](#positional-and-keyword-arguments).
 
 As a `@STAR` example, consider the following query:
 
@@ -635,15 +635,15 @@ FROM foo AS bar
 Note these aspects of the rendered query:
 - Each column is `CAST` to its data type in the table `foo` (e.g., `a` to `TEXT`)
 - Each column selection uses the alias `bar` (e.g., `"bar"."a"`)
-- Column `c` is not present because it was passed to `@STAR`'s `except` argument
+- Column `c` is not present because it was passed to `@STAR`'s `except_` argument
 - Each column alias is prefixed with `baz_` and suffixed with `_qux` (e.g., `"baz_a_qux"`)
 
 Now consider a more complex example that provides different prefixes to `a` and `b` than to `d` and includes an explicit column `my_column`:
 
 ```sql linenums="1"
 SELECT
-  @STAR(foo, bar, except=[c, d], 'ab_pre_'),
-  @STAR(foo, bar, except=[a, b, c], 'd_pre_'),
+  @STAR(foo, bar, except_=[c, d], 'ab_pre_'),
+  @STAR(foo, bar, except_=[a, b, c], 'd_pre_'),
   my_column
 FROM foo AS bar
 ```
@@ -661,7 +661,7 @@ FROM foo AS bar
 
 Note these aspects of the rendered query:
 - Columns `a` and `b` have the prefix `"ab_pre_"` , while column `d` has the prefix `"d_pre_"`
-- Column `c` is not present because it was passed to the `except` argument in both `@STAR` calls
+- Column `c` is not present because it was passed to the `except_` argument in both `@STAR` calls
 - `my_column` is present in the query
 
 ### @GENERATE_SURROGATE_KEY

sqlmesh/core/macros.py

@@ -183,35 +183,39 @@ def send(
         if not callable(func):
             raise SQLMeshError(f"Macro '{name}' does not exist.")
 
+        try:
+            # Bind the macro's actual parameters to its formal parameters
+            sig = inspect.signature(func)
+            bound = sig.bind(self, *args, **kwargs)
+            bound.apply_defaults()
+        except Exception as e:
+            print_exception(e, self.python_env)
+            raise MacroEvalError("Error trying to eval macro.") from e
+
         try:
             annotations = t.get_type_hints(func)
         except NameError:  # forward references aren't handled
             annotations = {}
 
+        # If the macro is annotated, we try coerce the actual parameters to the corresponding types
         if annotations:
-            spec = inspect.getfullargspec(func)
-            callargs = inspect.getcallargs(func, self, *args, **kwargs)
-            new_args: t.List[t.Any] = []
-
-            for arg, value in callargs.items():
+            for arg, value in bound.arguments.items():
                 typ = annotations.get(arg)
-
-                if value is self:
+                if not typ:
                     continue
-                if arg == spec.varargs:
-                    new_args.extend(self._coerce(v, typ) for v in value)
-                elif arg == spec.varkw:
-                    for k, v in value.items():
-                        kwargs[k] = self._coerce(v, typ)
-                elif arg in kwargs:
-                    kwargs[arg] = self._coerce(value, typ)
-                else:
-                    new_args.append(self._coerce(value, typ))
 
-            args = new_args  # type: ignore
+                # Changes to bound.arguments will reflect in bound.args and bound.kwargs
+                # https://docs.python.org/3/library/inspect.html#inspect.BoundArguments.arguments
+                param = sig.parameters[arg]
+                if param.kind is inspect.Parameter.VAR_POSITIONAL:
+                    bound.arguments[arg] = tuple(self._coerce(v, typ) for v in value)
+                elif param.kind is inspect.Parameter.VAR_KEYWORD:
+                    bound.arguments[arg] = {k: self._coerce(v, typ) for k, v in value.items()}
+                else:
+                    bound.arguments[arg] = self._coerce(value, typ)
 
         try:
-            return func(self, *args, **kwargs)
+            return func(*bound.args, **bound.kwargs)
         except Exception as e:
             print_exception(e, self.python_env)
             raise MacroEvalError("Error trying to eval macro.") from e
@@ -337,7 +341,8 @@ def evaluate(self, node: MacroFunc) -> exp.Expression | t.List[exp.Expression] |
                 else:
                     if kwargs:
                         raise MacroEvalError(
-                            f"Positional argument cannot follow keyword argument.\n  {func.sql(dialect=self.dialect)} at '{self._path}'"
+                            "Positional argument cannot follow keyword argument.\n  "
+                            f"{func.sql(dialect=self.dialect)} at '{self._path}'"
                         )
 
                     args.append(e)
@@ -479,7 +484,7 @@ def _coerce(self, expr: exp.Expression, typ: t.Any, strict: bool = False) -> t.A
                 return expr
             if issubclass(base, exp.Expression):
                 d = Dialect.get_or_raise(self.dialect)
-                into = base if base in d.parser().EXPRESSION_PARSERS else None
+                into = base if base in d.parser_class.EXPRESSION_PARSERS else None
                 if into is None:
                     if isinstance(expr, exp.Literal):
                         coerced = parse_one(expr.this)
@@ -809,7 +814,7 @@ def star(
         >>> from sqlglot import parse_one, exp
         >>> from sqlglot.schema import MappingSchema
         >>> from sqlmesh.core.macros import MacroEvaluator
-        >>> sql = "SELECT @STAR(foo, bar, [c], 'baz_') FROM foo AS bar"
+        >>> sql = "SELECT @STAR(foo, bar, except_ := [c], prefix := 'baz_') FROM foo AS bar"
         >>> MacroEvaluator(schema=MappingSchema({"foo": {"a": exp.DataType.build("string"), "b": exp.DataType.build("string"), "c": exp.DataType.build("string"), "d": exp.DataType.build("int")}})).transform(parse_one(sql)).sql()
         'SELECT CAST("bar"."a" AS TEXT) AS "baz_a", CAST("bar"."b" AS TEXT) AS "baz_b", CAST("bar"."d" AS INT) AS "baz_d" FROM foo AS bar'
     """

sqlmesh/dbt/relation.py

@@ -4,4 +4,4 @@
 if DBT_VERSION < (1, 8):
     from dbt.contracts.relation import *  # type: ignore  # noqa: F403
 else:
-    from dbt.adapters.contracts.relation import *  # noqa: F403
+    from dbt.adapters.contracts.relation import *  # type: ignore # noqa: F403

tests/core/test_macros.py

@@ -76,6 +76,18 @@ def suffix_idents_2(evaluator: MacroEvaluator, items: t.Tuple[str, ...], suffix:
     def stamped(evaluator, query: exp.Select) -> exp.Subquery:
         return query.select(exp.Literal.string("2024-01-01").as_("stamp")).subquery()
 
+    @macro()
+    def test_arg_resolution(evaluator, pos_only, /, a1, *, a2=1, **rest):
+        return 1
+
+    @macro()
+    def test_default_arg_coercion(
+        evaluator: MacroEvaluator,
+        a1: int = 1,
+        a2: int = exp.Literal.number(2),  # type: ignore
+    ):
+        return sum([a1, a2])
+
     return MacroEvaluator(
         "hive",
         {"test": Executable(name="test", payload="def test(_):\n    return 'test'")},
@@ -440,6 +452,11 @@ def test_ast_correctness(macro_evaluator):
             "SELECT 'a1-b1-c2-d:d1-e:e2'",
             {},
         ),
+        (
+            """select @TEST_DEFAULT_ARG_COERCION()""",
+            "SELECT 3",
+            {},
+        ),
     ],
 )
 def test_macro_functions(macro_evaluator: MacroEvaluator, assert_exp_eq, sql, expected, args):
@@ -537,6 +554,34 @@ def test_macro_coercion(macro_evaluator: MacroEvaluator, assert_exp_eq):
         )
 
 
-def test_positional_follows_kwargs(macro_evaluator: MacroEvaluator):
+def test_positional_follows_kwargs(macro_evaluator):
     with pytest.raises(MacroEvalError, match="Positional argument cannot follow"):
-        macro_evaluator.evaluate(parse_one("@repeated(x, multi := True, 3)"))  # type: ignore
+        macro_evaluator.evaluate(parse_one("@repeated(x, multi := True, 3)"))
+
+
+def test_macro_parameter_resolution(macro_evaluator):
+    with pytest.raises(MacroEvalError) as e:
+        macro_evaluator.evaluate(parse_one("@test_arg_resolution()"))
+    assert str(e.value.__cause__) == "missing a required argument: 'pos_only'"
+
+    with pytest.raises(MacroEvalError) as e:
+        macro_evaluator.evaluate(parse_one("@test_arg_resolution(a1 := 1)"))
+    assert str(e.value.__cause__) == "missing a required argument: 'pos_only'"
+
+    with pytest.raises(MacroEvalError) as e:
+        macro_evaluator.evaluate(parse_one("@test_arg_resolution(1)"))
+    assert str(e.value.__cause__) == "missing a required argument: 'a1'"
+
+    with pytest.raises(MacroEvalError) as e:
+        macro_evaluator.evaluate(parse_one("@test_arg_resolution(1, a2 := 2)"))
+    assert str(e.value.__cause__) == "missing a required argument: 'a1'"
+
+    with pytest.raises(MacroEvalError) as e:
+        macro_evaluator.evaluate(parse_one("@test_arg_resolution(pos_only := 1)"))
+    assert str(e.value.__cause__) == (
+        "'pos_only' parameter is positional only, but was passed as a keyword"
+    )
+
+    with pytest.raises(MacroEvalError) as e:
+        macro_evaluator.evaluate(parse_one("@test_arg_resolution(1, 2, 3)"))
+    assert str(e.value.__cause__) == "too many positional arguments"