diff --git a/docs/usage/bulk_ingest.rst b/docs/usage/bulk_ingest.rst index dae7970ea..d151238c3 100644 --- a/docs/usage/bulk_ingest.rst +++ b/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 diff --git a/sqlspec/adapters/asyncpg/driver.py b/sqlspec/adapters/asyncpg/driver.py index 1f37b66ae..ac51380e4 100644 --- a/sqlspec/adapters/asyncpg/driver.py +++ b/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, ...]": diff --git a/tests/integration/adapters/postgres/asyncpg/test_load_from_records.py b/tests/integration/adapters/postgres/asyncpg/test_load_from_records.py new file mode 100644 index 000000000..051ac4bb0 --- /dev/null +++ b/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}") diff --git a/tests/unit/adapters/test_asyncpg/test_load_from_records.py b/tests/unit/adapters/test_asyncpg/test_load_from_records.py new file mode 100644 index 000000000..d5e53fcbd --- /dev/null +++ b/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() diff --git a/tests/unit/storage/test_bridge.py b/tests/unit/storage/test_bridge.py index 4e3d0ac39..dcce37e5d 100644 --- a/tests/unit/storage/test_bridge.py +++ b/tests/unit/storage/test_bridge.py @@ -73,10 +73,12 @@ def test_encode_row_payload_jsonl_returns_line_delimited_bytes() -> None: class DummyAsyncpgConnection: def __init__(self) -> None: - self.calls: list[tuple[str, list[tuple[object, ...]], list[str]]] = [] + self.calls: list[tuple[str, str | None, list[tuple[object, ...]], list[str]]] = [] - async def copy_records_to_table(self, table: str, *, records: list[tuple[object, ...]], columns: list[str]) -> None: - self.calls.append((table, records, columns)) + async def copy_records_to_table( + self, table: str, *, records: list[tuple[object, ...]], columns: list[str], schema_name: str | None = None + ) -> None: + self.calls.append((table, schema_name, records, columns)) class DummyPsqlpyConnection: @@ -280,8 +282,9 @@ async def _fake_read(self, *_: object, **__: object) -> tuple[pa.Table, dict[str job = await driver.load_from_storage("public.ingest_target", "file://tmp/part-0.parquet", file_format="parquet") - assert driver.connection.calls[0][0] == "public.ingest_target" - assert driver.connection.calls[0][2] == ["id", "name"] + assert driver.connection.calls[0][0] == "ingest_target" + assert driver.connection.calls[0][1] == "public" + assert driver.connection.calls[0][3] == ["id", "name"] assert job.telemetry["rows_processed"] == arrow_table.num_rows assert job.telemetry["destination"] == "public.ingest_target" diff --git a/tools/scripts/bench_tuning.py b/tools/scripts/bench_tuning.py index 61e49f685..aa2702c48 100644 --- a/tools/scripts/bench_tuning.py +++ b/tools/scripts/bench_tuning.py @@ -48,6 +48,7 @@ class BenchmarkResult: min_seconds: float | None skipped: bool message: str = "" + rows_per_second: float | None = None @dataclass(frozen=True) @@ -61,14 +62,16 @@ class Scenario: runner: Callable[[BenchmarkOptions], list[BenchmarkResult]] -def _summarize(scenario: str, variant: str, samples: list[float]) -> BenchmarkResult: +def _summarize(scenario: str, variant: str, samples: list[float], row_count: int | None = None) -> BenchmarkResult: + median_seconds = statistics.median(samples) return BenchmarkResult( scenario=scenario, variant=variant, iterations=len(samples), - median_seconds=statistics.median(samples), + median_seconds=median_seconds, min_seconds=min(samples), skipped=False, + rows_per_second=row_count / median_seconds if row_count is not None else None, ) @@ -154,6 +157,63 @@ async def run() -> list[BenchmarkResult]: return asyncio.run(run()) +async def _run_asyncpg_record_copy(dsn: str, options: BenchmarkOptions) -> list[BenchmarkResult]: + import asyncpg + + from sqlspec.adapters.asyncpg import AsyncpgDriver + + batch_sizes = (1, 2, 10, 100, 1_000, 5_000) + connection = await asyncpg.connect(dsn=dsn) + driver = AsyncpgDriver(connection, driver_features={"storage_capabilities": {"arrow_import_enabled": True}}) + results: list[BenchmarkResult] = [] + try: + await connection.execute("CREATE TEMP TABLE sqlspec_bench_asyncpg_record_copy (id int PRIMARY KEY, value text)") + for batch_size in batch_sizes: + rows = [(index, f"value-{index}") for index in range(batch_size)] + + async def direct_copy(batch: list[tuple[int, str]] = rows) -> None: + await driver.load_from_records("sqlspec_bench_asyncpg_record_copy", batch, columns=["id", "value"]) + + async def execute_many(batch: list[tuple[int, str]] = rows) -> None: + await connection.executemany( + "INSERT INTO sqlspec_bench_asyncpg_record_copy (id, value) VALUES ($1, $2)", batch + ) + + for variant, operation in (("direct-copy", direct_copy), ("execute-many", execute_many)): + for _ in range(options.warmup): + await connection.execute("TRUNCATE TABLE sqlspec_bench_asyncpg_record_copy") + await operation() + samples: list[float] = [] + for _ in range(options.iterations): + await connection.execute("TRUNCATE TABLE sqlspec_bench_asyncpg_record_copy") + started = time.perf_counter() + await operation() + samples.append(time.perf_counter() - started) + results.append( + _summarize( + "asyncpg_record_copy", f"{variant}:batch_size={batch_size}", samples, row_count=batch_size + ) + ) + finally: + await connection.close() + return results + + +def run_asyncpg_record_copy(options: BenchmarkOptions) -> list[BenchmarkResult]: + """Compare direct-record COPY with parameterized executemany by batch size.""" + dsn = os.getenv("SQLSPEC_BENCH_ASYNCPG_DSN") + if not dsn: + return [ + _skipped( + "asyncpg_record_copy", + "direct-copy:batch_size=1", + options, + "set SQLSPEC_BENCH_ASYNCPG_DSN to run the AsyncPG record COPY crossover benchmark", + ) + ] + return asyncio.run(_run_asyncpg_record_copy(dsn, options)) + + def _run_oracle_variant(stmtcachesize: int, options: BenchmarkOptions) -> BenchmarkResult: import oracledb @@ -211,6 +271,13 @@ def run_arrow_odbc_fetch(options: BenchmarkOptions) -> list[BenchmarkResult]: SCENARIOS: dict[str, Scenario] = { + "asyncpg_record_copy": Scenario( + name="asyncpg_record_copy", + description="Compare direct-record COPY with executemany across batch sizes.", + requires="SQLSPEC_BENCH_ASYNCPG_DSN", + skip_message="Skipped unless SQLSPEC_BENCH_ASYNCPG_DSN points at a PostgreSQL database.", + runner=run_asyncpg_record_copy, + ), "asyncpg_stmt_cache": Scenario( name="asyncpg_stmt_cache", description="Compare asyncpg native statement cache enabled vs disabled on one pooled connection.", @@ -254,13 +321,17 @@ def _print_scenarios() -> None: def _print_results(results: list[BenchmarkResult]) -> None: - click.echo("scenario variant median_s min_s status") + click.echo("scenario variant median_s min_s rows/s status") for result in results: if result.skipped: - click.echo(f"{result.scenario:<20} {result.variant:<26} {'-':>8} {'-':>8} skipped: {result.message}") + click.echo( + f"{result.scenario:<20} {result.variant:<26} {'-':>8} {'-':>8} {'-':>12} skipped: {result.message}" + ) continue + rows_per_second = f"{result.rows_per_second:,.1f}" if result.rows_per_second is not None else "-" click.echo( - f"{result.scenario:<20} {result.variant:<26} {result.median_seconds:>8.6f} {result.min_seconds:>8.6f} ok" + f"{result.scenario:<20} {result.variant:<26} {result.median_seconds:>8.6f} " + f"{result.min_seconds:>8.6f} {rows_per_second:>12} ok" )