Fix: ensure session-scoped Snowflake warehouse is rolled back on failure (#4640)

Author: georgesittas

sqlmesh/core/engine_adapter/snowflake.py

@@ -93,8 +93,10 @@ def session(self, properties: SessionProperties) -> t.Iterator[None]:
             return
 
         self.execute(f"USE WAREHOUSE {warehouse_sql}")
-        yield
-        self.execute(f"USE WAREHOUSE {current_warehouse_sql}")
+        try:
+            yield
+        finally:
+            self.execute(f"USE WAREHOUSE {current_warehouse_sql}")
 
     @property
     def _current_warehouse(self) -> exp.Identifier:

tests/core/engine_adapter/test_snowflake.py

@@ -98,6 +98,7 @@ def test_session(
     adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter)
     adapter.cursor.fetchone.return_value = (current_warehouse,)
 
+    # Test normal execution
     with adapter.session({"warehouse": configured_warehouse}):
         pass
 
@@ -114,6 +115,27 @@ def test_session(
 
     assert to_sql_calls(adapter) == expected_calls
 
+    # Test exception handling - warehouse should still be reset
+    if should_change:
+        adapter.cursor.execute.reset_mock()
+        adapter.cursor.fetchone.return_value = (current_warehouse,)
+
+        try:
+            with adapter.session({"warehouse": configured_warehouse}):
+                adapter.execute("SELECT 1")
+                raise RuntimeError("Test exception")
+        except RuntimeError:
+            pass
+
+        expected_exception_calls = [
+            "SELECT CURRENT_WAREHOUSE()",
+            f"USE WAREHOUSE {configured_warehouse_exp}",
+            "SELECT 1",
+            f"USE WAREHOUSE {current_warehouse_exp}",
+        ]
+
+        assert to_sql_calls(adapter) == expected_exception_calls
+
 
 def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
     adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter)