diff --git a/CHANGELOG.md b/CHANGELOG.md index 5f3a9b1..87fa621 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -39,6 +39,12 @@ follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). from the registered Arrow data, one transaction) and `table_columns`. `ensure_table` accepts and ignores an upsert key: DuckLake has no constraints, so the key is guaranteed by the write path. +- `continuo-trino-adapter`: `append` and `upsert` through a stage table + shaped like the target, then one `INSERT ... SELECT` or one `MERGE` — one + Iceberg commit per write; the stage is dropped best-effort afterwards. + `table_columns` reads `information_schema`. `ensure_table` accepts an + upsert key (no constraint DDL) and rejects it with `format_version: 1`. + `load` now shares the stage helpers (no behavior change). ### Changed diff --git a/adapters/trino/continuo_trino_adapter/adapter.py b/adapters/trino/continuo_trino_adapter/adapter.py index ca983d6..32da131 100644 --- a/adapters/trino/continuo_trino_adapter/adapter.py +++ b/adapters/trino/continuo_trino_adapter/adapter.py @@ -3,7 +3,8 @@ One class covers both roles Continuo asks of an engine: validation DDL (``ensure_schema``/``drop_schema``, the empty builds and clones, ``check_binds``) and python-node data-plane I/O (``fetch`` executes one declared read and returns -an Arrow table; ``ensure_table``/``load`` build and replace a table's contents). +an Arrow table; ``ensure_table``/``load`` build and replace a table's contents; +``append`` and ``_upsert`` add or merge rows; ``table_columns`` reads a table's shape). Trino's namespace is catalog.schema: the catalog comes from TRINO_CATALOG and every statement is fully qualified, so the contract's ``schema`` argument stays @@ -46,6 +47,14 @@ ``finally``. If recovery fails, the private old relation is deliberately retained because it contains the original target data. +``append`` and ``_upsert`` write through the same kind of stage (created LIKE +the target, batch-inserted) and then touch the target with exactly one +statement — ``INSERT ... SELECT`` or ``MERGE INTO ... USING stage`` — so each +write is one Iceberg commit: all of its rows or none. The stage is dropped +afterwards, best-effort and only logged, so a committed write never reports +failure. A 0-row write creates no stage. ``MERGE`` needs Iceberg format +version 2 (Trino's default). + ``load`` assumes a single writer per table: Continuo's scheduler runs at most one Job per node at a time. Concurrent ``load()`` calls against the same table are unsupported because their target renames would race. Swap relations use a @@ -65,13 +74,19 @@ import os import uuid +from collections.abc import Callable, Sequence from typing import Any from urllib.parse import urlsplit, urlunsplit 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] @@ -246,6 +261,34 @@ def _sibling_location(location: str, name: str) -> str: ) +def _append_statement(target_ref: str, stage_ref: str, columns: Sequence[str]) -> str: + """The append's one write to the target: ``INSERT ... SELECT`` every staged row.""" + col_list = ", ".join(_quote(name) for name in columns) + return f"INSERT INTO {target_ref} ({col_list}) SELECT {col_list} FROM {stage_ref}" + + +def _merge_statement( + target_ref: str, stage_ref: str, columns: Sequence[str], primary_key: Sequence[str] +) -> str: + """The upsert's one write to the target: ``MERGE INTO target USING stage`` on the key. + + A matched row gets every non-key column from the stage; unmatched stage + rows are inserted. With every column in the key there is nothing to + overwrite, so there is no ``WHEN MATCHED`` branch. + """ + key = set(primary_key) + on = " AND ".join(f"t.{_quote(name)} = s.{_quote(name)}" for name in primary_key) + statement = f"MERGE INTO {target_ref} AS t USING {stage_ref} AS s ON ({on})" + updates = [name for name in columns if name not in key] + if updates: + statement += " WHEN MATCHED THEN UPDATE SET " + ", ".join( + f"{_quote(name)} = s.{_quote(name)}" for name in updates + ) + col_list = ", ".join(_quote(name) for name in columns) + values = ", ".join(f"s.{_quote(name)}" for name in columns) + return statement + f" WHEN NOT MATCHED THEN INSERT ({col_list}) VALUES ({values})" + + def _arrow_table_from_rows(colnames: list[str], rows: list[Any]) -> "pa.Table": """Build a column-wise Arrow table from cursor description names and fetched rows. @@ -281,8 +324,11 @@ def _arrow_table_from_rows(colnames: list[str], rows: list[Any]) -> "pa.Table": # Batch size for multi-row INSERT VALUES statements in load(). _INSERT_BATCH_SIZE = 500 +# Name prefix of the stage tables every write path stages its rows in. +_STAGE_PREFIX = "__continuo_stage_" + -class TrinoAdapter(WarehouseAdapter): +class TrinoAdapter(WarehouseAdapter, AppendCapable, UpsertCapable): """WarehouseAdapter speaking Trino (Iceberg connector) over the trino DBAPI.""" def __init__(self, conn: "trino.dbapi.Connection", catalog: str) -> None: @@ -465,6 +511,25 @@ def fetch(self, sql: str) -> "pa.Table": finally: cur.close() + def table_columns(self, schema: str, table: str) -> list[LiveColumn]: + """Names and nullability of ``schema.table`` from ``information_schema.columns``. + + Trino stores schema, table and column names lowercased, so the lookup + lowercases *schema* and *table*; callers compare column names + case-insensitively. + + Raises: + LookupError: If the table does not exist. + """ + rows = self._execute( + f"SELECT column_name, is_nullable FROM {_quote(self._catalog)}.information_schema.columns " + "WHERE table_schema = ? AND table_name = ? ORDER BY ordinal_position", + [schema.lower(), table.lower()], + ) + if not rows: + raise LookupError(f"table {schema}.{table} does not exist in catalog {self._catalog}") + return [LiveColumn(name=row[0], nullable=row[1] == "YES") for row in rows] + @classmethod def validate_config( cls, config: dict[str, Any] | None, column_names: list[str] @@ -506,6 +571,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*. @@ -520,8 +586,20 @@ def ensure_table( table behind (fail closed), then rendered into the CREATE TABLE's ``WITH`` clause. Required and keyword-only, matching the contract 0.6.0 port signature. + + *primary_key* declares nothing: Trino has no constraint syntax, so an + upsert node's key is guaranteed by ``upsert`` (the port's duplicate + check, then one MERGE). It does gate one config value: Trino's MERGE + needs Iceberg format version 2, so a key with ``format_version`` below + 2 is rejected before any DDL. """ properties = _table_properties(config) + format_version = (config or {}).get("format_version") + if primary_key and isinstance(format_version, int) and format_version < 2: + raise ValueError( + f"config 'format_version' {format_version} cannot hold an upsert table: " + "Trino's MERGE needs Iceberg format version 2 or later" + ) for col in columns: validate_column_type(col["type"]) @@ -532,6 +610,40 @@ def ensure_table( logger.info("ensuring table %s exists", ref) self._execute(f"CREATE TABLE IF NOT EXISTS {ref} ({col_defs}){properties}") + def _create_stage(self, schema: str, table: str, token: str) -> str: + """Create an empty stage table shaped exactly like ``schema.table``; return its reference. + + ``LIKE ... INCLUDING PROPERTIES`` copies the target's columns, NOT NULL + constraints and Iceberg properties, so a row the target would refuse + fails here, before the target is touched. The explicit location, a + fresh sibling of the target's, avoids Iceberg's "non-empty location" + rejection. The target must exist. + """ + target_ref = self._table_ref(schema, table) + stage_table = f"{_STAGE_PREFIX}{token}" + stage_ref = self._table_ref(schema, stage_table) + stage_location = _sibling_location(self._table_location(schema, table), stage_table) + logger.info("creating staging table %s like %s", stage_ref, target_ref) + self._execute( + f"CREATE TABLE {stage_ref} " + f"(LIKE {target_ref} INCLUDING PROPERTIES) " + f"WITH (location = {_sql_string(stage_location)})" + ) + return stage_ref + + def _drop_stage(self, stage_ref: str) -> None: + """Drop a stage table best-effort: a failure is logged, never raised. + + Callers drop the stage after their write; a write that already + committed must not report failure because its cleanup did not. + """ + try: + self._execute(f"DROP TABLE {stage_ref}") + except Exception: + logger.warning( + "best-effort cleanup of staging table %s failed", stage_ref, exc_info=True + ) + def _insert_batches(self, table_ref: str, columns: list[str], data: "pa.Table") -> None: """Insert *data*'s rows into *table_ref* in batches of ``_INSERT_BATCH_SIZE``.""" rows = data.to_pylist() @@ -558,22 +670,11 @@ def load(self, schema: str, table: str, data: "pa.Table") -> None: """ columns = data.schema.names target_ref = self._table_ref(schema, table) - load_token = uuid.uuid4().hex - stage_table = f"__continuo_stage_{load_token}" - old_table = f"__continuo_old_{load_token}" - stage_ref = self._table_ref(schema, stage_table) - old_ref = self._table_ref(schema, old_table) - target_location = self._table_location(schema, table) - stage_location = _sibling_location(target_location, stage_table) - stage_created = False + token = uuid.uuid4().hex + old_ref = self._table_ref(schema, f"__continuo_old_{token}") + stage_ref: str | None = None try: - logger.info("creating staging table %s like %s", stage_ref, target_ref) - self._execute( - f"CREATE TABLE {stage_ref} " - f"(LIKE {target_ref} INCLUDING PROPERTIES) " - f"WITH (location = {_sql_string(stage_location)})" - ) - stage_created = True + stage_ref = self._create_stage(schema, table, token) if data.num_rows: logger.info("loading %d row(s) into staging table %s", data.num_rows, stage_ref) self._insert_batches(stage_ref, columns, data) @@ -596,14 +697,72 @@ def load(self, schema: str, table: str, data: "pa.Table") -> None: ) raise else: - stage_created = False + stage_ref = None # it is the target now self._execute(f"DROP TABLE {old_ref}") finally: - if stage_created: - try: - self._execute(f"DROP TABLE {stage_ref}") - except Exception: - logger.warning("best-effort cleanup of staging table %s failed", stage_ref) + if stage_ref is not None: + self._drop_stage(stage_ref) + + def _write_via_stage( + self, schema: str, table: str, data: "pa.Table", final_statement: Callable[[str], str] + ) -> None: + """Stage *data* beside ``schema.table``, then write the target in one statement. + + Trino/Iceberg has no multi-statement transaction, but each statement is + one Iceberg commit: the rows are batch-inserted into a stage shaped + like the target (a refused row fails there, before the target is + touched), then ``final_statement(stage_ref)`` moves them all in one + commit. The stage is dropped afterwards, best-effort, so a committed + write never reports failure because of its cleanup. + """ + stage_ref = self._create_stage(schema, table, uuid.uuid4().hex) + try: + logger.info("loading %d row(s) into staging table %s", data.num_rows, stage_ref) + self._insert_batches(stage_ref, data.schema.names, data) + self._execute(final_statement(stage_ref)) + finally: + self._drop_stage(stage_ref) + + def append(self, schema: str, table: str, data: "pa.Table") -> None: + """Insert every row of *data* into ``schema.table`` as one Iceberg commit. + + Stage, then one ``INSERT ... SELECT`` (see ``_write_via_stage``). A + 0-row *data* creates no stage and writes nothing. + """ + if not data.num_rows: + logger.info("append of 0 rows into %s.%s: nothing to write", schema, table) + return + target_ref = self._table_ref(schema, table) + columns = data.schema.names + logger.info("appending %d row(s) into %s", data.num_rows, target_ref) + self._write_via_stage( + schema, table, data, lambda stage_ref: _append_statement(target_ref, stage_ref, columns) + ) + + def _upsert( + self, schema: str, table: str, data: "pa.Table", primary_key: Sequence[str] + ) -> None: + """Insert new keys and overwrite existing ones as one Iceberg commit. + + Stage, then one ``MERGE INTO target USING stage`` (see + ``_write_via_stage``). Keys arrive unique (``UpsertCapable.upsert``); + Trino itself fails only when two stage rows match one existing row. A + 0-row *data* creates no stage. + """ + if not data.num_rows: + logger.info("upsert of 0 rows into %s.%s: nothing to write", schema, table) + return + target_ref = self._table_ref(schema, table) + columns = data.schema.names + logger.info( + "upserting %d row(s) into %s on (%s)", data.num_rows, target_ref, ", ".join(primary_key) + ) + self._write_via_stage( + schema, + table, + data, + lambda stage_ref: _merge_statement(target_ref, stage_ref, columns, primary_key), + ) def close(self) -> None: """Release the underlying connection.""" diff --git a/adapters/trino/tests/test_adapter_trino_writes.py b/adapters/trino/tests/test_adapter_trino_writes.py new file mode 100644 index 0000000..877402d --- /dev/null +++ b/adapters/trino/tests/test_adapter_trino_writes.py @@ -0,0 +1,231 @@ +"""Write paths of the trino adapter over a scripted connection: statements in order. + +The live behavior (Iceberg commits, NOT NULL in the stage, MERGE semantics) is +covered by the integration suites; this file pins which statements run, in +what order, and what happens when one of them fails. +""" +import logging +import re + +import pyarrow as pa +import pytest + +from continuo_engine_contract.keys import DuplicateKeyError +from continuo_engine_contract.port import LiveColumn +import continuo_trino_adapter.adapter as adapter_module +from continuo_trino_adapter.adapter import TrinoAdapter + + +class _ScriptedCursor: + def __init__(self, conn: "_ScriptedConnection") -> None: + self._conn = conn + self._rows: list = [] + + def execute(self, statement, params=None): + self._conn.statements.append(statement) + self._conn.params.append(params) + for marker, error in self._conn.fail_on.items(): + if marker in statement: + raise error + self._rows = next( + (rows for marker, rows in self._conn.rows_for.items() if marker in statement), [] + ) + + def fetchall(self): + return list(self._rows) + + def close(self): + pass + + +class _ScriptedConnection: + """Records statements; answers and fails by substring of the statement.""" + + def __init__(self, rows_for=None, fail_on=None): + self.statements: list[str] = [] + self.params: list = [] + self.rows_for = {"$properties": [["s3://warehouse/s/t-0abc"]], **(rows_for or {})} + self.fail_on = dict(fail_on or {}) + + def cursor(self): + return _ScriptedCursor(self) + + +def _adapter(conn): + return TrinoAdapter(conn, "iceberg") + + +def _ref(statement, kind): + """The quoted reference of the __continuo__ table in *statement*.""" + return re.search(rf'"iceberg"\."s"\."__continuo_{kind}_[0-9a-f]{{32}}"', statement).group(0) + + +def _name(ref): + return ref.rsplit(".", 1)[1].strip('"') + + +LOCATION_QUERY = 'SELECT value FROM "iceberg"."s"."t$properties" WHERE key = \'location\'' + + +def test_load_stages_then_swaps_and_drops_the_old_table(): + conn = _ScriptedConnection() + _adapter(conn).load("s", "t", pa.table({"id": [1, 2]})) + s = conn.statements + stage, old = _ref(s[1], "stage"), _ref(s[3], "old") + assert s == [ + LOCATION_QUERY, + f'CREATE TABLE {stage} (LIKE "iceberg"."s"."t" INCLUDING PROPERTIES) ' + f"WITH (location = 's3://warehouse/s/{_name(stage)}')", + f'INSERT INTO {stage} ("id") VALUES (?), (?)', + f'ALTER TABLE "iceberg"."s"."t" RENAME TO {old}', + f'ALTER TABLE {stage} RENAME TO "iceberg"."s"."t"', + f"DROP TABLE {old}", + ] + + +def test_a_failed_stage_insert_on_load_drops_the_stage_and_never_swaps(): + conn = _ScriptedConnection(fail_on={"VALUES (?": RuntimeError("NOT NULL violated")}) + with pytest.raises(RuntimeError, match="NOT NULL"): + _adapter(conn).load("s", "t", pa.table({"id": [1]})) + stage = _ref(conn.statements[1], "stage") + assert conn.statements[-1] == f"DROP TABLE {stage}" + assert not any("RENAME" in statement for statement in conn.statements) + + +def test_table_columns_reads_information_schema_with_lowercased_names(): + conn = _ScriptedConnection(rows_for={"information_schema.columns": [["id", "NO"], ["label", "YES"]]}) + columns = _adapter(conn).table_columns("MixedSchema", "MixedTable") + assert columns == [LiveColumn("id", False), LiveColumn("label", True)] + assert conn.statements == [ + 'SELECT column_name, is_nullable FROM "iceberg".information_schema.columns ' + "WHERE table_schema = ? AND table_name = ? ORDER BY ordinal_position" + ] + assert conn.params == [["mixedschema", "mixedtable"]] + + +def test_table_columns_of_a_missing_table_raises_lookup_error(): + with pytest.raises(LookupError, match=r"s\.t"): + _adapter(_ScriptedConnection()).table_columns("s", "t") + + +def test_append_statement_names_and_quotes_every_column(): + assert adapter_module._append_statement( + '"iceberg"."s"."t"', '"iceberg"."s"."st"', ["order", "Rate", "pct%"] + ) == ( + 'INSERT INTO "iceberg"."s"."t" ("order", "Rate", "pct%") ' + 'SELECT "order", "Rate", "pct%" FROM "iceberg"."s"."st"' + ) + + +def test_append_stages_then_inserts_once_then_drops_the_stage(): + conn = _ScriptedConnection() + _adapter(conn).append("s", "t", pa.table({"id": [1, 2]})) + s = conn.statements + stage = _ref(s[1], "stage") + assert s == [ + LOCATION_QUERY, + f'CREATE TABLE {stage} (LIKE "iceberg"."s"."t" INCLUDING PROPERTIES) ' + f"WITH (location = 's3://warehouse/s/{_name(stage)}')", + f'INSERT INTO {stage} ("id") VALUES (?), (?)', + f'INSERT INTO "iceberg"."s"."t" ("id") SELECT "id" FROM {stage}', + f"DROP TABLE {stage}", + ] + + +def test_append_of_zero_rows_creates_no_stage(): + conn = _ScriptedConnection() + _adapter(conn).append("s", "t", pa.table({"id": pa.array([], pa.int64())})) + assert conn.statements == [] + + +def test_a_failed_stage_insert_never_touches_the_target_and_drops_the_stage(): + conn = _ScriptedConnection(fail_on={"VALUES (?": RuntimeError("NOT NULL violated")}) + with pytest.raises(RuntimeError, match="NOT NULL"): + _adapter(conn).append("s", "t", pa.table({"id": [1]})) + stage = _ref(conn.statements[1], "stage") + assert not any(st.startswith('INSERT INTO "iceberg"."s"."t"') for st in conn.statements) + assert conn.statements[-1] == f"DROP TABLE {stage}" + + +def test_a_failed_stage_drop_after_the_commit_is_only_logged(caplog): + conn = _ScriptedConnection( + fail_on={'DROP TABLE "iceberg"."s"."__continuo_stage_': RuntimeError("catalog hiccup")} + ) + with caplog.at_level(logging.WARNING, logger="continuo_trino_adapter"): + _adapter(conn).append("s", "t", pa.table({"id": [1]})) + assert any(st.startswith('INSERT INTO "iceberg"."s"."t"') for st in conn.statements) + assert "best-effort cleanup of staging table" in caplog.text + + +_FX = [ + {"name": "rate_date", "type": "DATE", "nullable": False}, + {"name": "base", "type": "VARCHAR(3)", "nullable": False}, +] + + +def test_ensure_table_with_a_key_emits_the_same_ddl_as_without(): + keyed, plain = _ScriptedConnection(), _ScriptedConnection() + _adapter(keyed).ensure_table("s", "t", _FX, config={}, primary_key=("rate_date", "base")) + _adapter(plain).ensure_table("s", "t", _FX, config={}) + assert keyed.statements == plain.statements + + +def test_an_upsert_key_on_a_format_version_1_table_is_rejected_before_any_ddl(): + conn = _ScriptedConnection() + with pytest.raises(ValueError, match="format_version"): + _adapter(conn).ensure_table( + "s", "t", _FX, config={"format_version": 1}, primary_key=("rate_date", "base") + ) + assert conn.statements == [] + + +def test_format_version_1_without_a_key_is_still_allowed(): + conn = _ScriptedConnection() + _adapter(conn).ensure_table("s", "t", _FX, config={"format_version": 1}) + assert conn.statements[-1].endswith("WITH (format_version = 1)") + + +def test_merge_statement_overwrites_non_key_columns_and_inserts_the_rest(): + assert adapter_module._merge_statement( + '"iceberg"."s"."t"', '"iceberg"."s"."st"', ["order", "Rate", "pct%"], ["order", "Rate"] + ) == ( + 'MERGE INTO "iceberg"."s"."t" AS t USING "iceberg"."s"."st" AS s ' + 'ON (t."order" = s."order" AND t."Rate" = s."Rate") ' + 'WHEN MATCHED THEN UPDATE SET "pct%" = s."pct%" ' + 'WHEN NOT MATCHED THEN INSERT ("order", "Rate", "pct%") ' + 'VALUES (s."order", s."Rate", s."pct%")' + ) + + +def test_merge_statement_with_every_column_in_the_key_has_no_matched_branch(): + assert adapter_module._merge_statement( + '"iceberg"."s"."t"', '"iceberg"."s"."st"', ["a", "b"], ["a", "b"] + ) == ( + 'MERGE INTO "iceberg"."s"."t" AS t USING "iceberg"."s"."st" AS s ' + 'ON (t."a" = s."a" AND t."b" = s."b") ' + 'WHEN NOT MATCHED THEN INSERT ("a", "b") VALUES (s."a", s."b")' + ) + + +def test_upsert_stages_then_merges_once_then_drops_the_stage(): + conn = _ScriptedConnection() + _adapter(conn).upsert("s", "t", pa.table({"id": [1, 2], "v": ["a", "b"]}), ["id"]) + s = conn.statements + stage = _ref(s[1], "stage") + assert s[2] == f'INSERT INTO {stage} ("id", "v") VALUES (?, ?), (?, ?)' + assert s[3] == adapter_module._merge_statement('"iceberg"."s"."t"', stage, ["id", "v"], ["id"]) + assert s[4] == f"DROP TABLE {stage}" + assert len(s) == 5 + + +def test_upsert_of_zero_rows_creates_no_stage(): + conn = _ScriptedConnection() + _adapter(conn).upsert("s", "t", pa.table({"id": pa.array([], pa.int64())}), ["id"]) + assert conn.statements == [] + + +def test_upsert_with_a_repeated_key_runs_no_statement(): + conn = _ScriptedConnection() + with pytest.raises(DuplicateKeyError): + _adapter(conn).upsert("s", "t", pa.table({"id": [1, 1]}), ["id"]) + assert conn.statements == [] diff --git a/adapters/trino/tests/test_integration_trino_conformance.py b/adapters/trino/tests/test_integration_trino_conformance.py new file mode 100644 index 0000000..be62be8 --- /dev/null +++ b/adapters/trino/tests/test_integration_trino_conformance.py @@ -0,0 +1,77 @@ +"""The engine-contract conformance suite against Trino + Iceberg (tests/smoke/trino-stack). + +Run the stack first: + docker compose -f tests/smoke/trino-stack/docker-compose.yml up -d --wait +""" +import os +import uuid + +import pytest + +import trino + +from continuo_engine_contract.testing.conformance import ( + ARROW, + COLUMNS, + AppendConformance, + ConformanceBase, + IntrospectionConformance, + UpsertConformance, + bulk_rows, + table_of, +) +from continuo_trino_adapter.adapter import TrinoAdapter + +pytestmark = pytest.mark.integration + + +@pytest.fixture +def adapter(): + conn = trino.dbapi.connect( + host="localhost", + port=int(os.environ.get("VR_IT_TRINO_PORT", "18080")), + user="continuo", + catalog="iceberg", + http_scheme="http", + ) + connected = TrinoAdapter(conn, "iceberg") + yield connected + connected.close() + + +@pytest.fixture +def schema(adapter): + name = f"conf_{uuid.uuid4().hex[:10]}" + yield name + adapter.drop_schema(name) + + +class TestTrinoIntrospection(IntrospectionConformance): + pass + + +class TestTrinoAppend(AppendConformance): + pass + + +def _stage_tables(adapter, schema): + rows = adapter._execute(f'SHOW TABLES FROM "iceberg"."{schema}"') + return [row[0] for row in rows if row[0].startswith("__continuo_stage_")] + + +def test_a_failed_append_leaves_no_stage_table(adapter, schema): + table = ConformanceBase().new_table(adapter, schema, COLUMNS) + with pytest.raises(Exception): + adapter.append(schema, table, table_of(ARROW, bulk_rows(1))) + assert _stage_tables(adapter, schema) == [] + + +class TestTrinoUpsert(UpsertConformance): + pass + + +def test_a_failed_upsert_leaves_no_stage_table(adapter, schema): + table = ConformanceBase().new_table(adapter, schema, COLUMNS, primary_key=["id"]) + with pytest.raises(Exception): + adapter.upsert(schema, table, table_of(ARROW, bulk_rows(1)), ["id"]) + assert _stage_tables(adapter, schema) == [] diff --git a/adapters/trino/tests/test_integration_trino_merge_probe.py b/adapters/trino/tests/test_integration_trino_merge_probe.py new file mode 100644 index 0000000..e1ec130 --- /dev/null +++ b/adapters/trino/tests/test_integration_trino_merge_probe.py @@ -0,0 +1,98 @@ +"""Probe (spec section 8): Trino 483 MERGE on Iceberg, as the trino adapter's _upsert runs it. + +The statement shape is the one _merge_statement emits: MERGE INTO target USING a +stage table, key equality in ON, every other column in WHEN MATCHED, and +WHEN NOT MATCHED alone when every column is in the key. Pinned against the +compose stack's Iceberg REST catalog. +""" +import os +import uuid + +import pytest + +import trino + +from continuo_trino_adapter.adapter import TrinoAdapter + +pytestmark = pytest.mark.integration + + +def _adapter() -> TrinoAdapter: + conn = trino.dbapi.connect( + host="localhost", + port=int(os.environ.get("VR_IT_TRINO_PORT", "18080")), + user="continuo", + catalog="iceberg", + http_scheme="http", + ) + return TrinoAdapter(conn, "iceberg") + + +@pytest.fixture +def probe(): + adapter = _adapter() + schema = f"merge_probe_{uuid.uuid4().hex[:8]}" + adapter._execute(f'CREATE SCHEMA "iceberg"."{schema}"') + yield adapter, schema + adapter._execute(f'DROP SCHEMA IF EXISTS "iceberg"."{schema}" CASCADE') + adapter.close() + + +def _target_and_stage(adapter, schema, target_properties=""): + target = f'"iceberg"."{schema}"."target"' + stage = f'"iceberg"."{schema}"."stage"' + adapter._execute(f'CREATE TABLE {target} ("id" BIGINT NOT NULL, "v" VARCHAR){target_properties}') + adapter._execute(f"INSERT INTO {target} VALUES (1, 'a'), (2, 'b')") + adapter._execute(f'CREATE TABLE {stage} ("id" BIGINT NOT NULL, "v" VARCHAR)') + return target, stage + + +def _merge(target, stage): + return ( + f'MERGE INTO {target} AS t USING {stage} AS s ON (t."id" = s."id") ' + 'WHEN MATCHED THEN UPDATE SET "v" = s."v" ' + 'WHEN NOT MATCHED THEN INSERT ("id", "v") VALUES (s."id", s."v")' + ) + + +def _rows(adapter, target): + return [tuple(row) for row in adapter._execute(f'SELECT "id", "v" FROM {target} ORDER BY "id"')] + + +def test_merge_updates_and_inserts_in_one_statement(probe): + adapter, schema = probe + target, stage = _target_and_stage(adapter, schema) + adapter._execute(f"INSERT INTO {stage} VALUES (2, 'B'), (3, 'c')") + adapter._execute(_merge(target, stage)) + assert _rows(adapter, target) == [(1, "a"), (2, "B"), (3, "c")] + + +def test_merge_with_only_when_not_matched_inserts_only_new_keys(probe): + adapter, schema = probe + target, stage = _target_and_stage(adapter, schema) + adapter._execute(f"INSERT INTO {stage} VALUES (2, 'b'), (3, 'c')") + insert_only = ( + f'MERGE INTO {target} AS t USING {stage} AS s ON (t."id" = s."id" AND t."v" = s."v") ' + 'WHEN NOT MATCHED THEN INSERT ("id", "v") VALUES (s."id", s."v")' + ) + adapter._execute(insert_only) + adapter._execute(insert_only) + assert _rows(adapter, target) == [(1, "a"), (2, "b"), (3, "c")] + + +def test_merge_on_format_version_1_fails_and_leaves_the_target_unchanged(probe): + adapter, schema = probe + target, stage = _target_and_stage(adapter, schema, " WITH (format_version = 1)") + adapter._execute(f"INSERT INTO {stage} VALUES (2, 'B')") + with pytest.raises(trino.exceptions.TrinoQueryError, match="format version"): + adapter._execute(_merge(target, stage)) + assert _rows(adapter, target) == [(1, "a"), (2, "b")] + + +def test_two_stage_rows_matching_one_target_row_fail(probe): + adapter, schema = probe + target, stage = _target_and_stage(adapter, schema) + adapter._execute(f"INSERT INTO {stage} VALUES (2, 'x'), (2, 'y')") + with pytest.raises(trino.exceptions.TrinoQueryError): + adapter._execute(_merge(target, stage)) + assert _rows(adapter, target) == [(1, "a"), (2, "b")] diff --git a/docs/boundary-contract.md b/docs/boundary-contract.md index 0d175cc..2bfe6af 100644 --- a/docs/boundary-contract.md +++ b/docs/boundary-contract.md @@ -138,7 +138,9 @@ parse time. - **trino (Iceberg)** → `partitioning` and `sorted_by` (each a non-empty list of non-empty strings — column names or Iceberg partition transforms like `day(event_ts)`), and `format` (one of - `PARQUET`/`ORC`/`AVRO`, case-insensitive). + `PARQUET`/`ORC`/`AVRO`, case-insensitive). An `upsert` node cannot set + `format_version: 1`: Trino's `MERGE` needs Iceberg format version 2 (the + default), and `ensure_table` rejects the combination before any DDL. - **duckdb (DuckLake)** → `partitioned_by` and `sorted_by`. `partitioned_by` is a non-empty list whose entries are a declared column name (identity partitioning) or `{column, transform, buckets}` with