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

Author: erindru

sqlmesh/core/engine_adapter/snowflake.py

@@ -3,6 +3,7 @@
 import contextlib
 import logging
 import typing as t
+import threading
 
 import pandas as pd
 from pandas.api.types import is_datetime64_any_dtype  # type: ignore
@@ -69,6 +70,10 @@ class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, ClusteredByMixi
     )
     MANAGED_TABLE_KIND = "DYNAMIC TABLE"
 
+    def __init__(self, *args: t.Any, **kwargs: t.Any):
+        super().__init__(*args, **kwargs)
+        self._snowpark_threadlocal = threading.local()
+
     @contextlib.contextmanager
     def session(self, properties: SessionProperties) -> t.Iterator[None]:
         warehouse = properties.get("warehouse")
@@ -104,9 +109,15 @@ 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()
+            # 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
+            if not hasattr(self._snowpark_threadlocal, "session"):
+                new_session = snowpark.Session.builder.configs(
+                    {"connection": self._connection_pool.get()}
+                ).create()
+                self._snowpark_threadlocal.session = new_session
+
+            return self._snowpark_threadlocal.session
         return None
 
     @property
@@ -584,3 +595,15 @@ 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 _cleanup_snowpark(self) -> None:
+        if hasattr(self._snowpark_threadlocal, "session") and (
+            session := self._snowpark_threadlocal.session
+        ):
+            session.close()
+            delattr(self._snowpark_threadlocal, "session")
+
+    def close(self) -> t.Any:
+        self._cleanup_snowpark()
+
+        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,43 @@ 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
+
+    @model(
+        name="my_model",
+        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()["my_model"].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
+
+    # todo: read table result