Add PostgreSQL bulk load via COPY FROM STDIN

.github/workflows/integration-tests-postgres.yml

@@ -10,6 +10,15 @@ on:
 jobs:
   integration-tests:
     runs-on: ubuntu-24.04
+    strategy:
+      matrix:
+        include:
+          - driver: psycopg2
+            package: psycopg2
+            db_string: postgresql+psycopg2://postgres:postgres@localhost:5432/postgres
+          - driver: psycopg
+            package: "psycopg[binary]"
+            db_string: postgresql+psycopg://postgres:postgres@localhost:5432/postgres
     services:
       postgres:
         image: postgres
@@ -25,23 +34,23 @@ jobs:
           --health-interval 10s
           --health-timeout 5s
           --health-retries 5
-          
+
     steps:
       - name: Check out repository code
         uses: actions/checkout@v4
-      
+
       - name: Set up Python 3.12
         uses: actions/setup-python@v5
         with:
           python-version: 3.12
-          
+
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install flake8 .[tests] psycopg2
-    
+          python -m pip install flake8 .[tests] "${{ matrix.package }}"
+
       - name: Test with pytest
         run: |
           pytest tests -x
         env:
-          DB_STRING: postgresql+psycopg2://postgres:postgres@localhost:5432/postgres
+          DB_STRING: ${{ matrix.db_string }}

src/xml2db/dialect/postgresql.py

@@ -1,3 +1,7 @@
+import csv
+import io
+from typing import Any
+
 from .base import DatabaseDialect
 
 
