From 0553ebaa204d9f8170065ddc613654aa7ce295b9 Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Wed, 21 Feb 2024 11:03:29 +0100 Subject: [PATCH 01/16] add regression test for old behavior - this should fail once closing sessions in stored procedures is a noop --- tests/unit/test_session.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 890cd7b2d27..71441aa722c 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -8,9 +8,9 @@ from unittest.mock import MagicMock import pytest - import snowflake.snowpark.session from snowflake.connector import ProgrammingError, SnowflakeConnection +from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages try: import pandas @@ -122,6 +122,22 @@ def test_close_exception(): session.close() +def test_close_exception_in_stored_procedure(): + fake_connection = mock.create_autospec(ServerConnection) + fake_connection._conn = mock.Mock() + fake_connection.is_closed = MagicMock(return_value=False) + with pytest.raises( + SnowparkSessionException, + match=SnowparkClientExceptionMessages.DONT_CLOSE_SESSION_IN_SP().message, + ): + 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() + + def test_resolve_import_path_ignore_import_path(tmp_path_factory): fake_connection = mock.create_autospec(ServerConnection) fake_connection._conn = mock.Mock() From 8b11568b62a1a74e10e3faf9c87d5c5d30ae0109 Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Wed, 21 Feb 2024 11:21:36 +0100 Subject: [PATCH 02/16] write a warning when closing a session inside a stored procedure - instead of throwing an error a warning is written when closing a session inside a stored procedure --- src/snowflake/snowpark/session.py | 2 +- tests/unit/test_session.py | 51 +++++++++++++++++++++++-------- 2 files changed, 40 insertions(+), 13 deletions(-) diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index 79a95dcf12b..a05e513dc7c 100644 --- a/src/snowflake/snowpark/session.py +++ b/src/snowflake/snowpark/session.py @@ -492,7 +492,7 @@ 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.") try: if self._conn.is_closed(): _logger.debug( diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 71441aa722c..2d5f67607a6 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -2,15 +2,16 @@ # 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 -from snowflake.snowpark._internal.error_message import SnowparkClientExceptionMessages try: import pandas @@ -122,20 +123,46 @@ def test_close_exception(): session.close() -def test_close_exception_in_stored_procedure(): +def test_close_exception_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) - with pytest.raises( - SnowparkSessionException, - match=SnowparkClientExceptionMessages.DONT_CLOSE_SESSION_IN_SP().message, - ): - 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() + 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_exception_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_exception_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): From 69ca98b0e3ce36a7c4a6db71cb8ae23fbf8d799d Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Wed, 21 Feb 2024 11:22:39 +0100 Subject: [PATCH 03/16] remove skip condition for session closing integration test when inside a stored procedure - skip should not be needed anymore due to the closing of a session being a no-op when inside a stored procedure --- tests/integ/test_session.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integ/test_session.py b/tests/integ/test_session.py index f38e0c688ea..ed2d114f5cd 100644 --- a/tests/integ/test_session.py +++ b/tests/integ/test_session.py @@ -570,7 +570,6 @@ 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.") def test_close_session_twice(db_parameters): new_session = Session.builder.configs(db_parameters).create() new_session.close() From 56459d8477e9dc90994a51299910dd6fb00d1246 Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Wed, 21 Feb 2024 11:28:41 +0100 Subject: [PATCH 04/16] rename test functions --- tests/unit/test_session.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 2d5f67607a6..25a00f727cf 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -123,7 +123,7 @@ def test_close_exception(): session.close() -def test_close_exception_in_stored_procedure_log_level_warning(caplog): +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() @@ -137,7 +137,7 @@ def test_close_exception_in_stored_procedure_log_level_warning(caplog): assert "Closing a session in a stored procedure is a no-op." in caplog.text -def test_close_exception_in_stored_procedure_log_level_info(caplog): +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() @@ -151,7 +151,7 @@ def test_close_exception_in_stored_procedure_log_level_info(caplog): assert "Closing a session in a stored procedure is a no-op." in caplog.text -def test_close_exception_in_stored_procedure_log_level_error(caplog): +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() From e371badb1cb057bc18bf4d4ce3944e35707c830f Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Wed, 21 Feb 2024 11:48:00 +0100 Subject: [PATCH 05/16] add no-op functionality when closing sessions in stored procedure --- src/snowflake/snowpark/session.py | 2 +- tests/unit/test_session.py | 26 ++++++++++++++++++++++---- 2 files changed, 23 insertions(+), 5 deletions(-) diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index a05e513dc7c..b7fc34f4435 100644 --- a/src/snowflake/snowpark/session.py +++ b/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 @@ -493,6 +492,7 @@ def close(self) -> None: """Close this session.""" if is_in_stored_procedure(): _logger.warning("Closing a session in a stored procedure is a no-op.") + return try: if self._conn.is_closed(): _logger.debug( diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 25a00f727cf..7f771bf5d7b 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -9,7 +9,6 @@ from unittest.mock import MagicMock import pytest - import snowflake.snowpark.session from snowflake.connector import ProgrammingError, SnowflakeConnection @@ -123,6 +122,27 @@ 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) @@ -221,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): From 72fef95249630a5f5917add94eb5becf040df014 Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Wed, 21 Feb 2024 14:05:12 +0100 Subject: [PATCH 06/16] update skif if reason - test cannot be run because creating a session in stored procedures is not allowed --- tests/integ/test_session.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integ/test_session.py b/tests/integ/test_session.py index ed2d114f5cd..fb715f47d4c 100644 --- a/tests/integ/test_session.py +++ b/tests/integ/test_session.py @@ -570,6 +570,7 @@ def test_use_secondary_roles(session): session.use_secondary_roles(current_role[1:-1]) +@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() From 4239b78847dc70fe718989104a4eb7601141404a Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Thu, 22 Feb 2024 09:03:17 +0100 Subject: [PATCH 07/16] rerun pre-commit for formatting --- src/snowflake/snowpark/session.py | 1 + tests/unit/test_session.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index b7fc34f4435..5191bc1897f 100644 --- a/src/snowflake/snowpark/session.py +++ b/src/snowflake/snowpark/session.py @@ -23,6 +23,7 @@ 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 diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 7f771bf5d7b..b3cf88d09d8 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -9,6 +9,7 @@ from unittest.mock import MagicMock import pytest + import snowflake.snowpark.session from snowflake.connector import ProgrammingError, SnowflakeConnection @@ -241,7 +242,9 @@ 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): From eaad7b537d9e1a27aae6dc1ca94f81057e325a3e Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Thu, 22 Feb 2024 09:15:24 +0100 Subject: [PATCH 08/16] combine log tests for different log levels into one using pytest parametrize --- tests/unit/test_session.py | 40 +++++++++----------------------------- 1 file changed, 9 insertions(+), 31 deletions(-) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index b3cf88d09d8..8f8408cd43a 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -144,36 +144,13 @@ def test_close_session_in_stored_procedure_no_op(): 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) +@pytest.mark.parametrize( + "warning_level, expected", + [(logging.WARNING, True), (logging.INFO, True), (logging.ERROR, False)], +) +def test_close_session_in_stored_procedure_log_level(caplog, warning_level, expected): + caplog.clear() + caplog.set_level(warning_level) fake_connection = mock.create_autospec(ServerConnection) fake_connection._conn = mock.Mock() fake_connection.is_closed = MagicMock(return_value=False) @@ -183,7 +160,8 @@ def test_close_session_in_stored_procedure_log_level_error(caplog): ) 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 + result = "Closing a session in a stored procedure is a no-op." in caplog.text + assert result == expected def test_resolve_import_path_ignore_import_path(tmp_path_factory): From 6ea9bcc055b31d8dc55c9f9504c012f803d10fe8 Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Fri, 23 Feb 2024 18:35:21 +0100 Subject: [PATCH 09/16] Change test_close_session_in_sp - closing session in stored procedure should not close session --- tests/integ/test_session.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/integ/test_session.py b/tests/integ/test_session.py index fb715f47d4c..623a97f3813 100644 --- a/tests/integ/test_session.py +++ b/tests/integ/test_session.py @@ -7,7 +7,6 @@ from functools import partial import pytest - import snowflake.connector from snowflake.connector.errors import ProgrammingError from snowflake.snowpark import Row, Session @@ -23,6 +22,7 @@ _get_active_session, _get_active_sessions, ) + from tests.utils import IS_IN_STORED_PROC, IS_IN_STORED_PROC_LOCALFS, TestFiles, Utils @@ -199,9 +199,8 @@ def test_close_session_in_sp(session): original_platform = internal_utils.PLATFORM internal_utils.PLATFORM = "XP" try: - with pytest.raises(SnowparkSessionException) as exec_info: - session.close() - assert exec_info.value.error_code == "1411" + session.close() + assert not session.connection.is_closed() finally: internal_utils.PLATFORM = original_platform From 51741afcbaefcd42937c38c8cbdcd149bd491a72 Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Fri, 23 Feb 2024 18:38:15 +0100 Subject: [PATCH 10/16] due to the changes this test does not need to be skipped --- tests/integ/test_session.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integ/test_session.py b/tests/integ/test_session.py index 623a97f3813..56377e14bf1 100644 --- a/tests/integ/test_session.py +++ b/tests/integ/test_session.py @@ -569,7 +569,6 @@ def test_use_secondary_roles(session): session.use_secondary_roles(current_role[1:-1]) -@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() From 4def5f67ee4d5c661d2a1db417f32dc94fc1ae7d Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Fri, 23 Feb 2024 19:32:45 +0100 Subject: [PATCH 11/16] rerun pre-commit --- tests/integ/test_session.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integ/test_session.py b/tests/integ/test_session.py index 56377e14bf1..9a473e8076e 100644 --- a/tests/integ/test_session.py +++ b/tests/integ/test_session.py @@ -7,6 +7,7 @@ from functools import partial import pytest + import snowflake.connector from snowflake.connector.errors import ProgrammingError from snowflake.snowpark import Row, Session @@ -22,7 +23,6 @@ _get_active_session, _get_active_sessions, ) - from tests.utils import IS_IN_STORED_PROC, IS_IN_STORED_PROC_LOCALFS, TestFiles, Utils From a58b62a555ef225742bc00d3b0f8f1b8a5b8df7d Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Sun, 25 Feb 2024 07:21:21 +0100 Subject: [PATCH 12/16] remove error message - not used anymore --- src/snowflake/snowpark/_internal/error_message.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/src/snowflake/snowpark/_internal/error_message.py b/src/snowflake/snowpark/_internal/error_message.py index 39d3d93078c..d89e8b40efd 100644 --- a/src/snowflake/snowpark/_internal/error_message.py +++ b/src/snowflake/snowpark/_internal/error_message.py @@ -423,13 +423,6 @@ def DONT_CREATE_SESSION_IN_SP() -> SnowparkSessionException: error_code="1410", ) - @staticmethod - def DONT_CLOSE_SESSION_IN_SP() -> SnowparkSessionException: - return SnowparkSessionException( - "In a stored procedure, you shouldn't close a session. The stored procedure manages the lifecycle of the provided session.", - error_code="1411", - ) - # General Error codes 15XX @staticmethod From a05151f7527ebb5077e317f949b6d83fdbccf274 Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Sun, 25 Feb 2024 07:24:27 +0100 Subject: [PATCH 13/16] readd skip needed since creation of sessions not allowed in SP --- tests/integ/test_session.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integ/test_session.py b/tests/integ/test_session.py index 9a473e8076e..c1f86f206e3 100644 --- a/tests/integ/test_session.py +++ b/tests/integ/test_session.py @@ -569,6 +569,7 @@ def test_use_secondary_roles(session): session.use_secondary_roles(current_role[1:-1]) +@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() From 12dcadbe71f9af19d8ad9cd3c8589a7fde1bf9a4 Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Sun, 25 Feb 2024 07:33:57 +0100 Subject: [PATCH 14/16] update CHANGELOG.md to describe changed session closing behavior in stored procedures --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index ba34f8ccafb..c4a3c4d2339 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,7 @@ ### Improvements - Added cleanup logic at interpreter shutdown to close all active sessions. +- Closing sessions within stored procedures now is a no-op logging a warning instead of raising an error. ### Bug Fixes From 1a578982bb5c1eea35c7ab6b27c36ff4fba36a72 Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Tue, 5 Mar 2024 11:46:09 +0100 Subject: [PATCH 15/16] skip closing session in __exit__ --- src/snowflake/snowpark/session.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index 5191bc1897f..1a3b88305e4 100644 --- a/src/snowflake/snowpark/session.py +++ b/src/snowflake/snowpark/session.py @@ -215,6 +215,8 @@ def _close_session_atexit(): This is the helper function to close all active sessions at interpreter shutdown. For example, when a jupyter notebook is shutting down, this will also close all active sessions and make sure send all telemetry to the server. """ + if is_in_stored_procedure(): + return with _session_management_lock: for session in _active_sessions.copy(): try: @@ -476,6 +478,8 @@ def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): + if not is_in_stored_procedure(): + self.close() self.close() def __str__(self): From 878528f109330d159a04e3253a0bef4e02a03482 Mon Sep 17 00:00:00 2001 From: zaramzamzam Date: Tue, 5 Mar 2024 11:52:31 +0100 Subject: [PATCH 16/16] remove closing of session when in stored procedure --- src/snowflake/snowpark/session.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index 1a3b88305e4..4aa83919316 100644 --- a/src/snowflake/snowpark/session.py +++ b/src/snowflake/snowpark/session.py @@ -480,7 +480,6 @@ def __enter__(self): def __exit__(self, exc_type, exc_val, exc_tb): if not is_in_stored_procedure(): self.close() - self.close() def __str__(self): return (