SNOW-2097824: cloudpickle callback function to support locally defined function (#3356)

Author: sfc-gh-aling

CHANGELOG.md

@@ -29,6 +29,7 @@
 
 - Fixed a bug in `DataFrameWriter.dbapi` (PrPr) that unicode or double-quoted column name in external database causes error because not quoted correctly.
 - Fixed a bug where named fields in nested OBJECT data could cause errors when containing spaces.
+- Fixed a bug in `DataFrameReader.dbapi` (PrPr) where the `create_connection` defined as local function was incompatible with multiprocessing.
 
 ### Snowpark Local Testing Updates
 

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

@@ -1,7 +1,8 @@
 #
 # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
 #
-
+import pickle
+import cloudpickle
 from enum import Enum
 
 from typing import List, Any, Iterator, Type, Callable, Optional
@@ -28,14 +29,26 @@ def __init__(
         session_init_statement: Optional[List[str]] = None,
         fetch_merge_count: Optional[int] = 1,
     ) -> None:
-        self.driver = driver_class(create_connection, dbms_type)
+        # we use cloudpickle to pickle the callback function so that local function and function defined in
+        # __main__ can be pickled and unpickled in subprocess
+        self.pickled_create_connection_callback = cloudpickle.dumps(
+            create_connection, protocol=pickle.HIGHEST_PROTOCOL
+        )
+        self.driver = None
+        self.driver_class = driver_class
+        self.dbms_type = dbms_type
         self.schema = schema
         self.fetch_size = fetch_size
         self.query_timeout = query_timeout
         self.session_init_statement = session_init_statement
         self.fetch_merge_count = fetch_merge_count
 
     def read(self, partition: str) -> Iterator[List[Any]]:
+        self.driver = self.driver_class(
+            cloudpickle.loads(self.pickled_create_connection_callback),
+            self.dbms_type,
+        )
+
         conn = self.driver.prepare_connection(
             self.driver.create_connection(), self.query_timeout
         )
@@ -74,4 +87,6 @@ def read(self, partition: str) -> Iterator[List[Any]]:
             conn.close()
 
     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
+        assert self.driver is not None
         return self.driver.data_source_data_to_pandas_df(data, self.schema)

tests/integ/test_data_source_api.py

@@ -6,8 +6,10 @@
 import logging
 import math
 import os
+import subprocess
 import tempfile
 import datetime
+from textwrap import dedent
 from unittest import mock
 from unittest.mock import patch, MagicMock, PropertyMock
 
@@ -929,3 +931,49 @@ def test_double_quoted_column_name_sql_server(session):
     assert df.collect() == [
         Row(Id=1, FullName="John Doe", Country="USA", Notes="Fake note")
     ]
+
+
+@pytest.mark.skipif(
+    IS_WINDOWS,
+    reason="sqlite3 file can not be shared across processes on windows",
+)
+def test_local_create_connection_function(session, db_parameters):
+    with tempfile.TemporaryDirectory() as temp_dir:
+        dbpath = os.path.join(temp_dir, "testsqlite3.db")
+        table_name, _, _, assert_data = sqlite3_db(dbpath)
+
+        # test local function definition
+        def local_create_connection():
+            import sqlite3
+
+            return sqlite3.connect(dbpath)
+
+        df = session.read.dbapi(
+            local_create_connection,
+            table=table_name,
+            custom_schema=SQLITE3_DB_CUSTOM_SCHEMA_STRING,
+        )
+        assert df.order_by("ID").collect() == assert_data
+
+        # test function is defined in the main
+        code = dedent(
+            f"""
+        if __name__ == "__main__":
+            import sqlite3
+            from snowflake.snowpark import Session
+
+            def local_create_connection():
+                return sqlite3.connect("{str(dbpath)}")
+
+            session = Session.builder.configs({str(db_parameters)}).create()
+            df = session.read.dbapi(
+                local_create_connection,
+                table='{table_name}',
+                custom_schema='{SQLITE3_DB_CUSTOM_SCHEMA_STRING}',
+            )
+            assert df.collect()
+            print("successful ingestion")
+        """
+        )
+        result = subprocess.run(["python", "-c", code], capture_output=True, text=True)
+        assert "successful ingestion" in result.stdout

tox.ini

@@ -7,8 +7,9 @@ omit = */src/snowflake/snowpark/modin/config/*
        */src/snowflake/snowpark/modin/plugin/docstrings/*
        */src/snowflake/snowpark/mock/*
        */src/snowflake/snowpark/_internal/data_source/datasource_typing.py
-       tests/integ/datasource/test_databricks.py
-       tests/integ/datasource/test_oracledb.py
+       tests/integ/datasource/*.py
+       tests/integ/test_data_source_api.py
+
 [coverage:run]
 relative_files = true
 branch = true