SNOW-2174093: improve error handling when closing cursor and connection in dbapi (#3487)

Author: sfc-gh-aling

CHANGELOG.md

@@ -23,6 +23,7 @@
 #### Bug Fixes
 
 - Fixed a bug caused by redundant validation when creating an iceberg table.
+- Fixed a bug in `DataFrameReader.dbapi` (PrPr) where closing the cursor or connection could unexpectedly raise an error and terminate the program.
 
 ### Snowpark Local Testing Updates
 

src/snowflake/snowpark/_internal/data_source/datasource_reader.py

@@ -83,8 +83,18 @@ def read(self, partition: str) -> Iterator[List[Any]]:
             else:
                 raise ValueError("fetch size cannot be smaller than 0")
         finally:
-            cursor.close()
-            conn.close()
+            try:
+                cursor.close()
+            except BaseException as exc:
+                logger.debug(
+                    f"Failed to close cursor after reading data due to error: {exc!r}"
+                )
+            try:
+                conn.close()
+            except BaseException as exc:
+                logger.debug(
+                    f"Failed to close connection after reading data due to error: {exc!r}"
+                )
 
     def data_source_data_to_pandas_df(self, data: List[Any]) -> "pd.DataFrame":
         # self.driver is guaranteed to be initialized in self.read() which is called prior to this method

src/snowflake/snowpark/_internal/data_source/drivers/base_driver.py

@@ -82,8 +82,20 @@ def infer_schema_from_description_with_error_control(
                 f" Please check the stack trace for more details."
             ) from exc
         finally:
-            cursor.close()
-            conn.close()
+            # Best effort to close cursor and connection; failures are non-critical and can be ignored.
+            try:
+                cursor.close()
+            except BaseException as exc:
+                logger.debug(
+                    f"Failed to close cursor after inferring schema from description due to error: {exc!r}"
+                )
+
+            try:
+                conn.close()
+            except BaseException as exc:
+                logger.debug(
+                    f"Failed to close connection after inferring schema from description due to error: {exc!r}"
+                )
 
     def udtf_ingestion(
         self,

tests/unit/test_data_source.py

@@ -0,0 +1,116 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+
+import pytest
+import re
+from unittest.mock import Mock, patch
+from snowflake.snowpark._internal.data_source.drivers.base_driver import BaseDriver
+from snowflake.snowpark._internal.data_source.datasource_reader import DataSourceReader
+from snowflake.snowpark._internal.data_source.utils import DBMS_TYPE
+from snowflake.snowpark.types import StructType, StructField, StringType
+
+
+@pytest.mark.parametrize(
+    "cursor_fails,conn_fails",
+    [
+        (True, False),  # cursor.close() fails
+        (False, True),  # connection.close() fails
+        (True, True),  # both fail
+    ],
+)
+def test_close_error_handling(cursor_fails, conn_fails):
+    """Test that errors during cursor/connection close are handled gracefully."""
+    # Setup mock driver
+    mock_create_connection = Mock()
+    driver = BaseDriver(mock_create_connection, DBMS_TYPE.UNKNOWN)
+
+    # Setup mocks
+    mock_conn = Mock()
+    mock_cursor = Mock()
+    mock_conn.cursor.return_value = mock_cursor
+    mock_create_connection.return_value = mock_conn
+
+    # Configure failures
+
+    if conn_fails:
+        mock_conn.close.side_effect = Exception("Connection close failed")
+    if cursor_fails:
+        mock_cursor.close.side_effect = Exception("Cursor close failed")
+
+    # Mock schema inference to succeed
+    expected_schema = StructType([StructField("test_col", StringType())])
+    driver.infer_schema_from_description = Mock(return_value=expected_schema)
+
+    # Test - should succeed despite close errors and log the failure
+    with patch(
+        "snowflake.snowpark._internal.data_source.drivers.base_driver.logger"
+    ) as mock_logger:
+        result = driver.infer_schema_from_description_with_error_control(
+            "test_table", False
+        )
+
+        assert result == expected_schema
+
+        # Use regex to match the log message flexibly
+        mock_logger.debug.assert_called()
+        args, kwargs = mock_logger.debug.call_args
+        assert re.search(r"Failed to close", args[0])
+
+
+@pytest.mark.parametrize(
+    "cursor_fails,conn_fails",
+    [
+        (True, False),  # cursor.close() fails
+        (False, True),  # connection.close() fails
+        (True, True),  # both fail
+    ],
+)
+def test_datasource_reader_close_error_handling(cursor_fails, conn_fails):
+    """Test that DataSourceReader handles cursor/connection close errors gracefully."""
+    # Setup mocks
+    mock_create_connection = Mock()
+    expected_schema = StructType([StructField("test_col", StringType())])
+
+    # Mock driver, connection, and cursor
+    mock_driver = Mock()
+    mock_conn = Mock()
+    mock_cursor = Mock()
+
+    # Configure the mock chain
+    mock_driver.prepare_connection.return_value = mock_conn
+    mock_driver.create_connection.return_value = mock_conn
+    mock_conn.cursor.return_value = mock_cursor
+    mock_cursor.fetchall.return_value = [("test_data",)]
+
+    # Configure failures
+    if cursor_fails:
+        mock_cursor.close.side_effect = Exception("Cursor close failed")
+    if conn_fails:
+        mock_conn.close.side_effect = Exception("Connection close failed")
+
+    # Create mock driver class that returns our mock driver
+    mock_driver_class = Mock(return_value=mock_driver)
+
+    # Create reader with the mock driver class
+    reader = DataSourceReader(
+        driver_class=mock_driver_class,
+        create_connection=mock_create_connection,
+        schema=expected_schema,
+        dbms_type=DBMS_TYPE.UNKNOWN,
+        fetch_size=0,  # Use 0 to trigger fetchall() path
+        query_timeout=0,
+        session_init_statement=None,
+        fetch_merge_count=1,
+    )
+
+    # Test - should succeed despite close errors and log the failure
+    with patch(
+        "snowflake.snowpark._internal.data_source.datasource_reader.logger"
+    ) as mock_logger:
+        # Consume the generator to trigger execution of the finally block
+        list(reader.read("SELECT * FROM test"))
+
+        mock_logger.debug.assert_called()
+        args, kwargs = mock_logger.debug.call_args
+        assert re.search(r"Failed to close", args[0])