Influxdata
/
influxdb
mirror of https://github.com/influxdata/influxdb.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

refactor(ingester): Replace test_util `DmlWrite` with `WriteOperation` 

This change replaces the test_util equality and write generation code
to use the new `IngestOp::Write(WriteOperation)` type, removing many
pointless conversions in tests.
pull/24376/head
Fraser Savage 2023-06-16 13:16:10 +01:00
parent 8908f5fc96
commit a3a4145774
No known key found for this signature in database
GPG Key ID: DE47C33CE8C5C446
10 changed files with 200 additions and 228 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
ingester/src/buffer_tree/namespace.rs
View File

@ -164,7 +164,7 @@ where
table_data table_data
.buffer_table_write( .buffer_table_write(
partitioned_data.sequence_number(), partitioned_data.sequence_number(),
partitioned_data.data(), partitioned_data.into_data(),
partition_key.clone(), partition_key.clone(),
) )
.await?; .await?;
@ -214,7 +214,6 @@ where
mod tests { mod tests {
use std::sync::Arc; use std::sync::Arc;
use dml::DmlOperation;
use metric::{Attributes, Metric}; use metric::{Attributes, Metric};
use super::*; use super::*;
@ -262,20 +261,17 @@ mod tests {
assert!(ns.table(ARBITRARY_TABLE_ID).is_none()); assert!(ns.table(ARBITRARY_TABLE_ID).is_none());
// Write some test data // Write some test data
ns.apply( ns.apply(IngestOp::Write(make_write_op(
DmlOperation::Write(make_write_op( &ARBITRARY_PARTITION_KEY,
&ARBITRARY_PARTITION_KEY, ARBITRARY_NAMESPACE_ID,
ARBITRARY_NAMESPACE_ID, &ARBITRARY_TABLE_NAME,
&ARBITRARY_TABLE_NAME, ARBITRARY_TABLE_ID,
ARBITRARY_TABLE_ID, 0,
0, &format!(
&format!( r#"{},city=Medford day="sun",temp=55 22"#,
r#"{},city=Medford day="sun",temp=55 22"#, &*ARBITRARY_TABLE_NAME
&*ARBITRARY_TABLE_NAME ),
), )))
))
.into(),
)
.await .await
.expect("buffer op should succeed"); .expect("buffer op should succeed");

251
ingester/src/buffer_tree/root.rs
View File

@ -232,7 +232,6 @@ mod tests {
use assert_matches::assert_matches; use assert_matches::assert_matches;
use data_types::{PartitionId, PartitionKey}; use data_types::{PartitionId, PartitionKey};
use datafusion::{assert_batches_eq, assert_batches_sorted_eq}; use datafusion::{assert_batches_eq, assert_batches_sorted_eq};
use dml::DmlOperation;
use futures::StreamExt; use futures::StreamExt;
use metric::{Attributes, Metric}; use metric::{Attributes, Metric};
@ -285,20 +284,17 @@ mod tests {
assert!(ns.table(ARBITRARY_TABLE_ID).is_none()); assert!(ns.table(ARBITRARY_TABLE_ID).is_none());
// Write some test data // Write some test data
ns.apply( ns.apply(IngestOp::Write(make_write_op(
DmlOperation::Write(make_write_op( &ARBITRARY_PARTITION_KEY,
&ARBITRARY_PARTITION_KEY, ARBITRARY_NAMESPACE_ID,
ARBITRARY_NAMESPACE_ID, &ARBITRARY_TABLE_NAME,
&ARBITRARY_TABLE_NAME, ARBITRARY_TABLE_ID,
ARBITRARY_TABLE_ID, 0,
0, &format!(
&format!( r#"{},city=Madrid day="sun",temp=55 22"#,
r#"{},city=Madrid day="sun",temp=55 22"#, &*ARBITRARY_TABLE_NAME
&*ARBITRARY_TABLE_NAME ),
), )))
))
.into(),
)
.await .await
.expect("buffer op should succeed"); .expect("buffer op should succeed");
@ -329,7 +325,7 @@ mod tests {
( (
$name:ident, $name:ident,
partitions = [$($partition:expr), +], // The set of PartitionData for the mock partition provider partitions = [$($partition:expr), +], // The set of PartitionData for the mock partition provider
writes = [$($write:expr), *], // The set of DmlWrite to apply() writes = [$($write:expr), *], // The set of WriteOperation to apply()
want = $want:expr // The expected results of querying ARBITRARY_NAMESPACE_ID and ARBITRARY_TABLE_ID want = $want:expr // The expected results of querying ARBITRARY_NAMESPACE_ID and ARBITRARY_TABLE_ID
) => { ) => {
paste::paste! { paste::paste! {
@ -352,9 +348,9 @@ mod tests {
Arc::new(metric::Registry::default()), Arc::new(metric::Registry::default()),
); );
// Write the provided DmlWrites // Write the provided WriteOperation
$( $(
buf.apply(DmlOperation::Write($write).into()) buf.apply(IngestOp::Write($write).into())
.await .await
.expect("failed to perform write"); .expect("failed to perform write");
)* )*
@ -629,39 +625,33 @@ mod tests {
); );
// Write data to partition p1, in the arbitrary table // Write data to partition p1, in the arbitrary table
buf.apply( buf.apply(IngestOp::Write(make_write_op(
DmlOperation::Write(make_write_op( &PartitionKey::from("p1"),
&PartitionKey::from("p1"), ARBITRARY_NAMESPACE_ID,
ARBITRARY_NAMESPACE_ID, &ARBITRARY_TABLE_NAME,
&ARBITRARY_TABLE_NAME, ARBITRARY_TABLE_ID,
ARBITRARY_TABLE_ID, 0,
0, &format!(
&format!( r#"{},region=Asturias temp=35 4242424242"#,
r#"{},region=Asturias temp=35 4242424242"#, &*ARBITRARY_TABLE_NAME
&*ARBITRARY_TABLE_NAME ),
), )))
))
.into(),
)
.await .await
.expect("failed to write initial data"); .expect("failed to write initial data");
// Write a duplicate record with the same series key & timestamp, but a // Write a duplicate record with the same series key & timestamp, but a
// different temp value. // different temp value.
buf.apply( buf.apply(IngestOp::Write(make_write_op(
DmlOperation::Write(make_write_op( &PartitionKey::from("p2"),
&PartitionKey::from("p2"), ARBITRARY_NAMESPACE_ID,
ARBITRARY_NAMESPACE_ID, &ARBITRARY_TABLE_NAME,
&ARBITRARY_TABLE_NAME, ARBITRARY_TABLE_ID,
ARBITRARY_TABLE_ID, 1,
1, &format!(
&format!( r#"{},region=Asturias temp=12 4242424242"#,
r#"{},region=Asturias temp=12 4242424242"#, &*ARBITRARY_TABLE_NAME
&*ARBITRARY_TABLE_NAME ),
), )))
))
.into(),
)
.await .await
.expect("failed to overwrite data"); .expect("failed to overwrite data");
@ -734,57 +724,48 @@ mod tests {
assert_eq!(buf.partitions().count(), 0); assert_eq!(buf.partitions().count(), 0);
// Write data to partition p1, in the arbitrary table // Write data to partition p1, in the arbitrary table
buf.apply( buf.apply(IngestOp::Write(make_write_op(
DmlOperation::Write(make_write_op( &PartitionKey::from("p1"),
&PartitionKey::from("p1"), ARBITRARY_NAMESPACE_ID,
ARBITRARY_NAMESPACE_ID, &ARBITRARY_TABLE_NAME,
&ARBITRARY_TABLE_NAME, ARBITRARY_TABLE_ID,
ARBITRARY_TABLE_ID, 0,
0, &format!(
&format!( r#"{},region=Asturias temp=35 4242424242"#,
r#"{},region=Asturias temp=35 4242424242"#, &*ARBITRARY_TABLE_NAME
&*ARBITRARY_TABLE_NAME ),
), )))
))
.into(),
)
.await .await
.expect("failed to write initial data"); .expect("failed to write initial data");
assert_eq!(buf.partitions().count(), 1); assert_eq!(buf.partitions().count(), 1);
// Write data to partition p2, in the arbitrary table // Write data to partition p2, in the arbitrary table
buf.apply( buf.apply(IngestOp::Write(make_write_op(
DmlOperation::Write(make_write_op( &PartitionKey::from("p2"),
&PartitionKey::from("p2"), ARBITRARY_NAMESPACE_ID,
ARBITRARY_NAMESPACE_ID, &ARBITRARY_TABLE_NAME,
&ARBITRARY_TABLE_NAME, ARBITRARY_TABLE_ID,
ARBITRARY_TABLE_ID, 0,
0, &format!(
&format!( r#"{},region=Asturias temp=35 4242424242"#,
r#"{},region=Asturias temp=35 4242424242"#, &*ARBITRARY_TABLE_NAME
&*ARBITRARY_TABLE_NAME ),
), )))
))
.into(),
)
.await .await
.expect("failed to write initial data"); .expect("failed to write initial data");
assert_eq!(buf.partitions().count(), 2); assert_eq!(buf.partitions().count(), 2);
// Write data to partition p3, in the second table // Write data to partition p3, in the second table
buf.apply( buf.apply(IngestOp::Write(make_write_op(
DmlOperation::Write(make_write_op( &PartitionKey::from("p3"),
&PartitionKey::from("p3"), ARBITRARY_NAMESPACE_ID,
ARBITRARY_NAMESPACE_ID, TABLE2_NAME,
TABLE2_NAME, TABLE2_ID,
TABLE2_ID, 0,
0, &format!(r#"{},region=Asturias temp=35 4242424242"#, TABLE2_NAME),
&format!(r#"{},region=Asturias temp=35 4242424242"#, TABLE2_NAME), )))
))
.into(),
)
.await .await
.expect("failed to write initial data"); .expect("failed to write initial data");
@ -832,20 +813,17 @@ mod tests {
}); });
// Write data to partition p1, in the arbitrary table // Write data to partition p1, in the arbitrary table
buf.apply( buf.apply(IngestOp::Write(make_write_op(
DmlOperation::Write(make_write_op( &PartitionKey::from("p1"),
&PartitionKey::from("p1"), ARBITRARY_NAMESPACE_ID,
ARBITRARY_NAMESPACE_ID, &ARBITRARY_TABLE_NAME,
&ARBITRARY_TABLE_NAME, ARBITRARY_TABLE_ID,
ARBITRARY_TABLE_ID, 0,
0, &format!(
&format!( r#"{},region=Asturias temp=35 4242424242"#,
r#"{},region=Asturias temp=35 4242424242"#, &*ARBITRARY_TABLE_NAME
&*ARBITRARY_TABLE_NAME ),
), )))
))
.into(),
)
.await .await
.expect("failed to write data"); .expect("failed to write data");
@ -911,20 +889,17 @@ mod tests {
); );
// Write data to partition p1, in the arbitrary table // Write data to partition p1, in the arbitrary table
buf.apply( buf.apply(IngestOp::Write(make_write_op(
DmlOperation::Write(make_write_op( &PartitionKey::from("p1"),
&PartitionKey::from("p1"), ARBITRARY_NAMESPACE_ID,
ARBITRARY_NAMESPACE_ID, &ARBITRARY_TABLE_NAME,
&ARBITRARY_TABLE_NAME, ARBITRARY_TABLE_ID,
ARBITRARY_TABLE_ID, 0,
0, &format!(
&format!( r#"{},region=Madrid temp=35 4242424242"#,
r#"{},region=Madrid temp=35 4242424242"#, &*ARBITRARY_TABLE_NAME
&*ARBITRARY_TABLE_NAME ),
), )))
))
.into(),
)
.await .await
.expect("failed to write initial data"); .expect("failed to write initial data");
@ -938,39 +913,33 @@ mod tests {
// Perform a write concurrent to the consumption of the query stream // Perform a write concurrent to the consumption of the query stream
// that creates a new partition (p2) in the same table. // that creates a new partition (p2) in the same table.
buf.apply( buf.apply(IngestOp::Write(make_write_op(
DmlOperation::Write(make_write_op( &PartitionKey::from("p2"),
&PartitionKey::from("p2"), ARBITRARY_NAMESPACE_ID,
ARBITRARY_NAMESPACE_ID, &ARBITRARY_TABLE_NAME,
&ARBITRARY_TABLE_NAME, ARBITRARY_TABLE_ID,
ARBITRARY_TABLE_ID, 1,
1, &format!(
&format!( r#"{},region=Asturias temp=20 4242424242"#,
r#"{},region=Asturias temp=20 4242424242"#, &*ARBITRARY_TABLE_NAME
&*ARBITRARY_TABLE_NAME ),
), )))
))
.into(),
)
.await .await
.expect("failed to perform concurrent write to new partition"); .expect("failed to perform concurrent write to new partition");
// Perform another write that hits the partition within the query // Perform another write that hits the partition within the query
// results snapshot (p1) before the partition is read. // results snapshot (p1) before the partition is read.
buf.apply( buf.apply(IngestOp::Write(make_write_op(
DmlOperation::Write(make_write_op( &PartitionKey::from("p1"),
&PartitionKey::from("p1"), ARBITRARY_NAMESPACE_ID,
ARBITRARY_NAMESPACE_ID, &ARBITRARY_TABLE_NAME,
&ARBITRARY_TABLE_NAME, ARBITRARY_TABLE_ID,
ARBITRARY_TABLE_ID, 2,
2, &format!(
&format!( r#"{},region=Murcia temp=30 4242424242"#,
r#"{},region=Murcia temp=30 4242424242"#, &*ARBITRARY_TABLE_NAME
&*ARBITRARY_TABLE_NAME ),
), )))
))
.into(),
)
.await .await
.expect("failed to perform concurrent write to existing partition"); .expect("failed to perform concurrent write to existing partition");

16
ingester/src/dml_payload/ingest_op.rs
View File

@ -40,6 +40,13 @@ impl IngestOp {
Self::Write(w) => w.span_context.as_ref(), Self::Write(w) => w.span_context.as_ref(),
} }
} }
/// Sets the tracing context associated with the [`IngestOp`]
pub fn set_span_context(&mut self, ctx: SpanContext) {
match self {
Self::Write(w) => w.span_context = Some(ctx),
}
}
} }
/// A decoded representation of the data contained by an RPC write /// A decoded representation of the data contained by an RPC write
@ -224,8 +231,13 @@ impl PartitionedData {
self.sequence_number self.sequence_number
} }
/// Consumes `self`, returning the data /// Returns a reference to the data
pub fn data(self) -> MutableBatch { pub fn data(&self) -> &MutableBatch {
&self.data
}
/// Consumes `self`, returning the owned data
pub fn into_data(self) -> MutableBatch {
self.data self.data
} }
} }

5
ingester/src/dml_sink/instrumentation.rs
View File

@ -74,7 +74,6 @@ mod tests {
use assert_matches::assert_matches; use assert_matches::assert_matches;
use data_types::{NamespaceId, PartitionId, PartitionKey, TableId}; use data_types::{NamespaceId, PartitionId, PartitionKey, TableId};
use dml::DmlOperation;
use lazy_static::lazy_static; use lazy_static::lazy_static;
use metric::Attributes; use metric::Attributes;
@ -123,14 +122,14 @@ mod tests {
let metrics = metric::Registry::default(); let metrics = metric::Registry::default();
let decorator = DmlSinkInstrumentation::new(LAYER_NAME, mock, &metrics); let decorator = DmlSinkInstrumentation::new(LAYER_NAME, mock, &metrics);
let op = DmlOperation::Write(make_write_op( let op = IngestOp::Write(make_write_op(
&PARTITION_KEY, &PARTITION_KEY,
NAMESPACE_ID, NAMESPACE_ID,
TABLE_NAME, TABLE_NAME,
TABLE_ID, TABLE_ID,
42, 42,
"banana-report,tag=1 v=2 42424242", "banana-report,tag=1 v=2 42424242",
)).into(); ));
// Call the decorator and assert the return value // Call the decorator and assert the return value
let got = decorator let got = decorator

25
ingester/src/dml_sink/tracing.rs
View File

@ -59,7 +59,6 @@ mod tests {
use assert_matches::assert_matches; use assert_matches::assert_matches;
use data_types::{NamespaceId, PartitionId, PartitionKey, TableId}; use data_types::{NamespaceId, PartitionId, PartitionKey, TableId};
use dml::{DmlMeta, DmlOperation};
use lazy_static::lazy_static; use lazy_static::lazy_static;
use trace::{ctx::SpanContext, span::SpanStatus, RingBufferTraceCollector, TraceCollector}; use trace::{ctx::SpanContext, span::SpanStatus, RingBufferTraceCollector, TraceCollector};
@ -120,7 +119,7 @@ mod tests {
let traces: Arc<dyn TraceCollector> = Arc::new(RingBufferTraceCollector::new(5)); let traces: Arc<dyn TraceCollector> = Arc::new(RingBufferTraceCollector::new(5));
let span = SpanContext::new(Arc::clone(&traces)); let span = SpanContext::new(Arc::clone(&traces));
let mut op = DmlOperation::Write(make_write_op( let mut op = IngestOp::Write(make_write_op(
&PARTITION_KEY, &PARTITION_KEY,
NAMESPACE_ID, NAMESPACE_ID,
TABLE_NAME, TABLE_NAME,
@ -130,17 +129,11 @@ mod tests {
)); ));
// Populate the metadata with a span context. // Populate the metadata with a span context.
let meta = op.meta(); op.set_span_context(span);
op.set_meta(DmlMeta::sequenced(
meta.sequence().unwrap(),
meta.producer_ts().unwrap(),
Some(span),
42,
));
// Drive the trace wrapper // Drive the trace wrapper
DmlSinkTracing::new(mock, "bananas") DmlSinkTracing::new(mock, "bananas")
.apply(op.into()) .apply(op)
.await .await
.expect("wrapper should not modify result"); .expect("wrapper should not modify result");
@ -156,7 +149,7 @@ mod tests {
let traces: Arc<dyn TraceCollector> = Arc::new(RingBufferTraceCollector::new(5)); let traces: Arc<dyn TraceCollector> = Arc::new(RingBufferTraceCollector::new(5));
let span = SpanContext::new(Arc::clone(&traces)); let span = SpanContext::new(Arc::clone(&traces));
let mut op = DmlOperation::Write(make_write_op( let mut op = IngestOp::Write(make_write_op(
&PARTITION_KEY, &PARTITION_KEY,
NAMESPACE_ID, NAMESPACE_ID,
TABLE_NAME, TABLE_NAME,
@ -166,17 +159,11 @@ mod tests {
)); ));
// Populate the metadata with a span context. // Populate the metadata with a span context.
let meta = op.meta(); op.set_span_context(span);
op.set_meta(DmlMeta::sequenced(
meta.sequence().unwrap(),
meta.producer_ts().unwrap(),
Some(span),
42,
));
// Drive the trace wrapper // Drive the trace wrapper
let got = DmlSinkTracing::new(mock, "bananas") let got = DmlSinkTracing::new(mock, "bananas")
.apply(op.into()) .apply(op)
.await .await
.expect_err("wrapper should not modify result"); .expect_err("wrapper should not modify result");
assert_matches!(got, DmlError::Wal(s) => { assert_matches!(got, DmlError::Wal(s) => {

20
ingester/src/init/wal_replay.rs
View File

@ -350,12 +350,12 @@ mod tests {
// Apply the first op through the decorator // Apply the first op through the decorator
wal_sink wal_sink
.apply(DmlOperation::Write(op1.clone()).into()) .apply(IngestOp::Write(op1.clone()))
.await .await
.expect("wal should not error"); .expect("wal should not error");
// And the second op // And the second op
wal_sink wal_sink
.apply(DmlOperation::Write(op2.clone()).into()) .apply(IngestOp::Write(op2.clone()))
.await .await
.expect("wal should not error"); .expect("wal should not error");
@ -364,7 +364,7 @@ mod tests {
// Write the third op // Write the third op
wal_sink wal_sink
.apply(DmlOperation::Write(op3.clone()).into()) .apply(IngestOp::Write(op3.clone()))
.await .await
.expect("wal should not error"); .expect("wal should not error");
@ -389,7 +389,13 @@ mod tests {
// Put at least one write into the buffer so it is a candidate for persistence // Put at least one write into the buffer so it is a candidate for persistence
partition partition
.buffer_write( .buffer_write(
op1.tables().next().unwrap().1.clone(), op1.tables()
.next()
.unwrap()
.1
.partitioned_data()
.data()
.clone(),
SequenceNumber::new(1), SequenceNumber::new(1),
) )
.unwrap(); .unwrap();
@ -414,9 +420,9 @@ mod tests {
IngestOp::Write(ref w2), IngestOp::Write(ref w2),
IngestOp::Write(ref w3) IngestOp::Write(ref w3)
] => { ] => {
assert_dml_writes_eq(w1.into(), op1); assert_dml_writes_eq(w1.clone(), op1);
assert_dml_writes_eq(w2.into(), op2); assert_dml_writes_eq(w2.clone(), op2);
assert_dml_writes_eq(w3.into(), op3); assert_dml_writes_eq(w3.clone(), op3);
} }
); );

21
ingester/src/persist/handle.rs
View File

@ -475,7 +475,6 @@ mod tests {
use std::{sync::Arc, task::Poll, time::Duration}; use std::{sync::Arc, task::Poll, time::Duration};
use assert_matches::assert_matches; use assert_matches::assert_matches;
use dml::DmlOperation;
use futures::Future; use futures::Future;
use iox_catalog::mem::MemCatalog; use iox_catalog::mem::MemCatalog;
use object_store::memory::InMemory; use object_store::memory::InMemory;
@ -492,6 +491,7 @@ mod tests {
post_write::mock::MockPostWriteObserver, BufferTree, post_write::mock::MockPostWriteObserver, BufferTree,
}, },
deferred_load::DeferredLoad, deferred_load::DeferredLoad,
dml_payload::IngestOp,
dml_sink::DmlSink, dml_sink::DmlSink,
ingest_state::IngestStateError, ingest_state::IngestStateError,
persist::{ persist::{
@ -523,17 +523,14 @@ mod tests {
); );
buffer_tree buffer_tree
.apply( .apply(IngestOp::Write(make_write_op(
DmlOperation::Write(make_write_op( &ARBITRARY_PARTITION_KEY,
&ARBITRARY_PARTITION_KEY, ARBITRARY_NAMESPACE_ID,
ARBITRARY_NAMESPACE_ID, &ARBITRARY_TABLE_NAME,
&ARBITRARY_TABLE_NAME, ARBITRARY_TABLE_ID,
ARBITRARY_TABLE_ID, 0,
0, &format!("{},good=yes level=1000 4242424242", &*ARBITRARY_TABLE_NAME),
&format!("{},good=yes level=1000 4242424242", &*ARBITRARY_TABLE_NAME), )))
))
.into(),
)
.await .await
.expect("failed to write partition test dataa"); .expect("failed to write partition test dataa");

6
ingester/src/persist/mod.rs
View File

@ -17,7 +17,6 @@ mod tests {
use assert_matches::assert_matches; use assert_matches::assert_matches;
use data_types::{CompactionLevel, ParquetFile}; use data_types::{CompactionLevel, ParquetFile};
use dml::DmlOperation;
use futures::TryStreamExt; use futures::TryStreamExt;
use iox_catalog::{ use iox_catalog::{
interface::{get_schema_by_id, Catalog, SoftDeletedRows}, interface::{get_schema_by_id, Catalog, SoftDeletedRows},
@ -41,6 +40,7 @@ mod tests {
post_write::mock::MockPostWriteObserver, post_write::mock::MockPostWriteObserver,
BufferTree, BufferTree,
}, },
dml_payload::IngestOp,
dml_sink::DmlSink, dml_sink::DmlSink,
ingest_state::IngestState, ingest_state::IngestState,
persist::handle::PersistHandle, persist::handle::PersistHandle,
@ -98,7 +98,7 @@ mod tests {
validate_or_insert_schema( validate_or_insert_schema(
write write
.tables() .tables()
.map(|(_id, data)| (&***ARBITRARY_TABLE_NAME, data)), .map(|(_id, data)| (&***ARBITRARY_TABLE_NAME, data.partitioned_data().data())),
&schema, &schema,
&mut *repos, &mut *repos,
) )
@ -108,7 +108,7 @@ mod tests {
drop(repos); // Don't you love this testing-only deadlock bug? #3859 drop(repos); // Don't you love this testing-only deadlock bug? #3859
// Apply the write // Apply the write
buf.apply(DmlOperation::Write(write).into()) buf.apply(IngestOp::Write(write))
.await .await
.expect("failed to apply write to buffer"); .expect("failed to apply write to buffer");

54
ingester/src/test_util.rs
View File

@ -1,7 +1,6 @@
use std::{collections::BTreeMap, sync::Arc, time::Duration}; use std::{collections::BTreeMap, sync::Arc, time::Duration};
use data_types::{NamespaceId, Partition, PartitionId, PartitionKey, SequenceNumber, TableId}; use data_types::{NamespaceId, Partition, PartitionId, PartitionKey, SequenceNumber, TableId};
use dml::{DmlMeta, DmlWrite};
use iox_catalog::{interface::Catalog, test_helpers::arbitrary_namespace}; use iox_catalog::{interface::Catalog, test_helpers::arbitrary_namespace};
use lazy_static::lazy_static; use lazy_static::lazy_static;
use mutable_batch_lp::lines_to_batches; use mutable_batch_lp::lines_to_batches;
@ -20,6 +19,7 @@ use crate::{
}, },
}, },
deferred_load::DeferredLoad, deferred_load::DeferredLoad,
dml_payload::{PartitionedData, TableData, WriteOperation},
}; };
pub(crate) const ARBITRARY_PARTITION_ID: PartitionId = PartitionId::new(1); pub(crate) const ARBITRARY_PARTITION_ID: PartitionId = PartitionId::new(1);
@ -259,7 +259,6 @@ macro_rules! make_partition_stream {
/// # Panics /// # Panics
/// ///
/// This method panics if `lines` contains data for more than one table. /// This method panics if `lines` contains data for more than one table.
// TODO(savage): Return a `dml_payload::WriteOperation` here.
#[track_caller] #[track_caller]
pub(crate) fn make_write_op( pub(crate) fn make_write_op(
partition_key: &PartitionKey, partition_key: &PartitionKey,
@ -268,7 +267,7 @@ pub(crate) fn make_write_op(
table_id: TableId, table_id: TableId,
sequence_number: i64, sequence_number: i64,
lines: &str, lines: &str,
) -> DmlWrite { ) -> WriteOperation {
let mut tables_by_name = lines_to_batches(lines, 0).expect("invalid LP"); let mut tables_by_name = lines_to_batches(lines, 0).expect("invalid LP");
assert_eq!( assert_eq!(
tables_by_name.len(), tables_by_name.len(),
@ -283,19 +282,18 @@ pub(crate) fn make_write_op(
.expect("table_name does not exist in LP"), .expect("table_name does not exist in LP"),
)] )]
.into_iter() .into_iter()
.map(|(table_id, mutable_batch)| {
(
table_id,
TableData::new(
table_id,
PartitionedData::new(SequenceNumber::new(sequence_number), mutable_batch),
),
)
})
.collect(); .collect();
DmlWrite::new( WriteOperation::new(namespace_id, tables_by_id, partition_key.clone(), None)
namespace_id,
tables_by_id,
partition_key.clone(),
DmlMeta::sequenced(
SequenceNumber::new(sequence_number),
iox_time::Time::MIN,
None,
42,
),
)
} }
pub(crate) async fn populate_catalog( pub(crate) async fn populate_catalog(
@ -318,17 +316,19 @@ pub(crate) async fn populate_catalog(
/// Assert `a` and `b` have identical metadata, and that when converting /// Assert `a` and `b` have identical metadata, and that when converting
/// them to Arrow batches they produces identical output. /// them to Arrow batches they produces identical output.
#[track_caller] #[track_caller]
pub(crate) fn assert_dml_writes_eq(a: DmlWrite, b: DmlWrite) { pub(crate) fn assert_dml_writes_eq(a: WriteOperation, b: WriteOperation) {
assert_eq!(a.namespace_id(), b.namespace_id(), "namespace"); assert_eq!(a.namespace(), b.namespace(), "namespace");
assert_eq!(a.table_count(), b.table_count(), "table count"); assert_eq!(a.tables().count(), b.tables().count(), "table count");
assert_eq!(a.min_timestamp(), b.min_timestamp(), "min timestamp");
assert_eq!(a.max_timestamp(), b.max_timestamp(), "max timestamp");
assert_eq!(a.partition_key(), b.partition_key(), "partition key"); assert_eq!(a.partition_key(), b.partition_key(), "partition key");
// Assert sequence numbers were reassigned // Assert sequence numbers were reassigned
let seq_a = a.meta().sequence(); for (a_table, b_table) in a.tables().zip(b.tables()) {
let seq_b = b.meta().sequence(); assert_eq!(
assert_eq!(seq_a, seq_b, "sequence numbers differ"); a_table.1.partitioned_data().sequence_number(),
b_table.1.partitioned_data().sequence_number(),
"sequence number"
);
}
let a = a.into_tables().collect::<BTreeMap<_, _>>(); let a = a.into_tables().collect::<BTreeMap<_, _>>();
let b = b.into_tables().collect::<BTreeMap<_, _>>(); let b = b.into_tables().collect::<BTreeMap<_, _>>();
@ -336,9 +336,15 @@ pub(crate) fn assert_dml_writes_eq(a: DmlWrite, b: DmlWrite) {
a.into_iter().zip(b.into_iter()).for_each(|(a, b)| { a.into_iter().zip(b.into_iter()).for_each(|(a, b)| {
assert_eq!(a.0, b.0, "table IDs differ - a table is missing!"); assert_eq!(a.0, b.0, "table IDs differ - a table is missing!");
assert_eq!( assert_eq!(
a.1.to_arrow(Projection::All) a.1.partitioned_data()
.data()
.clone()
.to_arrow(Projection::All)
.expect("failed projection for a"), .expect("failed projection for a"),
b.1.to_arrow(Projection::All) b.1.partitioned_data()
.data()
.clone()
.to_arrow(Projection::All)
.expect("failed projection for b"), .expect("failed projection for b"),
"table data differs" "table data differs"
); );

2
ingester/src/wal/wal_sink.rs
View File

@ -291,7 +291,7 @@ mod tests {
// Apply the op through the decorator, which should time out // Apply the op through the decorator, which should time out
let err = wal_sink let err = wal_sink
.apply(DmlOperation::Write(op.clone()).into()) .apply(IngestOp::Write(op.clone()))
.await .await
.expect_err("write should time out"); .expect_err("write should time out");