Fix(snowflake): Allow models that utilize Snowpark to execute concurrently (#4431)

Author: erindru

sqlmesh/core/engine_adapter/snowflake.py

@@ -68,6 +68,7 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, ClusteredByMixi
         },
     )
     MANAGED_TABLE_KIND = "DYNAMIC TABLE"
+    SNOWPARK = "snowpark"
 
     @contextlib.contextmanager
     def session(self, properties: SessionProperties) -> t.Iterator[None]:
@@ -104,9 +105,16 @@ def _current_warehouse(self) -> exp.Identifier:
     @property
     def snowpark(self) -> t.Optional[SnowparkSession]:
         if snowpark:
-            return snowpark.Session.builder.configs(
-                {"connection": self._connection_pool.get()}
-            ).getOrCreate()
+            if not self._connection_pool.get_attribute(self.SNOWPARK):
+                # Snowpark sessions are not thread safe so we create a session per thread to prevent them from interfering with each other
+                # The sessions are cleaned up when close() is called
+                new_session = snowpark.Session.builder.configs(
+                    {"connection": self._connection_pool.get()}
+                ).create()
+                self._connection_pool.set_attribute(self.SNOWPARK, new_session)
+
+            return self._connection_pool.get_attribute(self.SNOWPARK)
+
         return None
 
     @property
@@ -584,3 +592,10 @@ def _columns_to_types(
             return columns_to_types_from_dtypes(query_or_df.sample(n=1).to_pandas().dtypes.items())
 
         return super()._columns_to_types(query_or_df, columns_to_types)
+
+    def close(self) -> t.Any:
+        if snowpark_session := self._connection_pool.get_attribute(self.SNOWPARK):
+            snowpark_session.close()  # type: ignore
+            self._connection_pool.set_attribute(self.SNOWPARK, None)
+
+        return super().close()

tests/core/engine_adapter/integration/__init__.py

@@ -707,9 +707,7 @@ def cleanup(self, ctx: t.Optional[Context] = None):
                 schema_name=schema_name, ignore_if_not_exists=True, cascade=True
             )
 
-        if snowpark := self.engine_adapter.snowpark:
-            # ensure that the next test gets a fresh Snowpark session
-            snowpark.close()
+        self.engine_adapter.close()
 
     def upsert_sql_model(self, model_definition: str) -> t.Tuple[Context, SqlModel]:
         if not self._context:

tests/core/engine_adapter/integration/test_integration_snowflake.py

@@ -2,6 +2,7 @@
 import pytest
 from pytest import FixtureRequest
 from sqlglot import exp
+from pathlib import Path
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 from sqlglot.helper import seq_get
 from sqlmesh.core.engine_adapter import SnowflakeEngineAdapter
@@ -10,6 +11,9 @@
 from sqlmesh.core.model import SqlModel, load_sql_based_model
 from sqlmesh.core.plan import Plan
 from tests.core.engine_adapter.integration import TestContext
+from sqlmesh import model, ExecutionContext
+from sqlmesh.core.model import ModelKindName
+from datetime import datetime
 
 from tests.core.engine_adapter.integration import (
     TestContext,
@@ -19,7 +23,9 @@
 )
 
 
-@pytest.fixture(params=list(generate_pytest_params(ENGINES_BY_NAME["snowflake"])))
+@pytest.fixture(
+    params=list(generate_pytest_params(ENGINES_BY_NAME["snowflake"], show_variant_in_test_id=False))
+)
 def ctx(
     request: FixtureRequest,
     create_test_context: t.Callable[[IntegrationTestEngine, str, str], t.Iterable[TestContext]],
@@ -220,3 +226,48 @@ def test_create_iceberg_table(ctx: TestContext, engine_adapter: SnowflakeEngineA
     result = sqlmesh.plan(auto_apply=True)
 
     assert len(result.new_snapshots) == 2
+
+
+def test_snowpark_concurrency(ctx: TestContext) -> None:
+    from snowflake.snowpark import DataFrame
+
+    table = ctx.table("my_model")
+
+    # this model will insert 10 records in batches of 1, with 4 batches at a time running concurrently
+    @model(
+        name=table.sql(),
+        kind=dict(
+            name=ModelKindName.INCREMENTAL_BY_TIME_RANGE,
+            time_column="ds",
+            batch_size=1,
+            batch_concurrency=4,
+        ),
+        columns={"id": "int", "ds": "date"},
+        start="2020-01-01",
+        end="2020-01-10",
+    )
+    def execute(context: ExecutionContext, start: datetime, **kwargs) -> DataFrame:
+        if snowpark := context.snowpark:
+            return snowpark.create_dataframe([(start.day, start.date())], schema=["id", "ds"])
+
+        raise ValueError("Snowpark not present!")
+
+    m = model.get_registry()[table.sql().lower()].model(
+        module_path=Path("."), path=Path("."), dialect="snowflake"
+    )
+
+    sqlmesh = ctx.create_context()
+
+    # verify that we are actually running in multithreaded mode
+    assert sqlmesh.concurrent_tasks > 1
+    assert ctx.engine_adapter._multithreaded
+
+    sqlmesh.upsert_model(m)
+
+    plan = sqlmesh.plan(auto_apply=True)
+
+    assert len(plan.new_snapshots) == 1
+
+    query = exp.select("*").from_(table)
+    df = ctx.engine_adapter.fetchdf(query, quote_identifiers=True)
+    assert len(df) == 10