fix(#663): adapt psycopg JSON records for Arrow COPY

sqlspec/adapters/psycopg/_typing.py

@@ -105,6 +105,14 @@ def prepare_statement(
         kwargs: "dict[str, Any] | None" = None,
     ) -> "SQL": ...
 
+    def prepare_driver_parameters(
+        self,
+        parameters: Any,
+        statement_config: "StatementConfig",
+        is_many: bool = False,
+        prepared_statement: Any | None = None,
+    ) -> Any: ...
+
     def _compiled_sql(self, statement: "SQL", statement_config: "StatementConfig") -> "tuple[str, Any]": ...
 
 

sqlspec/adapters/psycopg/driver.py

@@ -65,6 +65,8 @@
 from sqlspec.utils.type_guards import is_readable, resolve_row_format
 
 if TYPE_CHECKING:
+    from collections import abc
+
     from sqlspec.adapters.psycopg._typing import PsycopgPipelineDriver
     from sqlspec.core import ArrowResult
     from sqlspec.driver import ExecutionResult
@@ -107,6 +109,15 @@ class PsycopgPipelineMixin:
 
     __slots__ = ()
 
+    def _prepare_records_for_arrow(
+        self, records: "abc.Sequence[abc.Mapping[str, Any]] | abc.Sequence[abc.Sequence[Any]]"
+    ) -> "abc.Sequence[abc.Mapping[str, Any]] | abc.Sequence[abc.Sequence[Any]]":
+        driver = cast("PsycopgPipelineDriver", self)
+        return cast(
+            "abc.Sequence[abc.Mapping[str, Any]] | abc.Sequence[abc.Sequence[Any]]",
+            driver.prepare_driver_parameters(records, driver.statement_config, is_many=True),
+        )
+
     def _prepare_pipeline_operations(self, stack: "StatementStack") -> "list[PreparedStackOperation] | None":
         prepared: list[PreparedStackOperation] = []
         for index, operation in enumerate(stack.operations):
@@ -485,14 +496,20 @@ def load_from_arrow(
             if exc_handler.pending_exception is not None:
                 raise exc_handler.pending_exception from None
         columns, records = self._arrow_table_to_rows(arrow_table)
+        prepared_records = cast(
+            "list[Any]",
+            self.prepare_driver_parameters(records, self.statement_config, is_many=True)
+            if records and self._arrow_rows_need_preparation(arrow_table)
+            else records,
+        )
         if records:
             copy_sql = build_copy_from_command(table, columns)
             exc_handler = self.handle_database_exceptions()
             with ExitStack() as stack:
                 stack.enter_context(exc_handler)
                 cursor = stack.enter_context(self.with_cursor(self.connection))
                 copy_ctx = stack.enter_context(cursor.copy(copy_sql))
-                for record in records:
+                for record in prepared_records:
                     copy_ctx.write_row(record)
             if exc_handler.pending_exception is not None:
                 raise exc_handler.pending_exception from None
@@ -1010,14 +1027,20 @@ async def load_from_arrow(
             if exc_handler.pending_exception is not None:
                 raise exc_handler.pending_exception from None
         columns, records = self._arrow_table_to_rows(arrow_table)
+        prepared_records = cast(
+            "list[Any]",
+            self.prepare_driver_parameters(records, self.statement_config, is_many=True)
+            if records and self._arrow_rows_need_preparation(arrow_table)
+            else records,
+        )
         if records:
             copy_sql = build_copy_from_command(table, columns)
             exc_handler = self.handle_database_exceptions()
             async with AsyncExitStack() as stack:
                 await stack.enter_async_context(exc_handler)
                 cursor = await stack.enter_async_context(self.with_cursor(self.connection))
                 copy_ctx = await stack.enter_async_context(cursor.copy(copy_sql))
-                for record in records:
+                for record in prepared_records:
                     await copy_ctx.write_row(record)
             if exc_handler.pending_exception is not None:
                 raise exc_handler.pending_exception from None

sqlspec/driver/_async.py

@@ -1689,7 +1689,8 @@ async def load_from_records(
         Returns:
             StorageBridgeJob with execution telemetry.
         """
-        arrow_table = self._records_to_arrow_table(records, columns)
+        prepared_records = self._prepare_records_for_arrow(records)
+        arrow_table = self._records_to_arrow_table(prepared_records, columns)
         return await self.load_from_arrow(table, arrow_table, overwrite=overwrite)
 
     def stage_artifact(self, request: "dict[str, Any]") -> "dict[str, Any]":

sqlspec/driver/_common.py

@@ -1910,6 +1910,12 @@ def _arrow_rows_need_preparation(table: "ArrowTable") -> bool:
         """Return whether Arrow rows may contain nested values needing preparation."""
         return arrow_table_needs_parameter_preparation(table)
 
+    def _prepare_records_for_arrow(
+        self, records: "abc.Sequence[abc.Mapping[str, Any]] | abc.Sequence[abc.Sequence[Any]]"
+    ) -> "abc.Sequence[abc.Mapping[str, Any]] | abc.Sequence[abc.Sequence[Any]]":
+        """Prepare records before Arrow schema inference."""
+        return records
+
     @staticmethod
     def _ingest_telemetry(table: "ArrowTable", *, format_label: str = "arrow") -> "StorageTelemetry":
         """Build telemetry dict from Arrow table statistics."""

sqlspec/driver/_sync.py

@@ -1628,7 +1628,8 @@ def load_from_records(
         Returns:
             StorageBridgeJob with execution telemetry.
         """
-        arrow_table = self._records_to_arrow_table(records, columns)
+        prepared_records = self._prepare_records_for_arrow(records)
+        arrow_table = self._records_to_arrow_table(prepared_records, columns)
         return self.load_from_arrow(table, arrow_table, overwrite=overwrite)
 
     def stage_artifact(self, request: "dict[str, Any]") -> "dict[str, Any]":

tests/integration/adapters/postgres/psycopg/test_load_from_records.py

@@ -0,0 +1,176 @@
+"""Integration tests for psycopg record loading."""
+
+from typing import TYPE_CHECKING
+
+import pyarrow as pa
+import pytest
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.psycopg import PsycopgAsyncConfig, PsycopgSyncConfig
+
+pytestmark = pytest.mark.xdist_group("postgres")
+
+_SYNC_TABLE = "test_psycopg_sync_load_from_records"
+_ASYNC_TABLE = "test_psycopg_async_load_from_records"
+_SYNC_ARROW_TABLE = "test_psycopg_sync_load_from_arrow"
+_ASYNC_ARROW_TABLE = "test_psycopg_async_load_from_arrow"
+_RECORDS = [
+    {
+        "id": 1,
+        "payload_json": {"name": "alpha", "items": [1, 2]},
+        "payload_jsonb": {"status": "ready", "details": {"attempt": 1}},
+        "metadata_jsonb": None,
+        "tags": ["python", "sql"],
+    },
+    {
+        "id": 2,
+        "payload_json": {"name": "beta", "items": []},
+        "payload_jsonb": {"status": "done", "details": {}},
+        "metadata_jsonb": {"worker": "two"},
+        "tags": [],
+    },
+]
+_ARROW_SCHEMA = pa.schema([
+    pa.field("id", pa.int64(), nullable=False),
+    pa.field("payload_json", pa.struct([pa.field("name", pa.string()), pa.field("items", pa.list_(pa.int64()))])),
+    pa.field(
+        "payload_jsonb",
+        pa.struct([pa.field("status", pa.string()), pa.field("details", pa.struct([pa.field("attempt", pa.int64())]))]),
+    ),
+    pa.field("metadata_jsonb", pa.struct([pa.field("worker", pa.string())])),
+    pa.field("tags", pa.list_(pa.string()), nullable=False),
+])
+_ARROW_RECORDS = [
+    {
+        "id": 1,
+        "payload_json": {"name": "alpha", "items": [1, 2]},
+        "payload_jsonb": {"status": "ready", "details": {"attempt": 1}},
+        "metadata_jsonb": None,
+        "tags": ["python", "sql"],
+    },
+    {
+        "id": 2,
+        "payload_json": {"name": "beta", "items": []},
+        "payload_jsonb": {"status": "done", "details": {"attempt": None}},
+        "metadata_jsonb": {"worker": "two"},
+        "tags": [],
+    },
+]
+
+
+def test_psycopg_sync_load_from_records_prepares_json_mappings(psycopg_sync_config: "PsycopgSyncConfig") -> None:
+    """Psycopg COPY should serialize mapping records for JSON columns."""
+    with psycopg_sync_config.provide_session() as session:
+        session.execute_script(f"DROP TABLE IF EXISTS {_SYNC_TABLE}")
+        session.execute_script(
+            f"""
+            CREATE TABLE {_SYNC_TABLE} (
+                id INTEGER PRIMARY KEY,
+                payload_json JSON NOT NULL,
+                payload_jsonb JSONB NOT NULL,
+                metadata_jsonb JSONB,
+                tags TEXT[] NOT NULL
+            )
+            """
+        )
+        session.commit()
+        try:
+            job = session.load_from_records(_SYNC_TABLE, _RECORDS)
+            rows = session.execute(f"SELECT * FROM {_SYNC_TABLE} ORDER BY id").get_data()
+
+            assert rows == _RECORDS
+            assert job.telemetry["rows_processed"] == 2
+        finally:
+            session.rollback()
+            session.execute_script(f"DROP TABLE IF EXISTS {_SYNC_TABLE}")
+            session.commit()
+
+
+async def test_psycopg_async_load_from_records_prepares_json_mappings(
+    psycopg_async_config: "PsycopgAsyncConfig",
+) -> None:
+    """Async psycopg COPY should serialize mapping records for JSON columns."""
+    async with psycopg_async_config.provide_session() as session:
+        await session.execute_script(f"DROP TABLE IF EXISTS {_ASYNC_TABLE}")
+        await session.execute_script(
+            f"""
+            CREATE TABLE {_ASYNC_TABLE} (
+                id INTEGER PRIMARY KEY,
+                payload_json JSON NOT NULL,
+                payload_jsonb JSONB NOT NULL,
+                metadata_jsonb JSONB,
+                tags TEXT[] NOT NULL
+            )
+            """
+        )
+        await session.commit()
+        try:
+            job = await session.load_from_records(_ASYNC_TABLE, _RECORDS)
+            rows = (await session.execute(f"SELECT * FROM {_ASYNC_TABLE} ORDER BY id")).get_data()
+
+            assert rows == _RECORDS
+            assert job.telemetry["rows_processed"] == 2
+        finally:
+            await session.rollback()
+            await session.execute_script(f"DROP TABLE IF EXISTS {_ASYNC_TABLE}")
+            await session.commit()
+
+
+def test_psycopg_sync_load_from_arrow_prepares_struct_mappings(psycopg_sync_config: "PsycopgSyncConfig") -> None:
+    """Psycopg COPY should adapt Arrow structs without changing primitive arrays."""
+    arrow_table = pa.Table.from_pylist(_ARROW_RECORDS, schema=_ARROW_SCHEMA)
+    with psycopg_sync_config.provide_session() as session:
+        session.execute_script(f"DROP TABLE IF EXISTS {_SYNC_ARROW_TABLE}")
+        session.execute_script(
+            f"""
+            CREATE TABLE {_SYNC_ARROW_TABLE} (
+                id INTEGER PRIMARY KEY,
+                payload_json JSON NOT NULL,
+                payload_jsonb JSONB NOT NULL,
+                metadata_jsonb JSONB,
+                tags TEXT[] NOT NULL
+            )
+            """
+        )
+        session.commit()
+        try:
+            job = session.load_from_arrow(_SYNC_ARROW_TABLE, arrow_table)
+            rows = session.execute(f"SELECT * FROM {_SYNC_ARROW_TABLE} ORDER BY id").get_data()
+
+            assert rows == _ARROW_RECORDS
+            assert job.telemetry["rows_processed"] == 2
+        finally:
+            session.rollback()
+            session.execute_script(f"DROP TABLE IF EXISTS {_SYNC_ARROW_TABLE}")
+            session.commit()
+
+
+async def test_psycopg_async_load_from_arrow_prepares_struct_mappings(
+    psycopg_async_config: "PsycopgAsyncConfig",
+) -> None:
+    """Async psycopg COPY should adapt Arrow structs without changing primitive arrays."""
+    arrow_table = pa.Table.from_pylist(_ARROW_RECORDS, schema=_ARROW_SCHEMA)
+    async with psycopg_async_config.provide_session() as session:
+        await session.execute_script(f"DROP TABLE IF EXISTS {_ASYNC_ARROW_TABLE}")
+        await session.execute_script(
+            f"""
+            CREATE TABLE {_ASYNC_ARROW_TABLE} (
+                id INTEGER PRIMARY KEY,
+                payload_json JSON NOT NULL,
+                payload_jsonb JSONB NOT NULL,
+                metadata_jsonb JSONB,
+                tags TEXT[] NOT NULL
+            )
+            """
+        )
+        await session.commit()
+        try:
+            job = await session.load_from_arrow(_ASYNC_ARROW_TABLE, arrow_table)
+            rows = (await session.execute(f"SELECT * FROM {_ASYNC_ARROW_TABLE} ORDER BY id")).get_data()
+
+            assert rows == _ARROW_RECORDS
+            assert job.telemetry["rows_processed"] == 2
+        finally:
+            await session.rollback()
+            await session.execute_script(f"DROP TABLE IF EXISTS {_ASYNC_ARROW_TABLE}")
+            await session.commit()