Repository: TobikoData/sqlmesh
Pull request #2983 (2024-08-04 23:50:45)

Fix!: Regression in Snowflake Python models using Snowpark DataFrame's
Prior to this change, attempting to create a Python model on Snowflake utilising a Snowpark DataFrame would produce the following error:

```
snowflake.connector.errors.ProgrammingError: SQL compilation error:  Object found is of type 'VIEW', not specified type 'TABLE'.
```

This is because for Snowpark, the DataFrame is made available via `CREATE TEMPORARY VIEW` but the cleanup logic in the Snowflake engine adapter assumed a Table (because Pandas DataFrames use Tables) and tried to issue `DROP TABLE` instead of `DROP VIEW`.

Now the cleanup logic will issue `DROP VIEW` if a Snowpark DataFrame is in use.

In order to add a test for this, I had to re-introduce a dependency on `snowflake-snowpark-python` (first introduced in #2666 and then removed in #2761)

Since Snowpark doesnt work on Python 3.12, I also added some guards to only install it on Python < 3.12 and also skip the test if Snowpark is not available


Patch:
diff --git a/docs/guides/multi_repo.md b/docs/guides/multi_repo.md
index 19e8d2d4..59ab149b 100644
--- a/docs/guides/multi_repo.md
+++ b/docs/guides/multi_repo.md
@@ -1,6 +1,8 @@
 # Multi-Repo guide
 
-Although mono repos are convenient and easy to use, sometimes your organization may choose to use multiple repos. SQLMesh provides native support for multiple repos and makes it easy to maintain data consistency and correctness even with multiple repos.
+Although mono repos are convenient and easy to use, sometimes your organization may choose to use multiple repos. 
+SQLMesh provides native support for multiple repos and makes it easy to maintain data consistency and correctness even with multiple repos.
+If you are wanting to separate your systems/data and provide isolation, checkout the [isolated systems guide](https://sqlmesh.readthedocs.io/en/stable/guides/isolated_systems/?h=isolated). 
 
 ## Bootstrapping multiple projects
 Setting up SQLMesh with multiple repos is quite simple. Copy the contents of this example [multi-repo project](https://github.com/TobikoData/sqlmesh/tree/main/examples/multi).
diff --git a/setup.py b/setup.py
index 0aef17b5..6e2a01fa 100644
--- a/setup.py
+++ b/setup.py
@@ -134,6 +134,8 @@ setup(
             # https://github.com/dbt-labs/dbt-snowflake/blob/main/dev-requirements.txt#L12
             "cryptography~=42.0.4",
             "snowflake-connector-python[pandas,secure-local-storage]",
+            # as at 2024-08-05, snowflake-snowpark-python is only available up to Python 3.11
+            "snowflake-snowpark-python; python_version<'3.12'",
         ],
         "trino": [
             "trino",
diff --git a/sqlmesh/core/engine_adapter/snowflake.py b/sqlmesh/core/engine_adapter/snowflake.py
index 38944d96..6b478a06 100644
--- a/sqlmesh/core/engine_adapter/snowflake.py
+++ b/sqlmesh/core/engine_adapter/snowflake.py
@@ -204,9 +204,11 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin):
             target_table or "pandas", quoted=False
         )  # write_pandas() re-quotes everything without checking if its already quoted
 
+        is_snowpark_dataframe = snowpark and isinstance(df, snowpark.dataframe.DataFrame)
+
         def query_factory() -> Query:
-            if snowpark and isinstance(df, snowpark.dataframe.DataFrame):
-                df.createOrReplaceTempView(temp_table.sql(dialect=self.dialect, identify=True))
+            if is_snowpark_dataframe:
+                df.createOrReplaceTempView(temp_table.sql(dialect=self.dialect, identify=True))  # type: ignore
             elif isinstance(df, pd.DataFrame):
                 from snowflake.connector.pandas_tools import write_pandas
 
@@ -258,13 +260,15 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin):
 
             return exp.select(*self._casted_columns(columns_to_types)).from_(temp_table)
 
+        def cleanup() -> None:
+            if is_snowpark_dataframe:
+                self.drop_view(temp_table)
+            else:
+                self.drop_table(temp_table)
+
         # the cleanup_func technically isnt needed because the temp table gets dropped when the session ends
         # but boy does it make our multi-adapter integration tests easier to write
-        return [
-            SourceQuery(
-                query_factory=query_factory, cleanup_func=lambda: self.drop_table(temp_table)
-            )
-        ]
+        return [SourceQuery(query_factory=query_factory, cleanup_func=cleanup)]
 
     def _fetch_native_df(
         self, query: t.Union[exp.Expression, str], quote_identifiers: bool = False
diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py
index 386809ff..7426095d 100644
--- a/sqlmesh/core/snapshot/evaluator.py
+++ b/sqlmesh/core/snapshot/evaluator.py
@@ -799,12 +799,11 @@ class SnapshotEvaluator:
         evaluation_strategy = _evaluation_strategy(snapshot, self.adapter)
 
         for is_table_deployable, table_name in table_names:
-            table = exp.to_table(table_name)
-            if table.db != snapshot.physical_schema:
-                raise SQLMeshError(
-                    f"Table '{table_name}' is not a part of the physical schema '{snapshot.physical_schema}' and so can't be dropped."
-                )
-            evaluation_strategy.delete(table_name, is_table_deployable=is_table_deployable)
+            evaluation_strategy.delete(
+                table_name,
+                is_table_deployable=is_table_deployable,
+                physical_schema=snapshot.physical_schema,
+            )
 
             if on_complete is not None:
                 on_complete(table_name)
@@ -1225,6 +1224,7 @@ class MaterializableStrategy(PromotableStrategy):
         self.adapter.alter_table(alter_expressions)
 
     def delete(self, name: str, **kwargs: t.Any) -> None:
+        _check_table_db_is_physical_schema(name, kwargs["physical_schema"])
         self.adapter.drop_table(name)
         logger.info("Dropped table '%s'", name)
 
@@ -1792,6 +1792,7 @@ class EngineManagedStrategy(MaterializableStrategy):
 
     def delete(self, name: str, **kwargs: t.Any) -> None:
         # a dev preview table is created as a normal table, so it needs to be dropped as a normal table
+        _check_table_db_is_physical_schema(name, kwargs["physical_schema"])
         if kwargs["is_table_deployable"]:
             self.adapter.drop_managed_table(name)
             logger.info("Dropped managed table '%s'", name)
@@ -1825,3 +1826,11 @@ def _check_destructive_schema_change(
         raise SQLMeshError(
             f"{warning_msg} To allow this, change the model's `on_destructive_change` setting to `warn` or `allow` or include it in the plan's `--allow-destructive-model` option."
         )
+
+
+def _check_table_db_is_physical_schema(table_name: str, physical_schema: str) -> None:
+    table = exp.to_table(table_name)
+    if table.db != physical_schema:
+        raise SQLMeshError(
+            f"Table '{table_name}' is not a part of the physical schema '{physical_schema}' and so can't be dropped."
+        )


Test patch:
diff --git a/tests/core/engine_adapter/test_snowflake.py b/tests/core/engine_adapter/test_snowflake.py
index 44b8ffe8..6dab3cf5 100644
--- a/tests/core/engine_adapter/test_snowflake.py
+++ b/tests/core/engine_adapter/test_snowflake.py
@@ -11,6 +11,7 @@ from sqlmesh.core.model import load_sql_based_model
 from sqlmesh.core.engine_adapter import SnowflakeEngineAdapter
 from sqlmesh.core.model.definition import SqlModel
 from sqlmesh.utils.errors import SQLMeshError
+from sqlmesh.utils import optional_import
 from tests.core.engine_adapter import to_sql_calls
 
 pytestmark = [pytest.mark.engine, pytest.mark.snowflake]
@@ -359,3 +360,40 @@ def test_set_current_schema(make_mocked_engine_adapter: t.Callable):
         'USE SCHEMA "FOO"."foo"',
         'USE SCHEMA "FOO"."fOo"',
     ]
+
+
+def test_replace_query_snowpark_dataframe(
+    mocker: MockerFixture, make_mocked_engine_adapter: t.Callable
+):
+    if not optional_import("snowflake.snowpark"):
+        pytest.skip("Snowpark not available in this environment")
+
+    from snowflake.snowpark.session import Session
+    from snowflake.snowpark.dataframe import DataFrame as SnowparkDataFrame
+
+    session = Session.builder.config("local_testing", True).create()
+    df: SnowparkDataFrame = session.create_dataframe([(1, "name")], schema=["ID", "NAME"])
+    assert isinstance(df, SnowparkDataFrame)
+
+    mocker.patch("sqlmesh.core.engine_adapter.base.random_id", return_value="e6wjkjj6")
+    spy = mocker.spy(df, "createOrReplaceTempView")
+
+    adapter: SnowflakeEngineAdapter = make_mocked_engine_adapter(SnowflakeEngineAdapter)
+    adapter._default_catalog = "foo"
+
+    adapter.replace_query(
+        table_name="foo",
+        query_or_df=df,
+        columns_to_types={"ID": exp.DataType.build("INT"), "NAME": exp.DataType.build("VARCHAR")},
+    )
+
+    # the Snowflake library generates "CREATE TEMPORARY VIEW" from a direct DataFrame call
+    # which doesnt pass through our EngineAdapter so we cant capture it
+    spy.assert_called()
+    assert "__temp_foo_e6wjkjj6" in spy.call_args[0][0]
+
+    # verify that DROP VIEW is called instead of DROP TABLE
+    assert to_sql_calls(adapter) == [
+        'CREATE OR REPLACE TABLE "foo" AS SELECT CAST("ID" AS INT) AS "ID", CAST("NAME" AS VARCHAR) AS "NAME" FROM (SELECT CAST("ID" AS INT) AS "ID", CAST("NAME" AS VARCHAR) AS "NAME" FROM "__temp_foo_e6wjkjj6") AS "_subquery"',
+        'DROP VIEW IF EXISTS "__temp_foo_e6wjkjj6"',
+    ]
diff --git a/tests/core/test_snapshot.py b/tests/core/test_snapshot.py
index 077b5448..544e299a 100644
--- a/tests/core/test_snapshot.py
+++ b/tests/core/test_snapshot.py
@@ -371,7 +371,7 @@ def test_missing_intervals_end_bounded_with_ignore_cron(make_snapshot):
 
 
 def test_missing_intervals_past_end_date_with_lookback(make_snapshot):
-    snapshot: Snapshot = make_snapshot(
+    snapshot: Snapshot = make_snapshot(  # type: ignore
         SqlModel(
             name="test_model",
             kind=IncrementalByTimeRangeKind(time_column=TimeColumn(column="ds"), lookback=2),
diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py
index 014effa0..cac20b35 100644
--- a/tests/core/test_snapshot_evaluator.py
+++ b/tests/core/test_snapshot_evaluator.py
@@ -32,8 +32,9 @@ from sqlmesh.core.model import (
     TimeColumn,
     ViewKind,
     load_sql_based_model,
+    ExternalModel,
 )
-from sqlmesh.core.model.kind import OnDestructiveChange
+from sqlmesh.core.model.kind import OnDestructiveChange, ExternalKind
 from sqlmesh.core.node import IntervalUnit
 from sqlmesh.core.snapshot import (
     DeployabilityIndex,
@@ -428,6 +429,29 @@ def test_cleanup(mocker: MockerFixture, adapter_mock, make_snapshot):
     )
 
 
+def test_cleanup_external_model(mocker: MockerFixture, adapter_mock, make_snapshot):
+    evaluator = SnapshotEvaluator(adapter_mock)
+
+    def create_and_cleanup_external_model(name: str, dev_table_only: bool):
+        model = ExternalModel(
+            name=name,
+            kind=ExternalKind(),
+        )
+
+        snapshot = make_snapshot(model)
+        snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+        snapshot.version = "test_version"
+
+        evaluator.promote([snapshot], EnvironmentNamingInfo(name="test_env"))
+        evaluator.cleanup(
+            [SnapshotTableCleanupTask(snapshot=snapshot.table_info, dev_table_only=dev_table_only)]
+        )
+        return snapshot
+
+    create_and_cleanup_external_model("catalog.test_schema.test_model", True)
+    adapter_mock.drop_table.assert_not_called()
+
+
 @pytest.mark.parametrize("view_exists", [True, False])
 def test_evaluate_materialized_view(
     mocker: MockerFixture, adapter_mock, make_snapshot, view_exists: bool