diff --git a/Cargo.lock b/Cargo.lock index 295f6afbb76ac..fef4e87c22b7f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4728,7 +4728,6 @@ dependencies = [ [[package]] name = "iceberg" version = "0.9.0" -source = "git+https://github.com/MaterializeInc/iceberg-rust.git?rev=dedd9231ee88ee979b648e14792878b40e74c20a#dedd9231ee88ee979b648e14792878b40e74c20a" dependencies = [ "anyhow", "apache-avro", @@ -4782,7 +4781,6 @@ dependencies = [ [[package]] name = "iceberg-catalog-rest" version = "0.9.0" -source = "git+https://github.com/MaterializeInc/iceberg-rust.git?rev=dedd9231ee88ee979b648e14792878b40e74c20a#dedd9231ee88ee979b648e14792878b40e74c20a" dependencies = [ "async-trait", "chrono", @@ -4802,7 +4800,6 @@ dependencies = [ [[package]] name = "iceberg-storage-opendal" version = "0.9.0" -source = "git+https://github.com/MaterializeInc/iceberg-rust.git?rev=dedd9231ee88ee979b648e14792878b40e74c20a#dedd9231ee88ee979b648e14792878b40e74c20a" dependencies = [ "anyhow", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index 080f0dfaa9227..ffd17d0d0e4c4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -668,9 +668,10 @@ async-compression = { git = "https://github.com/MaterializeInc/async-compression # Custom iceberg features for mz # All changes should go to the `mz_v0.9.0` branch. -iceberg = { git = "https://github.com/MaterializeInc/iceberg-rust.git", rev = "dedd9231ee88ee979b648e14792878b40e74c20a" } -iceberg-catalog-rest = { git = "https://github.com/MaterializeInc/iceberg-rust.git", rev = "dedd9231ee88ee979b648e14792878b40e74c20a" } -iceberg-storage-opendal = { git = "https://github.com/MaterializeInc/iceberg-rust.git", rev = "dedd9231ee88ee979b648e14792878b40e74c20a" } +# TODO: point back at a MaterializeInc/iceberg-rust rev before merging. +iceberg = { path = "../iceberg-rust/crates/iceberg" } +iceberg-catalog-rest = { path = "../iceberg-rust/crates/catalog/rest" } +iceberg-storage-opendal = { path = "../iceberg-rust/crates/storage/opendal" } # Custom duckdb crate to support mz needs # All changes should go to the `mz_changes` branch. diff --git a/src/storage/src/sink/iceberg.rs b/src/storage/src/sink/iceberg.rs index aed8ad817272f..a486956f8281c 100644 --- a/src/storage/src/sink/iceberg.rs +++ b/src/storage/src/sink/iceberg.rs @@ -105,7 +105,7 @@ use iceberg::spec::{ }; use iceberg::spec::{Schema, SchemaRef}; use iceberg::table::Table; -use iceberg::transaction::{ApplyTransactionAction, Transaction}; +use iceberg::transaction::{RowDeltaAction, TransactionAction}; use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder; use iceberg::writer::base_writer::equality_delete_writer::{ EqualityDeleteFileWriterBuilder, EqualityDeleteWriterConfig, @@ -120,7 +120,7 @@ use iceberg::writer::file_writer::location_generator::{ }; use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder; use iceberg::writer::{IcebergWriter, IcebergWriterBuilder}; -use iceberg::{Catalog, NamespaceIdent, TableCreation, TableIdent}; +use iceberg::{Catalog, NamespaceIdent, TableCommit, TableCreation, TableIdent}; use itertools::Itertools; use mz_arrow_util::builder::{ARROW_EXTENSION_NAME_KEY, ArrowBuilder}; use mz_interchange::avro::DiffPair; @@ -154,7 +154,7 @@ use timely::dataflow::channels::pact::{Exchange, Pipeline}; use timely::dataflow::operators::vec::{Broadcast, Map, ToStream}; use timely::dataflow::operators::{CapabilitySet, Concatenate}; use timely::progress::{Antichain, Timestamp as _}; -use tracing::debug; +use tracing::{debug, info, warn}; use crate::healthcheck::{HealthStatusMessage, HealthStatusUpdate, StatusNamespace}; use crate::metrics::sink::iceberg::IcebergSinkMetrics; @@ -777,11 +777,63 @@ async fn reload_table( } } +/// After attempting a commit, we expect to know whether it succeeded. +enum CommitState { + /// The Catalog told us whether the commit succeeded and gave us an updated table + /// (which we can safely use for another attempt). + Known(Table), + /// We couldn't get an answer from the Catalog. + /// We don't have an updated table. We don't know if the commit succeeded or failed. + /// Before making any more attempts, we must reload the table + /// and check for ourselves whether the commit is there. + Unresolved(Table), +} + +/// Check whether the most recent Materialize snapshot on the table belongs to another writer. +fn check_fencing( + last: Option<&(Antichain, u64)>, + sink_version: u64, + frontier: &Antichain, + conn_table: &str, +) -> Result<(), anyhow::Error> { + let Some((last_frontier, last_version)) = last else { + return Ok(()); + }; + if *last_version > sink_version { + anyhow::bail!( + "Iceberg table '{}' has been modified by another writer \ + with version {}. Current sink version: {}. \ + Frontiers may be out of sync, aborting to avoid data loss.", + conn_table, + last_version, + sink_version, + ); + } + + // Check if someone has already written this batch (or an even later batch). + // If it was us, we should already know. + // (Either we received a success response or we reloaded the table and checked.) + // So it must be another writer. + if PartialOrder::less_equal(frontier, last_frontier) { + anyhow::bail!( + "Iceberg table '{}' has been modified by another writer. \ + Current frontier: {:?}, last frontier: {:?}.", + conn_table, + frontier, + last_frontier, + ); + } + Ok(()) +} + /// Attempt a single commit of a batch of data files to an Iceberg table. -/// On conflict or failure, reloads the table and returns a retryable error. -/// On success, returns the updated table state. +/// +/// If a previous attempt left the outcome unknown, first reload the table to establish whether +/// that attempt succeeded. Try again only once we're sure it failed. +/// On conflict, reload the table and return a retryable error. +/// On success, return the updated table state. async fn try_commit_batch( - mut table: Table, + state: CommitState, snapshot_properties: Vec<(String, String)>, data_files: Vec, delete_files: Vec, @@ -793,10 +845,71 @@ async fn try_commit_batch( batch_lower: &Antichain, batch_upper: &Antichain, metrics: &IcebergSinkMetrics, -) -> (Table, RetryResult<(), anyhow::Error>) { - let tx = Transaction::new(&table); - let mut action = tx - .row_delta() +) -> (CommitState, RetryResult<(), anyhow::Error>) { + let table = match state { + CommitState::Known(table) => table, + CommitState::Unresolved(stale) => { + let reloaded = match reload_table( + catalog, + conn_namespace.to_string(), + conn_table.to_string(), + stale.clone(), + ) + .await + { + Ok(reloaded) => reloaded, + // We can't proceed until we've checked the table. + Err(e) => { + return ( + CommitState::Unresolved(stale), + RetryResult::RetryableErr(anyhow!(e)), + ); + } + }; + + let mut snapshots: Vec<_> = reloaded.metadata().snapshots().cloned().collect(); + match retrieve_upper_from_snapshots(&mut snapshots) { + // Our own commit for this batch is already on the table. + // It landed and we never saw the response. + Ok(Some((last_frontier, last_version))) + if last_version == sink_version && last_frontier == *frontier => + { + info!( + namespace = %conn_namespace, + table = %conn_table, + lower = %batch_lower.pretty(), + upper = %batch_upper.pretty(), + "iceberg commit with lost response found on reload, treating as success" + ); + return (CommitState::Known(reloaded), RetryResult::Ok(())); + } + // Our commit isn't there. + // We can retry only if no other writer has taken over the table in the meantime. + Ok(last) => { + if let Err(e) = check_fencing(last.as_ref(), sink_version, frontier, conn_table) + { + return (CommitState::Known(reloaded), RetryResult::FatalErr(e)); + } + info!( + namespace = %conn_namespace, + table = %conn_table, + lower = %batch_lower.pretty(), + upper = %batch_upper.pretty(), + "iceberg commit with lost response not found on reload, committing again" + ); + reloaded + } + Err(e) => { + return ( + CommitState::Unresolved(stale), + RetryResult::RetryableErr(anyhow!(e)), + ); + } + } + } + }; + + let mut action = RowDeltaAction::new() .set_snapshot_properties(snapshot_properties.into_iter().collect()) .with_check_duplicate(false); @@ -806,13 +919,18 @@ async fn try_commit_batch( .add_delete_files(delete_files); } - let tx = match action - .apply(tx) - .context("Failed to apply data file addition to iceberg table transaction") - { - Ok(tx) => tx, + // Build the commit's metadata updates and requirements against our own view of + // the table and send them to the catalog directly, instead of going through + // `Transaction::commit`. That path reloads the table and rebases the commit onto + // whatever it finds, retrying conflicts internally, so it would silently commit + // over another writer. Generated this way, the requirements pin the table state + // we know, and any interleaved write surfaces as a commit conflict below, where + // `check_fencing` decides whether retrying is safe. + let mut action_commit = match Arc::new(action).commit(&table).await { + Ok(action_commit) => action_commit, Err(e) => { - match reload_table( + // Nothing was sent to the catalog, so this commit definitely didn't happen. + let reloaded = match reload_table( catalog, conn_namespace.to_string(), conn_table.to_string(), @@ -820,26 +938,32 @@ async fn try_commit_batch( ) .await { - Ok(reloaded) => table = reloaded, + Ok(reloaded) => reloaded, Err(reload_err) => { - return (table, RetryResult::RetryableErr(anyhow!(reload_err))); + return ( + CommitState::Known(table), + RetryResult::RetryableErr(anyhow!(reload_err)), + ); } - } + }; return ( - table, - RetryResult::RetryableErr(anyhow!( - "Failed to apply data file addition to iceberg table transaction: {}", - e - )), + CommitState::Known(reloaded), + RetryResult::RetryableErr(anyhow!("Failed to build iceberg table commit: {}", e)), ); } }; - let new_table = tx.commit(catalog).await; + let table_commit = TableCommit::builder() + .ident(table.identifier().clone()) + .updates(action_commit.take_updates()) + .requirements(action_commit.take_requirements()) + .build(); + + let new_table = catalog.update_table(table_commit).await; match new_table { Err(e) if matches!(e.kind(), ErrorKind::CatalogCommitConflicts) => { metrics.commit_conflicts.inc(); - match reload_table( + let table = match reload_table( catalog, conn_namespace.to_string(), conn_table.to_string(), @@ -847,9 +971,12 @@ async fn try_commit_batch( ) .await { - Ok(reloaded) => table = reloaded, + Ok(reloaded) => reloaded, Err(e) => { - return (table, RetryResult::RetryableErr(anyhow!(e))); + return ( + CommitState::Known(table), + RetryResult::RetryableErr(anyhow!(e)), + ); } }; @@ -858,41 +985,20 @@ async fn try_commit_batch( let last = match last { Ok(val) => val, Err(e) => { - return (table, RetryResult::RetryableErr(anyhow!(e))); + return ( + CommitState::Known(table), + RetryResult::RetryableErr(anyhow!(e)), + ); } }; // Check if another writer has advanced the frontier beyond ours (fencing check) - if let Some((last_frontier, last_version)) = last { - if last_version > sink_version { - return ( - table, - RetryResult::FatalErr(anyhow!( - "Iceberg table '{}' has been modified by another writer \ - with version {}. Current sink version: {}. \ - Frontiers may be out of sync, aborting to avoid data loss.", - conn_table, - last_version, - sink_version, - )), - ); - } - if PartialOrder::less_equal(frontier, &last_frontier) { - return ( - table, - RetryResult::FatalErr(anyhow!( - "Iceberg table '{}' has been modified by another writer. \ - Current frontier: {:?}, last frontier: {:?}.", - conn_table, - frontier, - last_frontier, - )), - ); - } + if let Err(e) = check_fencing(last.as_ref(), sink_version, frontier, conn_table) { + return (CommitState::Known(table), RetryResult::FatalErr(e)); } ( - table, + CommitState::Known(table), RetryResult::RetryableErr(anyhow!( "Commit conflict detected when committing batch [{}, {}) \ to Iceberg table '{}.{}'. Retrying...", @@ -903,11 +1009,29 @@ async fn try_commit_batch( )), ) } + // The catalog may have applied this commit before the success response was lost. + // Mark the outcome as unknown so the next attempt starts by reading the table to see what happened. + Err(e) if matches!(e.kind(), ErrorKind::Unexpected) => { + metrics.commit_failures.inc(); + warn!( + namespace = %conn_namespace, + table = %conn_table, + lower = %batch_lower.pretty(), + upper = %batch_upper.pretty(), + error = %e, + "iceberg commit outcome unknown, will reload the table to check before retrying" + ); + ( + CommitState::Unresolved(table), + RetryResult::RetryableErr(anyhow!(e)), + ) + } + // All other errors are definite: retrying will not change the outcome. Err(e) => { metrics.commit_failures.inc(); - (table, RetryResult::RetryableErr(anyhow!(e))) + (CommitState::Known(table), RetryResult::FatalErr(anyhow!(e))) } - Ok(new_table) => (new_table, RetryResult::Ok(())), + Ok(new_table) => (CommitState::Known(new_table), RetryResult::Ok(())), } } @@ -2538,9 +2662,9 @@ fn commit_to_iceberg<'scope>( ("mz-sink-version".to_string(), sink_version.to_string()), ]; - let (table_state, commit_result) = Retry::default() + let (commit_state, commit_result) = Retry::default() .max_tries(5) - .retry_async_with_state(table, |_, table| { + .retry_async_with_state(CommitState::Known(table), |_, commit_state| { let snapshot_properties = snapshot_properties.clone(); let data_files = data_files.clone(); let delete_files = delete_files.clone(); @@ -2553,7 +2677,7 @@ fn commit_to_iceberg<'scope>( let batch_upper = batch.1.clone(); async move { try_commit_batch( - table, + commit_state, snapshot_properties, data_files, delete_files, @@ -2576,12 +2700,21 @@ fn commit_to_iceberg<'scope>( connection.namespace, connection.table ) }); - table = table_state; let duration = instant.elapsed(); metrics .commit_duration_seconds .observe(duration.as_secs_f64()); commit_result?; + table = match commit_state { + CommitState::Known(table) => table, + CommitState::Unresolved(_) => { + anyhow::bail!( + "invariant: unresolved commit state but the commit succeeded, Iceberg table '{}.{}'", + connection.namespace, + connection.table + ) + } + }; debug!( ?sink_id, diff --git a/test/iceberg/fenced-writer-setup.td b/test/iceberg/fenced-writer-setup.td new file mode 100644 index 0000000000000..1c2e4f852b791 --- /dev/null +++ b/test/iceberg/fenced-writer-setup.td @@ -0,0 +1,40 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Use of this software is governed by the Business Source License +# included in the LICENSE file at the root of this repository. +# +# As of the Change Date specified in that file, in accordance with +# the Business Source License, use of this software will be governed +# by the Apache License, Version 2.0. + +> CREATE SECRET access_key_secret AS '${arg.s3-access-key}' + +> CREATE CONNECTION aws_conn TO AWS ( + ACCESS KEY ID = 'tduser', + SECRET ACCESS KEY = SECRET access_key_secret, + ENDPOINT = '${arg.aws-endpoint}', + REGION = 'us-east-1' + ); + +> CREATE CONNECTION polaris TO ICEBERG CATALOG ( + CATALOG TYPE = 'REST', + URL = 'http://polaris:8181/api/catalog', + CREDENTIAL = 'root:root', + WAREHOUSE = 'default_catalog', + SCOPE = 'PRINCIPAL_ROLE:ALL' + ); + +> CREATE TABLE fence_src (id INT, val TEXT); + +> INSERT INTO fence_src VALUES (1, 'a'), (2, 'b'), (3, 'c'); + +> CREATE SINK fence_sink + FROM fence_src + INTO ICEBERG CATALOG CONNECTION polaris ( + NAMESPACE 'default_namespace', + TABLE 'fence_table' + ) + USING AWS CONNECTION aws_conn + KEY (id) NOT ENFORCED + MODE UPSERT + WITH (COMMIT INTERVAL '2s'); diff --git a/test/iceberg/idempotent-retry-setup.td b/test/iceberg/idempotent-retry-setup.td new file mode 100644 index 0000000000000..8ad4c2279251c --- /dev/null +++ b/test/iceberg/idempotent-retry-setup.td @@ -0,0 +1,40 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Use of this software is governed by the Business Source License +# included in the LICENSE file at the root of this repository. +# +# As of the Change Date specified in that file, in accordance with +# the Business Source License, use of this software will be governed +# by the Apache License, Version 2.0. + +> CREATE SECRET access_key_secret AS '${arg.s3-access-key}' + +> CREATE CONNECTION aws_conn TO AWS ( + ACCESS KEY ID = 'tduser', + SECRET ACCESS KEY = SECRET access_key_secret, + ENDPOINT = '${arg.aws-endpoint}', + REGION = 'us-east-1' + ); + +> CREATE CONNECTION polaris TO ICEBERG CATALOG ( + CATALOG TYPE = 'REST', + URL = 'http://polaris-proxy:8181/api/catalog', + CREDENTIAL = 'root:root', + WAREHOUSE = 'default_catalog', + SCOPE = 'PRINCIPAL_ROLE:ALL' + ); + +> CREATE TABLE retry_src (id INT, val TEXT); + +> INSERT INTO retry_src VALUES (1, 'a'), (2, 'b'), (3, 'c'); + +> CREATE SINK retry_sink + FROM retry_src + INTO ICEBERG CATALOG CONNECTION polaris ( + NAMESPACE 'default_namespace', + TABLE 'retry_table' + ) + USING AWS CONNECTION aws_conn + KEY (id) NOT ENFORCED + MODE UPSERT + WITH (COMMIT INTERVAL '2s'); diff --git a/test/iceberg/idempotent-retry-verify.td b/test/iceberg/idempotent-retry-verify.td new file mode 100644 index 0000000000000..c54f765eca88f --- /dev/null +++ b/test/iceberg/idempotent-retry-verify.td @@ -0,0 +1,16 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Use of this software is governed by the Business Source License +# included in the LICENSE file at the root of this repository. +# +# As of the Change Date specified in that file, in accordance with +# the Business Source License, use of this software will be governed +# by the Apache License, Version 2.0. + +$ duckdb-execute name=iceberg +CREATE SECRET s3_secret (TYPE S3, KEY_ID 'tduser', SECRET '${arg.s3-access-key}', ENDPOINT '${arg.aws-endpoint}', URL_STYLE 'path', USE_SSL false, REGION 'minio'); +SET unsafe_enable_version_guessing = true; + +$ duckdb-query name=iceberg +SELECT COUNT(*) FROM iceberg_scan('s3://test-bucket/default_namespace/retry_table') +13 diff --git a/test/iceberg/mzcompose.py b/test/iceberg/mzcompose.py index d56b1fc512f23..0e7ed777fbb2d 100644 --- a/test/iceberg/mzcompose.py +++ b/test/iceberg/mzcompose.py @@ -13,11 +13,15 @@ import time import urllib.error import urllib.request +from collections.abc import Callable -from materialize.mzcompose.composition import Composition, Service +from materialize.mzcompose.composition import Composition +from materialize.mzcompose.composition import Service as UpService from materialize.mzcompose.helpers.iceberg import ( + get_polaris_access_token, setup_polaris_for_iceberg, ) +from materialize.mzcompose.service import Service from materialize.mzcompose.services.materialized import Materialized from materialize.mzcompose.services.minio import Mc, Minio from materialize.mzcompose.services.mz import Mz @@ -31,6 +35,21 @@ Minio(), PolarisBootstrap(), Polaris(), + Service( + "polaris-proxy", + { + "image": "python:3.11-slim", + "command": ["python", "-u", "polaris_proxy.py"], + "working_dir": "/workdir", + "volumes": [".:/workdir"], + "ports": [8181], + "environment": [ + "UPSTREAM_HOST=polaris", + "UPSTREAM_PORT=8181", + "PROXY_PORT=8181", + ], + }, + ), Materialized( depends_on=["minio"], sanity_restart=False, @@ -49,8 +68,8 @@ def _setup(c: Composition) -> str: c.up( "postgres", "materialized", - Service("polaris-bootstrap", idle=True), - Service("polaris", idle=True), + UpService("polaris-bootstrap", idle=True), + UpService("polaris", idle=True), ) _, key = setup_polaris_for_iceberg(c) return key @@ -343,6 +362,238 @@ def modify_table_loop() -> None: ) +def workflow_idempotent_retry(c: Composition) -> None: + """Regression test: dropping a single catalog commit response must not + fence the sink off or cause duplicate row commits.""" + key = _setup(c) + c.invoke("up", "--detach", "--wait", "--no-recreate", "polaris-proxy") + + c.run_testdrive_files( + f"--var=s3-access-key={key}", + "--var=aws-endpoint=minio:9000", + "idempotent-retry-setup.td", + ) + + proxy_base = f"http://localhost:{c.port('polaris-proxy', 8181)}" + + def proxy_post(path: str) -> None: + with urllib.request.urlopen( + urllib.request.Request(f"{proxy_base}{path}", data=b"", method="POST") + ) as resp: + resp.read() + + def proxy_status() -> dict[str, int]: + with urllib.request.urlopen(f"{proxy_base}/__control/status") as resp: + return json.loads(resp.read()) + + def await_condition(what: str, timeout: float, check: Callable[[], bool]) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + if check(): + return + time.sleep(0.5) + raise AssertionError(f"timed out waiting for {what}") + + # Arm the drop only once the sink has a commit through the proxy, so the + # dropped commit is a steady-state one rather than table bootstrap. + await_condition( + "first sink commit", + timeout=60, + check=lambda: proxy_status()["commits_ok"] >= 1, + ) + proxy_post("/__control/drop_next_commit") + + for i in range(10): + c.sql(f"INSERT INTO retry_src VALUES ({i + 4}, 'row_{i + 4}')") + time.sleep(1) + + await_condition( + "dropped commit response", + timeout=60, + check=lambda: proxy_status()["commits_dropped"] >= 1, + ) + + def messages_committed() -> int: + rows = c.sql_query( + "SELECT COALESCE(SUM(messages_committed), 0) " + "FROM mz_internal.mz_sink_statistics st " + "JOIN mz_sinks s ON st.id = s.id " + "WHERE s.name = 'retry_sink'" + ) + return int(rows[0][0]) + + # 3 initial rows + 10 inserted rows, all committed despite the dropped response. + await_condition( + "all 13 rows committed", + timeout=120, + check=lambda: messages_committed() >= 13, + ) + + status_rows = c.sql_query( + "SELECT s.status, COALESCE(s.error, '') " + "FROM mz_internal.mz_sink_statuses s " + "JOIN mz_sinks ON s.id = mz_sinks.id " + "WHERE mz_sinks.name = 'retry_sink'" + ) + assert status_rows, "retry_sink not found in mz_sink_statuses" + status, error = status_rows[0] + assert status == "running", f"retry_sink is {status!r} (error={error!r})" + + c.run_testdrive_files( + f"--var=s3-access-key={key}", + "--var=aws-endpoint=minio:9000", + "idempotent-retry-verify.td", + ) + + +def workflow_fenced_writer(c: Composition) -> None: + """Regression test: once a snapshot with a newer mz-sink-version is on the + table, the running sink must stop committing. + + Currently fails: iceberg-rust's commit path reloads the table and rebases + onto the newest snapshot before every attempt, so the fenced sink never + sees a conflict and commits right over the newer writer. Its snapshot then + becomes the latest one, so even the startup fencing check of a later + restart no longer notices the newer writer.""" + key = _setup(c) + + c.run_testdrive_files( + f"--var=s3-access-key={key}", + "--var=aws-endpoint=minio:9000", + "fenced-writer-setup.td", + ) + + token = get_polaris_access_token(c) + table_url = ( + f"http://localhost:{c.port('polaris', 8181)}" + "/api/catalog/v1/default_catalog/namespaces/default_namespace/tables/fence_table" + ) + headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"} + + def load_metadata() -> dict | None: + req = urllib.request.Request(table_url, headers=headers) + try: + with urllib.request.urlopen(req) as resp: + return json.loads(resp.read())["metadata"] + except urllib.error.HTTPError as e: + if e.code == 404: + return None + raise + + def await_condition(what: str, timeout: float, check: Callable[[], bool]) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + if check(): + return + time.sleep(0.5) + raise AssertionError(f"timed out waiting for {what}") + + def first_commit_done() -> bool: + meta = load_metadata() + return meta is not None and meta.get("current-snapshot-id") not in (None, -1) + + await_condition("first sink commit", timeout=60, check=first_commit_done) + + def forge_newer_writer_snapshot() -> int: + """Commit a snapshot that claims mz-sink-version 999, as a sink with a + newer version would. Reuses the current snapshot's manifest list so the + table stays readable without writing new files. Returns the forged + snapshot's sequence number.""" + meta = load_metadata() + assert meta is not None + current_id = meta["refs"]["main"]["snapshot-id"] + current = next(s for s in meta["snapshots"] if s["snapshot-id"] == current_id) + forged_seq = meta["last-sequence-number"] + 1 + body = json.dumps( + { + "requirements": [ + { + "type": "assert-ref-snapshot-id", + "ref": "main", + "snapshot-id": current_id, + } + ], + "updates": [ + { + "action": "add-snapshot", + "snapshot": { + "snapshot-id": current_id + 1, + "parent-snapshot-id": current_id, + "sequence-number": forged_seq, + "timestamp-ms": int(time.time() * 1000), + "manifest-list": current["manifest-list"], + "schema-id": meta["current-schema-id"], + "summary": { + "operation": "append", + "mz-sink-id": "u0", + "mz-frontier": current["summary"]["mz-frontier"], + "mz-sink-version": "999", + }, + }, + }, + { + "action": "set-snapshot-ref", + "ref-name": "main", + "type": "branch", + "snapshot-id": current_id + 1, + }, + ], + } + ).encode() + req = urllib.request.Request( + table_url, data=body, headers=headers, method="POST" + ) + with urllib.request.urlopen(req) as resp: + resp.read() + return forged_seq + + # The sink may commit between reading the metadata and posting the forged + # snapshot, failing our requirement. Retry on conflict with fresh metadata. + for attempt in range(5): + try: + forged_seq = forge_newer_writer_snapshot() + break + except urllib.error.HTTPError as e: + if e.code != 409 or attempt == 4: + raise + else: + raise AssertionError("unreachable") + + # Give the fenced sink something to commit. + for i in range(5): + c.sql(f"INSERT INTO fence_src VALUES ({i + 4}, 'row_{i + 4}')") + time.sleep(1) + + deadline = time.time() + 90 + while True: + meta = load_metadata() + assert meta is not None + for snapshot in meta["snapshots"]: + if ( + snapshot["sequence-number"] > forged_seq + and snapshot["summary"].get("mz-sink-version") != "999" + ): + raise AssertionError( + "fenced sink committed past the newer writer: " + f"snapshot {snapshot['snapshot-id']} " + f"summary {snapshot['summary']}" + ) + + status_rows = c.sql_query( + "SELECT s.status, COALESCE(s.error, '') " + "FROM mz_internal.mz_sink_statuses s " + "JOIN mz_sinks ON s.id = mz_sinks.id " + "WHERE mz_sinks.name = 'fence_sink'" + ) + assert status_rows, "fence_sink not found in mz_sink_statuses" + status, error = status_rows[0] + if status != "running" and ("another writer" in error or "Fenced off" in error): + return + if time.time() > deadline: + raise AssertionError(f"sink not fenced: status={status!r} error={error!r}") + time.sleep(1) + + def workflow_large_upsert_batch(c: Composition) -> None: """Regression test for database-issues#11326: DeltaWriter seen_rows eviction caused equality deletes within the same snapshot, which diff --git a/test/iceberg/polaris_proxy.py b/test/iceberg/polaris_proxy.py new file mode 100644 index 0000000000000..1ce5e27843e10 --- /dev/null +++ b/test/iceberg/polaris_proxy.py @@ -0,0 +1,162 @@ +# Copyright Materialize, Inc. and contributors. All rights reserved. +# +# Use of this software is governed by the Business Source License +# included in the LICENSE file at the root of this repository. +# +# As of the Change Date specified in that file, in accordance with +# the Business Source License, use of this software will be governed +# by the Apache License, Version 2.0. + +import http.server +import json +import os +import socketserver +import sys +import threading +import time +import urllib.error +import urllib.request + +UPSTREAM_HOST = os.environ.get("UPSTREAM_HOST", "polaris") +UPSTREAM_PORT = int(os.environ.get("UPSTREAM_PORT", "8181")) + +_lock = threading.Lock() +_drop_armed = False +# Successful (2xx upstream) table commits seen, including ones whose response we dropped. +_commits_ok = 0 +_commits_dropped = 0 + + +def _log(msg: str) -> None: + ts = time.strftime("%H:%M:%S") + sys.stderr.write(f"[polaris-proxy {ts}] {msg}\n") + sys.stderr.flush() + + +class Handler(http.server.BaseHTTPRequestHandler): + def log_message(self, fmt, *args): # type: ignore[override] + return + + def _handle_control(self) -> bool: + global _drop_armed + if not self.path.startswith("/__control/"): + return False + if self.path == "/__control/drop_next_commit" and self.command == "POST": + with _lock: + _drop_armed = True + self.send_response(200) + self.send_header("Content-Length", "6") + self.end_headers() + self.wfile.write(b"armed\n") + return True + if self.path == "/__control/status" and self.command == "GET": + with _lock: + payload = json.dumps( + {"commits_ok": _commits_ok, "commits_dropped": _commits_dropped} + ).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + return True + self.send_response(404) + self.send_header("Content-Length", "0") + self.end_headers() + return True + + def _is_table_commit_endpoint(self) -> bool: + parts = self.path.split("?", 1)[0].split("/") + if "tables" not in parts: + return False + idx = parts.index("tables") + return idx < len(parts) - 1 and parts[idx + 1] != "" + + def _forward(self, body: bytes | None) -> None: + global _drop_armed, _commits_ok, _commits_dropped + upstream_url = f"http://{UPSTREAM_HOST}:{UPSTREAM_PORT}{self.path}" + headers = { + k: v + for k, v in self.headers.items() + if k.lower() not in ("host", "transfer-encoding", "content-length") + } + req = urllib.request.Request( + upstream_url, data=body, headers=headers, method=self.command + ) + try: + resp = urllib.request.urlopen(req) + resp_body = resp.read() + status = resp.status + resp_headers = list(resp.headers.items()) + except urllib.error.HTTPError as e: + resp_body = e.read() + status = e.code + resp_headers = list(e.headers.items()) + + should_drop = False + if ( + self.command == "POST" + and self._is_table_commit_endpoint() + and 200 <= status < 300 + ): + with _lock: + _commits_ok += 1 + if _drop_armed: + _drop_armed = False + should_drop = True + _commits_dropped += 1 + + if should_drop: + _log(f"dropping response for {self.command} {self.path}") + self.send_response(502) + self.send_header("Content-Length", "0") + self.end_headers() + return + + self.send_response(status) + for k, v in resp_headers: + if k.lower() not in ("transfer-encoding", "connection", "content-length"): + self.send_header(k, v) + self.send_header("Content-Length", str(len(resp_body))) + self.end_headers() + if resp_body: + self.wfile.write(resp_body) + + def _do_with_body(self) -> None: + if self._handle_control(): + return + length = int(self.headers.get("Content-Length", 0)) + body = self.rfile.read(length) if length else None + self._forward(body) + + def _do_without_body(self) -> None: + if self._handle_control(): + return + self._forward(None) + + def do_GET(self): + self._do_without_body() + + def do_HEAD(self): + self._do_without_body() + + def do_DELETE(self): + self._do_with_body() + + def do_POST(self): + self._do_with_body() + + def do_PUT(self): + self._do_with_body() + + +class ThreadedHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer): + daemon_threads = True + allow_reuse_address = True + + +if __name__ == "__main__": + port = int(os.environ.get("PROXY_PORT", "8181")) + server = ThreadedHTTPServer(("0.0.0.0", port), Handler) + _log(f"listening on :{port}, upstream={UPSTREAM_HOST}:{UPSTREAM_PORT}") + server.serve_forever()