SNOW-1955847: DBAPI Postgres Support

CHANGELOG.md

@@ -4,6 +4,10 @@
 
 ### Snowpark Python API Updates
 
+#### New Features
+
+- Added PostgreSQL support to `DataFrameReader.dbapi` (PrPr) for both Parquet and UDTF-based ingestion.
+
 #### Improvements
 
 - Invoking snowflake system procedures does not invoke an additional `describe procedure` call to check the return type of the procedure.

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

@@ -55,7 +55,7 @@ def read(self, partition: str) -> Iterator[List[Any]]:
                 yield result
             elif self.fetch_size > 0:
                 cap_size = self.fetch_merge_count * self.fetch_size
-                cursor = cursor.execute(partition)
+                cursor.execute(partition)
                 batch = []
                 while True:
                     rows = cursor.fetchmany(self.fetch_size)

src/snowflake/snowpark/_internal/data_source/dbms_dialects/__init__.py

@@ -7,6 +7,7 @@
     "SqlServerDialect",
     "OracledbDialect",
     "DatabricksDialect",
+    "PostgresDialect",
 ]
 
 from snowflake.snowpark._internal.data_source.dbms_dialects.base_dialect import (
@@ -24,3 +25,6 @@
 from snowflake.snowpark._internal.data_source.dbms_dialects.databricks_dialect import (
     DatabricksDialect,
 )
+from snowflake.snowpark._internal.data_source.dbms_dialects.postgresql_dialect import (
+    PostgresDialect,
+)

src/snowflake/snowpark/_internal/data_source/dbms_dialects/postgresql_dialect.py

@@ -0,0 +1,8 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+from snowflake.snowpark._internal.data_source.dbms_dialects import BaseDialect
+
+
+class PostgresDialect(BaseDialect):
+    pass

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

@@ -8,6 +8,7 @@
     "SqliteDriver",
     "PyodbcDriver",
     "DatabricksDriver",
+    "Psycopg2Driver",
 ]
 
 from snowflake.snowpark._internal.data_source.drivers.base_driver import BaseDriver
@@ -19,3 +20,6 @@
 from snowflake.snowpark._internal.data_source.drivers.databricks_driver import (
     DatabricksDriver,
 )
+from snowflake.snowpark._internal.data_source.drivers.psycopg2_driver import (
+    Psycopg2Driver,
+)

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

@@ -40,8 +40,8 @@ def to_snow_type(self, schema: List[Any]) -> StructType:
             f"{self.__class__.__name__} has not implemented to_snow_type function"
         )
 
+    @staticmethod
     def prepare_connection(
-        self,
         conn: "Connection",
         query_timeout: int = 0,
     ) -> "Connection":
@@ -81,7 +81,7 @@ def udtf_ingestion(
         udtf_name = f"data_source_udtf_{generate_random_alphanumeric(5)}"
         start = time.time()
         session.udtf.register(
-            self.udtf_class_builder(fetch_size=fetch_size),
+            self.udtf_class_builder(fetch_size=fetch_size, schema=schema),
             name=udtf_name,
             output_schema=StructType(
                 [
@@ -104,7 +104,9 @@ def udtf_ingestion(
         ]
         return res.select(cols)
 
-    def udtf_class_builder(self, fetch_size: int = 1000) -> type:
+    def udtf_class_builder(
+        self, fetch_size: int = 1000, schema: StructType = None
+    ) -> type:
         create_connection = self.create_connection
 
         class UDTFIngestion:

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

@@ -109,8 +109,8 @@ def to_snow_type(self, schema: List[Any]) -> StructType:
 
         return StructType(fields)
 
+    @staticmethod
     def prepare_connection(
-        self,
         conn: "Connection",
         query_timeout: int = 0,
     ) -> "Connection":
@@ -119,7 +119,9 @@ def prepare_connection(
             conn.outputtypehandler = output_type_handler
         return conn
 
-    def udtf_class_builder(self, fetch_size: int = 1000) -> type:
+    def udtf_class_builder(
+        self, fetch_size: int = 1000, schema: StructType = None
+    ) -> type:
         create_connection = self.create_connection
 
         def oracledb_output_type_handler(cursor, metadata):