SNOW-2301201: use server cursor to fetch data (#3726)

Author: sfc-gh-aling

CHANGELOG.md

@@ -48,6 +48,7 @@
 
 - Enhanced error handling in `DataFrameReader.dbapi` thread-based ingestion to prevent unnecessary operations, which improves resource efficiency.
 - Bumped cloudpickle dependency to also support `cloudpickle==3.1.1` in addition to previous versions.
+- Improved `DataFrameReader.dbapi` (PuPr) ingestion performance for PostgreSQL and MySQL by using server side cursor to fetch data.
 
 ### Snowpark pandas API Updates
 

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

@@ -52,8 +52,8 @@ def read(self, partition: str) -> Iterator[List[Any]]:
         conn = self.driver.prepare_connection(
             self.driver.create_connection(), self.query_timeout
         )
-        cursor = conn.cursor()
         try:
+            cursor = conn.cursor()
             if self.session_init_statement:
                 for statement in self.session_init_statement:
                     try:
@@ -62,6 +62,9 @@ def read(self, partition: str) -> Iterator[List[Any]]:
                         raise SnowparkDataframeReaderException(
                             f"Failed to execute session init statement: '{statement}' due to exception '{exc!r}'"
                         )
+            # use server side cursor to fetch data if supported by the driver
+            # some drivers do not support execute twice on server side cursor (e.g. psycopg2)
+            cursor = self.driver.get_server_cursor_if_supported(conn)
             if self.fetch_size == 0:
                 cursor.execute(partition)
                 result = cursor.fetchall()

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

@@ -260,3 +260,20 @@ def to_result_snowpark_df_udtf(
             for field in schema.fields
         ]
         return res_df.select(cols, _emit_ast=_emit_ast)
+
+    def get_server_cursor_if_supported(self, conn: "Connection") -> "Cursor":
+        """
+        This method is used to get a server cursor if the driver and the DBMS supports it.
+        It can be overridden by the driver to return a server cursor if supported.
+        Otherwise, it will return the default cursor supported by the driver and the DBMS.
+
+        - databricks-sql-connector: no concept of client/server cursor, no need to override
+        - python-oracledb: default to the server cursor, no need to override
+        - psycopg2: default to the client cursor which needs to be overridden to return the server cursor
+        - pymysql: default to the client cursor which needs to be overridden to return the server cursor
+
+        TODO:
+        - pyodbc: This is a Python wrapper on top of ODBC drivers, the ODBC driver and the DBMS may or may not support server cursor
+         and if they do support, the way to get the server cursor may vary across different DBMS. we need to document pyodbc.
+        """
+        return conn.cursor()

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

@@ -7,6 +7,7 @@
 
 from snowflake.snowpark._internal.data_source.datasource_typing import Connection
 from snowflake.snowpark._internal.data_source.drivers import BaseDriver
+from snowflake.snowpark._internal.utils import generate_random_alphanumeric
 from snowflake.snowpark.functions import to_variant, parse_json, column
 from snowflake.snowpark.types import (
     StructType,
@@ -28,6 +29,9 @@
 if TYPE_CHECKING:
     from snowflake.snowpark.session import Session  # pragma: no cover
     from snowflake.snowpark.dataframe import DataFrame  # pragma: no cover
+    from snowflake.snowpark._internal.data_source.datasource_typing import (
+        Cursor,
+    )  # pragma: no cover
 
 
 logger = logging.getLogger(__name__)
@@ -277,7 +281,9 @@ def prepare_connection_in_udtf(
         class UDTFIngestion:
             def process(self, query: str):
                 conn = prepare_connection_in_udtf(create_connection())
-                cursor = conn.cursor()
+                cursor = conn.cursor(
+                    f"SNOWPARK_CURSOR_{generate_random_alphanumeric(5)}"
+                )
                 cursor.execute(query)
                 while True:
                     rows = cursor.fetchmany(fetch_size)
@@ -286,3 +292,6 @@ def process(self, query: str):
                     yield from rows
 
         return UDTFIngestion
+
+    def get_server_cursor_if_supported(self, conn: "Connection") -> "Cursor":
+        return conn.cursor(f"SNOWPARK_CURSOR_{generate_random_alphanumeric(5)}")

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

@@ -192,9 +192,10 @@ def udtf_class_builder(
 
         class UDTFIngestion:
             def process(self, query: str):
+                import pymysql
 
                 conn = create_connection()
-                cursor = conn.cursor()
+                cursor = pymysql.cursors.SSCursor(conn)
                 cursor.execute(query)
                 while True:
                     rows = cursor.fetchmany(fetch_size)
@@ -256,3 +257,8 @@ def to_result_snowpark_df_udtf(
             else:
                 cols.append(res_df[field.name].cast(field.datatype).alias(field.name))
         return res_df.select(cols, _emit_ast=_emit_ast)
+
+    def get_server_cursor_if_supported(self, conn: "Connection") -> "Cursor":
+        import pymysql
+
+        return pymysql.cursors.SSCursor(conn)

tests/integ/datasource/test_mysql.py

@@ -287,3 +287,12 @@ def test_pymysql_driver_udtf_class_builder():
     # Verify we got data with the right structure (2 columns)
     assert len(column_result_rows) > 0
     assert len(column_result_rows[0]) == 2  # Two columns
+
+
+def test_server_side_cursor():
+    conn = create_connection_mysql()
+    driver = PymysqlDriver(create_connection_mysql, DBMS_TYPE.MYSQL_DB)
+    cursor = driver.get_server_cursor_if_supported(conn)
+    assert isinstance(cursor, pymysql.cursors.SSCursor)
+    cursor.close()
+    conn.close()

tests/integ/datasource/test_postgres.py

@@ -464,3 +464,12 @@ def test_unit_generate_select_query():
         'SELECT TO_JSON("jsonb_col")::TEXT AS jsonb_col FROM test_table'
     )
     assert jsonb_query == expected_jsonb_query
+
+
+def test_server_side_cursor(session):
+    conn = create_postgres_connection()
+    driver = Psycopg2Driver(create_postgres_connection, DBMS_TYPE.POSTGRES_DB)
+    cursor = driver.get_server_cursor_if_supported(conn)
+    assert cursor.name is not None  # Server-side cursor should have a name
+    cursor.close()
+    conn.close()

tests/unit/test_data_source.py

@@ -81,6 +81,7 @@ def test_datasource_reader_close_error_handling(cursor_fails, conn_fails):
     mock_driver.prepare_connection.return_value = mock_conn
     mock_driver.create_connection.return_value = mock_conn
     mock_conn.cursor.return_value = mock_cursor
+    mock_driver.get_server_cursor_if_supported.return_value = mock_cursor
     mock_cursor.fetchall.return_value = [("test_data",)]
 
     # Configure failures