Alter closing sessions inside a stored procedure to write a warning instead of throwing an error

src/snowflake/snowpark/session.py

@@ -23,7 +23,6 @@
 
 import cloudpickle
 import pkg_resources
-
 from snowflake.connector import ProgrammingError, SnowflakeConnection
 from snowflake.connector.options import installed_pandas, pandas
 from snowflake.connector.pandas_tools import write_pandas
@@ -492,7 +491,8 @@ def _generate_new_action_id(self) -> int:
     def close(self) -> None:
         """Close this session."""
         if is_in_stored_procedure():
-            raise SnowparkClientExceptionMessages.DONT_CLOSE_SESSION_IN_SP()
+            _logger.warning("Closing a session in a stored procedure is a no-op.")
+            return
         try:
             if self._conn.is_closed():
                 _logger.debug(

tests/integ/test_session.py

@@ -570,7 +570,7 @@ def test_use_secondary_roles(session):
     session.use_secondary_roles(current_role[1:-1])
 
 
-@pytest.mark.skipif(IS_IN_STORED_PROC, reason="SP doesn't allow to close a session.")
+@pytest.mark.skipif(IS_IN_STORED_PROC, reason="Can't create a session in SP")
 def test_close_session_twice(db_parameters):
     new_session = Session.builder.configs(db_parameters).create()
     new_session.close()

tests/unit/test_session.py

@@ -2,13 +2,13 @@
 # Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
 #
 import json
+import logging
 import os
 from typing import Optional
 from unittest import mock
 from unittest.mock import MagicMock
 
 import pytest
-
 import snowflake.snowpark.session
 from snowflake.connector import ProgrammingError, SnowflakeConnection
 
@@ -122,6 +122,69 @@ def test_close_exception():
         session.close()
 
 
+def test_close_session_in_stored_procedure_no_op():
+    fake_connection = mock.create_autospec(ServerConnection)
+    fake_connection._conn = mock.Mock()
+    fake_connection.is_closed = MagicMock(return_value=False)
+    session = Session(fake_connection)
+    with mock.patch.object(
+        snowflake.snowpark.session, "is_in_stored_procedure"
+    ) as mock_fn, mock.patch.object(
+        session._conn, "close"
+    ) as mock_close, mock.patch.object(
+        session, "cancel_all"
+    ) as mock_cancel_all, mock.patch.object(
+        snowflake.snowpark.session, "_remove_session"
+    ) as mock_remove:
+        mock_fn.return_value = True
+        session.close()
+        mock_cancel_all.assert_not_called()
+        mock_close.assert_not_called()
+        mock_remove.assert_not_called()
+
+
+def test_close_session_in_stored_procedure_log_level_warning(caplog):
+    caplog.set_level(logging.WARNING)
+    fake_connection = mock.create_autospec(ServerConnection)
+    fake_connection._conn = mock.Mock()
+    fake_connection.is_closed = MagicMock(return_value=False)
+    session = Session(fake_connection)
+    with mock.patch.object(
+        snowflake.snowpark.session, "is_in_stored_procedure"
+    ) as mock_fn:
+        mock_fn.return_value = True
+        session.close()
+    assert "Closing a session in a stored procedure is a no-op." in caplog.text
+
+
+def test_close_session_in_stored_procedure_log_level_info(caplog):
+    caplog.set_level(logging.WARNING)
+    fake_connection = mock.create_autospec(ServerConnection)
+    fake_connection._conn = mock.Mock()
+    fake_connection.is_closed = MagicMock(return_value=False)
+    session = Session(fake_connection)
+    with mock.patch.object(
+        snowflake.snowpark.session, "is_in_stored_procedure"
+    ) as mock_fn:
+        mock_fn.return_value = True
+        session.close()
+    assert "Closing a session in a stored procedure is a no-op." in caplog.text
+
+
+def test_close_session_in_stored_procedure_log_level_error(caplog):
+    caplog.set_level(logging.ERROR)
+    fake_connection = mock.create_autospec(ServerConnection)
+    fake_connection._conn = mock.Mock()
+    fake_connection.is_closed = MagicMock(return_value=False)
+    session = Session(fake_connection)
+    with mock.patch.object(
+        snowflake.snowpark.session, "is_in_stored_procedure"
+    ) as mock_fn:
+        mock_fn.return_value = True
+        session.close()
+    assert "Closing a session in a stored procedure is a no-op." not in caplog.text
+
+
 def test_resolve_import_path_ignore_import_path(tmp_path_factory):
     fake_connection = mock.create_autospec(ServerConnection)
     fake_connection._conn = mock.Mock()
@@ -178,9 +241,7 @@ def test_resolve_package_terms_not_accepted():
     def get_information_schema_packages(table_name: str):
         if table_name == "information_schema.packages":
             result = MagicMock()
-            result.filter().group_by().agg()._internal_collect_with_tag.return_value = (
-                []
-            )
+            result.filter().group_by().agg()._internal_collect_with_tag.return_value = []
             return result
 
     def run_query(sql: str):