Repository: TobikoData/sqlmesh
Pull request #2353 (2024-03-27 02:05:41)

Feat: generate CTE fixtures too when using the create_test command
Addresses the seventh item in https://github.com/TobikoData/sqlmesh/issues/1637. Was wondering whether we want to make the CTE generation optional. What do people think about this?

EDIT: went ahead and made it optional for the time being.


Patch:
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index 4bd8e96c..81c58813 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -478,6 +478,13 @@ def dag(ctx: click.Context, file: str, select_model: t.List[str]) -> None:
     "name",
     help="The name of the test that will be created. By default, it's inferred based on the model's name.",
 )
+@click.option(
+    "--include-ctes",
+    "include_ctes",
+    is_flag=True,
+    default=False,
+    help="When true, CTE fixtures will also be generated.",
+)
 @click.pass_obj
 @error_handler
 def create_test(
@@ -486,8 +493,9 @@ def create_test(
     queries: t.List[t.Tuple[str, str]],
     overwrite: bool = False,
     variables: t.Optional[t.List[t.Tuple[str, str]]] = None,
-    name: t.Optional[str] = None,
     path: t.Optional[str] = None,
+    name: t.Optional[str] = None,
+    include_ctes: bool = False,
 ) -> None:
     """Generate a unit test fixture for a given model."""
     obj.create_test(
@@ -495,8 +503,9 @@ def create_test(
         input_queries=dict(queries),
         overwrite=overwrite,
         variables=dict(variables) if variables else None,
-        name=name,
         path=path,
+        name=name,
+        include_ctes=include_ctes,
     )
 
 
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 2e624f12..c012f66d 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -1281,6 +1281,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         variables: t.Optional[t.Dict[str, str]] = None,
         path: t.Optional[str] = None,
         name: t.Optional[str] = None,
+        include_ctes: bool = False,
     ) -> None:
         """Generate a unit test fixture for a given model.
 
@@ -1295,6 +1296,7 @@ class GenericContext(BaseContext, t.Generic[C]):
                 By default, the fixture will be created under the test directory and the file name
                 will be inferred from the test's name.
             name: The name of the test. This is inferred from the model name by default.
+            include_ctes: When true, CTE fixtures will also be generated.
         """
         input_queries = {
             # The get_model here has two purposes: return normalized names & check for missing deps
@@ -1313,6 +1315,7 @@ class GenericContext(BaseContext, t.Generic[C]):
             variables=variables,
             path=path,
             name=name,
+            include_ctes=include_ctes,
         )
 
     def test(
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index 1b98d4c1..3ee0908e 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -750,6 +750,11 @@ class SQLMeshMagics(Magics):
         type=str,
         help="The name of the test that will be created. By default, it's inferred based on the model's name.",
     )
+    @argument(
+        "--include-ctes",
+        action="store_true",
+        help="When true, CTE fixtures will also be generated.",
+    )
     @line_magic
     @pass_sqlmesh_context
     def create_test(self, context: Context, line: str) -> None:
@@ -762,8 +767,9 @@ class SQLMeshMagics(Magics):
             input_queries={k: v.strip('"') for k, v in dict(zip(queries, queries)).items()},
             overwrite=args.overwrite,
             variables=dict(zip(variables, variables)) if variables else None,
-            name=args.name,
             path=args.path,
+            name=args.name,
+            include_ctes=args.include_ctes,
         )
 
     @magic_arguments()


Test patch:
diff --git a/sqlmesh/core/test/definition.py b/sqlmesh/core/test/definition.py
index 0eb5125c..267bdfcf 100644
--- a/sqlmesh/core/test/definition.py
+++ b/sqlmesh/core/test/definition.py
@@ -399,6 +399,7 @@ def generate_test(
     variables: t.Optional[t.Dict[str, str]] = None,
     path: t.Optional[str] = None,
     name: t.Optional[str] = None,
+    include_ctes: bool = False,
 ) -> None:
     """Generate a unit test fixture for a given model.
 
@@ -417,6 +418,7 @@ def generate_test(
             By default, the fixture will be created under the test directory and the file name
             will be inferred from the test's name.
         name: The name of the test. This is inferred from the model name by default.
+        include_ctes: When true, CTE fixtures will also be generated.
     """
     test_name = name or f"test_{model.view_name}"
     path = path or f"{test_name}.yaml"
