diff --git a/.github/workflows/integration-tests-postgres.yml b/.github/workflows/integration-tests-postgres.yml index 5598f96..702fe16 100644 --- a/.github/workflows/integration-tests-postgres.yml +++ b/.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 }} diff --git a/src/xml2db/dialect/postgresql.py b/src/xml2db/dialect/postgresql.py index 23a9cfe..d909e30 100644 --- a/src/xml2db/dialect/postgresql.py +++ b/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")) diff --git a/tests/test_bulk_insert_postgres.py b/tests/test_bulk_insert_postgres.py new file mode 100644 index 0000000..60c534a --- /dev/null +++ b/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)