Repository: TobikoData/sqlmesh
Pull request #2041 (2024-01-29 19:16:27)

Feat: Macro typing coercion
A spike at allowing users to explicitly define types for macro arguments. This makes the interface a little bit nicer since you know exactly what you are gonna get. The implementation is a proof of concept at where we may add this functionality.

Thoughts?

z3z1ma: Test failures unrelated to my change AFAICT
tobymao: can you update the documentation to indicate this new behavior of type annotations?
tobymao: one last small suggestion.

can we get a lot more tests, basically covering the gamut of coercisions and docs? thanks @z3z1ma !
z3z1ma: cool thanks for help @tobymao you rock. i will get to the tests and docs tomorrow
z3z1ma: @tobymao I added tests that hit every code path I think, I might add another 1 or 2. But otherwise going to focus on the docs.

Patch:
diff --git a/examples/sushi/macros/macros.py b/examples/sushi/macros/macros.py
index abf0ebb9..f4c2b6a8 100644
--- a/examples/sushi/macros/macros.py
+++ b/examples/sushi/macros/macros.py
@@ -1,8 +1,9 @@
 from macros.utils import between  # type: ignore
+from sqlglot import exp
 
 from sqlmesh import macro
 
 
 @macro()
-def incremental_by_ds(evaluator, column):
+def incremental_by_ds(evaluator, column: exp.Column):
     return between(evaluator, column, evaluator.locals["start_date"], evaluator.locals["end_date"])
diff --git a/sqlmesh/core/engine_adapter/mssql.py b/sqlmesh/core/engine_adapter/mssql.py
index 875b8fa9..1988441e 100644
--- a/sqlmesh/core/engine_adapter/mssql.py
+++ b/sqlmesh/core/engine_adapter/mssql.py
@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import typing as t
 
+import numpy as np
 import pandas as pd
 from pandas.api.types import is_datetime64_any_dtype  # type: ignore
 from sqlglot import exp
