diff --git a/sqlspec/adapters/psycopg/_typing.py b/sqlspec/adapters/psycopg/_typing.py index 568f8fe01..9af670d66 100644 --- a/sqlspec/adapters/psycopg/_typing.py +++ b/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]": ... diff --git a/sqlspec/adapters/psycopg/driver.py b/sqlspec/adapters/psycopg/driver.py index 7b41666f6..4f9d6a79b 100644 --- a/sqlspec/adapters/psycopg/driver.py +++ b/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,6 +496,12 @@ 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() @@ -492,7 +509,7 @@ def load_from_arrow( 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,6 +1027,12 @@ 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() @@ -1017,7 +1040,7 @@ async def load_from_arrow( 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 diff --git a/sqlspec/driver/_async.py b/sqlspec/driver/_async.py index 5a6e7a6fa..4f80499c6 100644 --- a/sqlspec/driver/_async.py +++ b/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]": diff --git a/sqlspec/driver/_common.py b/sqlspec/driver/_common.py index 1924f8ed5..70e557d79 100644 --- a/sqlspec/driver/_common.py +++ b/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.""" diff --git a/sqlspec/driver/_sync.py b/sqlspec/driver/_sync.py index bd78a0e89..97ce30afe 100644 --- a/sqlspec/driver/_sync.py +++ b/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]": diff --git a/tests/integration/adapters/postgres/psycopg/test_load_from_records.py b/tests/integration/adapters/postgres/psycopg/test_load_from_records.py new file mode 100644 index 000000000..7d11a76e4 --- /dev/null +++ b/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()