SNOW-2201361:fix dbapi does not work in python stored proc (#3541)

Author: sfc-gh-yuwang

CHANGELOG.md

@@ -4,6 +4,10 @@
 
 ### Snowpark Python API Updates
 
+#### Bug Fixes
+
+- Fixed a bug in `DataFrameReader.dbapi` (PrPr) that `dbapi` fail in python stored procedure with process exit with code 1.
+
 #### New Features
 
 - Added support for the following AI-powered functions in `functions.py`:

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

@@ -1,7 +1,7 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-
+import os
 import queue
 import time
 import traceback
@@ -11,7 +11,7 @@
 from threading import BoundedSemaphore
 from io import BytesIO
 from enum import Enum
-from typing import Any, Tuple, Optional, Callable, Dict, Union
+from typing import Any, Tuple, Optional, Callable, Dict, Union, Set
 import logging
 from snowflake.snowpark._internal.data_source.dbms_dialects import (
     Sqlite3Dialect,
@@ -309,6 +309,7 @@ def _retry_run(func: Callable, *args, **kwargs) -> Any:
 def worker_process(
     partition_queue: Union[mp.Queue, queue.Queue],
     parquet_queue: Union[mp.Queue, queue.Queue],
+    process_or_thread_error_indicator: Union[mp.Queue, queue.Queue],
     reader,
     stop_event: threading.Event = None,
 ):
@@ -326,6 +327,8 @@ def worker_process(
                 stop_event,
             )
         except queue.Empty:
+            # indicate whether a process is exit gracefully
+            process_or_thread_error_indicator.put(os.getpid())
             # No more work available, exit gracefully
             break
         except Exception as e:
@@ -358,9 +361,22 @@ def process_completed_futures(thread_futures):
                 raise
 
 
+def _drain_process_status_queue(
+    process_or_thread_error_indicator: Union[mp.Queue, queue.Queue],
+) -> Set:
+    result = set()
+    while True:
+        try:
+            result.add(process_or_thread_error_indicator.get(block=False))
+        except queue.Empty:
+            break
+    return result
+
+
 def process_parquet_queue_with_threads(
     session: "snowflake.snowpark.Session",
     parquet_queue: Union[mp.Queue, queue.Queue],
+    process_or_thread_error_indicator: Union[mp.Queue, queue.Queue],
     workers: list,
     total_partitions: int,
     snowflake_stage_name: str,
@@ -382,6 +398,7 @@ def process_parquet_queue_with_threads(
     Args:
         session: Snowflake session for database operations
         parquet_queue: Multiprocessing queue containing parquet data
+        process_or_thread_error_indicator: Multiprocessing queue containing process exit information
         workers: List of worker processes or thread futures to monitor
         total_partitions: Total number of partitions expected
         snowflake_stage_name: Name of the Snowflake stage for uploads
@@ -395,6 +412,7 @@ def process_parquet_queue_with_threads(
     """
 
     completed_partitions = set()
+    gracefully_exited_processes = set()
     # process parquet_queue may produce more data than the threads can handle,
     # so we use semaphore to limit the number of threads
     backpressure_semaphore = BoundedSemaphore(value=_MAX_WORKER_SCALE * max_workers)
@@ -448,10 +466,18 @@ def process_parquet_queue_with_threads(
                 if fetch_with_process:
                     # Check if any processes have failed
                     for i, process in enumerate(workers):
-                        if not process.is_alive() and process.exitcode != 0:
-                            raise SnowparkDataframeReaderException(
-                                f"Partition {i} data fetching process failed with exit code {process.exitcode}"
+                        if not process.is_alive():
+                            gracefully_exited_processes = (
+                                gracefully_exited_processes.union(
+                                    _drain_process_status_queue(
+                                        process_or_thread_error_indicator
+                                    )
+                                )
                             )
+                            if process.pid not in gracefully_exited_processes:
+                                raise SnowparkDataframeReaderException(
+                                    f"Partition {i} data fetching process failed with exit code {process.exitcode} or failed silently"
+                                )
                 else:
                     # Check if any threads have failed
                     for i, future in enumerate(workers):
@@ -469,11 +495,18 @@ def process_parquet_queue_with_threads(
 
     if fetch_with_process:
         # Wait for all processes to complete
-        for idx, process in enumerate(workers):
+        for process in workers:
             process.join()
-            if process.exitcode != 0:
+        # empty parquet queue to get all signals after each process ends
+        gracefully_exited_processes = gracefully_exited_processes.union(
+            _drain_process_status_queue(process_or_thread_error_indicator)
+        )
+
+        # check if any process fails
+        for idx, process in enumerate(workers):
+            if process.pid not in gracefully_exited_processes:
                 raise SnowparkDataframeReaderException(
-                    f"Partition {idx} data fetching process failed with exit code {process.exitcode}"
+                    f"Partition {idx} data fetching process failed with exit code {process.exitcode} or failed silently"
                 )
     else:
         # Wait for all threads to complete

src/snowflake/snowpark/dataframe_reader.py

@@ -1491,6 +1491,7 @@ def create_oracledb_connection():
             # Determine the number of processes or threads to use
             max_workers = max_workers or os.cpu_count()
             queue_class = mp.Queue if fetch_with_process else queue.Queue
+            process_or_thread_error_indicator = queue_class()
             # a queue of partitions to be processed, this is filled by the partitioner before starting the workers
             partition_queue = queue_class()
             # a queue of parquet BytesIO objects to be uploaded
@@ -1509,7 +1510,12 @@ def create_oracledb_connection():
                 for _worker_id in range(max_workers):
                     process = mp.Process(
                         target=worker_process,
-                        args=(partition_queue, parquet_queue, partitioner.reader()),
+                        args=(
+                            partition_queue,
+                            parquet_queue,
+                            process_or_thread_error_indicator,
+                            partitioner.reader(),
+                        ),
                     )
                     process.start()
                     workers.append(process)
@@ -1523,6 +1529,7 @@ def create_oracledb_connection():
                         worker_process,
                         partition_queue,
                         parquet_queue,
+                        process_or_thread_error_indicator,
                         partitioner.reader(),
                         data_fetching_thread_stop_event,
                     )
@@ -1533,6 +1540,7 @@ def create_oracledb_connection():
             process_parquet_queue_with_threads(
                 session=self._session,
                 parquet_queue=parquet_queue,
+                process_or_thread_error_indicator=process_or_thread_error_indicator,
                 workers=workers,
                 total_partitions=len(partitioned_queries),
                 snowflake_stage_name=snowflake_stage_name,

tests/integ/test_data_source_api.py

@@ -1160,13 +1160,18 @@ def test_worker_process_unit(fetch_with_process):
             multiprocessing.Queue() if fetch_with_process else queue.Queue()
         )
         parquet_queue = multiprocessing.Queue() if fetch_with_process else queue.Queue()
+        process_or_thread_error_indicator = (
+            multiprocessing.Queue() if fetch_with_process else queue.Queue()
+        )
 
         # Set up partition_queue to return test data, then raise queue.Empty
         partition_queue.put((0, f"SELECT * FROM {table_name} WHERE id <= 3"))
         partition_queue.put((1, f"SELECT * FROM {table_name} WHERE id > 3"))
 
         # Call the worker_process function directly (using real sqlite3 operations)
-        worker_process(partition_queue, parquet_queue, reader)
+        worker_process(
+            partition_queue, parquet_queue, process_or_thread_error_indicator, reader
+        )
 
         expected_order = [
             "data_partition0_fetch0.parquet",
@@ -1191,7 +1196,9 @@ def test_worker_process_unit(fetch_with_process):
 
         # check error handling
         partition_queue.put((0, "SELECT * FROM NON_EXISTING_TABLE"))
-        worker_process(partition_queue, parquet_queue, reader)
+        worker_process(
+            partition_queue, parquet_queue, process_or_thread_error_indicator, reader
+        )
         error_signal, error_instance = parquet_queue.get()
         assert error_signal == PARTITION_TASK_ERROR_SIGNAL
         assert isinstance(
@@ -1520,13 +1527,15 @@ def test_thread_worker_exception(exception, match_message):
 
     # Create test parameters
     parquet_queue = queue.Queue()
+    process_or_thread_error_indicator = queue.Queue()
     workers = [mock_future]  # Single worker that will fail
     total_partitions = 0  # No partitions to complete
 
     with pytest.raises(SnowparkDataframeReaderException, match=match_message):
         process_parquet_queue_with_threads(
             session=mock_session,
             parquet_queue=parquet_queue,
+            process_or_thread_error_indicator=process_or_thread_error_indicator,
             workers=workers,
             total_partitions=total_partitions,
             snowflake_stage_name="test_stage",
@@ -1553,6 +1562,7 @@ def test_process_worker_non_zero_exitcode():
 
     # Create test parameters
     parquet_queue = queue.Queue()
+    process_or_thread_error_indicator = queue.Queue()
     workers = [mock_process]  # Single worker that will fail
     total_partitions = 0  # No partitions to complete
 
@@ -1563,6 +1573,7 @@ def test_process_worker_non_zero_exitcode():
         process_parquet_queue_with_threads(
             session=mock_session,
             parquet_queue=parquet_queue,
+            process_or_thread_error_indicator=process_or_thread_error_indicator,
             workers=workers,
             total_partitions=total_partitions,
             snowflake_stage_name="test_stage",
@@ -1589,6 +1600,7 @@ def test_queue_empty_process_failure():
 
     # Create empty queue to trigger queue.Empty exception
     parquet_queue = queue.Queue()
+    process_or_thread_error_indicator = queue.Queue()
     workers = [mock_process]
     total_partitions = 1  # Set to 1 so the loop continues
 
@@ -1599,6 +1611,7 @@ def test_queue_empty_process_failure():
         process_parquet_queue_with_threads(
             session=mock_session,
             parquet_queue=parquet_queue,
+            process_or_thread_error_indicator=process_or_thread_error_indicator,
             workers=workers,
             total_partitions=total_partitions,
             snowflake_stage_name="test_stage",
@@ -1636,13 +1649,15 @@ def test_queue_empty_thread_failure(exception, match_message):
 
     # Create empty queue to trigger queue.Empty exception
     parquet_queue = queue.Queue()
+    process_or_thread_error_indicator = queue.Queue()
     workers = [mock_future]
     total_partitions = 1  # Set to 1 so the loop continues
 
     with pytest.raises(SnowparkDataframeReaderException, match=match_message):
         process_parquet_queue_with_threads(
             session=mock_session,
             parquet_queue=parquet_queue,
+            process_or_thread_error_indicator=process_or_thread_error_indicator,
             workers=workers,
             total_partitions=total_partitions,
             snowflake_stage_name="test_stage",