SNOW-2372025: support passing connection parameters to dbapi (#3845)

Author: sfc-gh-aling

CHANGELOG.md

@@ -7,6 +7,7 @@
 #### New Features
 
 - Added a new function `service` in `snowflake.snowpark.functions` that allows users to create a callable representing a Snowpark Container Services (SPCS) service.
+- Added `connection_parameters` parameter to `DataFrameReader.dbapi()` (PuPr) method to allow passing keyword arguments to the `create_connection` callable.
 - Added support for `Session.begin_transaction`, `Session.commit` and `Session.rollback`.
 - Added support for the following functions in `functions.py`:
     - Geospatial functions:

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

@@ -37,7 +37,7 @@
 class DataSourcePartitioner:
     def __init__(
         self,
-        create_connection: Callable[[], "Connection"],
+        create_connection: Callable[..., "Connection"],
         table_or_query: str,
         is_query: bool,
         column: Optional[str] = None,
@@ -50,6 +50,7 @@ def __init__(
         predicates: Optional[List[str]] = None,
         session_init_statement: Optional[List[str]] = None,
         fetch_merge_count: Optional[int] = 1,
+        connection_parameters: Optional[dict] = None,
     ) -> None:
         self.create_connection = create_connection
         self.table_or_query = table_or_query
@@ -64,14 +65,21 @@ def __init__(
         self.predicates = predicates
         self.session_init_statement = session_init_statement
         self.fetch_merge_count = fetch_merge_count
-        conn = create_connection()
+        self.connection_parameters = connection_parameters
+        conn = (
+            create_connection(**connection_parameters)
+            if connection_parameters
+            else create_connection()
+        )
         dbms_type, driver_type = detect_dbms(conn)
         self.driver_type = driver_type
         self.dbms_type = dbms_type
         self.dialect_class = DBMS_MAP.get(dbms_type, BaseDialect)
         self.driver_class = DRIVER_MAP.get(driver_type, BaseDriver)
         self.dialect = self.dialect_class()
-        self.driver = self.driver_class(create_connection, dbms_type)
+        self.driver = self.driver_class(
+            create_connection, dbms_type, connection_parameters
+        )
 
         self._query_input_alias = (
             f"SNOWPARK_DBAPI_QUERY_INPUT_ALIAS_{generate_random_alphanumeric(5).upper()}"
@@ -89,6 +97,7 @@ def reader(self) -> DataSourceReader:
             self.query_timeout,
             self.session_init_statement,
             self.fetch_merge_count,
+            self.connection_parameters,
         )
 
     @cached_property

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

@@ -21,19 +21,23 @@ class DataSourceReader:
     def __init__(
         self,
         driver_class: Type[BaseDriver],
-        create_connection: Callable[[], "Connection"],
+        create_connection: Callable[..., "Connection"],
         schema: StructType,
         dbms_type: Enum,
         fetch_size: Optional[int] = 0,
         query_timeout: Optional[int] = 0,
         session_init_statement: Optional[List[str]] = None,
         fetch_merge_count: Optional[int] = 1,
+        connection_parameters: Optional[dict] = None,
     ) -> None:
         # 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.pickled_connection_parameters = cloudpickle.dumps(
+            connection_parameters, protocol=pickle.HIGHEST_PROTOCOL
+        )
         self.driver = None
         self.driver_class = driver_class
         self.dbms_type = dbms_type
@@ -44,14 +48,19 @@ def __init__(
         self.fetch_merge_count = fetch_merge_count
 
     def read(self, partition: str) -> Iterator[List[Any]]:
+        connection_parameters = cloudpickle.loads(self.pickled_connection_parameters)
         self.driver = self.driver_class(
             cloudpickle.loads(self.pickled_create_connection_callback),
             self.dbms_type,
+            connection_parameters,
         )
 
-        conn = self.driver.prepare_connection(
-            self.driver.create_connection(), self.query_timeout
+        create_conn_result = (
+            self.driver.create_connection(**connection_parameters)
+            if connection_parameters
+            else self.driver.create_connection()
         )
+        conn = self.driver.prepare_connection(create_conn_result, self.query_timeout)
         try:
             cursor = conn.cursor()
             if self.session_init_statement:

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

@@ -43,13 +43,21 @@
 class BaseDriver:
     def __init__(
         self,
-        create_connection: Callable[[], "Connection"],
+        create_connection: Callable[..., "Connection"],
         dbms_type: Enum,
+        connection_parameters: Optional[dict] = None,
     ) -> None:
         self.create_connection = create_connection
         self.dbms_type = dbms_type
+        self.connection_parameters = connection_parameters
         self.raw_schema = None
 
+    def _call_create_connection(self) -> "Connection":
+        """Call create_connection with connection_parameters if provided."""
+        if self.connection_parameters:
+            return self.create_connection(**self.connection_parameters)
+        return self.create_connection()
+
     def to_snow_type(self, schema: List[Any]) -> StructType:
         raise NotImplementedError(
             f"{self.__class__.__name__} has not implemented to_snow_type function"
@@ -100,7 +108,7 @@ def infer_schema_from_description(
     def infer_schema_from_description_with_error_control(
         self, table_or_query: str, is_query: bool, query_input_alias: str
     ) -> StructType:
-        conn = self.create_connection()
+        conn = self._call_create_connection()
         cursor = conn.cursor()
         try:
             return self.infer_schema_from_description(
@@ -184,10 +192,16 @@ def udtf_class_builder(
     ) -> type:
         create_connection = self.create_connection
         prepare_connection = self.prepare_connection
+        connection_parameters = self.connection_parameters
 
         class UDTFIngestion:
             def process(self, query: str):
-                conn = prepare_connection(create_connection(), query_timeout)
+                conn_result = (
+                    create_connection(**connection_parameters)
+                    if connection_parameters
+                    else create_connection()
+                )
+                conn = prepare_connection(conn_result, query_timeout)
                 cursor = conn.cursor()
                 if session_init_statement is not None:
                     for statement in session_init_statement:

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

@@ -88,10 +88,15 @@ def udtf_class_builder(
         query_timeout: int = 0,
     ) -> type:
         create_connection = self.create_connection
+        connection_parameters = self.connection_parameters
 
         class UDTFIngestion:
             def process(self, query: str):
-                conn = create_connection()
+                conn = (
+                    create_connection(**connection_parameters)
+                    if connection_parameters
+                    else create_connection()
+                )
                 cursor = conn.cursor()
                 if session_init_statement is not None:
                     for statement in session_init_statement:

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

@@ -143,6 +143,7 @@ def udtf_class_builder(
         query_timeout: int = 0,
     ) -> type:
         create_connection = self.create_connection
+        connection_parameters = self.connection_parameters
 
         def oracledb_output_type_handler(cursor, metadata):
             from oracledb import (
@@ -166,7 +167,11 @@ def convert_to_hex(value):
 
         class UDTFIngestion:
             def process(self, query: str):
-                conn = create_connection()
+                conn = (
+                    create_connection(**connection_parameters)
+                    if connection_parameters
+                    else create_connection()
+                )
                 if query_timeout > 0:
                     conn.call_timeout = query_timeout * 1000
                 if conn.outputtypehandler is None:

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

@@ -3,7 +3,7 @@
 #
 import logging
 from enum import Enum
-from typing import Callable, List, Any, TYPE_CHECKING
+from typing import List, Any, TYPE_CHECKING
 
 from snowflake.snowpark._internal.data_source.datasource_typing import Connection
 from snowflake.snowpark._internal.data_source.drivers import BaseDriver
@@ -167,11 +167,6 @@ class Psycopg2TypeCode(Enum):
 
 
 class Psycopg2Driver(BaseDriver):
-    def __init__(
-        self, create_connection: Callable[[], "Connection"], dbms_type: Enum
-    ) -> None:
-        super().__init__(create_connection, dbms_type)
-
     def to_snow_type(self, schema: List[Any]) -> StructType:
         # The psycopg2 spec is defined in the following links:
         # https://www.psycopg.org/docs/cursor.html#cursor.description
@@ -272,6 +267,7 @@ def udtf_class_builder(
         query_timeout: int = 0,
     ) -> type:
         create_connection = self.create_connection
+        connection_parameters = self.connection_parameters
 
         # TODO: SNOW-2101485 use class method to prepare connection
         # ideally we should use the same function as prepare_connection
@@ -291,7 +287,12 @@ def prepare_connection_in_udtf(
 
         class UDTFIngestion:
             def process(self, query: str):
-                conn = prepare_connection_in_udtf(create_connection(), query_timeout)
+                conn_result = (
+                    create_connection(**connection_parameters)
+                    if connection_parameters
+                    else create_connection()
+                )
+                conn = prepare_connection_in_udtf(conn_result, query_timeout)
                 cursor = conn.cursor(
                     f"SNOWPARK_CURSOR_{generate_random_alphanumeric(5)}"
                 )

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

@@ -203,12 +203,17 @@ def udtf_class_builder(
         query_timeout: int = 0,
     ) -> type:
         create_connection = self.create_connection
+        connection_parameters = self.connection_parameters
 
         class UDTFIngestion:
             def process(self, query: str):
                 import pymysql
 
-                conn = create_connection()
+                conn = (
+                    create_connection(**connection_parameters)
+                    if connection_parameters
+                    else create_connection()
+                )
                 cursor = pymysql.cursors.SSCursor(conn)
                 if session_init_statement is not None:
                     for statement in session_init_statement:

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

@@ -86,6 +86,7 @@ def udtf_class_builder(
     ) -> type:
         create_connection = self.create_connection
         prepare_connection = self.prepare_connection
+        connection_parameters = self.connection_parameters
 
         def binary_converter(value):
             return value.hex() if value is not None else None
@@ -94,7 +95,12 @@ class UDTFIngestion:
             def process(self, query: str):
                 import pyodbc
 
-                conn = prepare_connection(create_connection(), query_timeout)
+                conn_result = (
+                    create_connection(**connection_parameters)
+                    if connection_parameters
+                    else create_connection()
+                )
+                conn = prepare_connection(conn_result, query_timeout)
                 if (
                     conn.get_output_converter(pyodbc.SQL_BINARY) is None
                     and conn.get_output_converter(pyodbc.SQL_VARBINARY) is None

src/snowflake/snowpark/dataframe_reader.py

@@ -1681,7 +1681,7 @@ def jdbc(
     @publicapi
     def dbapi(
         self,
-        create_connection: Callable[[], "Connection"],
+        create_connection: Callable[..., "Connection"],
         *,
         table: Optional[str] = None,
         query: Optional[str] = None,
@@ -1698,6 +1698,7 @@ def dbapi(
         udtf_configs: Optional[dict] = None,
         fetch_merge_count: int = 1,
         fetch_with_process: bool = False,
+        connection_parameters: Optional[dict] = None,
         _emit_ast: bool = True,
     ) -> DataFrame:
         """
@@ -1718,7 +1719,9 @@ def dbapi(
         column name with table in external data source.
 
         Args:
-            create_connection: A callable that takes no arguments and returns a DB-API compatible database connection.
+            create_connection: A callable that returns a DB-API compatible database connection.
+                The callable can optionally accept keyword arguments via `**kwargs`.
+                If connection_parameters is provided, those will be passed as keyword arguments to this callable.
                 The callable must be picklable, as it will be passed to and executed in child processes.
             table: The name of the table in the external data source.
                 This parameter cannot be used together with the `query` parameter.
@@ -1777,6 +1780,11 @@ def dbapi(
                 like Parquet file generation. When using multiprocessing, guard your script with
                 `if __name__ == "__main__":` and call `multiprocessing.freeze_support()` on Windows if needed.
                 This parameter has no effect in UDFT ingestion.
+            connection_parameters: Optional dictionary of parameters to pass to the create_connection callable.
+                If provided, these parameters will be unpacked and passed as keyword arguments
+                to create_connection(`**connection_parameters`).
+                This allows for flexible connection configuration without hardcoding values in the callable.
+                Example: {"timeout": 30, "isolation_level": "READ_UNCOMMITTED"}
 
         Example::
             .. code-block:: python
@@ -1796,8 +1804,27 @@ def create_oracledb_connection():
                     connection = oracledb.connect(...)
                     return connection
 
-                if __name__ == "__main__":
-                    df = session.read.dbapi(create_oracledb_connection, table=..., fetch_with_process=True)
+                df = session.read.dbapi(create_oracledb_connection, table=..., fetch_with_process=True)
+
+        Example::
+            .. code-block:: python
+
+                import sqlite3
+                def create_sqlite_connection(timeout=5.0, isolation_level=None, **kwargs):
+                    connection = sqlite3.connect(
+                        database=":memory:",
+                        timeout=timeout,
+                        isolation_level=isolation_level
+                    )
+                    return connection
+
+                connection_params = {"timeout": 30.0, "isolation_level": "DEFERRED"}
+                df = session.read.dbapi(
+                    create_sqlite_connection,
+                    table=...,
+                    connection_parameters=connection_params
+                )
+
         """
         if (not table and not query) or (table and query):
             raise SnowparkDataframeReaderException(
@@ -1831,6 +1858,7 @@ def create_oracledb_connection():
             predicates,
             session_init_statement,
             fetch_merge_count,
+            connection_parameters,
         )
         struct_schema = partitioner.schema
         partitioned_queries = partitioner.partitions