diff --git a/CHANGELOG.md b/CHANGELOG.md index fc26b00..8884755 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,11 @@ follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). `UpsertConformance`). Adapter integration suites subclass the reusable classes and run them against their real engine, providing `adapter` and `schema` fixtures. Install with `continuo-engine-contract[testing]`. +- `continuo-postgres-adapter`: `append` (batched inserts in one transaction), + `upsert` (`INSERT ... ON CONFLICT (key) DO UPDATE`, `DO NOTHING` when every + column is in the key) and `table_columns`; `ensure_table` declares an + upsert key as the table's `PRIMARY KEY`. `load` now shares the insert path + (no behavior change). ### Changed diff --git a/adapters/postgres/continuo_postgres_adapter/adapter.py b/adapters/postgres/continuo_postgres_adapter/adapter.py index 6e27e55..2f45770 100644 --- a/adapters/postgres/continuo_postgres_adapter/adapter.py +++ b/adapters/postgres/continuo_postgres_adapter/adapter.py @@ -8,8 +8,10 @@ without scanning data. Schema creation is serialized on a session advisory lock because parallel root validation nodes race on ``CREATE SCHEMA``. * **Python-node data-plane I/O** — ``fetch`` executes one declared read and - returns an Arrow table; ``ensure_table``/``load`` build and atomically - replace a table's contents. + returns an Arrow table; ``ensure_table`` builds the table (an upsert key as + its ``PRIMARY KEY``); ``load`` atomically replaces its contents, ``append`` + adds rows and ``_upsert`` inserts or overwrites by key, all through one + insert helper; ``table_columns`` reads its shape. The connection runs with ``autocommit = True``: every method owns its transaction explicitly rather than relying on psycopg2's implicit @@ -17,22 +19,29 @@ ``build_empty_from_sql``, ``clone_empty_from_prod``) needs nothing further — under autocommit each statement commits as it runs. A method that must apply several statements atomically (``build_empty_from_columns``, ``check_binds``, -``ensure_table``, ``load``) wraps them in its own explicit +``ensure_table``, ``load``, ``append``, ``_upsert``) wraps them in its own explicit ``BEGIN``/``COMMIT``, with a ``ROLLBACK`` on the error path so a mid-sequence failure leaves no partial change and no aborted transaction behind for the next call on the same connection. """ +import contextlib import hashlib import logging import os +from collections.abc import Iterator, Sequence from typing import Any import psycopg2 # type: ignore[import-untyped] import pyarrow as pa # type: ignore[import-untyped] from continuo_engine_contract.config import ensure_known_keys # type: ignore[import-untyped] -from continuo_engine_contract.port import WarehouseAdapter # type: ignore[import-untyped] +from continuo_engine_contract.port import ( # type: ignore[import-untyped] + AppendCapable, + LiveColumn, + UpsertCapable, + WarehouseAdapter, +) from continuo_engine_contract.sql import ensure_single_read # type: ignore[import-untyped] from continuo_engine_contract.types import validate_column_type # type: ignore[import-untyped] from psycopg2 import errors as pg_errors # type: ignore[import-untyped] @@ -305,6 +314,65 @@ def _column_ddl(columns: list[dict[str, Any]]) -> "pg_sql.Composed": return pg_sql.SQL(", ").join(col_defs) +def _primary_key_ddl(primary_key: Sequence[str]) -> "pg_sql.Composed": + """``PRIMARY KEY (...)`` for an upsert node's table: enforced and indexed by postgres.""" + return pg_sql.SQL("PRIMARY KEY ({})").format( + pg_sql.SQL(", ").join(pg_sql.Identifier(name) for name in primary_key) + ) + + +# Rows per INSERT statement on every write path (execute_values' page_size). +_INSERT_PAGE_SIZE = 1000 + + +def _insert_rows( + cur: Any, + schema: str, + table: str, + data: "pa.Table", + suffix: "pg_sql.Composable | None" = None, +) -> None: + """Insert every row of *data* into ``schema.table`` inside the caller's transaction. + + Batched INSERTs (``execute_values``, ``_INSERT_PAGE_SIZE`` rows per + statement) that name *data*'s columns, so the table's own column order + does not matter. *suffix* (an ``ON CONFLICT`` clause) follows the VALUES + list. ``execute_values`` reparses every percent token after the composable + identifiers have rendered, so literal percents in identifiers are escaped + and only its one values-list token is left unescaped. The single insert + path behind ``load``, ``append`` and ``_upsert``. + """ + columns = data.schema.names + values = [tuple(row[c] for c in columns) for row in data.to_pylist()] + prefix = pg_sql.SQL("INSERT INTO {}.{} ({})").format( + pg_sql.Identifier(schema), + pg_sql.Identifier(table), + pg_sql.SQL(", ").join(pg_sql.Identifier(c) for c in columns), + ) + statement = prefix.as_string(cur).replace("%", "%%") + " VALUES %s" + if suffix is not None: + statement += " " + suffix.as_string(cur).replace("%", "%%") + execute_values(cur, statement, values, page_size=_INSERT_PAGE_SIZE) + + +def _on_conflict(columns: Sequence[str], primary_key: Sequence[str]) -> "pg_sql.Composed": + """The upsert's ``ON CONFLICT`` clause. + + Overwrites every non-key column from ``EXCLUDED``; with every column in + the key there is nothing to overwrite, so ``DO NOTHING``. + """ + key = set(primary_key) + target = pg_sql.SQL(", ").join(pg_sql.Identifier(name) for name in primary_key) + updates = [name for name in columns if name not in key] + if not updates: + return pg_sql.SQL("ON CONFLICT ({}) DO NOTHING").format(target) + assignments = pg_sql.SQL(", ").join( + pg_sql.SQL("{} = EXCLUDED.{}").format(pg_sql.Identifier(name), pg_sql.Identifier(name)) + for name in updates + ) + return pg_sql.SQL("ON CONFLICT ({}) DO UPDATE SET {}").format(target, assignments) + + def _arrow_table_from_rows(colnames: list[str], rows: list[tuple[Any, ...]]) -> "pa.Table": """Build a column-wise Arrow table from cursor description names and fetched rows. @@ -337,7 +405,7 @@ def _arrow_table_from_rows(colnames: list[str], rows: list[tuple[Any, ...]]) -> return pa.table({name: pa.array(values) for name, values in zip(colnames, by_column)}) -class PostgresAdapter(WarehouseAdapter): +class PostgresAdapter(WarehouseAdapter, AppendCapable, UpsertCapable): """WarehouseAdapter speaking postgres over a psycopg2 connection.""" def __init__(self, conn: "psycopg2.extensions.connection") -> None: @@ -571,6 +639,26 @@ def fetch(self, sql: str) -> "pa.Table": raise return table + def table_columns(self, schema: str, table: str) -> list[LiveColumn]: + """Names and nullability of ``schema.table`` from ``information_schema.columns``. + + Names are matched exactly: this adapter creates every schema, table and + column with quoted identifiers, so the catalog holds them as written. + + Raises: + LookupError: If the table does not exist (or this role cannot see it). + """ + with self._conn.cursor() as cur: + cur.execute( + "SELECT column_name, is_nullable FROM information_schema.columns " + "WHERE table_schema = %s AND table_name = %s ORDER BY ordinal_position", + (schema, table), + ) + rows = cur.fetchall() + if not rows: + raise LookupError(f"table {schema}.{table} does not exist") + return [LiveColumn(name=name, nullable=is_nullable == "YES") for name, is_nullable in rows] + @classmethod def validate_config( cls, config: dict[str, Any] | None, column_names: list[str] @@ -609,6 +697,7 @@ def ensure_table( columns: list[dict[str, Any]], *, config: dict[str, Any], + primary_key: Sequence[str] = (), ) -> None: """CREATE TABLE IF NOT EXISTS with typed DDL compiled from *columns*. @@ -620,17 +709,33 @@ def ensure_table( keyword-only, matching the contract 0.6.0 port signature. The create and every index either all land or none do, in the same explicit BEGIN/COMMIT/ROLLBACK shape build_empty_from_columns uses. + + *primary_key* (an upsert node's key, in the declared column spelling) + becomes the created table's ``PRIMARY KEY``, which ``_upsert``'s + ``ON CONFLICT`` needs. Like the indexes it applies only when this call + creates the table: an existing table without the constraint makes the + upsert fail at ``ON CONFLICT`` rather than collect duplicate keys. """ - indexes = _validated_indexes(config, table, [col["name"] for col in columns]) + declared = [col["name"] for col in columns] + indexes = _validated_indexes(config, table, declared) for col in columns: validate_column_type(col["type"]) + undeclared = [name for name in primary_key if name not in declared] + if undeclared: + raise ValueError( + f"primary key on undeclared column(s) {undeclared!r}; " + f"declared columns: {declared!r}" + ) self.ensure_schema(schema) + elements = _column_ddl(columns) + if primary_key: + elements = pg_sql.SQL(", ").join([elements, _primary_key_ddl(primary_key)]) stmt = pg_sql.SQL("CREATE TABLE IF NOT EXISTS {}.{} ({})").format( pg_sql.Identifier(schema), pg_sql.Identifier(table), - _column_ddl(columns), + elements, ) with self._conn.cursor() as cur: cur.execute("BEGIN") @@ -650,51 +755,86 @@ def ensure_table( logger.error("ensure_table rollback failed: %s", rollback_exc) raise - def load(self, schema: str, table: str, data: "pa.Table") -> None: - """Atomically replace ``schema.table``'s contents with *data*. + @contextlib.contextmanager + def _transaction(self, operation: str) -> Iterator[Any]: + """Run the body on one cursor inside an explicit ``BEGIN``/``COMMIT``. - One explicit transaction: TRUNCATE then batched inserts - (``execute_values``, page_size 1000) in the Arrow table's column - order; commits on success, rolls back and re-raises on any error. A - 0-row *data* is just a TRUNCATE. The transaction is opened explicitly - (``BEGIN``) rather than relying on an implicit one, matching every - other multi-statement method on this class under ``autocommit = True``. + Any error rolls back and re-raises. A rollback that fails as well is + only logged, so the caller sees the original error and the session is + never left in an aborted transaction. """ - columns = data.schema.names with self._conn.cursor() as cur: cur.execute("BEGIN") try: - cur.execute( - pg_sql.SQL("TRUNCATE {}.{}").format( - pg_sql.Identifier(schema), pg_sql.Identifier(table) - ) - ) - if data.num_rows: - rows = data.to_pylist() - values = [tuple(row[c] for c in columns) for row in rows] - insert_prefix = pg_sql.SQL("INSERT INTO {}.{} ({})").format( - pg_sql.Identifier(schema), - pg_sql.Identifier(table), - pg_sql.SQL(", ").join(pg_sql.Identifier(c) for c in columns), - ) - # execute_values reparses every percent token after composable - # identifiers have rendered. Escape literal identifier percents, - # then append its one values-list expansion token unescaped. - insert_stmt = ( - insert_prefix.as_string(cur).replace("%", "%%") + " VALUES %s" - ) - logger.info( - "loading %d row(s) into %s.%s", data.num_rows, schema, table - ) - execute_values(cur, insert_stmt, values, page_size=1000) + yield cur cur.execute("COMMIT") except Exception: try: cur.execute("ROLLBACK") except psycopg2.Error as rollback_exc: - logger.error("load rollback failed: %s", rollback_exc) + logger.error("%s rollback failed: %s", operation, rollback_exc) raise + def load(self, schema: str, table: str, data: "pa.Table") -> None: + """Atomically replace ``schema.table``'s contents with *data*. + + One explicit transaction: TRUNCATE then batched inserts + (``_insert_rows``, page size 1000) that name the Arrow table's columns; + commits on success, rolls back and re-raises on any error. A 0-row + *data* is just a TRUNCATE. The transaction is opened explicitly + (``BEGIN``) rather than relying on an implicit one, matching every + other multi-statement method on this class under ``autocommit = True``. + """ + with self._transaction("load") as cur: + cur.execute( + pg_sql.SQL("TRUNCATE {}.{}").format( + pg_sql.Identifier(schema), pg_sql.Identifier(table) + ) + ) + if data.num_rows: + logger.info("loading %d row(s) into %s.%s", data.num_rows, schema, table) + _insert_rows(cur, schema, table, data) + + def append(self, schema: str, table: str, data: "pa.Table") -> None: + """Atomically insert every row of *data* into ``schema.table``; earlier rows stay. + + One explicit transaction of batched inserts that name *data*'s columns + (the same insert path as ``load``, without the TRUNCATE). A 0-row + *data* opens no transaction at all. + """ + if not data.num_rows: + logger.info("append of 0 rows into %s.%s: nothing to write", schema, table) + return + with self._transaction("append") as cur: + logger.info("appending %d row(s) into %s.%s", data.num_rows, schema, table) + _insert_rows(cur, schema, table, data) + + def _upsert( + self, schema: str, table: str, data: "pa.Table", primary_key: Sequence[str] + ) -> None: + """Insert new keys and overwrite existing ones in one transaction. + + Batched ``INSERT ... ON CONFLICT (key) DO UPDATE`` (``DO NOTHING`` when + every column is in the key) through the shared insert path. Needs the + table's ``PRIMARY KEY`` (see ``ensure_table``): postgres rejects + ``ON CONFLICT`` on columns without a matching constraint, so a table + created without it fails here instead of collecting duplicates. Keys + arrive unique (``UpsertCapable.upsert``), which also keeps one + statement from touching a row twice. A 0-row *data* opens no + transaction. + """ + if not data.num_rows: + logger.info("upsert of 0 rows into %s.%s: nothing to write", schema, table) + return + with self._transaction("upsert") as cur: + logger.info( + "upserting %d row(s) into %s.%s on (%s)", + data.num_rows, schema, table, ", ".join(primary_key), + ) + _insert_rows( + cur, schema, table, data, suffix=_on_conflict(data.schema.names, primary_key) + ) + def close(self) -> None: """Release the underlying connection.""" self._conn.close() diff --git a/adapters/postgres/tests/test_adapter_runtime_postgres.py b/adapters/postgres/tests/test_adapter_runtime_postgres.py index 7cf72a9..3ec4845 100644 --- a/adapters/postgres/tests/test_adapter_runtime_postgres.py +++ b/adapters/postgres/tests/test_adapter_runtime_postgres.py @@ -549,3 +549,89 @@ def test_config_validation_survives_a_second_recognized_top_level_key(monkeypatc conn = _FakeConnection() PostgresAdapter(conn).ensure_table("s", "t", _ONE_COL, config={"fillfactor": 90}) assert conn.committed # DDL ran; no KeyError on the way in + + +def test_load_of_zero_rows_truncates_in_one_explicit_transaction(): + """load's transaction shape, pinned across the insert-path refactor.""" + conn = _FakeConnection() + PostgresAdapter(conn).load("s", "t", pa.table({"id": pa.array([], pa.int64())})) + (cur,) = conn.cursors + assert cur.statements == [ + "BEGIN", + pg_sql.SQL("TRUNCATE {}.{}").format(pg_sql.Identifier("s"), pg_sql.Identifier("t")), + "COMMIT", + ] + + +def test_append_of_zero_rows_opens_no_transaction(): + conn = _FakeConnection() + PostgresAdapter(conn).append("s", "t", pa.table({"id": pa.array([], pa.int64())})) + assert conn.cursors == [] + + +_KEYED = [ + {"name": "rate_date", "type": "DATE", "nullable": False}, + {"name": "base", "type": "VARCHAR(3)", "nullable": False}, + {"name": "rate", "type": "DECIMAL(12,6)", "nullable": True}, +] + + +def test_ensure_table_with_a_primary_key_declares_it_in_the_create_table(): + conn = _FakeConnection() + PostgresAdapter(conn).ensure_table( + "s", "t", _KEYED, config={}, primary_key=("rate_date", "base") + ) + expected = pg_sql.SQL("CREATE TABLE IF NOT EXISTS {}.{} ({})").format( + pg_sql.Identifier("s"), + pg_sql.Identifier("t"), + pg_sql.SQL(", ").join([ + adapter_module._column_ddl(_KEYED), + pg_sql.SQL("PRIMARY KEY ({})").format( + pg_sql.SQL(", ").join([pg_sql.Identifier("rate_date"), pg_sql.Identifier("base")]) + ), + ]), + ) + assert conn.cursors[-1].statements[1] == expected + + +def test_ensure_table_without_a_primary_key_emits_todays_create_table(): + conn = _FakeConnection() + PostgresAdapter(conn).ensure_table("s", "t", _KEYED, config={}) + expected = pg_sql.SQL("CREATE TABLE IF NOT EXISTS {}.{} ({})").format( + pg_sql.Identifier("s"), pg_sql.Identifier("t"), adapter_module._column_ddl(_KEYED) + ) + assert conn.cursors[-1].statements[1] == expected + + +def test_ensure_table_rejects_a_primary_key_on_an_undeclared_column_and_emits_no_ddl(): + conn = _FakeConnection() + with pytest.raises(ValueError, match=r"primary key on undeclared column\(s\) \['quote'\]"): + PostgresAdapter(conn).ensure_table( + "s", "t", _KEYED, config={}, primary_key=("rate_date", "quote") + ) + assert conn.cursors == [] + + +def test_on_conflict_overwrites_every_non_key_column(): + assert adapter_module._on_conflict(["rate_date", "base", "rate"], ["rate_date", "base"]) == ( + pg_sql.SQL("ON CONFLICT ({}) DO UPDATE SET {}").format( + pg_sql.SQL(", ").join([pg_sql.Identifier("rate_date"), pg_sql.Identifier("base")]), + pg_sql.SQL(", ").join([ + pg_sql.SQL("{} = EXCLUDED.{}").format( + pg_sql.Identifier("rate"), pg_sql.Identifier("rate") + ) + ]), + ) + ) + + +def test_on_conflict_does_nothing_when_every_column_is_in_the_key(): + assert adapter_module._on_conflict(["a", "b"], ["a", "b"]) == pg_sql.SQL( + "ON CONFLICT ({}) DO NOTHING" + ).format(pg_sql.SQL(", ").join([pg_sql.Identifier("a"), pg_sql.Identifier("b")])) + + +def test_upsert_of_zero_rows_opens_no_transaction(): + conn = _FakeConnection() + PostgresAdapter(conn).upsert("s", "t", pa.table({"id": pa.array([], pa.int64())}), ["id"]) + assert conn.cursors == [] diff --git a/adapters/postgres/tests/test_integration_postgres_conformance.py b/adapters/postgres/tests/test_integration_postgres_conformance.py new file mode 100644 index 0000000..c74fb17 --- /dev/null +++ b/adapters/postgres/tests/test_integration_postgres_conformance.py @@ -0,0 +1,103 @@ +"""The engine-contract conformance suite against postgres:16 (tests/smoke/postgres-stack). + +Run the stack first: + docker compose -f tests/smoke/postgres-stack/docker-compose.yml up -d --wait +""" +import datetime +import os +import uuid + +import psycopg2 +import pytest + +from continuo_engine_contract.testing.conformance import ( + ARROW, + COLUMNS, + FX_COLUMNS, + NAMES, + AppendConformance, + ConformanceBase, + IntrospectionConformance, + UpsertConformance, + table_of, +) +from continuo_postgres_adapter.adapter import PostgresAdapter + +pytestmark = pytest.mark.integration + +PG = dict( + host="localhost", + port=os.environ.get("VR_IT_PG_PORT", "15499"), + dbname="warehouse", + user="continuo", + password="continuo", +) + + +@pytest.fixture +def adapter(): + connected = PostgresAdapter(psycopg2.connect(**PG)) + yield connected + connected.close() + + +@pytest.fixture +def schema(adapter): + name = f"conf_{uuid.uuid4().hex[:10]}" + yield name + adapter.drop_schema(name) + + +class TestPostgresIntrospection(IntrospectionConformance): + pass + + +class TestPostgresAppend(AppendConformance): + pass + + +class TestPostgresUpsert(UpsertConformance): + pass + + +def _primary_key_columns(schema: str, table: str) -> list[str]: + conn = psycopg2.connect(**PG) + try: + with conn.cursor() as cur: + cur.execute( + "SELECT kcu.column_name FROM information_schema.table_constraints tc " + "JOIN information_schema.key_column_usage kcu " + "ON kcu.constraint_name = tc.constraint_name " + "AND kcu.table_schema = tc.table_schema AND kcu.table_name = tc.table_name " + "WHERE tc.table_schema = %s AND tc.table_name = %s " + "AND tc.constraint_type = 'PRIMARY KEY' ORDER BY kcu.ordinal_position", + (schema, table), + ) + return [row[0] for row in cur.fetchall()] + finally: + conn.close() + + +def test_ensure_table_with_a_primary_key_creates_an_enforced_constraint(adapter, schema): + """Review Focus 4: the key exists as soon as the first run created the table.""" + adapter.ensure_table(schema, "keyed", FX_COLUMNS, config={}, primary_key=["rate_date", "base"]) + assert _primary_key_columns(schema, "keyed") == ["rate_date", "base"] + conn = psycopg2.connect(**PG) + conn.autocommit = True + try: + with conn.cursor() as cur: + insert = f'INSERT INTO "{schema}"."keyed" VALUES (%s, %s, %s)' + cur.execute(insert, (datetime.date(2026, 10, 8), "EUR", None)) + with pytest.raises(psycopg2.errors.UniqueViolation): + cur.execute(insert, (datetime.date(2026, 10, 8), "EUR", None)) + finally: + conn.close() + + +def test_upsert_into_a_table_without_its_primary_key_fails_and_writes_nothing(adapter, schema): + """ON CONFLICT needs the constraint; postgres refuses instead of collecting duplicates.""" + adapter.ensure_table(schema, "unkeyed", COLUMNS, config={}) + adapter.append(schema, "unkeyed", table_of(ARROW, [(1, "a", 10)])) + with pytest.raises(psycopg2.errors.InvalidColumnReference): + adapter.upsert(schema, "unkeyed", table_of(ARROW, [(1, "b", 20)]), ["id"]) + ConformanceBase().expect_rows(adapter, schema, "unkeyed", NAMES, [(1, "a", 10)])