perf(#664): bypass Arrow for asyncpg record COPY

docs/usage/bulk_ingest.rst

@@ -18,9 +18,10 @@ Three methods cover the common shapes. They share return type
   a staged artifact (a local path or cloud URI) into a table.
 - ``load_from_records(table, records, *, columns=None, overwrite=False)`` --
   load in-memory rows. ``records`` may be mappings (columns derived from the
-  keys) or positional sequences (``columns`` required). Records are normalized
-  and routed through the adapter's native ``load_from_arrow`` path, so every
-  adapter that supports Arrow ingest supports records too.
+  keys) or positional sequences (``columns`` required). Adapters normally
+  normalize records through their native Arrow ingest path. AsyncPG sends
+  validated record tuples directly to binary ``COPY``; callers that pass an
+  actual Arrow input still use ``load_from_arrow`` unchanged.
 
 .. code-block:: python
 

sqlspec/adapters/asyncpg/driver.py

@@ -2,9 +2,13 @@
 
 import re
 from collections import OrderedDict
+from collections.abc import Mapping
 from io import BytesIO
 from typing import TYPE_CHECKING, Any, Final, cast
 
+from sqlglot import exp, parse_one
+from sqlglot.errors import ParseError
+
 from sqlspec.adapters.asyncpg._typing import AsyncpgCursor, AsyncpgPostgresError, AsyncpgSessionContext
 from sqlspec.adapters.asyncpg.core import (
     PREPARED_STATEMENT_CACHE_SIZE,
@@ -36,7 +40,7 @@
     StackExecutionObserver,
     describe_stack_statement,
 )
-from sqlspec.exceptions import SQLSpecError, StackExecutionError
+from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError, StackExecutionError
 from sqlspec.utils.logging import get_logger
 from sqlspec.utils.text import normalize_identifier, quote_identifier
 from sqlspec.utils.type_guards import has_sqlstate
