SNOW-1955847: DBAPI Postgres Support (#3351)

Author: sfc-gh-aling

.github/workflows/parameters/parameters_dbapi.py.gpg

@@ -210,6 +210,8 @@ jobs:
           TOX_PARALLEL_NO_SPINNER: 1
         shell: bash
       - name: Run data source tests
+        # psycopg2 is not supported on macos 3.9
+        if: ${{ !(matrix.os == 'macos-latest' && matrix.python-version == '3.9') }}
         run: python -m tox -e datasource
         env:
           PYTHON_VERSION: ${{ matrix.python-version }}

.github/workflows/precommit.yml

@@ -6,8 +6,9 @@
 
 #### New Features
 
-- Added support for ingestion with Snowflake UDTF to databricks in `DataFrameReader.dbapi` (PrPr).
-- Added support for Mysql in `DataFrameWriter.dbapi` (PrPr).
+- Added support for MySQL in `DataFrameWriter.dbapi` (PrPr) for both Parquet and UDTF-based ingestion.
+- Added support for PostgreSQL in `DataFrameReader.dbapi` (PrPr) for both Parquet and UDTF-based ingestion.
+- Added support for Databricks in `DataFrameWriter.dbapi` (PrPr) for UDTF-based ingestion.
 
 #### Bug Fixes
 

CHANGELOG.md

@@ -7,6 +7,7 @@
     "SqlServerDialect",
     "OracledbDialect",
     "DatabricksDialect",
+    "PostgresDialect",
     "MysqlDialect",
 ]
 
@@ -25,6 +26,9 @@
 from snowflake.snowpark._internal.data_source.dbms_dialects.databricks_dialect import (
     DatabricksDialect,
 )
+from snowflake.snowpark._internal.data_source.dbms_dialects.postgresql_dialect import (
+    PostgresDialect,
+)
 from snowflake.snowpark._internal.data_source.dbms_dialects.mysql_dialect import (
     MysqlDialect,
 )

scripts/parameters.py.gpg

@@ -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], is_query: bool
+    ) -> 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/dbms_dialects/__init__.py

@@ -8,6 +8,7 @@
     "SqliteDriver",
     "PyodbcDriver",
     "DatabricksDriver",
+    "Psycopg2Driver",
     "PymysqlDriver",
 ]
 
@@ -20,4 +21,7 @@
 from snowflake.snowpark._internal.data_source.drivers.databricks_driver import (
     DatabricksDriver,
 )
+from snowflake.snowpark._internal.data_source.drivers.psycopg2_driver import (
+    Psycopg2Driver,
+)
 from snowflake.snowpark._internal.data_source.drivers.pymsql_driver import PymysqlDriver

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

@@ -51,8 +51,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":
@@ -100,7 +100,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(
                 [
@@ -119,7 +119,9 @@ def udtf_ingestion(
         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:

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

@@ -61,7 +61,9 @@ def to_snow_type(self, schema: List[Any]) -> StructType:
             all_columns.append(StructField(column_name, data_type, True))
         return StructType(all_columns)
 
-    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/base_driver.py

@@ -103,8 +103,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":
@@ -113,7 +113,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):