@@ -11,3 +15,85 @@ class PostgreSQLDialect(DatabaseDialect):
     """
 
     MAX_IDENTIFIER_LENGTH: int = 63
+
+    def bulk_insert(self, conn: Any, table: Any, records: list) -> None:
+        """Bulk-insert records via PostgreSQL's ``COPY FROM STDIN``.
+
+        Builds an in-memory CSV payload and streams it to the server using
+        the driver's native COPY protocol.  Supported drivers:
+
+        - **psycopg2** — uses ``cursor.copy_expert()``.
+        - **psycopg** (psycopg3) — uses ``cursor.copy()``.
+
+        Falls back to the base-class parameterised executemany for any other
+        driver.
+
+        Args:
+            conn: A SQLAlchemy ``Connection`` already within a transaction.
+            table: The SQLAlchemy ``Table`` object to insert into.
+            records: A list of dicts mapping column keys to Python values.
+        """
+        if not records:
+            return
+
+        driver = conn.dialect.driver
+        if driver not in ("psycopg2", "psycopg"):
+            super().bulk_insert(conn, table, records)
+            return
+
+        col_by_key = {col.key: col for col in table.columns}
+        col_keys = [k for k in records[0] if k in col_by_key]
+
+        # Python-side scalar defaults absent from records (e.g. default=False).
+        # executemany applies these automatically; our COPY path must do it manually.
+        extra_defaults: dict = {}
+        for col in table.columns:
+            if col.key not in records[0] and col.key in col_by_key:
+                d = col.default
+                if d is not None and d.is_scalar:
+                    extra_defaults[col.key] = d.arg
+
+        all_col_keys = col_keys + list(extra_defaults.keys())
+
+        buf = io.StringIO()
+        writer = csv.writer(buf)
+        writer.writerow(all_col_keys)
+        for record in records:
+            row = []
+            for key in all_col_keys:
+                v = record.get(key) if key in col_keys else extra_defaults[key]
+                if v is None:
+                    row.append("")
+                elif isinstance(v, bytes):
+                    # PostgreSQL bytea hex format: \xDEADBEEF
+                    row.append("\\x" + v.hex())
+                elif isinstance(v, bool):
+                    # bool must precede the general str() path: bool subclasses int,
+                    # so str(True) → '1' which PostgreSQL COPY rejects for boolean.
+                    row.append("true" if v else "false")
+                else:
+                    # str() on datetime → "YYYY-MM-DD HH:MM:SS[.f][±HH:MM]",
+                    # which PostgreSQL's text input parser accepts.
+                    row.append(str(v))
+            writer.writerow(row)
+        buf.seek(0)
+
+        col_names = ", ".join(f'"{col_by_key[k].name}"' for k in all_col_keys)
+        full_name = (
+            f'"{table.schema}"."{table.name}"'
+            if table.schema
+            else f'"{table.name}"'
+        )
+        copy_sql = (
+            f"COPY {full_name} ({col_names}) FROM STDIN "
+            f"WITH (FORMAT CSV, HEADER, NULL '')"
+        )
+
+        raw_conn = conn.connection.dbapi_connection
+        if driver == "psycopg2":
+            cur = raw_conn.cursor()
+            cur.copy_expert(copy_sql, buf)
+        else:  # psycopg3
+            cur = raw_conn.cursor()
+            with cur.copy(copy_sql) as copy:
+                copy.write(buf.read().encode("utf-8"))

tests/test_bulk_insert_postgres.py

@@ -0,0 +1,257 @@
+"""Integration tests for PostgreSQLDialect.bulk_insert with psycopg2 and psycopg3."""
+import datetime
+import os
+
+import pytest
+from sqlalchemy import (
+    BigInteger,
+    Boolean,
+    Column,
+    DateTime,
+    Double,
+    Integer,
+    LargeBinary,
+    MetaData,
+    SmallInteger,
+    String,
+    Table,
+    create_engine,
+    select,
+)
+from sqlalchemy.engine import make_url
+
+from xml2db.dialect.postgresql import PostgreSQLDialect
+
+
+def _make_pg_engine(driver: str):
+    """Return a PostgreSQL engine using the given driver, or skip."""
+    db_string = os.getenv("DB_STRING", "")
+    if not db_string:
+        pytest.skip("DB_STRING not set")
+    url = make_url(db_string)
+    if url.get_dialect().name != "postgresql":
+        pytest.skip("DB_STRING is not a PostgreSQL connection")
+    if driver == "psycopg2":
+        pytest.importorskip("psycopg2")
+    else:
+        pytest.importorskip("psycopg")
+    url = url.set(drivername=f"postgresql+{driver}")
+    return create_engine(url)
+
+
+@pytest.fixture(params=["psycopg2", "psycopg"])
+def pg_engine(request):
+    engine = _make_pg_engine(request.param)
+    yield engine
+    engine.dispose()
+
+
+def _make_table(engine, name, *extra_cols):
+    meta = MetaData()
+    table = Table(
+        name,
+        meta,
+        Column("id", Integer, key="id"),
+        Column("label", String(200), key="label"),
+        *extra_cols,
+    )
+    meta.create_all(engine)
+    return table, meta
+
+
+def _roundtrip(engine, table, records):
+    dialect = PostgreSQLDialect()
+    with engine.begin() as conn:
+        dialect.bulk_insert(conn, table, records)
+    with engine.connect() as conn:
+        return conn.execute(select(table).order_by(table.c.id)).mappings().all()
+
+
+def _drop(meta, engine):
+    meta.drop_all(engine)
+
+
+# ---------------------------------------------------------------------------
+# Basic types
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.dbtest
+def test_pg_bulk_insert_basic(pg_engine):
+    table, meta = _make_table(pg_engine, "pg_bi_basic")
+    try:
+        records = [{"id": 1, "label": "hello"}, {"id": 2, "label": None}]
+        rows = _roundtrip(pg_engine, table, records)
+        assert len(rows) == 2
+        assert rows[0]["id"] == 1
+        assert rows[0]["label"] == "hello"
+        assert rows[1]["label"] is None
+    finally:
+        _drop(meta, pg_engine)
+
+
+@pytest.mark.dbtest
+def test_pg_bulk_insert_empty(pg_engine):
+    table, meta = _make_table(pg_engine, "pg_bi_empty")
+    try:
+        dialect = PostgreSQLDialect()
+        with pg_engine.begin() as conn:
+            dialect.bulk_insert(conn, table, [])
+        with pg_engine.connect() as conn:
+            count = conn.execute(select(table)).rowcount
+        assert count == 0
+    finally:
+        _drop(meta, pg_engine)
+
+
+# ---------------------------------------------------------------------------
+# Numeric types
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.dbtest
+def test_pg_bulk_insert_numeric_types(pg_engine):
+    meta = MetaData()
+    table = Table(
+        "pg_bi_numeric",
+        meta,
+        Column("i", Integer, key="i"),
+        Column("bi", BigInteger, key="bi"),
+        Column("si", SmallInteger, key="si"),
+        Column("d", Double, key="d"),
+    )
+    meta.create_all(pg_engine)
+    try:
+        records = [{"i": 1, "bi": 10**15, "si": 32767, "d": 3.14}]
+        dialect = PostgreSQLDialect()
+        with pg_engine.begin() as conn:
+            dialect.bulk_insert(conn, table, records)
+        with pg_engine.connect() as conn:
+            rows = conn.execute(select(table)).mappings().all()
+        assert rows[0]["i"] == 1
+        assert rows[0]["bi"] == 10**15
+        assert rows[0]["si"] == 32767
+        assert abs(rows[0]["d"] - 3.14) < 1e-9
+    finally:
+        meta.drop_all(pg_engine)
+
+
+# ---------------------------------------------------------------------------
+# Boolean
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.dbtest
+def test_pg_bulk_insert_boolean(pg_engine):
+    meta = MetaData()
+    table = Table(
+        "pg_bi_bool",
+        meta,
+        Column("id", Integer, key="id"),
+        Column("flag", Boolean, key="flag"),
+    )
+    meta.create_all(pg_engine)
+    try:
+        records = [
+            {"id": 1, "flag": True},
+            {"id": 2, "flag": False},
+            {"id": 3, "flag": None},
+        ]
+        dialect = PostgreSQLDialect()
+        with pg_engine.begin() as conn:
+            dialect.bulk_insert(conn, table, records)
+        with pg_engine.connect() as conn:
+            rows = conn.execute(select(table).order_by(table.c.id)).mappings().all()
+        assert rows[0]["flag"] is True
+        assert rows[1]["flag"] is False
+        assert rows[2]["flag"] is None
+    finally:
+        meta.drop_all(pg_engine)
+
+
+# ---------------------------------------------------------------------------
+# DateTime
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.dbtest
+def test_pg_bulk_insert_datetime(pg_engine):
+    meta = MetaData()
+    table = Table(
+        "pg_bi_dt",
+        meta,
+        Column("id", Integer, key="id"),
+        Column("ts", DateTime(timezone=True), key="ts"),
+    )
+    meta.create_all(pg_engine)
+    try:
+        dt = datetime.datetime(2023, 9, 27, 14, 35, 54, 274602,
+                               tzinfo=datetime.timezone.utc)
+        records = [{"id": 1, "ts": dt}, {"id": 2, "ts": None}]
+        dialect = PostgreSQLDialect()
+        with pg_engine.begin() as conn:
+            dialect.bulk_insert(conn, table, records)
+        with pg_engine.connect() as conn:
+            rows = conn.execute(select(table).order_by(table.c.id)).mappings().all()
+        assert rows[0]["ts"] is not None
+        assert rows[1]["ts"] is None
+    finally:
+        meta.drop_all(pg_engine)
+
+
+# ---------------------------------------------------------------------------
+# Binary (bytea)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.dbtest
+def test_pg_bulk_insert_binary(pg_engine):
+    meta = MetaData()
+    table = Table(
+        "pg_bi_binary",
+        meta,
+        Column("id", Integer, key="id"),
+        Column("data", LargeBinary, key="data"),
+    )
+    meta.create_all(pg_engine)
+    try:
+        payload = b"\xde\xad\xbe\xef" * 8
+        records = [{"id": 1, "data": payload}, {"id": 2, "data": None}]
+        dialect = PostgreSQLDialect()
+        with pg_engine.begin() as conn:
+            dialect.bulk_insert(conn, table, records)
+        with pg_engine.connect() as conn:
+            rows = conn.execute(select(table).order_by(table.c.id)).mappings().all()
+        assert bytes(rows[0]["data"]) == payload
+        assert rows[1]["data"] is None
+    finally:
+        meta.drop_all(pg_engine)
+
+
+# ---------------------------------------------------------------------------
+# Python-side scalar column defaults
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.dbtest
+def test_pg_bulk_insert_scalar_column_default(pg_engine):
+    meta = MetaData()
+    table = Table(
+        "pg_bi_default",
+        meta,
+        Column("id", Integer, key="id"),
+        Column("flag", Boolean, default=False, key="flag"),
+    )
+    meta.create_all(pg_engine)
+    try:
+        # Records do NOT contain 'flag'; the default must be applied.
+        records = [{"id": 1}, {"id": 2}]
+        dialect = PostgreSQLDialect()
+        with pg_engine.begin() as conn:
+            dialect.bulk_insert(conn, table, records)
+        with pg_engine.connect() as conn:
+            rows = conn.execute(select(table).order_by(table.c.id)).mappings().all()
+        assert rows[0]["flag"] is False
+        assert rows[1]["flag"] is False
+    finally:
+        meta.drop_all(pg_engine)