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
 
 - Added support for reading XML files with namespaces using `rowTag` and `stripNamespaces` options.

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

@@ -68,7 +68,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,40 @@
+#
+# Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved.
+#
+from typing import List
+
+from snowflake.snowpark._internal.data_source.dbms_dialects import BaseDialect
+from snowflake.snowpark._internal.data_source.drivers.psycopg2_driver import (
+    Psycopg2TypeCode,
+)
+from snowflake.snowpark.types import StructType
+
+
+class PostgresDialect(BaseDialect):
+    @staticmethod
+    def generate_select_query(
+        table_or_query: str, schema: StructType, raw_schema: List[tuple]
+    ) -> str:
+        cols = []
+        for _field, raw_field in zip(schema.fields, raw_schema):
+            # databricks-sql-connector returns list of tuples for MapType
+            # here we push down to-dict conversion to Databricks
+            type_code = raw_field[1]
+            if type_code in (
+                Psycopg2TypeCode.JSONB.value,
+                Psycopg2TypeCode.JSON.value,
+            ):
+                cols.append(f"""TO_JSON("{raw_field[0]}")::TEXT AS {raw_field[0]}""")
+            elif type_code == Psycopg2TypeCode.CASHOID.value:
+                cols.append(
+                    f"""CASE WHEN "{raw_field[0]}" IS NULL THEN NULL ELSE FORMAT('"%s"', "{raw_field[0]}"::TEXT) END AS {raw_field[0]}"""
+                )
+            elif type_code == Psycopg2TypeCode.BYTEAOID.value:
+                cols.append(f"""ENCODE("{raw_field[0]}", 'HEX') AS {raw_field[0]}""")
+            elif type_code == Psycopg2TypeCode.TIMETZOID.value:
+                cols.append(f""""{raw_field[0]}"::TIME AS {raw_field[0]}""")
+            elif type_code == Psycopg2TypeCode.INTERVALOID.value:
+                cols.append(f""""{raw_field[0]}"::TEXT AS {raw_field[0]}""")
+            else:
+                cols.append(f'"{raw_field[0]}"')
+        return f"""SELECT {", ".join(cols)} FROM {table_or_query}"""

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":
@@ -75,13 +75,14 @@ def udtf_ingestion(
         fetch_size: int = 1000,
         imports: Optional[List[str]] = None,
         packages: Optional[List[str]] = None,
+        _emit_ast: bool = True,
     ) -> "snowflake.snowpark.DataFrame":
         from snowflake.snowpark._internal.data_source.utils import UDTF_PACKAGE_MAP
 
         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(
                 [
@@ -97,14 +98,12 @@ def udtf_ingestion(
         call_udtf_sql = f"""
             select * from {partition_table}, table({udtf_name}({PARTITION_TABLE_COLUMN_NAME}))
             """
-        res = session.sql(call_udtf_sql)
-        cols = [
-            res[field.name].cast(field.datatype).alias(field.name)
-            for field in schema.fields
-        ]
-        return res.select(cols)
+        res = session.sql(call_udtf_sql, _emit_ast=_emit_ast)
+        return self.to_result_snowpark_df_udtf(res, schema, _emit_ast=_emit_ast)
 
-    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:
@@ -169,3 +168,15 @@ def to_result_snowpark_df(
         session: "Session", table_name, schema, _emit_ast: bool = True
     ) -> "DataFrame":
         return session.table(table_name, _emit_ast=_emit_ast)
+
+    @staticmethod
+    def to_result_snowpark_df_udtf(
+        res_df: "DataFrame",
+        schema: StructType,
+        _emit_ast: bool = True,
+    ):
+        cols = [
+            res_df[field.name].cast(field.datatype).alias(field.name)
+            for field in schema.fields
+        ]
+        return res_df.select(cols, _emit_ast=_emit_ast)

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):