@@ -333,20 +337,85 @@ async def load_from_arrow(
         """Load Arrow data into a PostgreSQL table via COPY."""
         self._require_capability("arrow_import_enabled")
         arrow_table = self._coerce_arrow_table(source)
+        table_name, schema_name, quoted_target = self._copy_target(table)
         if overwrite:
             try:
-                await self.connection.execute(f"TRUNCATE TABLE {table}")
+                await self.connection.execute(f"TRUNCATE TABLE {quoted_target}")
             except AsyncpgPostgresError as exc:
                 msg = f"Failed to truncate table '{table}': {exc}"
                 raise SQLSpecError(msg) from exc
         columns, records = self._arrow_table_to_rows(arrow_table)
         if records:
-            await self.connection.copy_records_to_table(table, records=records, columns=columns)
+            await self.connection.copy_records_to_table(
+                table_name, records=records, columns=columns, schema_name=schema_name
+            )
         telemetry_payload = self._ingest_telemetry(arrow_table)
         telemetry_payload["destination"] = table
         self._attach_partition_telemetry(telemetry_payload, partitioner)
         return self._storage_job(telemetry_payload, telemetry)
 
+    async def load_from_records(
+        self,
+        table: str,
+        records: "Sequence[Mapping[str, Any]] | Sequence[Sequence[Any]]",
+        *,
+        columns: "list[str] | None" = None,
+        overwrite: bool = False,
+    ) -> "StorageBridgeJob":
+        """Load mapping or positional records directly with binary COPY."""
+        self._require_capability("arrow_import_enabled")
+        materialized = list(records)
+        if not materialized:
+            msg = "load_from_records requires at least one record."
+            raise ImproperConfigurationError(msg)
+
+        first = materialized[0]
+        if isinstance(first, Mapping):
+            resolved_columns = columns if columns is not None else list(first.keys())
+            expected_keys = set(resolved_columns)
+            copy_rows: list[tuple[Any, ...]] = []
+            for record in materialized:
+                if not isinstance(record, Mapping):
+                    msg = "load_from_records mapping records must all be mappings."
+                    raise ImproperConfigurationError(msg)
+                if set(record.keys()) != expected_keys:
+                    msg = "load_from_records mapping records must all share the same keys."
+                    raise ImproperConfigurationError(msg)
+                copy_rows.append(tuple(record[column] for column in resolved_columns))
+        else:
+            if columns is None:
+                msg = "load_from_records requires columns when records are positional sequences."
+                raise ImproperConfigurationError(msg)
+            resolved_columns = columns
+            copy_rows = []
+            for record in materialized:
+                if isinstance(record, Mapping):
+                    msg = "load_from_records positional records must all have the same shape."
+                    raise ImproperConfigurationError(msg)
+                row = tuple(record)
+                if len(row) != len(resolved_columns):
+                    msg = "load_from_records positional records must match the number of columns."
+                    raise ImproperConfigurationError(msg)
+                copy_rows.append(row)
+
+        table_name, schema_name, quoted_target = self._copy_target(table)
+        if overwrite:
+            try:
+                await self.connection.execute(f"TRUNCATE TABLE {quoted_target}")
+            except AsyncpgPostgresError as exc:
+                msg = f"Failed to truncate table '{table}': {exc}"
+                raise SQLSpecError(msg) from exc
+        await self.connection.copy_records_to_table(
+            table_name, records=copy_rows, columns=resolved_columns, schema_name=schema_name
+        )
+        telemetry_payload: StorageTelemetry = {
+            "bytes_processed": 0,
+            "destination": table,
+            "format": "records",
+            "rows_processed": len(copy_rows),
+        }
+        return self._storage_job(telemetry_payload)
+
     async def load_from_storage(
         self,
         table: str,
@@ -391,6 +460,43 @@ def resolve_rowcount(self, cursor: "AsyncpgConnection") -> int:
         """Resolve rowcount from asyncpg status for the direct execution path."""
         return parse_status(cursor)
 
+    @staticmethod
+    def _copy_target(table: str) -> "tuple[str, str | None, str]":
+        try:
+            parsed = parse_one(table, read="postgres", into=exp.Table)
+        except ParseError as exc:
+            msg = "AsyncPG COPY targets must be unqualified or schema-qualified table names."
+            raise ImproperConfigurationError(msg) from exc
+        schema_expression = parsed.args.get("db")
+        schema_identifier = schema_expression if isinstance(schema_expression, exp.Identifier) else None
+        if (
+            parsed.catalog
+            or not isinstance(parsed.this, exp.Identifier)
+            or not parsed.this.name
+            or parsed.this.name == "*"
+            or (schema_expression is not None and schema_identifier is None)
+            or (schema_identifier is not None and not schema_identifier.name)
+        ):
+            msg = "AsyncPG COPY targets must be unqualified or schema-qualified table names."
+            raise ImproperConfigurationError(msg)
+        table_identifier = parsed.this
+        table_name = (
+            table_identifier.name
+            if table_identifier.quoted
+            else normalize_identifier(table_identifier.name, "postgres")
+        )
+        schema_name = None
+        if schema_identifier is not None:
+            schema_name = (
+                schema_identifier.name
+                if schema_identifier.quoted
+                else normalize_identifier(schema_identifier.name, "postgres")
+            )
+        quoted_target = quote_identifier(table_name)
+        if schema_name is not None:
+            quoted_target = f"{quote_identifier(schema_name)}.{quoted_target}"
+        return table_name, schema_name, quoted_target
+
     async def _execute_stack_native(
         self, stack: "StatementStack", *, continue_on_error: bool
     ) -> "tuple[StackResult, ...]":

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

@@ -0,0 +1,61 @@
+"""AsyncPG direct record COPY integration tests."""
+
+from typing import Any
+from uuid import uuid4
+
+import pyarrow as pa
+import pytest
+
+from sqlspec.adapters.asyncpg import AsyncpgDriver
+
+pytestmark = pytest.mark.xdist_group("postgres")
+
+
+async def test_load_from_records_preserves_native_postgres_values(asyncpg_async_driver: AsyncpgDriver) -> None:
+    table = '"sqlspec.bulk"."Record.Target"'
+    await asyncpg_async_driver.execute_script(
+        'CREATE SCHEMA IF NOT EXISTS "sqlspec.bulk"; '
+        'DROP TABLE IF EXISTS "sqlspec.bulk"."Record.Target"; '
+        'CREATE TABLE "sqlspec.bulk"."Record.Target" '
+        "(id UUID PRIMARY KEY, payload JSONB, labels TEXT[], note TEXT)"
+    )
+    first_id = uuid4()
+    second_id = uuid4()
+    try:
+        first = {"kind": "nested", "metadata": {"enabled": True}, "items": [1, {"name": "two"}]}
+        second = {"kind": "different", "other": [False, None]}
+        mapping_records: list[dict[str, Any]] = [
+            {"id": first_id, "payload": first, "labels": ["a", "b"], "note": None},
+            {"note": "second", "labels": [], "payload": second, "id": second_id},
+        ]
+        job = await asyncpg_async_driver.load_from_records(table, mapping_records)
+
+        rows = await asyncpg_async_driver.select(
+            'SELECT id, payload, labels, note FROM "sqlspec.bulk"."Record.Target" ORDER BY note NULLS FIRST'
+        )
+        assert [row["id"] for row in rows] == [first_id, second_id]
+        assert rows[0]["payload"] == first
+        assert rows[0]["labels"] == ["a", "b"]
+        assert rows[1]["payload"] == second
+        assert job.telemetry["rows_processed"] == 2
+
+        await asyncpg_async_driver.load_from_records(
+            table,
+            [(first_id, {"replacement": True}, ["updated"], "overwritten")],
+            columns=["id", "payload", "labels", "note"],
+            overwrite=True,
+        )
+        assert await asyncpg_async_driver.select_value('SELECT count(*) FROM "sqlspec.bulk"."Record.Target"') == 1
+    finally:
+        await asyncpg_async_driver.execute_script('DROP SCHEMA IF EXISTS "sqlspec.bulk" CASCADE')
+
+
+async def test_load_from_arrow_keeps_arrow_copy_path(asyncpg_async_driver: AsyncpgDriver) -> None:
+    table = "asyncpg_arrow_path"
+    await asyncpg_async_driver.execute_script(f"DROP TABLE IF EXISTS {table}; CREATE TABLE {table} (id INTEGER)")
+    try:
+        job = await asyncpg_async_driver.load_from_arrow(table, pa.table({"id": [1, 2]}), overwrite=True)
+        assert job.telemetry["format"] == "arrow"
+        assert await asyncpg_async_driver.select_value(f"SELECT count(*) FROM {table}") == 2
+    finally:
+        await asyncpg_async_driver.execute_script(f"DROP TABLE IF EXISTS {table}")

tests/unit/adapters/test_asyncpg/test_load_from_records.py

@@ -0,0 +1,168 @@
+"""AsyncPG direct record COPY tests."""
+
+from typing import Any
+from unittest.mock import AsyncMock
+
+import pytest
+from asyncpg import PostgresError
+
+from sqlspec.adapters.asyncpg import AsyncpgDriver
+from sqlspec.adapters.cockroach_asyncpg import CockroachAsyncpgDriver
+from sqlspec.exceptions import ImproperConfigurationError, SQLSpecError, StorageCapabilityError
+
+DRIVER_FEATURES = {"storage_capabilities": {"arrow_import_enabled": True}}
+
+
+@pytest.fixture
+def connection() -> AsyncMock:
+    return AsyncMock()
+
+
+@pytest.mark.parametrize("driver_type", [AsyncpgDriver, CockroachAsyncpgDriver])
+async def test_load_from_records_copies_mapping_rows_directly(
+    driver_type: type[AsyncpgDriver], connection: AsyncMock
+) -> None:
+    driver = driver_type(connection, driver_features=DRIVER_FEATURES)  # type: ignore[arg-type]
+    payload = {"nested": {"enabled": True}, "items": [1, 2]}
+    mapping_records: list[dict[str, Any]] = [{"id": 1, "payload": payload}, {"payload": None, "id": 2}]
+
+    job = await driver.load_from_records('"audit.schema"."Event.Log"', mapping_records, columns=["payload", "id"])
+
+    connection.copy_records_to_table.assert_awaited_once_with(
+        "Event.Log", records=[(payload, 1), (None, 2)], columns=["payload", "id"], schema_name="audit.schema"
+    )
+    assert job.status == "completed"
+    assert job.telemetry == {
+        "bytes_processed": 0,
+        "destination": '"audit.schema"."Event.Log"',
+        "format": "records",
+        "rows_processed": 2,
+    }
+
+
+async def test_load_from_records_preserves_first_mapping_key_order(connection: AsyncMock) -> None:
+    driver = AsyncpgDriver(connection, driver_features=DRIVER_FEATURES)  # type: ignore[arg-type]
+
+    await driver.load_from_records("events", [{"b": 2, "a": 1}, {"a": 3, "b": 4}])
+
+    connection.copy_records_to_table.assert_awaited_once_with(
+        "events", records=[(2, 1), (4, 3)], columns=["b", "a"], schema_name=None
+    )
+
+
+async def test_load_from_records_truncates_quoted_target_before_copy(connection: AsyncMock) -> None:
+    driver = AsyncpgDriver(connection, driver_features=DRIVER_FEATURES)  # type: ignore[arg-type]
+
+    await driver.load_from_records('"audit.schema"."Event.Log"', [(1,)], columns=["id"], overwrite=True)
+
+    assert connection.method_calls[0].args == ('TRUNCATE TABLE "audit.schema"."Event.Log"',)
+    assert connection.method_calls[1].args[0] == "Event.Log"
+
+
+async def test_load_from_records_reports_truncate_failure(connection: AsyncMock) -> None:
+    driver = AsyncpgDriver(connection, driver_features=DRIVER_FEATURES)  # type: ignore[arg-type]
+    connection.execute.side_effect = PostgresError("truncate failed")
+
+    with pytest.raises(SQLSpecError, match="Failed to truncate table 'events': truncate failed"):
+        await driver.load_from_records("events", [(1,)], columns=["id"], overwrite=True)
+
+    connection.copy_records_to_table.assert_not_awaited()
+
+
+async def test_load_from_records_checks_capability_before_input_or_database(connection: AsyncMock) -> None:
+    driver = AsyncpgDriver(  # type: ignore[arg-type]
+        connection, driver_features={"storage_capabilities": {"arrow_import_enabled": False}}
+    )
+
+    with pytest.raises(StorageCapabilityError, match="Arrow import"):
+        await driver.load_from_records("", [])
+
+    connection.execute.assert_not_awaited()
+    connection.copy_records_to_table.assert_not_awaited()
+
+
+@pytest.mark.parametrize(
+    "target",
+    [
+        "",
+        " ",
+        "*",
+        '"*"',
+        '""',
+        'public.""',
+        "events.*",
+        "events()",
+        "events().target",
+        "catalog.public.events",
+        "events alias",
+        "events,other",
+    ],
+)
+async def test_load_from_records_rejects_invalid_copy_targets_before_database_calls(
+    connection: AsyncMock, target: str
+) -> None:
+    driver = AsyncpgDriver(connection, driver_features=DRIVER_FEATURES)  # type: ignore[arg-type]
+
+    with pytest.raises(ImproperConfigurationError, match="COPY targets"):
+        await driver.load_from_records(target, [(1,)], columns=["id"])
+
+    connection.execute.assert_not_awaited()
+    connection.copy_records_to_table.assert_not_awaited()
+
+
+@pytest.mark.parametrize(
+    "target,expected_table,expected_schema",
+    [
+        ("Public.Events", "events", "public"),
+        ('"Public"."Events"', "Events", "Public"),
+        ('"audit.schema"."Event.Log"', "Event.Log", "audit.schema"),
+    ],
+)
+async def test_load_from_records_normalizes_copy_target_identifiers(
+    connection: AsyncMock, target: str, expected_table: str, expected_schema: str
+) -> None:
+    driver = AsyncpgDriver(connection, driver_features=DRIVER_FEATURES)  # type: ignore[arg-type]
+
+    await driver.load_from_records(target, [(1,)], columns=["id"])
+
+    call = connection.copy_records_to_table.await_args
+    assert call.args[0] == expected_table
+    assert call.kwargs["schema_name"] == expected_schema
+
+
+@pytest.mark.parametrize(
+    "records,columns,match",
+    [
+        ([], None, "at least one record"),
+        ([(1,)], None, "requires columns"),
+        ([(1, 2)], ["id"], "number of columns"),
+        ([{"id": 1}, (2,)], None, "all be mappings"),
+        ([{"id": 1}, {"id": 2, "extra": 3}], None, "share the same keys"),
+        ([(1,), {"id": 2}], ["id"], "same shape"),
+    ],
+)
+async def test_load_from_records_validates_before_database_calls(
+    connection: AsyncMock, records: list[object], columns: list[str] | None, match: str
+) -> None:
+    driver = AsyncpgDriver(connection, driver_features=DRIVER_FEATURES)  # type: ignore[arg-type]
+
+    with pytest.raises(ImproperConfigurationError, match=match):
+        await driver.load_from_records("events", records, columns=columns)  # type: ignore[arg-type]
+
+    connection.execute.assert_not_awaited()
+    connection.copy_records_to_table.assert_not_awaited()
+
+
+async def test_load_from_records_does_not_use_arrow_normalization(
+    connection: AsyncMock, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    driver = AsyncpgDriver(connection, driver_features=DRIVER_FEATURES)  # type: ignore[arg-type]
+
+    def fail_if_called(*args: object) -> None:
+        raise AssertionError("Arrow normalization was called")
+
+    monkeypatch.setattr(AsyncpgDriver, "_records_to_arrow_table", fail_if_called)
+
+    await driver.load_from_records("events", [(1,)], columns=["id"])
+
+    connection.copy_records_to_table.assert_awaited_once()