@@ -174,7 +175,7 @@ class MSSQLEngineAdapter(
                 columns_to_types_create = columns_to_types.copy()
                 self._convert_df_datetime(df, columns_to_types_create)
                 self.create_table(temp_table, columns_to_types_create)
-                rows: t.List[t.Tuple[t.Any, ...]] = list(df.itertuples(index=False, name=None))  # type: ignore
+                rows: t.List[t.Tuple[t.Any, ...]] = list(df.replace({np.nan: None}).itertuples(index=False, name=None))  # type: ignore
                 conn = self._connection_pool.get()
                 conn.bulk_copy(temp_table.sql(dialect=self.dialect), rows)
             return exp.select(*self._casted_columns(columns_to_types)).from_(temp_table)  # type: ignore
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index 04258443..fc4d5048 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -1,13 +1,15 @@
 from __future__ import annotations
 
+import inspect
+import logging
 import typing as t
 from enum import Enum
-from functools import reduce
+from functools import reduce, wraps
 from string import Template
 
 import sqlglot
 from jinja2 import Environment
-from sqlglot import Generator, exp
+from sqlglot import Generator, exp, parse_one
 from sqlglot.executor.env import ENV
 from sqlglot.executor.python import Python
 from sqlglot.helper import csv, ensure_collection
@@ -15,6 +17,7 @@ from sqlglot.schema import MappingSchema
 
 from sqlmesh.core.dialect import (
     SQLMESH_MACRO_PREFIX,
+    Dialect,
     MacroDef,
     MacroFunc,
     MacroSQL,
@@ -33,6 +36,8 @@ if t.TYPE_CHECKING:
     from sqlmesh.core.engine_adapter import EngineAdapter
     from sqlmesh.core.snapshot import Snapshot
 
+logger = logging.getLogger(__name__)
+
 
 class RuntimeStage(Enum):
     LOADING = "loading"
@@ -349,6 +354,75 @@ class MacroEvaluator:
             )
         return self.locals["engine_adapter"]
 
+    def _coerce(self, expr: exp.Expression, typ: t.Any, strict: bool = False) -> t.Any:
+        """Coerces the given expression to the specified type on a best-effort basis."""
+        base_err_msg = f"Failed to coerce expression '{expr}' to type '{typ}'."
+        try:
+            if typ is None:
+                return expr
+            base = t.get_origin(typ) or typ
+            # We need to handle t.Union first since we cannot use isinstance with it
+            if base is t.Union:
+                for branch in t.get_args(typ):
+                    try:
+                        return self._coerce(expr, branch, True)
+                    except Exception:
+                        pass
+                raise SQLMeshError(base_err_msg)
+            if isinstance(expr, base):
+                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
+                if into is None:
+                    if isinstance(expr, exp.Literal):
+                        coerced = parse_one(expr.this)
+                    else:
+                        raise SQLMeshError(
+                            f"{base_err_msg} Coercion to {base} requires a literal expression."
+                        )
+                else:
+                    coerced = parse_one(
+                        expr.this if isinstance(expr, exp.Literal) else expr.sql(), into=into
+                    )
+                if isinstance(coerced, base):
+                    return coerced
+                else:
+                    raise SQLMeshError(base_err_msg)
+            if base in (int, float, str) and isinstance(expr, exp.Literal):
+                return base(expr.this)
+            if base is bool and isinstance(expr, exp.Boolean):
+                return expr.this
+            if base is str and isinstance(expr, exp.Expression):
+                return expr.sql(self.dialect)
+            if base is tuple and isinstance(expr, (exp.Tuple, exp.Array)):
+                generic = t.get_args(typ)
+                if not generic:
+                    return tuple(expr.expressions)
+                if generic[-1] is ...:
+                    return tuple(self._coerce(expr, generic[0]) for expr in expr.expressions)
+                elif len(generic) == len(expr.expressions):
+                    return tuple(
+                        self._coerce(expr, generic[i]) for i, expr in enumerate(expr.expressions)
+                    )
+                raise SQLMeshError(f"{base_err_msg} Expected {len(generic)} items.")
+            if base is list and isinstance(expr, (exp.Array, exp.Tuple)):
+                generic = t.get_args(typ)
+                if not generic:
+                    return expr.expressions
+                return [self._coerce(expr, generic[0]) for expr in expr.expressions]
+            raise SQLMeshError(base_err_msg)
+        except Exception:
+            if strict:
+                raise
+            logger.warning(
+                "Coercion of expression '%s' to type '%s' failed. Using non coerced expression.",
+                expr,
+                typ,
+                exc_info=True,
+            )
+            return expr
+
 
 class macro(registry_decorator):
     """Specifies a function is a macro and registers it the global MACROS registry.
@@ -372,7 +446,31 @@ class macro(registry_decorator):
     def __call__(
         self, func: t.Callable[..., DECORATOR_RETURN_TYPE]
     ) -> t.Callable[..., DECORATOR_RETURN_TYPE]:
-        wrapper = super().__call__(func)
+        @wraps(func)
+        def _typed_func(
+            evaluator: MacroEvaluator, *args_: t.Any, **kwargs_: t.Any
+        ) -> DECORATOR_RETURN_TYPE:
+            spec = inspect.getfullargspec(func)
+            annotations = t.get_type_hints(func)
+            kwargs = inspect.getcallargs(func, evaluator, *args_, **kwargs_)
+            for param, value in kwargs.items():
+                coercible_type = annotations.get(param)
+                if not coercible_type:
+                    continue
+                kwargs[param] = evaluator._coerce(value, coercible_type)
+            args = [kwargs.pop(k) for k in spec.args if k in kwargs]
+            if spec.varargs:
+                args.extend(kwargs.pop(spec.varargs, []))
+            return func(*args, **kwargs)
+
+        try:
+            annotated = any(t.get_type_hints(func).keys() - {"return"})
+        except TypeError:
+            annotated = False
+
+        wrapper = super().__call__(
+            func if not annotated else t.cast(t.Callable[..., DECORATOR_RETURN_TYPE], _typed_func)
+        )
 
         # This is useful to identify macros at runtime
         setattr(wrapper, "__sqlmesh_macro__", True)
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index ee0d9ae8..32a2efa5 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -1170,7 +1170,9 @@ class SeedModel(_SqlBasedModel):
             for column in date_or_time_columns:
                 df[column] = pd.to_datetime(df[column])
             df[bool_columns] = df[bool_columns].apply(lambda i: str_to_bool(str(i)))
-            df[string_columns] = df[string_columns].astype(str)
+            df[string_columns] = df[string_columns].mask(
+                cond=lambda x: x.notna(), other=df[string_columns].astype(str)  # type: ignore
+            )
             yield df
 
     def text_diff(self, other: Node) -> str:


Test patch:
diff --git a/tests/core/engine_adapter/test_integration.py b/tests/core/engine_adapter/test_integration.py
index 6a149bee..ee377034 100644
--- a/tests/core/engine_adapter/test_integration.py
+++ b/tests/core/engine_adapter/test_integration.py
@@ -7,6 +7,7 @@ import sys
 import typing as t
 from datetime import timedelta
 
+import numpy as np
 import pandas as pd
 import pytest
 from sqlglot import exp, parse_one
@@ -174,7 +175,7 @@ class TestContext:
     def get_current_data(self, table: exp.Table) -> pd.DataFrame:
         df = self.engine_adapter.fetchdf(exp.select("*").from_(table), quote_identifiers=True)
         if self.dialect == "snowflake" and "id" in df.columns:
-            df["id"] = df["id"].astype("int")
+            df["id"] = df["id"].apply(lambda x: x if pd.isna(x) else int(x))
         return df
 
     def compare_with_current(self, table: exp.Table, expected: pd.DataFrame) -> None:
@@ -820,6 +821,34 @@ def test_drop_schema(ctx: TestContext):
     assert len(results.views) == 0
 
 
+def test_nan_roundtrip(ctx: TestContext):
+    if ctx.test_type != "df":
+        pytest.skip("NaN roundtrip test only relevant for dataframes.")
+    ctx.engine_adapter.DEFAULT_BATCH_SIZE = sys.maxsize
+    ctx.init()
+    table = ctx.table("test_table")
+    # Initial Load
+    input_data = pd.DataFrame(
+        [
+            {"id": 1, "ds": "2022-01-01"},
+            {"id": 2, "ds": "2022-01-02"},
+            {"id": np.nan, "ds": np.nan},
+        ]
+    )
+    ctx.engine_adapter.create_table(table, ctx.columns_to_types)
+    ctx.engine_adapter.replace_query(
+        table,
+        ctx.input_data(input_data),
+        columns_to_types=ctx.columns_to_types,
+    )
+    results = ctx.get_metadata_results()
+    assert not results.views
+    assert not results.materialized_views
+    assert len(results.tables) == len(results.non_temp_tables) == 1
+    assert results.non_temp_tables[0] == table.name
+    ctx.compare_with_current(table, input_data)
+
+
 def test_replace_query(ctx: TestContext):
     ctx.engine_adapter.DEFAULT_BATCH_SIZE = sys.maxsize
     ctx.init()
diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py
index 74508c6b..cadf8db3 100644
--- a/tests/core/test_macros.py
+++ b/tests/core/test_macros.py
@@ -27,6 +27,41 @@ def macro_evaluator() -> MacroEvaluator:
     def noop(evaluator: MacroEvaluator):
         return None
 
+    @macro()
+    def bitshift_square(evaluator: MacroEvaluator, x: int, y: int) -> int:
+        return (x >> y) ** 2
+
+    @macro()
+    def prefix_db(evaluator: MacroEvaluator, table: exp.Table, prefix: str) -> exp.Table:
+        table.set("db", prefix + table.db)
+        return table
+
+    @macro()
+    def repeated(evaluator: MacroEvaluator, expr: str, times: int = 2, multi: bool = False):
+        if multi is True:
+            return (expr,) * times
+        return expr * times
+
+    @macro()
+    def split(evaluator: MacroEvaluator, string: str, sep: str = ","):
+        return string.split(sep)
+
+    @macro()
+    def cte_tag_name(evaluator: MacroEvaluator, with_: exp.Select):
+        for cte in with_.find_all(exp.CTE):
+            name = cte.alias_or_name
+            for query in cte.find_all(exp.Select):
+                query.select(exp.Literal.string(name).as_("source"), copy=False)
+        return with_
+
+    @macro()
+    def suffix_idents(evaluator: MacroEvaluator, items: t.List[str], suffix: str):
+        return [item + suffix for item in items]
+
+    @macro()
+    def suffix_idents_2(evaluator: MacroEvaluator, items: t.Tuple[str, ...], suffix: str):
+        return [item + suffix for item in items]
+
     return MacroEvaluator(
         "hive",
         {"test": Executable(name="test", payload=f"def test(_):\n    return 'test'")},
@@ -292,12 +327,131 @@ def test_ast_correctness(macro_evaluator):
             "SELECT d",
             {},
         ),
+        (
+            """@BITSHIFT_SQUARE(50, '3')""",
+            "36",
+            {},
+        ),
+        (
+            """@PREFIX_DB(my.schema.table, 'dev_')""",
+            "my.dev_schema.table",
+            {},
+        ),
+        (
+            """select @REPEATED(test, 3)""",
+            "SELECT testtesttest",
+            {},
+        ),
+        (
+            """select @REPEATED(test, 3, true)""",
+            "SELECT test, test, test",
+            {},
+        ),
+        (
+            """select @SPLIT('a,b,c')""",
+            "SELECT a, b, c",
+            {},
+        ),
+        (
+            """@CTE_TAG_NAME(WITH step1 AS (SELECT 1) SELECT * FROM step1)""",
+            "WITH step1 AS (SELECT 1, 'step1' AS source) SELECT * FROM step1",
+            {},
+        ),
+        (
+            """@CTE_TAG_NAME('WITH step1 AS (SELECT 1) SELECT * FROM step1')""",
+            "WITH step1 AS (SELECT 1, 'step1' AS source) SELECT * FROM step1",
+            {},
+        ),
+        (
+            """SELECT @SUFFIX_IDENTS(['a', 'b', 'c'], 'z')""",
+            "SELECT az, bz, cz",
+            {},
+        ),
+        (
+            """SELECT @SUFFIX_IDENTS_2(['a', 'b', 'c'], 'z')""",
+            "SELECT az, bz, cz",
+            {},
+        ),
     ],
 )
-def test_macro_functions(macro_evaluator, assert_exp_eq, sql, expected, args):
+def test_macro_functions(macro_evaluator: MacroEvaluator, assert_exp_eq, sql, expected, args):
     macro_evaluator.locals = args or {}
     assert_exp_eq(macro_evaluator.transform(parse_one(sql)), expected)
 
 
-def test_macro_returns_none(macro_evaluator):
+def test_macro_returns_none(macro_evaluator: MacroEvaluator):
     assert macro_evaluator.transform(parse_one("@NOOP()")) is None
+
+
+def test_macro_coercion(macro_evaluator: MacroEvaluator, assert_exp_eq):
+    coerce = macro_evaluator._coerce
+    assert coerce(exp.Literal.number(1), int) == 1
+    assert coerce(exp.Literal.number(1.1), float) == 1.1
+    assert coerce(exp.Literal.string("Hi mom"), str) == "Hi mom"
+    assert coerce(exp.true(), bool) is True
+
+    # Coercing a string literal to a column should return a column with the same name
+    assert_exp_eq(coerce(exp.Literal.string("order"), exp.Column), exp.column("order"))
+    # Not possible to coerce this string literal Cast to an exp.Column node -- so it should just return the input
+    assert_exp_eq(
+        coerce(exp.Literal.string("order::date"), exp.Column), exp.Literal.string("order::date")
+    )
+    # This however, is correctly coercible since it's a cast
+    assert_exp_eq(
+        coerce(exp.Literal.string("order::date"), exp.Cast), exp.cast(exp.column("order"), "DATE")
+    )
+
+    # Here we resolve ambiguity via the user type hint
+    assert_exp_eq(coerce(exp.Literal.string("order"), exp.Identifier), exp.to_identifier("order"))
+    assert_exp_eq(coerce(exp.Literal.string("order"), exp.Table), exp.table_("order"))
+
+    # Resolve a union type hint by choosing the first one that works
+    assert_exp_eq(
+        coerce(exp.Literal.string("order::date"), t.Union[exp.Column, exp.Cast]),
+        exp.cast(exp.column("order"), "DATE"),
+    )
+
+    # Simply ask for a string, and always get a string
+    assert coerce(exp.column("order"), str) == "order"
+    assert (
+        coerce(parse_one("SELECT x FROM UNNEST(y) AS x WHERE 1 = 1"), str)
+        == "SELECT x FROM UNNEST(y) AS x WHERE 1 = 1"
+    )
+
+    # From a string literal to a Select should parse the string literal, and the inverse operation works as well
+    assert_exp_eq(
+        coerce(exp.Literal.string("SELECT 1 FROM a"), exp.Select), parse_one("SELECT 1 FROM a")
+    )
+    assert coerce(parse_one("SELECT 1 FROM a"), str) == "SELECT 1 FROM a"
+
+    # Get a list of exp directly instead of an exp.Array
+    assert coerce(parse_one("[1, 2, 3]"), list) == [
+        exp.Literal.number(1),
+        exp.Literal.number(2),
+        exp.Literal.number(3),
+    ]
+
+    # Generics work as well, recursively resolving inner types
+    assert coerce(parse_one("[1, 2, 3]"), t.List[int]) == [1, 2, 3]
+    assert coerce(parse_one("[1, 2, 3]"), t.Tuple[int, int, float]) == (1, 2, 3.0)
+    assert coerce(parse_one("[1, 2, 3]"), t.Tuple[int, ...]) == (1, 2, 3)
+    assert coerce(parse_one("[1, 2, 3]"), t.Tuple[int, str, float]) == (1, "2", 3.0)
+    assert coerce(
+        parse_one("[1, 2, [3]]"), t.Tuple[int, str, t.Union[float, t.Tuple[float, ...]]]
+    ) == (1, "2", (3.0,))
+
+    # Using exp.Expression will always return the input expression
+    assert coerce(parse_one("order", into=exp.Column), exp.Expression) == exp.column("order")
+    assert coerce(exp.Literal.string("OK"), exp.Expression) == exp.Literal.string("OK")
+
+    # Strict flag allows raising errors and is used when recursively coercing expressions
+    # otherwise, in general, we want to be lenient and just warn the user when something is not possible
+    with pytest.raises(SQLMeshError):
+        coerce(exp.Literal.string("order"), exp.Select, strict=True)
+
+    with pytest.raises(SQLMeshError):
+        _ = coerce(
+            exp.Literal.string("order::date"),
+            t.Union[exp.Column, exp.Identifier],
+            strict=True,
+        )
diff --git a/tests/core/test_selector.py b/tests/core/test_selector.py
index b3740b0e..6d4fc537 100644
--- a/tests/core/test_selector.py
+++ b/tests/core/test_selector.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import typing as t
+from pathlib import Path
 from unittest.mock import call
 
 import pytest
@@ -373,6 +374,51 @@ def test_expand_model_selections(
     assert selector.expand_model_selections(selections) == output
 
 
+@pytest.mark.parametrize(
+    "expressions, expected_fqns",
+    [
+        (["git:main"], {'"test_model_a"', '"test_model_c"'}),
+        (["git:main & +*model_c"], {'"test_model_c"'}),
+    ],
+)
+def test_expand_git_selection(
+    mocker: MockerFixture, expressions: t.List[str], expected_fqns: t.Set[str]
+):
+    models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
+
+    model_a = SqlModel(name="test_model_a", query=d.parse_one("SELECT 1 AS a"))
+    model_a._path = Path("/path/to/test_model_a.sql")
+    models[model_a.fqn] = model_a
+
+    model_b = SqlModel(name="test_model_b", query=d.parse_one("SELECT 2 AS b"))
+    model_b._path = Path("/path/to/test_model_b.sql")
+    models[model_b.fqn] = model_b
+
+    model_c = SqlModel(name="test_model_c", query=d.parse_one("SELECT 3 AS c"))
+    model_c._path = Path("/path/to/test_model_c.sql")
+    models[model_c.fqn] = model_c
+
+    model_d = SqlModel(
+        name="test_model_d",
+        query=d.parse_one("SELECT c FROM test_model_c"),
+        depends_on={"test_model_c"},
+    )
+    model_d._path = Path("/path/to/test_model_d.sql")
+    models[model_d.fqn] = model_d
+
+    git_client_mock = mocker.Mock()
+    git_client_mock.list_untracked_files.return_value = []
+    git_client_mock.list_changed_files.return_value = [model_a._path, model_c._path]
+
+    selector = Selector(mocker.Mock(), models)
+    selector._git_client = git_client_mock
+
+    assert selector.expand_model_selections(expressions) == expected_fqns
+
+    git_client_mock.list_changed_files.assert_called_once_with(target_branch="main")
+    git_client_mock.list_untracked_files.assert_called_once()
+
+
 def test_select_models_with_external_parent(mocker: MockerFixture):
     default_catalog = "test_catalog"
     added_model = SqlModel(