skade-sql/tests/merge_into.rs
//! `MERGE INTO` through skade-sql — the row-level Iceberg write path.
//!
//! # The red these guards were written against
//!
//! DataFusion 55 plans `MERGE INTO` end to end and then stops at execution: no
//! `TableProvider` in this tree implemented `merge_into`, so the statement came
//! back as
//!
//! ```text
//! Error during planning: MERGE INTO not supported for Base table
//! caused by: MERGE INTO operation on table 'tgt'
//! ```
//!
//! Every `#[tokio::test]` below was run against that state first (worktree at
//! `origin/main`, commit `b4613a9`) and every one of them failed with that
//! message. They are the record of closing it, not decoration on top.
//!
//! # LAW 2 — the guards that can fail
//!
//! - [`the_count_is_the_warehouses_not_the_inputs`] pins the count in BOTH
//! directions: one merge where affected < source rows and one where
//! affected > source rows. An implementation that echoed the input row count
//! passes neither.
//! - [`a_merge_that_matches_nothing_rewrites_no_files`] asserts the snapshot id
//! is *unchanged* — not "the rows are still right", which a full rewrite also
//! satisfies.
//! - [`a_failed_commit_leaves_the_table_exactly_as_it_was`] induces a real
//! commit failure and asserts head + contents + history.
//! - the refusal guards assert the message NAMES the unsupported shape; a
//! silent no-op fails them.
use std::sync::Arc;
use anyhow::Result;
use arrow_array::{Array, Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use datafusion::prelude::SessionContext;
use skade::Warehouse;
use skade_sql::WarehouseSql;
fn tgt_schema() -> Schema {
Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, true),
Field::new("qty", DataType::Int64, true),
])
}
fn tgt_batch(rows: &[(i64, &str, i64)]) -> Result<RecordBatch> {
Ok(RecordBatch::try_new(Arc::new(tgt_schema()), vec![
Arc::new(Int64Array::from(
rows.iter().map(|r| r.0).collect::<Vec<_>>(),
)),
Arc::new(StringArray::from(
rows.iter().map(|r| r.1).collect::<Vec<_>>(),
)),
Arc::new(Int64Array::from(
rows.iter().map(|r| r.2).collect::<Vec<_>>(),
)),
])?)
}
fn src_batch(rows: &[(i64, &str, i64)]) -> Result<RecordBatch> {
let schema = Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, true),
Field::new("qty", DataType::Int64, true),
]);
Ok(RecordBatch::try_new(Arc::new(schema), vec![
Arc::new(Int64Array::from(
rows.iter().map(|r| r.0).collect::<Vec<_>>(),
)),
Arc::new(StringArray::from(
rows.iter().map(|r| r.1).collect::<Vec<_>>(),
)),
Arc::new(Int64Array::from(
rows.iter().map(|r| r.2).collect::<Vec<_>>(),
)),
])?)
}
/// A warehouse with `tgt` holding `rows`, written as ONE append per slice so a
/// caller can control how many data files the merge has to consider.
async fn warehouse_with(
dir: &std::path::Path,
appends: &[&[(i64, &str, i64)]],
) -> Result<Warehouse> {
let wh = skade::open(dir.join("lake")).await?;
let mut t = wh.create_table("tgt", &tgt_schema()).await?;
for rows in appends {
t.append(&[tgt_batch(rows)?]).await?;
}
Ok(wh)
}
/// A session with `tgt` (writable, from the catalog) plus `src` as an in-memory
/// table.
async fn session_with_source(wh: &Warehouse, src: RecordBatch) -> Result<SessionContext> {
let ctx = wh.session().await?;
ctx.register_batch("src", src)?;
Ok(ctx)
}
/// The single `count` value a DML statement returns.
async fn run_count(ctx: &SessionContext, sql: &str) -> Result<u64> {
let batches = ctx.sql(sql).await?.collect().await?;
let b = batches
.iter()
.find(|b| b.num_rows() > 0)
.ok_or_else(|| anyhow::anyhow!("no count row returned by `{sql}`"))?;
let col = b
.column_by_name("count")
.ok_or_else(|| anyhow::anyhow!("no `count` column: {:?}", b.schema()))?;
Ok(col
.as_any()
.downcast_ref::<arrow_array::UInt64Array>()
.ok_or_else(|| anyhow::anyhow!("count is not UInt64"))?
.value(0))
}
/// `tgt`'s rows, sorted, as `(id, name, qty)` strings.
async fn rows(wh: &Warehouse) -> Result<Vec<(i64, String, String)>> {
let batches = wh.sql("SELECT id, name, qty FROM tgt ORDER BY id").await?;
let mut out = Vec::new();
for b in batches {
let id = b
.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.expect("id");
for i in 0..b.num_rows() {
let cell = |c: usize| {
let a = b.column(c);
if a.is_null(i) {
"NULL".to_string()
} else {
skade::arrow_cast::display::array_value_to_string(a, i).unwrap()
}
};
out.push((id.value(i), cell(1), cell(2)));
}
}
out.sort();
Ok(out)
}
async fn head(wh: &Warehouse) -> Result<i64> {
Ok(wh
.table("tgt")
.await?
.current_snapshot_id()
.expect("tgt has a snapshot"))
}
// ───────────────────────────── the four clause kinds ────────────────────────
/// `WHEN MATCHED THEN UPDATE` — the shape lane FR could plan but not execute.
#[tokio::test]
async fn when_matched_then_update_rewrites_only_the_matched_rows() -> Result<()> {
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[&[(1, "a", 10), (2, "b", 20), (3, "c", 30)]]).await?;
let ctx = session_with_source(&wh, src_batch(&[(2, "B", 200)])?).await?;
let n = run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id = src.id \
WHEN MATCHED THEN UPDATE SET name = src.name, qty = src.qty",
)
.await?;
assert_eq!(n, 1, "one target row matched, so one row was affected");
assert_eq!(rows(&wh).await?, vec![
(1, "a".into(), "10".into()),
(2, "B".into(), "200".into()),
(3, "c".into(), "30".into()),
]);
Ok(())
}
/// `WHEN MATCHED THEN DELETE` — must work, not silently no-op.
#[tokio::test]
async fn when_matched_then_delete_removes_the_row() -> Result<()> {
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[&[(1, "a", 10), (2, "b", 20), (3, "c", 30)]]).await?;
let ctx = session_with_source(&wh, src_batch(&[(1, "x", 0), (3, "x", 0)])?).await?;
let n = run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id = src.id WHEN MATCHED THEN DELETE",
)
.await?;
assert_eq!(n, 2);
assert_eq!(rows(&wh).await?, vec![(2, "b".into(), "20".into())]);
Ok(())
}
/// `WHEN NOT MATCHED THEN INSERT` — the source rows with no target row.
#[tokio::test]
async fn when_not_matched_then_insert_adds_the_new_rows() -> Result<()> {
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[&[(1, "a", 10)]]).await?;
let ctx = session_with_source(&wh, src_batch(&[(1, "A", 11), (2, "b", 20)])?).await?;
let n = run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id = src.id \
WHEN MATCHED THEN UPDATE SET name = src.name, qty = src.qty \
WHEN NOT MATCHED THEN INSERT (id, name, qty) VALUES (src.id, src.name, src.qty)",
)
.await?;
assert_eq!(n, 2, "one updated + one inserted");
assert_eq!(rows(&wh).await?, vec![
(1, "A".into(), "11".into()),
(2, "b".into(), "20".into()),
]);
Ok(())
}
/// `WHEN NOT MATCHED BY SOURCE THEN DELETE` — the clause that touches target
/// rows the source never mentions. It either works or refuses by name; here it
/// works.
#[tokio::test]
async fn when_not_matched_by_source_then_delete_prunes_the_target() -> Result<()> {
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[&[(1, "a", 10), (2, "b", 20), (3, "c", 30)]]).await?;
let ctx = session_with_source(&wh, src_batch(&[(2, "B", 22)])?).await?;
let n = run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id = src.id \
WHEN MATCHED THEN UPDATE SET name = src.name \
WHEN NOT MATCHED BY SOURCE THEN DELETE",
)
.await?;
assert_eq!(n, 3, "one updated + two deleted");
assert_eq!(rows(&wh).await?, vec![(2, "B".into(), "20".into())]);
Ok(())
}
/// `WHEN NOT MATCHED BY SOURCE THEN UPDATE`.
#[tokio::test]
async fn when_not_matched_by_source_then_update_rewrites_the_untouched_rows() -> Result<()> {
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[&[(1, "a", 10), (2, "b", 20)]]).await?;
let ctx = session_with_source(&wh, src_batch(&[(2, "B", 22)])?).await?;
let n = run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id = src.id \
WHEN NOT MATCHED BY SOURCE THEN UPDATE SET name = 'stale'",
)
.await?;
assert_eq!(n, 1);
assert_eq!(rows(&wh).await?, vec![
(1, "stale".into(), "10".into()),
(2, "b".into(), "20".into()),
]);
Ok(())
}
/// A clause `AND` predicate narrows which matched rows fire.
#[tokio::test]
async fn a_clause_predicate_narrows_the_matched_rows() -> Result<()> {
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[&[(1, "a", 10), (2, "b", 20)]]).await?;
let ctx = session_with_source(&wh, src_batch(&[(1, "A", 1), (2, "B", 2)])?).await?;
let n = run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id = src.id \
WHEN MATCHED AND tgt.qty > 15 THEN UPDATE SET name = src.name",
)
.await?;
assert_eq!(n, 1, "only id=2 has qty > 15");
assert_eq!(rows(&wh).await?, vec![
(1, "a".into(), "10".into()),
(2, "B".into(), "20".into()),
]);
Ok(())
}
// ───────────────────────────── LAW 2 guards ─────────────────────────────────
/// THE COUNT COMES FROM THE COMMIT, NOT THE INPUT.
///
/// Two merges on the same shape of table, one whose affected count is BELOW the
/// source row count and one ABOVE it. An implementation that returned
/// `source.num_rows()` fails the first; one that returned the target row count
/// fails both; one that returned the rewritten-file row count fails the second.
#[tokio::test]
async fn the_count_is_the_warehouses_not_the_inputs() -> Result<()> {
// 4 source rows, 2 of them match → 2.
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[&[(1, "a", 10), (2, "b", 20), (3, "c", 30)]]).await?;
let ctx =
session_with_source(&wh, src_batch(&[(1, "A", 1), (2, "B", 2), (8, "H", 8), (9, "I", 9)])?)
.await?;
let below = run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id = src.id \
WHEN MATCHED THEN UPDATE SET name = src.name",
)
.await?;
assert_eq!(
below, 2,
"4 source rows, 3 target rows, 2 matched — the count is 2, and it is neither input"
);
// 1 source row, 3 target rows deleted for not matching it → 3.
let tmp2 = tempfile::tempdir()?;
let wh2 =
warehouse_with(tmp2.path(), &[&[(1, "a", 10), (2, "b", 20), (3, "c", 30), (4, "d", 40)]])
.await?;
let ctx2 = session_with_source(&wh2, src_batch(&[(4, "D", 44)])?).await?;
let above = run_count(
&ctx2,
"MERGE INTO tgt USING src ON tgt.id = src.id \
WHEN NOT MATCHED BY SOURCE THEN DELETE",
)
.await?;
assert_eq!(above, 3, "1 source row, 3 rows affected — the count is above the input");
Ok(())
}
/// A MERGE THAT MATCHES NOTHING MUST NOT REWRITE FILES.
///
/// The assertion is on the SNAPSHOT ID, not on the rows: a full rewrite would
/// leave the rows correct too. No commit at all is the correct answer, and the
/// only one that leaves the head where it was.
#[tokio::test]
async fn a_merge_that_matches_nothing_rewrites_no_files() -> Result<()> {
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[
&[(1, "a", 10), (2, "b", 20)],
&[(3, "c", 30), (4, "d", 40)],
])
.await?;
let before = head(&wh).await?;
let ctx = session_with_source(&wh, src_batch(&[(77, "z", 0)])?).await?;
let n = run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id = src.id \
WHEN MATCHED THEN UPDATE SET name = src.name",
)
.await?;
assert_eq!(n, 0);
assert_eq!(
head(&wh).await?,
before,
"no row matched, so no data file was rewritten and no snapshot was committed"
);
assert_eq!(rows(&wh).await?.len(), 4);
Ok(())
}
/// A merge that matches rows in ONE of two data files rewrites only that file.
///
/// Measured on the manifest: the untouched file's path must still be live in
/// the new snapshot, and the touched file's path must not.
#[tokio::test]
async fn only_the_files_holding_matched_keys_are_rewritten() -> Result<()> {
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[
&[(1, "a", 10), (2, "b", 20)],
&[(3, "c", 30), (4, "d", 40)],
])
.await?;
let before = live_data_files(&wh).await?;
assert_eq!(before.len(), 2, "two appends, two data files");
let ctx = session_with_source(&wh, src_batch(&[(3, "C", 33)])?).await?;
let n = run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id = src.id \
WHEN MATCHED THEN UPDATE SET name = src.name",
)
.await?;
assert_eq!(n, 1);
let after = live_data_files(&wh).await?;
let survivors: Vec<&String> = after.iter().filter(|p| before.contains(p)).collect();
assert_eq!(
survivors.len(),
1,
"exactly one original data file is carried forward untouched; got {after:?} from {before:?}"
);
assert_eq!(after.len(), 2, "one carried + one rewritten");
Ok(())
}
/// A FAILED COMMIT LEAVES THE TABLE EXACTLY AS IT WAS.
///
/// The failure is induced at the level the merge actually commits at: an
/// overwrite naming a data file that is not live. `existing_manifest` refuses
/// it, so the catalog is never touched.
#[tokio::test]
async fn a_failed_commit_leaves_the_table_exactly_as_it_was() -> Result<()> {
use iceberg::transaction::{ApplyTransactionAction, Transaction};
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[&[(1, "a", 10), (2, "b", 20)]]).await?;
let before_head = head(&wh).await?;
let before_rows = rows(&wh).await?;
let before_files = live_data_files(&wh).await?;
let t = wh.table("tgt").await?;
let tx = Transaction::new(t.inner());
let action = tx
.overwrite_files()
.delete_files(["/nowhere/not-a-live-file.parquet".to_string()]);
let err = action
.apply(tx)?
.commit(wh.catalog().as_ref())
.await
.expect_err("an overwrite naming a file that is not live must fail");
let msg = err.to_string();
assert!(
msg.contains("not live in the current snapshot"),
"the failure must name what was wrong, got: {msg}"
);
assert_eq!(head(&wh).await?, before_head, "the head did not move");
assert_eq!(rows(&wh).await?, before_rows, "the rows are unchanged");
assert_eq!(live_data_files(&wh).await?, before_files);
Ok(())
}
/// The pre-merge snapshot is still readable after a successful merge.
#[tokio::test]
async fn the_pre_merge_snapshot_still_reads_back() -> Result<()> {
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[&[(1, "a", 10), (2, "b", 20)]]).await?;
let before = head(&wh).await?;
let ctx = session_with_source(&wh, src_batch(&[(1, "A", 11)])?).await?;
run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id = src.id \
WHEN MATCHED THEN UPDATE SET name = src.name, qty = src.qty",
)
.await?;
assert_ne!(head(&wh).await?, before, "the merge moved the head");
let ident = wh.table("tgt").await?.ident().clone();
let old = wh.catalog().load_table_at(&ident, before).await?;
let batches = scan_to_batches(&old).await?;
let n: usize = batches.iter().map(arrow_array::RecordBatch::num_rows).sum();
assert_eq!(n, 2, "the pre-merge snapshot still has its own two rows");
let names: Vec<String> = batches
.iter()
.flat_map(|b| {
let a = b
.column_by_name("name")
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap()
.clone();
(0..a.len()).map(move |i| a.value(i).to_string()).collect::<Vec<_>>()
})
.collect();
assert!(
names.contains(&"a".to_string()),
"history still holds the PRE-merge value, got {names:?}"
);
Ok(())
}
// ───────────────────────────── refusals, by name ────────────────────────────
/// A non-equality `ON` is refused, and the message says why.
#[tokio::test]
async fn a_non_equality_on_condition_is_refused_by_name() -> Result<()> {
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[&[(1, "a", 10)]]).await?;
let before = head(&wh).await?;
let ctx = session_with_source(&wh, src_batch(&[(1, "A", 1)])?).await?;
let err = run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id > src.id WHEN MATCHED THEN DELETE",
)
.await
.expect_err("a non-equi ON must be refused");
let msg = err.to_string();
assert!(
msg.contains("conjunction of") && msg.contains("equalities"),
"the refusal must name the shape it wanted, got: {msg}"
);
assert_eq!(head(&wh).await?, before, "and it changed nothing");
Ok(())
}
/// The clause/action combinations that name no row to act on are refused BY
/// NAME — and the refusal comes from wherever the shape is first recognised.
///
/// `WHEN NOT MATCHED THEN DELETE` never reaches the provider: sqlparser itself
/// rejects it. `WHEN NOT MATCHED BY SOURCE THEN INSERT` is the one that gets
/// past the parser, and `compile_clauses` refuses it. Both are asserted here so
/// the guard fails if EITHER layer ever starts accepting the shape and doing
/// nothing with it.
#[tokio::test]
async fn clause_action_combinations_that_name_no_row_are_refused_by_name() -> Result<()> {
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[&[(1, "a", 10)]]).await?;
let before = head(&wh).await?;
let ctx = session_with_source(&wh, src_batch(&[(2, "B", 2)])?).await?;
let parser = run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id = src.id WHEN NOT MATCHED THEN DELETE",
)
.await
.expect_err("WHEN NOT MATCHED THEN DELETE must be refused")
.to_string();
assert!(
parser.contains("DELETE is not allowed in a NOT MATCHED merge clause"),
"the parser must name the clause, got: {parser}"
);
let provider = run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id = src.id \
WHEN NOT MATCHED BY SOURCE THEN INSERT (id, name, qty) VALUES (9, 'z', 9)",
)
.await
.expect_err("WHEN NOT MATCHED BY SOURCE THEN INSERT must be refused")
.to_string();
assert!(
provider.contains("NOT MATCHED BY SOURCE") || provider.contains("not allowed"),
"the refusal must name the clause, got: {provider}"
);
assert_eq!(head(&wh).await?, before, "and neither changed anything");
Ok(())
}
/// A target row matched by two source rows is ambiguous. Refused, not
/// arbitrarily resolved.
#[tokio::test]
async fn a_doubly_matched_target_row_is_refused_by_name() -> Result<()> {
let tmp = tempfile::tempdir()?;
let wh = warehouse_with(tmp.path(), &[&[(1, "a", 10)]]).await?;
let before = head(&wh).await?;
let ctx = session_with_source(&wh, src_batch(&[(1, "A", 1), (1, "Z", 9)])?).await?;
let err = run_count(
&ctx,
"MERGE INTO tgt USING src ON tgt.id = src.id \
WHEN MATCHED THEN UPDATE SET name = src.name",
)
.await
.expect_err("two source rows for one target row must be refused");
let msg = err.to_string();
assert!(
msg.contains("more than \n one source row")
|| msg.contains("more than"),
"the refusal must say the target row was matched twice, got: {msg}"
);
assert_eq!(head(&wh).await?, before);
Ok(())
}
// ───────────────────────────── helpers ──────────────────────────────────────
/// Every live data-file path in `tgt`'s current snapshot, from the manifests.
async fn live_data_files(wh: &Warehouse) -> Result<Vec<String>> {
use futures::StreamExt;
let t = wh.table("tgt").await?;
let scan = t.inner().scan().build()?;
let mut stream = scan.plan_files().await?;
let mut out = Vec::new();
while let Some(task) = stream.next().await {
out.push(task?.data_file_path().to_string());
}
out.sort();
out.dedup();
Ok(out)
}
/// Every row of one iceberg table handle, through iceberg's own scan.
async fn scan_to_batches(t: &iceberg::table::Table) -> Result<Vec<RecordBatch>> {
use futures::StreamExt;
let mut stream = t.scan().build()?.to_arrow().await?;
let mut out = Vec::new();
while let Some(b) = stream.next().await {
out.push(b?);
}
Ok(out)
}