@@ -471,6 +473,26 @@ def generate_test(
             engine_adapter=test_engine_adapter,
             table_mapping=mapping,
         )
+
+        if include_ctes:
+            ctes = {}
+            previous_ctes: t.List[exp.CTE] = []
+            for cte in model_query.ctes:
+                cte_query = cte.this
+                for prev in previous_ctes:
+                    cte_query = cte_query.with_(prev.alias, prev.this)
+
+                cte_output = t.cast(SqlModelTest, test)._execute(cte_query)
+                ctes[cte.alias] = pandas_timestamp_to_pydatetime(
+                    cte_output.apply(lambda col: col.map(_normalize_dataframe)),
+                    cte_query.named_selects,
+                ).to_dict(orient="records")
+
+                previous_ctes.append(cte)
+
+            if ctes:
+                outputs["ctes"] = ctes
+
         output = t.cast(SqlModelTest, test)._execute(model_query)
     else:
         output = t.cast(PythonModelTest, test)._execute_model()
diff --git a/tests/core/test_test.py b/tests/core/test_test.py
index 6607cd83..7c6c6af1 100644
--- a/tests/core/test_test.py
+++ b/tests/core/test_test.py
@@ -5,6 +5,7 @@ import typing as t
 from pathlib import Path
 
 import pytest
+from sqlglot import exp
 
 from sqlmesh.cli.example_project import init_example_project
 from sqlmesh.core import constants as c
@@ -712,8 +713,16 @@ def test_test_generation(tmp_path: Path) -> None:
         default_connection=DuckDBConnectionConfig(),
         model_defaults=ModelDefaultsConfig(dialect="duckdb"),
     )
-
     context = Context(paths=tmp_path, config=config)
+
+    query = context.get_model("sqlmesh_example.full_model").render_query()
+    assert isinstance(query, exp.Query)
+
+    context.upsert_model(
+        "sqlmesh_example.full_model",
+        query=exp.select(*query.named_selects).from_("cte").with_("cte", as_=query),
+    )
+
     context.plan(auto_apply=True)
 
     input_queries = {
@@ -733,12 +742,14 @@ def test_test_generation(tmp_path: Path) -> None:
     assert len(test) == 1
     assert "test_example_full_model" in test
     assert "vars" not in test["test_example_full_model"]
+    assert "ctes" not in test["test_example_full_model"]["outputs"]
 
     context.create_test(
         "sqlmesh_example.full_model",
         input_queries=input_queries,
         overwrite=True,
         variables={"start": "2020-01-01", "end": "2024-01-01"},
+        include_ctes=True,
     )
 
     test = load_yaml(context.path / c.TESTS / "test_full_model.yaml")
@@ -747,17 +758,23 @@ def test_test_generation(tmp_path: Path) -> None:
     assert "test_full_model" in test
     assert "vars" in test["test_full_model"]
     assert test["test_full_model"]["vars"] == {"start": "2020-01-01", "end": "2024-01-01"}
+    assert "ctes" in test["test_full_model"]["outputs"]
+    assert "cte" in test["test_full_model"]["outputs"]["ctes"]
 
     result = context.test()
     _check_successful_or_raise(result)
 
     context.create_test(
-        "sqlmesh_example.full_model", input_queries=input_queries, name="new_name", path="foo/bar"
+        "sqlmesh_example.full_model",
+        input_queries=input_queries,
+        name="new_name",
+        path="foo/bar",
     )
 
     test = load_yaml(context.path / c.TESTS / "foo/bar.yaml")
     assert len(test) == 1
     assert "new_name" in test
+    assert "ctes" not in test["new_name"]["outputs"]
 
 
 def test_source_func() -> None: