#![cfg(test)]
use std::sync::{Arc, Mutex};
use arrow_array::{Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use async_trait::async_trait;
use tempfile::tempdir;
use crate::{
error::Result,
listener::Listener,
mapping::{MappingSpec, SyncMode},
mutation::{GraphMutation, SinkReport, VertexDelete, VertexSet},
registry::TransformRegistry,
sink::GraphSink,
source::{TableSource, TableSpec},
state::{SyncState, TableState},
transform::{
EdgeTransform, EndpointSpec, TableContext, Transform, VertexTransform,
deletes_from_change_rows, is_changelog_batch, partition_change_rows,
},
};
fn users_batch() -> RecordBatch {
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("user_id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
Field::new("city", DataType::Utf8, false),
Field::new("company_id", DataType::Int64, false),
])),
vec![
Arc::new(Int64Array::from(vec![1, 2])),
Arc::new(StringArray::from(vec!["alice", "bob"])),
Arc::new(StringArray::from(vec!["oslo", "bergen"])),
Arc::new(Int64Array::from(vec![10, 10])),
],
)
.unwrap()
}
const MAPPING_YAML: &str = r#"
version: 1
tables:
- table: lake.social.users
mode: incremental
pipeline:
- vertex:
label: Person
id_column: user_id
properties: [name, city]
rename: { city: location }
- fk_edges:
type: WORKS_AT
src: { label: Person, column: user_id }
dst: { label: Company, column: company_id, id_column: company_id }
properties: [city]
rename: { city: based_in }
"#;
#[test]
fn vertex_transform_projects_and_renames() {
let t = VertexTransform {
label: "Person".into(),
id_column: "user_id".into(),
properties: vec!["name".into(), "city".into()],
rename: [("city".to_string(), "location".to_string())].into(),
};
let ctx = TableContext { table: "t".into() };
let m = t.apply(&ctx, &[users_batch()]).unwrap();
assert_eq!(m.vertices.len(), 1);
assert!(m.edges.is_empty());
let set = &m.vertices[0];
assert_eq!(set.label, "Person");
assert_eq!(set.id_column, "user_id");
let names: Vec<_> = set
.batch
.schema()
.fields()
.iter()
.map(|f| f.name().clone())
.collect();
assert_eq!(names, vec!["user_id", "name", "location"]);
assert_eq!(set.batch.num_rows(), 2);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"vertex_transform_projects_and_renames",
set.label == "Person"
&& set.batch.num_rows() == 2
&& names == ["user_id", "name", "location"],
&format!(
"label={}, cols={:?}, rows={}",
set.label,
names,
set.batch.num_rows()
),
);
}
#[test]
fn vertex_transform_missing_column_errors() {
let t = VertexTransform {
label: "Person".into(),
id_column: "nope".into(),
properties: vec![],
rename: Default::default(),
};
let ctx = TableContext { table: "t".into() };
let res = t.apply(&ctx, &[users_batch()]);
assert!(res.is_err());
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"vertex_transform_missing_column_errors",
res.is_err(),
"missing id column → Err",
);
}
#[test]
fn edge_transform_maps_endpoints_to_src_dst() {
let t = EdgeTransform {
edge_type: "WORKS_AT".into(),
src: EndpointSpec {
label: "Person".into(),
column: "user_id".into(),
id_column: None,
},
dst: EndpointSpec {
label: "Company".into(),
column: "company_id".into(),
id_column: None,
},
properties: vec![],
rename: Default::default(),
key: Vec::new(),
replace: None,
};
let ctx = TableContext { table: "t".into() };
let m = t.apply(&ctx, &[users_batch()]).unwrap();
assert_eq!(m.edges.len(), 1);
let set = &m.edges[0];
assert_eq!(set.src_id_column, "user_id");
assert_eq!(set.dst_id_column, "company_id");
let names: Vec<_> = set
.batch
.schema()
.fields()
.iter()
.map(|f| f.name().clone())
.collect();
assert_eq!(names, vec!["src", "dst"]);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"edge_transform_maps_endpoints_to_src_dst",
set.src_id_column == "user_id"
&& set.dst_id_column == "company_id"
&& names == ["src", "dst"],
&format!(
"src_id={}, dst_id={}, cols={:?}",
set.src_id_column, set.dst_id_column, names
),
);
}
#[test]
fn edge_transform_projects_properties_alongside_endpoints() {
let t = EdgeTransform {
edge_type: "WORKS_AT".into(),
src: EndpointSpec {
label: "Person".into(),
column: "user_id".into(),
id_column: None,
},
dst: EndpointSpec {
label: "Company".into(),
column: "company_id".into(),
id_column: None,
},
properties: vec!["name".into(), "city".into()],
rename: [("city".to_string(), "location".to_string())].into(),
key: Vec::new(),
replace: None,
};
let ctx = TableContext { table: "t".into() };
let m = t.apply(&ctx, &[users_batch()]).unwrap();
assert_eq!(m.edges.len(), 1);
let set = &m.edges[0];
let names: Vec<_> = set
.batch
.schema()
.fields()
.iter()
.map(|f| f.name().clone())
.collect();
assert_eq!(names, vec!["src", "dst", "name", "location"]);
assert_eq!(set.batch.num_rows(), 2);
let rows = crate::sink::rows_as_cypher_maps(&set.batch).unwrap();
assert_eq!(
rows[0],
"{src: 1, dst: 10, props: {`name`: 'alice', `location`: 'oslo'}}"
);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"edge_transform_projects_properties_alongside_endpoints",
names == ["src", "dst", "name", "location"] && set.batch.num_rows() == 2,
&format!(
"cols={:?}, rows={}, cypher={}",
names,
set.batch.num_rows(),
rows[0]
),
);
}
#[test]
fn edge_transform_self_key_emits_both_src_and_dst() {
let t = EdgeTransform {
edge_type: "SAME_AS".into(),
src: EndpointSpec {
label: "Person".into(),
column: "user_id".into(),
id_column: None,
},
dst: EndpointSpec {
label: "Person".into(),
column: "user_id".into(),
id_column: None,
},
properties: vec![],
rename: Default::default(),
key: Vec::new(),
replace: None,
};
let ctx = TableContext { table: "t".into() };
let m = t.apply(&ctx, &[users_batch()]).unwrap();
let set = &m.edges[0];
let names: Vec<_> = set
.batch
.schema()
.fields()
.iter()
.map(|f| f.name().clone())
.collect();
let pass = names == ["src", "dst"] && set.batch.num_rows() == 2;
assert!(
pass,
"self-key edge must keep both endpoints, got {names:?}"
);
let rows = crate::sink::rows_as_cypher_maps(&set.batch).unwrap();
assert_eq!(rows[0], "{src: 1, dst: 1, props: {}}");
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"edge_transform_self_key_emits_both_src_and_dst",
pass,
&format!("cols={names:?}, cypher={}", rows[0]),
);
}
#[test]
fn edge_transform_property_colliding_with_endpoint_errors_clearly() {
let t = EdgeTransform {
edge_type: "WORKS_AT".into(),
src: EndpointSpec {
label: "Person".into(),
column: "user_id".into(),
id_column: None,
},
dst: EndpointSpec {
label: "Company".into(),
column: "company_id".into(),
id_column: None,
},
properties: vec!["name".into()],
rename: [("name".to_string(), "src".to_string())].into(),
key: Vec::new(),
replace: None,
};
let ctx = TableContext { table: "t".into() };
let res = t.apply(&ctx, &[users_batch()]);
let is_dup = matches!(
res,
Err(crate::error::BifrostError::DuplicateOutputColumn(ref n)) if n == "src"
);
assert!(is_dup, "expected DuplicateOutputColumn(src), got {res:?}");
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"edge_transform_property_colliding_with_endpoint_errors_clearly",
is_dup,
"property→'src' collision → DuplicateOutputColumn",
);
}
#[test]
fn vertex_transform_id_also_in_properties_dedups() {
let t = VertexTransform {
label: "Person".into(),
id_column: "user_id".into(),
properties: vec!["user_id".into(), "name".into()],
rename: Default::default(),
};
let ctx = TableContext { table: "t".into() };
let m = t.apply(&ctx, &[users_batch()]).unwrap();
let set = &m.vertices[0];
let names: Vec<_> = set
.batch
.schema()
.fields()
.iter()
.map(|f| f.name().clone())
.collect();
let pass = names == ["user_id", "name"];
assert!(
pass,
"id_column listed in properties must dedup, got {names:?}"
);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"vertex_transform_id_also_in_properties_dedups",
pass,
&format!("cols={names:?}"),
);
}
#[test]
fn mapping_yaml_parses_and_compiles() {
let spec = MappingSpec::from_yaml(MAPPING_YAML).unwrap();
assert_eq!(spec.version, 1);
assert_eq!(spec.tables.len(), 1);
assert_eq!(spec.tables[0].mode, SyncMode::Incremental);
let transforms = spec.tables[0].compile(&TransformRegistry::new()).unwrap();
assert_eq!(transforms.len(), 2);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"mapping_yaml_parses_and_compiles",
spec.version == 1 && spec.tables.len() == 1 && transforms.len() == 2,
&format!(
"v{}, {} table(s), {} transform(s)",
spec.version,
spec.tables.len(),
transforms.len()
),
);
}
#[test]
fn mapping_edge_properties_flow_into_edge_set() {
let spec = MappingSpec::from_yaml(MAPPING_YAML).unwrap();
let transforms = spec.tables[0].compile(&TransformRegistry::new()).unwrap();
let ctx = TableContext {
table: "lake.social.users".into(),
};
let edge_set = transforms
.iter()
.filter_map(|t| {
t.apply(&ctx, &[users_batch()])
.unwrap()
.edges
.into_iter()
.next()
})
.next()
.expect("the fk_edges step yields an EdgeSet");
assert_eq!(edge_set.edge_type, "WORKS_AT");
let names: Vec<_> = edge_set
.batch
.schema()
.fields()
.iter()
.map(|f| f.name().clone())
.collect();
assert_eq!(names, vec!["src", "dst", "based_in"]);
let rows = crate::sink::rows_as_cypher_maps(&edge_set.batch).unwrap();
assert_eq!(rows[0], "{src: 1, dst: 10, props: {`based_in`: 'oslo'}}");
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"mapping_edge_properties_flow_into_edge_set",
edge_set.edge_type == "WORKS_AT" && names == ["src", "dst", "based_in"],
&format!(
"type={}, cols={:?}, cypher={}",
edge_set.edge_type, names, rows[0]
),
);
}
#[test]
fn mapping_custom_step_requires_registration() {
let yaml = r#"
version: 1
tables:
- table: t
pipeline:
- custom: missing
"#;
let spec = MappingSpec::from_yaml(yaml).unwrap();
assert!(spec.tables[0].compile(&TransformRegistry::new()).is_err());
struct Nop;
impl Transform for Nop {
fn apply(&self, _: &TableContext, _: &[RecordBatch]) -> Result<GraphMutation> {
Ok(GraphMutation::default())
}
}
let mut registry = TransformRegistry::new();
registry.register("missing", Arc::new(Nop));
let after = spec.tables[0].compile(®istry).is_ok();
assert!(after);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"mapping_custom_step_requires_registration",
after,
"unregistered custom step errs; registered compiles",
);
}
#[tokio::test]
async fn sync_state_round_trips_through_file() {
let dir = tempdir().unwrap();
let path = dir.path().join("state.json");
{
let mut state = SyncState::open(&path).unwrap();
assert_eq!(state.get("t"), TableState::default());
state
.advance(
"t",
TableState {
snapshot_id: Some(42),
..Default::default()
},
)
.await
.unwrap();
}
let reopened = SyncState::open(&path).unwrap();
assert_eq!(reopened.get("t").snapshot_id, Some(42));
assert!(reopened.get("t").synced_at.is_some());
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"sync_state_round_trips_through_file",
reopened.get("t").snapshot_id == Some(42) && reopened.get("t").synced_at.is_some(),
&format!("reopened snapshot_id={:?}", reopened.get("t").snapshot_id),
);
}
struct MockSource {
snapshot: i64,
}
#[async_trait]
impl TableSource for MockSource {
async fn fetch_new(
&self,
_spec: &TableSpec,
state: &TableState,
) -> Result<Option<(Vec<RecordBatch>, TableState)>> {
if state.snapshot_id == Some(self.snapshot) {
return Ok(None);
}
Ok(Some((
vec![users_batch()],
TableState {
snapshot_id: Some(self.snapshot),
..state.clone()
},
)))
}
}
#[derive(Clone, Default)]
struct RecordingSink {
mutations: Arc<Mutex<Vec<GraphMutation>>>,
}
#[async_trait]
impl GraphSink for RecordingSink {
async fn apply(&mut self, mutation: &GraphMutation) -> Result<SinkReport> {
self.mutations.lock().unwrap().push(mutation.clone());
Ok(SinkReport {
vertices_written: mutation.vertices.iter().map(|v| v.batch.num_rows()).sum(),
edges_written: mutation.edges.iter().map(|e| e.batch.num_rows()).sum(),
deleted: mutation.deletes.iter().map(|d| d.batch.num_rows()).sum(),
})
}
}
#[tokio::test]
async fn listener_syncs_then_goes_idle() {
let dir = tempdir().unwrap();
let spec = MappingSpec::from_yaml(MAPPING_YAML).unwrap();
let sink = RecordingSink::default();
let mutations = sink.mutations.clone();
let mut listener = Listener::new(
spec,
TransformRegistry::new(),
MockSource { snapshot: 7 },
sink,
SyncState::open(dir.path().join("state.json")).unwrap(),
);
let round = listener.sync_once().await.unwrap();
assert_eq!(round.failed, vec![]);
assert_eq!(round.synced.len(), 1);
let (table, rows, report) = &round.synced[0];
assert_eq!(table, "lake.social.users");
assert_eq!(*rows, 2);
assert_eq!(report.vertices_written, 2);
assert_eq!(report.edges_written, 2);
assert_eq!(mutations.lock().unwrap().len(), 1);
let round = listener.sync_once().await.unwrap();
assert!(round.synced.is_empty());
assert!(round.failed.is_empty());
assert_eq!(mutations.lock().unwrap().len(), 1);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"listener_syncs_then_goes_idle",
round.synced.is_empty() && mutations.lock().unwrap().len() == 1,
&format!(
"1 sync round (2v/2e) then idle; {} mutation(s)",
mutations.lock().unwrap().len()
),
);
}
#[tokio::test]
async fn listener_metrics_accumulate_across_rounds() {
let dir = tempdir().unwrap();
let spec = MappingSpec::from_yaml(MAPPING_YAML).unwrap();
let mut listener = Listener::new(
spec,
TransformRegistry::new(),
MockSource { snapshot: 7 },
RecordingSink::default(),
SyncState::open(dir.path().join("state.json")).unwrap(),
);
listener.sync_once().await.unwrap();
let m = listener.metrics().clone();
assert_eq!(m.rounds_total, 1);
assert_eq!(m.rounds_failed, 0);
assert_eq!(m.tables_synced, 1);
assert_eq!(m.rows_ingested, 2);
assert_eq!(m.vertices_written, 2);
assert_eq!(m.edges_written, 2);
assert_eq!(m.assets_emitted, 1);
listener.sync_once().await.unwrap();
let m = listener.metrics();
assert_eq!(m.rounds_total, 2);
assert_eq!(m.rounds_failed, 0);
assert_eq!(m.tables_synced, 1);
assert_eq!(m.rows_ingested, 2);
let text = crate::render_prometheus(listener.metrics());
assert!(text.contains("bifrost_rounds_total 2"));
assert!(text.contains("bifrost_rows_ingested 2"));
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"listener_metrics_accumulate_across_rounds",
text.contains("bifrost_rounds_total 2") && text.contains("bifrost_rows_ingested 2"),
"prometheus reflects rounds_total=2, rows_ingested=2",
);
}
#[tokio::test]
async fn listener_metrics_count_a_fully_failed_round() {
struct FailSource;
#[async_trait]
impl TableSource for FailSource {
async fn fetch_new(
&self,
_spec: &TableSpec,
_state: &TableState,
) -> Result<Option<(Vec<RecordBatch>, TableState)>> {
Err(crate::error::BifrostError::Other("source down".into()))
}
}
let dir = tempdir().unwrap();
let spec = MappingSpec::from_yaml(MAPPING_YAML).unwrap();
let mut listener = Listener::new(
spec,
TransformRegistry::new(),
FailSource,
RecordingSink::default(),
SyncState::open(dir.path().join("state.json")).unwrap(),
);
let round = listener.sync_once().await.unwrap();
assert!(round.is_failure(), "all tables failed → no-progress round");
let m = listener.metrics();
assert_eq!(m.rounds_total, 1);
assert_eq!(m.rounds_failed, 1);
assert_eq!(m.tables_synced, 0);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"listener_metrics_count_a_fully_failed_round",
round.is_failure() && m.rounds_failed == 1 && m.tables_synced == 0,
&format!(
"rounds_failed={}, tables_synced={}",
m.rounds_failed, m.tables_synced
),
);
}
mod file_source {
use super::*;
use crate::source::FileSource;
fn drop_parquet(root: &std::path::Path, table: &str, name: &str, batch: &RecordBatch) {
let dir = root.join(table);
std::fs::create_dir_all(&dir).unwrap();
graphar::io::write_chunk(
dir.join(name),
std::slice::from_ref(batch),
&graphar::FileType::Parquet,
)
.unwrap();
}
fn spec(table: &str, mode: SyncMode) -> TableSpec {
TableSpec::new(table, mode)
}
#[tokio::test]
async fn missing_directory_yields_nothing() {
let dir = tempdir().unwrap();
let src = FileSource::new(dir.path());
let got = src
.fetch_new(
&spec("lake.absent", SyncMode::Incremental),
&TableState::default(),
)
.await
.unwrap();
assert!(got.is_none());
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"file_source::missing_directory_yields_nothing",
got.is_none(),
"absent table dir → None",
);
}
#[tokio::test]
async fn reads_files_then_records_them_done() {
let dir = tempdir().unwrap();
drop_parquet(dir.path(), "lake.users", "chunk0.parquet", &users_batch());
let src = FileSource::new(dir.path());
let (batches, state) = src
.fetch_new(
&spec("lake.users", SyncMode::Incremental),
&TableState::default(),
)
.await
.unwrap()
.expect("first drop produces rows");
assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
assert_eq!(state.done_files, vec!["chunk0.parquet".to_string()]);
let idle = src
.fetch_new(&spec("lake.users", SyncMode::Incremental), &state)
.await
.unwrap();
assert!(idle.is_none(), "no new files → no round");
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"file_source::reads_files_then_records_them_done",
state.done_files == vec!["chunk0.parquet".to_string()] && idle.is_none(),
&format!("done_files={:?}, idle round skipped", state.done_files),
);
}
#[tokio::test]
async fn incremental_ships_only_new_files() {
let dir = tempdir().unwrap();
drop_parquet(dir.path(), "lake.users", "chunk0.parquet", &users_batch());
let src = FileSource::new(dir.path());
let (_, state) = src
.fetch_new(
&spec("lake.users", SyncMode::Incremental),
&TableState::default(),
)
.await
.unwrap()
.unwrap();
drop_parquet(dir.path(), "lake.users", "chunk1.parquet", &users_batch());
let (batches, state2) = src
.fetch_new(&spec("lake.users", SyncMode::Incremental), &state)
.await
.unwrap()
.expect("new file produces a round");
assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 2);
assert_eq!(
state2.done_files,
vec!["chunk0.parquet".to_string(), "chunk1.parquet".to_string()]
);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"file_source::incremental_ships_only_new_files",
batches.iter().map(|b| b.num_rows()).sum::<usize>() == 2
&& state2.done_files.len() == 2,
&format!(
"read {} new row(s), {} done files",
batches.iter().map(|b| b.num_rows()).sum::<usize>(),
state2.done_files.len()
),
);
}
#[tokio::test]
async fn full_mode_rereads_every_file() {
let dir = tempdir().unwrap();
drop_parquet(dir.path(), "lake.users", "chunk0.parquet", &users_batch());
drop_parquet(dir.path(), "lake.users", "chunk1.parquet", &users_batch());
let src = FileSource::new(dir.path());
let prior = TableState {
done_files: vec!["chunk0.parquet".into(), "chunk1.parquet".into()],
..Default::default()
};
let (batches, _) = src
.fetch_new(&spec("lake.users", SyncMode::Full), &prior)
.await
.unwrap()
.expect("full mode always reads");
let total = batches.iter().map(|b| b.num_rows()).sum::<usize>();
assert_eq!(total, 4);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"file_source::full_mode_rereads_every_file",
total == 4,
&format!("full mode re-read {total} rows from 2 files"),
);
}
#[tokio::test]
async fn non_chunk_files_are_ignored() {
let dir = tempdir().unwrap();
drop_parquet(dir.path(), "lake.users", "chunk0.parquet", &users_batch());
std::fs::write(dir.path().join("lake.users").join("_SUCCESS"), b"").unwrap();
let src = FileSource::new(dir.path());
let (_, state) = src
.fetch_new(
&spec("lake.users", SyncMode::Incremental),
&TableState::default(),
)
.await
.unwrap()
.unwrap();
assert_eq!(state.done_files, vec!["chunk0.parquet".to_string()]);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"file_source::non_chunk_files_are_ignored",
state.done_files == vec!["chunk0.parquet".to_string()],
"stray _SUCCESS ignored; only chunk0.parquet read",
);
}
}
mod airflow {
use crate::airflow::{AssetEmitter, AssetEvent, asset_event_body, emit_round, round_to_events};
use crate::error::Result;
use crate::listener::RoundReport;
use crate::mutation::SinkReport;
use async_trait::async_trait;
use std::sync::{Arc, Mutex};
fn report() -> RoundReport {
RoundReport {
synced: vec![
(
"lake.users".into(),
3,
SinkReport {
vertices_written: 3,
edges_written: 2,
deleted: 0,
},
),
(
"lake.companies".into(),
1,
SinkReport {
vertices_written: 1,
edges_written: 0,
deleted: 0,
},
),
],
failed: vec![("lake.broken".into(), "boom".into())],
}
}
#[test]
fn round_to_events_one_per_synced_table() {
let events = round_to_events("graph://social/", &report());
assert_eq!(events.len(), 2, "failed table emits nothing");
assert_eq!(events[0].uri, "graph://social/lake.users");
assert_eq!(events[0].extra["rows"], 3);
assert_eq!(events[0].extra["vertices_written"], 3);
assert_eq!(events[0].extra["edges_written"], 2);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"airflow::round_to_events_one_per_synced_table",
events.len() == 2 && events[0].uri == "graph://social/lake.users",
&format!("{} event(s); failed table emits none", events.len()),
);
}
#[test]
fn empty_round_emits_no_events() {
assert!(round_to_events("p", &RoundReport::default()).is_empty());
}
#[test]
fn asset_event_body_matches_airflow_shape() {
let ev = AssetEvent {
uri: "graph://social/lake.users".into(),
extra: serde_json::json!({ "rows": 3 }),
};
let body = asset_event_body(&ev);
assert_eq!(body["asset_uri"], "graph://social/lake.users");
assert_eq!(body["extra"]["rows"], 3);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"airflow::asset_event_body_matches_airflow_shape",
body["asset_uri"] == "graph://social/lake.users" && body["extra"]["rows"] == 3,
"AIP-82 body carries asset_uri + extra.rows",
);
}
#[derive(Default)]
struct RecordingEmitter {
seen: Arc<Mutex<Vec<AssetEvent>>>,
}
#[async_trait]
impl AssetEmitter for RecordingEmitter {
async fn emit(&self, events: &[AssetEvent]) -> Result<()> {
self.seen.lock().unwrap().extend_from_slice(events);
Ok(())
}
}
#[tokio::test]
async fn emit_round_delivers_each_event() {
let emitter = RecordingEmitter::default();
emit_round(&emitter, "graph://social/", &report())
.await
.unwrap();
let seen = emitter.seen.lock().unwrap();
assert_eq!(seen.len(), 2);
assert_eq!(seen[1].uri, "graph://social/lake.companies");
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"airflow::emit_round_delivers_each_event",
seen.len() == 2 && seen[1].uri == "graph://social/lake.companies",
&format!("emitter saw {} event(s)", seen.len()),
);
}
#[tokio::test]
async fn emit_round_skips_empty() {
let emitter = RecordingEmitter::default();
emit_round(&emitter, "p", &RoundReport::default())
.await
.unwrap();
assert!(emitter.seen.lock().unwrap().is_empty());
}
}
mod cdc {
use super::*;
use crate::mapping::SyncMode;
fn change_batch(change_types: &[&str]) -> RecordBatch {
let n = change_types.len();
let ids: Vec<i64> = (0..n as i64).collect();
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("user_id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
Field::new("company_id", DataType::Int64, false),
Field::new(SyncMode::CHANGE_TYPE_COLUMN, DataType::Utf8, false),
])),
vec![
Arc::new(Int64Array::from(ids)),
Arc::new(StringArray::from(vec!["x"; n])),
Arc::new(Int64Array::from(vec![10i64; n])),
Arc::new(StringArray::from(change_types.to_vec())),
],
)
.unwrap()
}
fn col_names(batch: &RecordBatch) -> Vec<String> {
batch
.schema()
.fields()
.iter()
.map(|f| f.name().clone())
.collect()
}
#[test]
fn detects_changelog_batches() {
let cdc = is_changelog_batch(&change_batch(&["INSERT"]));
let plain = is_changelog_batch(&users_batch());
assert!(cdc);
assert!(!plain);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"cdc::detects_changelog_batches",
cdc && !plain,
"change batch detected as CDC; plain batch not",
);
}
#[test]
fn partition_splits_upserts_from_deletes_and_strips_sentinel() {
let batch = change_batch(&["INSERT", "DELETE", "UPDATE_AFTER", "UPDATE_BEFORE"]);
let split = partition_change_rows(&[batch]).unwrap();
let up_rows: usize = split.upserts.iter().map(|b| b.num_rows()).sum();
let del_rows: usize = split.deletes.iter().map(|b| b.num_rows()).sum();
assert_eq!(up_rows, 2);
assert_eq!(del_rows, 1);
for b in split.upserts.iter().chain(split.deletes.iter()) {
assert!(!col_names(b).contains(&SyncMode::CHANGE_TYPE_COLUMN.to_string()));
assert!(col_names(b).contains(&"user_id".to_string()));
}
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"cdc::partition_splits_upserts_from_deletes_and_strips_sentinel",
up_rows == 2 && del_rows == 1,
&format!("{up_rows} upsert rows, {del_rows} delete rows; sentinel stripped"),
);
}
#[test]
fn partition_passes_non_cdc_batches_through_as_upserts() {
let split = partition_change_rows(&[users_batch()]).unwrap();
assert_eq!(split.upserts.len(), 1);
assert!(split.deletes.is_empty());
assert_eq!(col_names(&split.upserts[0]), col_names(&users_batch()));
}
#[test]
fn partition_rejects_unknown_change_type() {
let res = partition_change_rows(&[change_batch(&["WIGGLE"])]);
assert!(res.is_err());
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"cdc::partition_rejects_unknown_change_type",
res.is_err(),
"unknown _change_type 'WIGGLE' → Err",
);
}
#[test]
fn deletes_from_change_rows_reinterprets_vertices_as_deletes() {
let vt = VertexTransform {
label: "Person".into(),
id_column: "user_id".into(),
properties: vec!["name".into()],
rename: Default::default(),
};
let transforms: Vec<Arc<dyn Transform>> = vec![Arc::new(vt)];
let ctx = TableContext {
table: "lake.users".into(),
};
let split = partition_change_rows(&[change_batch(&["DELETE", "INSERT"])]).unwrap();
let deletes = deletes_from_change_rows(&transforms, &ctx, &split.deletes).unwrap();
assert_eq!(deletes.len(), 1);
assert_eq!(deletes[0].label, "Person");
assert_eq!(deletes[0].id_column, "user_id");
assert_eq!(deletes[0].batch.num_rows(), 1); crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"cdc::deletes_from_change_rows_reinterprets_vertices_as_deletes",
deletes.len() == 1 && deletes[0].label == "Person" && deletes[0].batch.num_rows() == 1,
&format!(
"{} delete set(s), {} row (only DELETE)",
deletes.len(),
deletes[0].batch.num_rows()
),
);
}
#[test]
fn graph_mutation_carries_and_merges_deletes() {
let del = VertexDelete {
label: "Person".into(),
id_column: "user_id".into(),
batch: users_batch(),
};
let mut m = GraphMutation {
edge_deletes: Vec::new(),
deletes: vec![del.clone()],
..Default::default()
};
assert!(!m.is_empty());
m.merge(GraphMutation {
edge_deletes: Vec::new(),
vertices: vec![VertexSet {
label: "Person".into(),
id_column: "user_id".into(),
batch: users_batch(),
}],
deletes: vec![del],
..Default::default()
});
assert_eq!(m.deletes.len(), 2);
assert_eq!(m.vertices.len(), 1);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"cdc::graph_mutation_carries_and_merges_deletes",
m.deletes.len() == 2 && m.vertices.len() == 1,
&format!(
"merged: {} delete set(s), {} vertex set(s)",
m.deletes.len(),
m.vertices.len()
),
);
}
fn update_pair_change_batch() -> RecordBatch {
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("user_id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
Field::new(SyncMode::CHANGE_TYPE_COLUMN, DataType::Utf8, false),
])),
vec![
Arc::new(Int64Array::from(vec![1, 1])),
Arc::new(StringArray::from(vec!["old", "new"])),
Arc::new(StringArray::from(vec!["UPDATE_BEFORE", "UPDATE_AFTER"])),
],
)
.unwrap()
}
#[derive(Default)]
struct InMemoryGraph {
nodes: std::collections::HashMap<(String, i64), String>,
}
impl InMemoryGraph {
fn ids(batch: &RecordBatch, col: &str) -> Vec<i64> {
batch
.column(batch.schema().index_of(col).unwrap())
.as_any()
.downcast_ref::<Int64Array>()
.unwrap()
.values()
.to_vec()
}
fn names(batch: &RecordBatch) -> Vec<String> {
let c = batch
.column(batch.schema().index_of("name").unwrap())
.as_any()
.downcast_ref::<StringArray>()
.unwrap();
(0..batch.num_rows())
.map(|r| c.value(r).to_string())
.collect()
}
}
#[async_trait]
impl GraphSink for InMemoryGraph {
async fn apply(&mut self, mutation: &GraphMutation) -> Result<SinkReport> {
let mut report = SinkReport::default();
for set in &mutation.vertices {
let ids = Self::ids(&set.batch, &set.id_column);
let names = Self::names(&set.batch);
for (id, name) in ids.into_iter().zip(names) {
self.nodes.insert((set.label.clone(), id), name);
report.vertices_written += 1;
}
}
for set in &mutation.deletes {
for id in Self::ids(&set.batch, &set.id_column) {
self.nodes.remove(&(set.label.clone(), id));
report.deleted += 1;
}
}
Ok(report)
}
}
fn person_transforms() -> Vec<Arc<dyn Transform>> {
vec![Arc::new(VertexTransform {
label: "Person".into(),
id_column: "user_id".into(),
properties: vec!["name".into()],
rename: Default::default(),
})]
}
#[tokio::test]
async fn cdc_update_pair_keeps_node_with_new_values() {
let transforms = person_transforms();
let ctx = TableContext {
table: "lake.users".into(),
};
let change = partition_change_rows(&[update_pair_change_batch()]).unwrap();
assert!(
change.deletes.is_empty(),
"UPDATE_BEFORE must not become a delete"
);
let up_rows: usize = change.upserts.iter().map(|b| b.num_rows()).sum();
assert_eq!(up_rows, 1);
let mut mutation = GraphMutation::default();
for t in &transforms {
mutation.merge(t.apply(&ctx, &change.upserts).unwrap());
}
mutation.deletes = deletes_from_change_rows(&transforms, &ctx, &change.deletes).unwrap();
let mut sink = InMemoryGraph::default();
let report = sink.apply(&mutation).await.unwrap();
let survived = sink.nodes.get(&("Person".to_string(), 1)).cloned();
assert_eq!(
survived.as_deref(),
Some("new"),
"updated node must survive with its new value, not be deleted"
);
assert_eq!(report.deleted, 0, "an UPDATE must delete nothing");
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"cdc::cdc_update_pair_keeps_node_with_new_values",
survived.as_deref() == Some("new") && report.deleted == 0,
&format!(
"after UPDATE: Person(1).name={survived:?}, deleted={}",
report.deleted
),
);
}
#[tokio::test]
async fn cdc_delete_still_removes_node() {
let transforms = person_transforms();
let ctx = TableContext {
table: "lake.users".into(),
};
let seed = RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("user_id", DataType::Int64, false),
Field::new("name", DataType::Utf8, false),
Field::new(SyncMode::CHANGE_TYPE_COLUMN, DataType::Utf8, false),
])),
vec![
Arc::new(Int64Array::from(vec![1, 1])),
Arc::new(StringArray::from(vec!["x", "x"])),
Arc::new(StringArray::from(vec!["INSERT", "DELETE"])),
],
)
.unwrap();
let change = partition_change_rows(&[seed]).unwrap();
assert_eq!(
change.deletes.iter().map(|b| b.num_rows()).sum::<usize>(),
1
);
let mut mutation = GraphMutation::default();
for t in &transforms {
mutation.merge(t.apply(&ctx, &change.upserts).unwrap());
}
mutation.deletes = deletes_from_change_rows(&transforms, &ctx, &change.deletes).unwrap();
let mut sink = InMemoryGraph::default();
let report = sink.apply(&mutation).await.unwrap();
assert!(
sink.nodes.get(&("Person".to_string(), 1)).is_none(),
"INSERT+DELETE in one round resolves to gone"
);
assert_eq!(report.deleted, 1);
}
}
struct FellBackSource;
#[async_trait]
impl TableSource for FellBackSource {
fn cdc_capability(&self, _table: &str) -> Option<crate::cdc::CdcCapability> {
Some(crate::cdc::CdcCapability {
chosen: crate::cdc::CdcKind::IcebergChangelogView,
server_version: Some("4.2.0".into()),
server_major_minor: Some((4, 2)),
native_expected: true,
native_usable: false,
fell_back: true,
reason: "CHANGES did not resolve — connector implements no Changelog".into(),
})
}
async fn fetch_new(
&self,
_spec: &TableSpec,
state: &TableState,
) -> Result<Option<(Vec<RecordBatch>, TableState)>> {
if state.snapshot_id == Some(9) {
return Ok(None);
}
Ok(Some((
vec![users_batch()],
TableState {
snapshot_id: Some(9),
..state.clone()
},
)))
}
}
#[tokio::test]
async fn cdc_round_records_the_backend_and_its_fallback_in_state_json() {
let dir = tempdir().unwrap();
let path = dir.path().join("state.json");
let yaml = MAPPING_YAML.replace("mode: incremental", "mode: cdc");
assert!(
yaml.contains("mode: cdc"),
"the fixture must exercise CDC mode"
);
let spec = MappingSpec::from_yaml(&yaml).unwrap();
let mut listener = Listener::new(
spec,
TransformRegistry::new(),
FellBackSource,
RecordingSink::default(),
SyncState::open(&path).unwrap(),
);
let round = listener.sync_once().await.unwrap();
assert_eq!(round.failed, vec![], "the round itself succeeds");
assert_eq!(round.synced.len(), 1);
let raw = std::fs::read_to_string(&path).unwrap();
let json: serde_json::Value = serde_json::from_str(&raw).unwrap();
let cdc = &json["lake.social.users"]["cdc"];
assert!(
!cdc.is_null(),
"state_json must carry the CDC verdict:\n{raw}"
);
assert_eq!(cdc["chosen"], "iceberg_changelog_view");
assert_eq!(cdc["server_version"], "4.2.0");
assert_eq!(cdc["native_expected"], true);
assert_eq!(cdc["native_usable"], false, "the degrade is recorded");
assert_eq!(cdc["fell_back"], true, "and it is NOT silent");
assert!(
cdc["reason"].as_str().unwrap().contains("Changelog"),
"the engine's own reason survives: {cdc}",
);
let back = SyncState::open(&path).unwrap();
let cap = back
.get("lake.social.users")
.cdc
.expect("verdict persisted");
assert!(cap.fell_back);
assert_eq!(cap.chosen, crate::cdc::CdcKind::IcebergChangelogView);
crate::testmatrix::emit(
crate::testmatrix::COMPONENT,
"cdc::cdc_round_records_the_backend_and_its_fallback_in_state_json",
cap.fell_back,
&format!(
"state_json records backend={} native_expected=true native_usable=false fell_back=true",
cap.chosen.as_str()
),
);
}