Chore: Improve plan error handling when --select-model is used and remote snapshot cant be rendered

Author: erindru

sqlmesh/core/context.py

@@ -1450,12 +1450,18 @@ def plan_builder(
 
         models_override: t.Optional[UniqueKeyDict[str, Model]] = None
         if select_models:
-            models_override = model_selector.select_models(
-                select_models,
-                environment,
-                fallback_env_name=create_from or c.PROD,
-                ensure_finalized_snapshots=self.config.plan.use_finalized_state,
-            )
+            try:
+                models_override = model_selector.select_models(
+                    select_models,
+                    environment,
+                    fallback_env_name=create_from or c.PROD,
+                    ensure_finalized_snapshots=self.config.plan.use_finalized_state,
+                )
+            except SQLMeshError as e:
+                logger.exception(e)  # ensure the full stack trace is logged
+                raise PlanError(
+                    f"{e}\nCheck the SQLMesh log file for the full stack trace.\nIf the model has been fixed locally, please ensure that the --select-model expression includes it."
+                )
             if not backfill_models:
                 # Only backfill selected models unless explicitly specified.
                 backfill_models = model_selector.expand_model_selections(select_models)

sqlmesh/core/selector.py

@@ -16,6 +16,7 @@
 from sqlmesh.utils import UniqueKeyDict
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.git import GitClient
+from sqlmesh.utils.errors import SQLMeshError
 
 
 if t.TYPE_CHECKING:
@@ -111,7 +112,16 @@ def select_models(
         def get_model(fqn: str) -> t.Optional[Model]:
             if fqn not in all_selected_models and fqn in env_models:
                 # Unselected modified or added model.
-                return env_models[fqn]
+                model_from_env = env_models[fqn]
+                try:
+                    # this triggers a render_query() which can throw an exception
+                    model_from_env.depends_on
+                    return model_from_env
+                except Exception as e:
+                    raise SQLMeshError(
+                        f"Model '{model_from_env.name}' sourced from state cannot be rendered "
+                        f"in the local environment due to:\n> {str(e)}"
+                    ) from e
             if fqn in all_selected_models and fqn in self._models:
                 # Selected modified or removed model.
                 return self._models[fqn]

tests/cli/test_integration_cli.py

@@ -40,9 +40,8 @@ def _invoke(sqlmesh_args: t.List[str], **kwargs: t.Any) -> subprocess.CompletedP
     return _invoke
 
 
-def test_load_snapshots_that_reference_nonexistent_python_libraries(
-    invoke_cli: InvokeCliType, tmp_path: Path
-) -> None:
+@pytest.fixture
+def duckdb_example_project(tmp_path: Path) -> Path:
     init_example_project(tmp_path, dialect="duckdb")
     config_path = tmp_path / "config.yaml"
 
@@ -54,6 +53,36 @@ def test_load_snapshots_that_reference_nonexistent_python_libraries(
     }
     config_path.write_text(yaml.dump(config_dict))
 
+    return tmp_path
+
+
+@pytest.fixture
+def last_log_file_contents(tmp_path: Path) -> t.Callable[[], str]:
+    def _fetch() -> str:
+        log_file = sorted(list((tmp_path / "logs").iterdir()))[-1]
+        return log_file.read_text()
+
+    return _fetch
+
+
+def test_load_snapshots_that_reference_nonexistent_python_libraries(
+    invoke_cli: InvokeCliType,
+    duckdb_example_project: Path,
+    last_log_file_contents: t.Callable[[], str],
+) -> None:
+    """
+    Scenario:
+     - A model is created using a macro that is imported from an external package
+     - That model is applied + snapshot committed to state
+     - The external package is removed locally and the import macro import is changed to an inline definition
+
+    Outcome:
+     - `sqlmesh plan` should not exit with an ImportError when it tries to render the query of the snapshot stored in state
+     - Instead, it should log a warning and proceed with applying the new model version
+    """
+
+    project_path = duckdb_example_project
+
     # simulate a 3rd party library that provides a macro
     site_packages = site.getsitepackages()[0]
     sqlmesh_test_macros_package_path = Path(site_packages) / "sqlmesh_test_macros"
@@ -67,11 +96,11 @@ def do_something(evaluator):
 """)
 
     # reference the macro from site-packages
-    (tmp_path / "macros" / "__init__.py").write_text("""
+    (project_path / "macros" / "__init__.py").write_text("""
 from sqlmesh_test_macros.macros import do_something
 """)
 
-    (tmp_path / "models" / "example.sql").write_text("""
+    (project_path / "models" / "example.sql").write_text("""
 MODEL (
     name example.test_model,
     kind FULL
@@ -99,7 +128,7 @@ def do_something(evaluator):
     shutil.rmtree(sqlmesh_test_macros_package_path)
 
     # Move the macro inline so its no longer being loaded from a library but still exists with the same signature
-    (tmp_path / "macros" / "__init__.py").write_text("""
+    (project_path / "macros" / "__init__.py").write_text("""
 from sqlmesh import macro
 
 @macro()
@@ -120,8 +149,7 @@ def do_something(evaluator):
     assert "Physical layer updated" in result.stdout
     assert "Virtual layer updated" in result.stdout
 
-    log_file = sorted(list((tmp_path / "logs").iterdir()))[-1]
-    log_file_contents = log_file.read_text()
+    log_file_contents = last_log_file_contents()
     assert "ModuleNotFoundError: No module named 'sqlmesh_test_macros'" in log_file_contents
     assert (
         "ERROR - Failed to cache optimized query for model 'example.test_model'"
@@ -131,3 +159,195 @@ def do_something(evaluator):
         'ERROR - Failed to cache snapshot SnapshotId<"db"."example"."test_model"'
         in log_file_contents
     )
+
+
+def test_model_selector_snapshot_references_nonexistent_python_libraries(
+    invoke_cli: InvokeCliType,
+    duckdb_example_project: Path,
+    last_log_file_contents: t.Callable[[], str],
+) -> None:
+    """
+    Scenario:
+     - A model is created using a macro that is imported from an external package
+     - That model is applied + snapshot committed to state
+     - The external package is removed locally and the import macro import is changed to an inline definition
+     - Thus, local version of the model can be rendered but the remote version in state cannot
+
+    Outcome:
+     - `sqlmesh plan --select-model <this model>` should work as it picks up the local version
+     - `sqlmesh plan --select-model <some other model> should exit with an error, because the plan needs a valid DAG and the remote version is no longer valid locally
+    """
+    project_path = duckdb_example_project
+
+    # simulate a 3rd party library that provides a macro
+    site_packages = site.getsitepackages()[0]
+    sqlmesh_test_macros_package_path = Path(site_packages) / "sqlmesh_test_macros"
+    sqlmesh_test_macros_package_path.mkdir()
+    (sqlmesh_test_macros_package_path / "macros.py").write_text("""
+from sqlmesh import macro
+
+@macro()
+def do_something(evaluator):
+    return "'value from site-packages'"
+""")
+
+    # reference the macro from site-packages
+    (project_path / "macros" / "__init__.py").write_text("""
+from sqlmesh_test_macros.macros import do_something
+""")
+
+    (project_path / "models" / "example.sql").write_text("""
+MODEL (
+    name sqlmesh_example.test_model,
+    kind FULL
+);
+
+select @do_something() as a
+""")
+
+    result = invoke_cli(["plan", "--no-prompts", "--auto-apply", "--skip-tests"])
+
+    assert result.returncode == 0
+    assert "Physical layer updated" in result.stdout
+    assert "Virtual layer updated" in result.stdout
+
+    # clear cache to ensure we are forced to reload everything
+    assert invoke_cli(["clean"]).returncode == 0
+
+    # deleting this removes the 'do_something()' macro used by the version of the snapshot stored in state
+    # when loading the old snapshot from state in the local python env, this will create an ImportError
+    shutil.rmtree(sqlmesh_test_macros_package_path)
+
+    # Move the macro inline so its no longer being loaded from a library but still exists with the same signature
+    (project_path / "macros" / "__init__.py").write_text("""
+from sqlmesh import macro
+
+@macro()
+def do_something(evaluator):
+    return "'some value not from site-packages'"
+""")
+
+    # the invalid snapshot is in state but is not preventing a plan
+    result = invoke_cli(
+        [
+            "plan",
+            "--no-prompts",
+            "--skip-tests",
+        ],
+        input="n",  # for the apply backfill (y/n) prompt
+    )
+    assert result.returncode == 0
+    assert "Apply - Backfill Tables [y/n]:" in result.stdout
+    assert "Physical layer updated" not in result.stdout
+
+    # the invalid snapshot in state should not prevent a plan if --select-model is used on it (since the local version can be rendered)
+    result = invoke_cli(
+        ["plan", "--select-model", "sqlmesh_example.test_model", "--no-prompts", "--skip-tests"],
+        input="n",  # for the apply backfill (y/n) prompt
+    )
+    assert result.returncode == 0
+    assert "ModuleNotFoundError" not in result.stdout
+    assert "sqlmesh_example.test_model" in result.stdout
+    assert "Apply - Backfill Tables" in result.stdout
+
+    # the invalid snapshot in state should prevent a plan if --select-model is used on another model
+    # (since this says to SQLMesh "source everything from state except this selected model" and the plan DAG must be valid to run the plan)
+    result = invoke_cli(
+        [
+            "plan",
+            "--select-model",
+            "sqlmesh_example.full_model",
+            "--no-prompts",
+            "--skip-tests",
+        ],
+        input="n",  # for the apply backfill (y/n) prompt
+    )
+    assert result.returncode == 1
+    assert (
+        "Model 'sqlmesh_example.test_model' sourced from state cannot be rendered in the local environment"
+        in result.stdout
+    )
+    assert "No module named 'sqlmesh_test_macros'" in result.stdout
+    assert (
+        "If the model has been fixed locally, please ensure that the --select-model expression includes it"
+        in result.stdout
+    )
+
+    # verify the full stack trace was logged
+    log_file_contents = last_log_file_contents()
+    assert "ModuleNotFoundError: No module named 'sqlmesh_test_macros'" in log_file_contents
+    assert (
+        "The above exception was the direct cause of the following exception:" in log_file_contents
+    )
+
+
+def test_model_selector_tags_picks_up_both_remote_and_local(
+    invoke_cli: InvokeCliType, duckdb_example_project: Path
+) -> None:
+    """
+    Scenario:
+     - A model that has already been applied to prod (so exists in state) has a tag added locally
+     - A new model is created locally that has the same tag
+
+    Outcome:
+     - `sqlmesh plan --select-model tag:<tag>` should include both models
+    """
+    project_path = duckdb_example_project
+
+    # default state of full_model
+    (project_path / "models" / "full_model.sql").write_text("""
+    MODEL (
+        name sqlmesh_example.full_model,
+        kind FULL,
+        cron '@daily',
+        grain item_id,
+        audits (assert_positive_order_ids),
+    );
+
+    SELECT
+        item_id,
+        COUNT(DISTINCT id) AS num_orders
+    FROM sqlmesh_example.incremental_model
+    GROUP BY item_id
+    """)
+
+    # apply plan - starting point
+    result = invoke_cli(["plan", "--no-prompts", "--auto-apply", "--skip-tests"])
+
+    assert result.returncode == 0
+    assert "Physical layer updated" in result.stdout
+    assert "Virtual layer updated" in result.stdout
+
+    # add a new model locally with tag:a
+    (project_path / "models" / "new_model.sql").write_text("""
+    MODEL (
+        name sqlmesh_example.new_model,
+        kind full,
+        tags (a)
+    );
+
+    SELECT 1;
+    """)
+
+    # update full_model with tag:a
+    (project_path / "models" / "full_model.sql").write_text("""
+    MODEL (
+        name sqlmesh_example.full_model,
+        kind FULL,
+        tags (a)
+    );
+
+    SELECT
+        item_id,
+        COUNT(DISTINCT id) AS num_orders
+    FROM sqlmesh_example.incremental_model
+    GROUP BY item_id
+    """)
+
+    result = invoke_cli(
+        ["plan", "--select-model", "tag:a", "--no-prompts", "--skip-tests"],
+        input="n",  # for the apply backfill (y/n) prompt
+    )
+    assert result.returncode == 0
+    assert "sqlmesh_example.full_model" in result.stdout  # metadata update: tags
+    assert "sqlmesh_example.new_model" in result.stdout  # added