Fix: ensure partial works when a query is used to produce data (#4583)

Author: georgesittas

sqlmesh/core/test/definition.py

@@ -559,7 +559,10 @@ def _create_df(
     ) -> pd.DataFrame:
         query = values.get("query")
         if query:
-            return self._execute(self._add_missing_columns(query, columns))
+            if not partial:
+                query = self._add_missing_columns(query, columns)
+
+            return self._execute(query)
 
         rows = values["rows"]
         if columns:

tests/core/test_test.py

@@ -1474,6 +1474,33 @@ def test_generate_input_data_using_sql(mocker: MockerFixture, tmp_path: Path) ->
         ).run()
     )
 
+    _check_successful_or_raise(
+        _create_test(
+            body=load_yaml(
+                """
+test_example_full_model_partial:
+  model: sqlmesh_example.full_model
+  inputs:
+    sqlmesh_example.incremental_model:
+      rows:
+      - id: 1
+        item_id: 1
+      - id: 2
+        item_id: 1
+      - id: 3
+        item_id: 2
+  outputs:
+    query:
+      partial: true
+      query: "SELECT 2 AS num_orders UNION ALL SELECT 1 AS num_orders"
+                """
+            ),
+            test_name="test_example_full_model_partial",
+            model=context.get_model("sqlmesh_example.full_model"),
+            context=context,
+        ).run()
+    )
+
     mocker.patch("sqlmesh.core.test.definition.random_id", return_value="jzngz56a")
     test = _create_test(
         body=load_yaml(