mz_testdrive/action/kafka/

verify_data.rs

// Copyright Materialize, Inc. and contributors. All rights reserved.
//
// Use of this software is governed by the Business Source License
// included in the LICENSE file.
//
// 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.

use std::collections::BTreeSet;
use std::fmt::Debug;
use std::time::{Duration, Instant};
use std::{cmp, str};

use anyhow::{Context, anyhow, bail, ensure};
use mz_kafka_util::client::get_partitions;
use mz_ore::task;
use mz_postgres_util::{Sql, query_one, sql};
use rdkafka::consumer::{BaseConsumer, CommitMode, Consumer, StreamConsumer};
use rdkafka::error::KafkaError;
use rdkafka::message::{Headers, Message};
use rdkafka::types::RDKafkaErrorCode;
use rdkafka::{Offset, TopicPartitionList};
use regex::Regex;
use tokio::pin;
use tokio_stream::StreamExt;

use crate::action::{ControlFlow, State};
use crate::format::avro::{self, DebugValue};
use crate::parser::BuiltinCommand;

#[derive(Debug, Clone, Copy)]
enum Format {
    Avro,
    Json,
    Bytes,
    Text,
}

impl TryFrom<&str> for Format {
    type Error = anyhow::Error;

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        match value {
            "avro" => Ok(Format::Avro),
            "json" => Ok(Format::Json),
            "bytes" => Ok(Format::Bytes),
            "text" => Ok(Format::Text),
            f => bail!("unknown format: {}", f),
        }
    }
}

#[derive(Debug)]
struct RecordFormat {
    key: Format,
    value: Format,
    requires_key: bool,
}

#[allow(dead_code)]
#[derive(Debug, Clone)]
enum DecodedValue {
    Avro(DebugValue),
    Json(serde_json::Value),
    Bytes(Vec<u8>),
    Text(String),
}

enum Topic {
    FromSink(String),
    Named(String),
}

#[derive(Debug, Clone)]
struct Record<A> {
    headers: Vec<String>,
    key: Option<A>,
    value: Option<A>,
    partition: Option<i32>,
}

async fn get_topic(sink: &str, topic_field: &str, state: &State) -> Result<String, anyhow::Error> {
    let query = sql!(
        "SELECT {} FROM mz_sinks JOIN mz_kafka_sinks \
         ON mz_sinks.id = mz_kafka_sinks.id \
         JOIN mz_schemas s ON s.id = mz_sinks.schema_id \
         LEFT JOIN mz_databases d ON d.id = s.database_id \
         WHERE d.name = $1 \
         AND s.name = $2 \
         AND mz_sinks.name = $3",
        Sql::ident(topic_field)
    );
    let sink_fields: Vec<&str> = sink.split('.').collect();
    let result = query_one(
        &state.materialize.pgclient,
        query,
        &[&sink_fields[0], &sink_fields[1], &sink_fields[2]],
    )
    .await
    .context("retrieving topic name")?
    .get(topic_field);
    Ok(result)
}

pub async fn run_verify_data(
    mut cmd: BuiltinCommand,
    state: &mut State,
) -> Result<ControlFlow, anyhow::Error> {
    let mut format = if let Some(format_str) = cmd.args.opt_string("format") {
        // If just a single format is provided, the user should specify `key=true` if they expect a
        // in each message.  However for format=avro, we will conveniently set this based
        // on the presence of the key schema in the registry, so this argument is not required.
        let requires_key: bool = cmd.args.opt_bool("key")?.unwrap_or(false);
        let format_type = format_str.as_str().try_into()?;
        RecordFormat {
            key: format_type,
            value: format_type,
            requires_key,
        }
    } else {
        let key_format = cmd.args.string("key-format")?.as_str().try_into()?;
        let value_format = cmd.args.string("value-format")?.as_str().try_into()?;
        RecordFormat {
            key: key_format,
            value: value_format,
            requires_key: true,
        }
    };

    let source = match (cmd.args.opt_string("sink"), cmd.args.opt_string("topic")) {
        (Some(sink), None) => Topic::FromSink(sink),
        (None, Some(topic)) => Topic::Named(topic),
        (Some(_), Some(_)) => bail!("Can't provide both `source` and `topic` to kafka-verify-data"),
        (None, None) => bail!("kafka-verify-data expects either `source` or `topic`"),
    };

    let sort_messages = cmd.args.opt_bool("sort-messages")?.unwrap_or(false);

    let header_keys: Vec<_> = cmd
        .args
        .opt_string("headers")
        .map(|s| s.split(',').map(str::to_owned).collect())
        .unwrap_or_default();

    let expected_messages = cmd.input;
    if expected_messages.len() == 0 {
        // verify with 0 messages doesn't check that no messages have been written -
        // it 'verifies' 0 messages and trivially returns true
        bail!("kafka-verify-data requires a non-empty list of expected messages");
    }
    let partial_search = cmd.args.opt_parse("partial-search")?;
    // When false, the topic is exempt from the end-of-file exhaustion check.
    // Set it on topics that legitimately keep records this command does not
    // verify: topics fed by more than one sink, topics that double as a sink's
    // progress topic, or topics whose later records are checked by other means.
    let exhaustive = cmd.args.opt_bool("exhaustive")?.unwrap_or(true);
    let debug_print_only = cmd.args.opt_bool("debug-print-only")?.unwrap_or(false);
    // When set, decode Avro from AWS Glue framing (18-byte header) and resolve the
    // writer schema from Glue rather than Confluent Schema Registry.
    let glue = cmd.args.opt_bool("glue")?.unwrap_or(false);
    cmd.args.done()?;

    let topic: String = match &source {
        Topic::FromSink(sink) => get_topic(sink, "topic", state).await?,
        Topic::Named(name) => name.clone(),
    };

    println!("Verifying results in Kafka topic {}", topic);

    let mut config = state.kafka_config.clone();
    config.set("enable.auto.commit", "false");
    config.set("enable.auto.offset.store", "false");

    let consumer: StreamConsumer = config.create().context("creating kafka consumer")?;
    consumer
        .subscribe(&[&topic])
        .context("subscribing to kafka topic")?;

    // NOTE: Consumption stops after exactly as many messages as expected. The
    // consumer group's committed offsets make consecutive kafka-verify-data
    // commands resume where the previous one stopped, so tests can verify a
    // topic in chunks. The end-of-file check below verifies that no records
    // remain after the final non-partial verification.
    let (mut stream_messages_remaining, stream_timeout) = match partial_search {
        Some(size) => (size, state.timeout),
        None => (expected_messages.len(), Duration::from_secs(15)),
    };

    let timeout = cmp::max(state.timeout, stream_timeout);

    let message_stream = consumer.stream().timeout(timeout);
    pin!(message_stream);

    // Collect all messages that arrive without timing out. If we trip
    // the timeout, suppress the error and return what we have. This
    // is nicer than returning "timeout expired", as the user will
    // instead get an error message about the expected messages that
    // were missing.
    let mut actual_bytes = vec![];

    let start = std::time::Instant::now();
    let mut topic_created = false;

    while stream_messages_remaining > 0 {
        match message_stream.next().await {
            Some(Ok(message)) => {
                let message = match message {
                    // We create topics after creating sinks, so we permit
                    // retries here while waiting for the topic to get created.
                    Err(KafkaError::MessageConsumption(
                        RDKafkaErrorCode::UnknownTopicOrPartition,
                    )) if start.elapsed() < timeout && !topic_created => {
                        println!("waiting for Kafka topic creation...");
                        continue;
                    }
                    e => e?,
                };

                stream_messages_remaining -= 1;
                topic_created = true;

                consumer
                    .store_offset_from_message(&message)
                    .context("storing message offset")?;

                let mut headers = vec![];
                for header_key in &header_keys {
                    // Expect a unique header with the given key and a UTF8-formatted body.
                    let hs = message.headers().context("expected headers for message")?;
                    let mut hs = hs.iter().filter(|i| i.key == header_key);
                    let h = hs.next();
                    if hs.next().is_some() {
                        bail!("expected at most one header with key {header_key}");
                    }
                    match h {
                        None => headers.push("<missing>".into()),
                        Some(h) => {
                            let value = str::from_utf8(h.value.unwrap_or(b"<null>"))?;
                            headers.push(value.into());
                        }
                    }
                }

                actual_bytes.push(Record {
                    headers,
                    key: message.key().map(|b| b.to_owned()),
                    value: message.payload().map(|b| b.to_owned()),
                    partition: Some(message.partition()),
                });
            }
            Some(Err(e)) => {
                println!("Received error from Kafka stream consumer: {}", e);
                break;
            }
            None => {
                break;
            }
        }
    }

    let (key_schema, value_schema) = if glue {
        // A Glue sink frames each record with an 18-byte header carrying the
        // schema-version UUID. Resolve the writer schema for each side from Glue
        // using the UUID on the first record of that side. All records in a
        // verify batch share the sink's schema, so one lookup per side suffices.
        resolve_glue_schemas(state, &actual_bytes, &mut format).await?
    } else {
        let key_schema = if let Format::Avro = format.key {
            // A missing key subject means the topic is unkeyed. Any other error
            // must propagate: swallowing it would silently disable key
            // verification.
            let schema = match state
                .ccsr_client
                .get_schema_by_subject(&format!("{}-key", topic))
                .await
            {
                Ok(key_schema) => {
                    Some(avro::parse_schema(&key_schema.raw, &[]).context("parsing avro schema")?)
                }
                Err(
                    mz_ccsr::GetBySubjectError::SubjectNotFound
                    | mz_ccsr::GetBySubjectError::VersionNotFound(_),
                ) => None,
                Err(e) => return Err(anyhow::Error::from(e).context("fetching key schema")),
            };
            // for avro, we can determine if a key is required based on the presence of the key schema
            // rather than requiring the user to specify the key=true flag
            if schema.is_some() {
                format.requires_key = true;
            }
            schema
        } else {
            None
        };
        let value_schema = if let Format::Avro = format.value {
            let val_schema = state
                .ccsr_client
                .get_schema_by_subject(&format!("{}-value", topic))
                .await
                .context("fetching schema")?
                .raw;
            Some(avro::parse_schema(&val_schema, &[]).context("parsing avro schema")?)
        } else {
            None
        };
        (key_schema, value_schema)
    };

    // The Glue path resolves schemas from the received records, so an empty batch
    // leaves the schema unset. Parsing the (required, non-empty) expected messages
    // would then unwrap a `None` schema and panic. Turn that into the normal
    // "sink produced no output" diagnostic instead.
    if glue {
        if matches!(format.value, Format::Avro) && value_schema.is_none() {
            bail!(
                "kafka-verify-data glue=true: no records received to resolve the \
                 value schema from Glue (did the sink produce no output?)"
            );
        }
        if format.requires_key && matches!(format.key, Format::Avro) && key_schema.is_none() {
            bail!(
                "kafka-verify-data glue=true: no keyed records received to resolve \
                 the key schema from Glue (did the sink produce no output?)"
            );
        }
    }

    let mut actual_messages =
        decode_messages(actual_bytes, &key_schema, &value_schema, &format, glue)?;

    if sort_messages {
        actual_messages.sort_by_key(|r| format!("{:?}", r));
    }

    if debug_print_only {
        bail!(
            "records in sink:\n{}",
            actual_messages
                .into_iter()
                .map(|a| format!("{:#?}", a))
                .collect::<Vec<_>>()
                .join("\n")
        );
    }

    let expected = parse_expected_messages(
        expected_messages,
        key_schema,
        value_schema,
        &format,
        &header_keys,
    )?;

    verify_with_partial_search(
        &expected,
        &actual_messages,
        &state.regex,
        &state.regex_replacement,
        partial_search.is_some(),
    )?;

    consumer
        .commit_consumer_state(CommitMode::Sync)
        .context("committing verified message offsets")?;

    if partial_search.is_some() || !exhaustive {
        state.kafka_verify_topics.remove(&topic);
    } else {
        state.kafka_verify_topics.insert(topic);
    }

    Ok(ControlFlow::Continue)
}

/// Verifies that no records remain after each topic's final exact verification.
pub async fn verify_topics_exhausted(state: &State) -> Result<(), anyhow::Error> {
    for topic in &state.kafka_verify_topics {
        let mut config = state.kafka_config.clone();
        config.set("enable.auto.commit", "false");
        config.set("enable.auto.offset.store", "false");
        config.set("enable.partition.eof", "true");

        let topic = topic.clone();
        let timeout = state.timeout;
        task::spawn_blocking(
            {
                let topic = topic.clone();
                move || format!("kafka_verify_topic_exhausted:{topic}")
            },
            move || verify_topic_exhausted(config, &topic, timeout),
        )
        .await?;
    }
    Ok(())
}

fn verify_topic_exhausted(
    config: rdkafka::ClientConfig,
    topic: &str,
    timeout: Duration,
) -> Result<(), anyhow::Error> {
    let consumer: BaseConsumer = config.create().context("creating kafka consumer")?;
    let deadline = Instant::now() + timeout;
    let partitions = get_partitions(consumer.client(), topic, remaining(deadline)?)?;

    let mut requested_offsets = TopicPartitionList::with_capacity(partitions.len());
    for partition in &partitions {
        requested_offsets.add_partition(topic, *partition);
    }
    let committed_offsets = consumer
        .committed_offsets(requested_offsets, remaining(deadline)?)
        .context("fetching committed message offsets")?;

    // Resume from each partition's committed offset, where the last
    // kafka-verify-data on this topic stopped, and confirm that only the
    // end-of-partition marker remains. `pending` tracks the partitions we have
    // not yet drained.
    let mut assignment = TopicPartitionList::with_capacity(partitions.len());
    let mut pending: BTreeSet<i32> = BTreeSet::new();
    for partition in partitions {
        let committed = committed_offsets
            .find_partition(topic, partition)
            .context("missing committed offset for topic partition")?
            .offset();
        let (low, high) = consumer
            .fetch_watermarks(topic, partition, remaining(deadline)?)
            .with_context(|| {
                format!("fetching watermarks for Kafka topic {topic} partition {partition}")
            })?;
        let start = verification_start_offset(committed, low, high)?;
        assignment.add_partition_offset(topic, partition, Offset::Offset(start))?;
        pending.insert(partition);
    }

    consumer
        .assign(&assignment)
        .context("assigning Kafka topic partitions")?;

    // `enable.partition.eof` makes librdkafka emit a `PartitionEOF` once a
    // partition is drained to its end. Wait for that marker on every partition
    // rather than comparing consumer positions to high watermarks: a position
    // never advances on a partition that begins already drained (so empty
    // partitions would hang forever), and with `read_committed` the high
    // watermark stays ahead of the last readable record whenever transaction
    // control markers trail the data.
    while !pending.is_empty() {
        match consumer.poll(remaining(deadline)?) {
            Some(Ok(message)) => {
                bail!(
                    "extra record after final kafka-verify-data for topic {topic}, partition {}, offset {}",
                    message.partition(),
                    message.offset(),
                );
            }
            Some(Err(KafkaError::PartitionEOF(partition))) => {
                pending.remove(&partition);
            }
            Some(Err(e)) => {
                return Err(e).context("reading final Kafka topic offsets");
            }
            None => {
                bail!("timed out verifying the end of Kafka topic {topic}");
            }
        }
    }

    Ok(())
}

fn verification_start_offset(committed: Offset, low: i64, high: i64) -> Result<i64, anyhow::Error> {
    match committed {
        Offset::Offset(offset) if (low..=high).contains(&offset) => Ok(offset),
        Offset::Offset(_) | Offset::Invalid => Ok(low),
        offset => bail!("unexpected committed Kafka offset {offset:?}"),
    }
}

fn remaining(deadline: Instant) -> Result<Duration, anyhow::Error> {
    deadline
        .checked_duration_since(Instant::now())
        .filter(|remaining| !remaining.is_zero())
        .context("timed out verifying final Kafka topic offsets")
}

/// Resolve the writer schema for each record side from AWS Glue.
///
/// Reads the schema-version UUID from the Glue header of the first record on each
/// side and fetches that version's definition from Glue. Returns
/// `(key_schema, value_schema)`, mirroring the Confluent resolution path, and
/// sets `format.requires_key` when a key schema is present. All records in a
/// verify batch share the sink's schema, so one lookup per side is enough.
async fn resolve_glue_schemas(
    state: &State,
    records: &[Record<Vec<u8>>],
    format: &mut RecordFormat,
) -> Result<(Option<mz_avro::Schema>, Option<mz_avro::Schema>), anyhow::Error> {
    let client = aws_sdk_glue::Client::new(&state.aws_config);

    async fn schema_for(
        client: &aws_sdk_glue::Client,
        payload: Option<&Vec<u8>>,
    ) -> Result<Option<mz_avro::Schema>, anyhow::Error> {
        let Some(payload) = payload else {
            return Ok(None);
        };
        let (schema_version_id, _) = mz_interchange::glue::extract_avro_header(payload)?;
        let resp = client
            .get_schema_version()
            .schema_version_id(schema_version_id.to_string())
            .send()
            .await
            .with_context(|| format!("fetching Glue schema version {schema_version_id}"))?;
        let definition = resp
            .schema_definition()
            .ok_or_else(|| anyhow!("Glue schema version {schema_version_id} has no definition"))?;
        Ok(Some(
            avro::parse_schema(definition, &[]).context("parsing Glue avro schema")?,
        ))
    }

    let value_schema = if let Format::Avro = format.value {
        schema_for(&client, records.iter().find_map(|r| r.value.as_ref())).await?
    } else {
        None
    };
    let key_schema = if let Format::Avro = format.key {
        let schema = schema_for(&client, records.iter().find_map(|r| r.key.as_ref())).await?;
        if schema.is_some() {
            format.requires_key = true;
        }
        schema
    } else {
        None
    };
    Ok((key_schema, value_schema))
}

/// Expect and split out `n` whitespace-delimited headers before the main contents of the 'expect' row.
fn split_headers(input: &str, n_headers: usize) -> anyhow::Result<(Vec<String>, &str)> {
    let whitespace = Regex::new("\\s+").expect("building known-valid regex");
    let mut parts = whitespace.splitn(input, n_headers + 1);
    let mut headers = Vec::with_capacity(n_headers);
    for _ in 0..n_headers {
        headers.push(
            parts
                .next()
                .context("expected another header in the input")?
                .to_string(),
        )
    }
    let rest = parts
        .next()
        .context("expected some contents after any message headers")?;

    ensure!(
        parts.next().is_none(),
        "more than n+1 elements from a call to splitn(_, n+1)"
    );

    Ok((headers, rest))
}

/// Decode Avro `bytes` with `schema`, stripping Glue framing when `glue` is set
/// and Confluent framing otherwise.
fn decode_avro(
    schema: &mz_avro::Schema,
    bytes: &[u8],
    glue: bool,
) -> Result<avro::Value, anyhow::Error> {
    if glue {
        avro::from_glue_bytes(schema, bytes)
    } else {
        avro::from_confluent_bytes(schema, bytes)
    }
}

fn decode_messages(
    actual_bytes: Vec<Record<Vec<u8>>>,
    key_schema: &Option<mz_avro::Schema>,
    value_schema: &Option<mz_avro::Schema>,
    format: &RecordFormat,
    glue: bool,
) -> Result<Vec<Record<DecodedValue>>, anyhow::Error> {
    let mut actual_messages = vec![];

    for record in actual_bytes {
        let Record { key, value, .. } = record;
        let key = if format.requires_key {
            match (key, format.key) {
                (Some(bytes), Format::Avro) => Some(DecodedValue::Avro(DebugValue(decode_avro(
                    key_schema.as_ref().unwrap(),
                    &bytes,
                    glue,
                )?))),
                (Some(bytes), Format::Json) => Some(DecodedValue::Json(
                    serde_json::from_slice(&bytes).context("decoding json")?,
                )),
                (Some(bytes), Format::Bytes) => Some(DecodedValue::Bytes(bytes)),
                (Some(bytes), Format::Text) => Some(DecodedValue::Text(String::from_utf8(bytes)?)),
                (None, _) if format.requires_key => bail!("empty message key"),
                (None, _) => None,
            }
        } else {
            None
        };

        let value = match (value, format.value) {
            (Some(bytes), Format::Avro) => Some(DecodedValue::Avro(DebugValue(decode_avro(
                value_schema.as_ref().unwrap(),
                &bytes,
                glue,
            )?))),
            (Some(bytes), Format::Json) => Some(DecodedValue::Json(
                serde_json::from_slice(&bytes).context("decoding json")?,
            )),
            (Some(bytes), Format::Bytes) => Some(DecodedValue::Bytes(bytes)),
            (Some(bytes), Format::Text) => Some(DecodedValue::Text(String::from_utf8(bytes)?)),
            (None, _) => None,
        };

        actual_messages.push(Record {
            headers: record.headers.clone(),
            key,
            value,
            partition: record.partition,
        });
    }

    Ok(actual_messages)
}

fn parse_expected_messages(
    expected_messages: Vec<String>,
    key_schema: Option<mz_avro::Schema>,
    value_schema: Option<mz_avro::Schema>,
    format: &RecordFormat,
    header_keys: &[String],
) -> Result<Vec<Record<DecodedValue>>, anyhow::Error> {
    let mut expected = vec![];

    for msg in expected_messages {
        let (headers, content) = split_headers(&msg, header_keys.len())?;
        let mut content = content.as_bytes();
        let mut deserializer = serde_json::Deserializer::from_reader(&mut content).into_iter();

        let key = if format.requires_key {
            let key: serde_json::Value = deserializer
                .next()
                .context("key missing in input line")?
                .context("parsing json")?;

            Some(match format.key {
                Format::Avro => DecodedValue::Avro(DebugValue(avro::from_json(
                    &key,
                    key_schema.as_ref().unwrap().top_node(),
                )?)),
                Format::Json => DecodedValue::Json(key),
                Format::Bytes => {
                    unimplemented!("bytes format not yet supported in tests")
                }
                Format::Text => DecodedValue::Text(
                    key.as_str()
                        .map(|s| s.to_string())
                        .unwrap_or_else(|| key.to_string()),
                ),
            })
        } else {
            None
        };

        let value = match deserializer.next().transpose().context("parsing json")? {
            None => None,
            Some(value) if value.as_str() == Some("<null>") => None,
            Some(value) => match format.value {
                Format::Avro => Some(DecodedValue::Avro(DebugValue(avro::from_json(
                    &value,
                    value_schema.as_ref().unwrap().top_node(),
                )?))),
                Format::Json => Some(DecodedValue::Json(value)),
                Format::Bytes => {
                    unimplemented!("bytes format not yet supported in tests")
                }
                // Take JSON strings verbatim, like the key path above. The
                // actual message is decoded as raw text, so keeping the JSON
                // quotes would make string values never match.
                Format::Text => Some(DecodedValue::Text(
                    value
                        .as_str()
                        .map(|s| s.to_string())
                        .unwrap_or_else(|| value.to_string()),
                )),
            },
        };

        let content =
            str::from_utf8(content).context("internal error: contents were previously a string")?;
        let partition = match content.trim().split_once("=") {
            None if content.trim() != "" => bail!("unexpected cruft at end of line: {content}"),
            None => None,
            Some((label, partition)) => {
                if label != "partition" {
                    bail!("partition expectation has unexpected label: {label}")
                }
                Some(partition.parse().context("parsing expected partition")?)
            }
        };

        expected.push(Record {
            headers,
            key,
            value,
            partition,
        });
    }

    Ok(expected)
}

fn verify_with_partial_search<A>(
    expected: &[Record<A>],
    actual: &[Record<A>],
    regex: &Option<Regex>,
    regex_replacement: &String,
    partial_search: bool,
) -> Result<(), anyhow::Error>
where
    A: Debug + Clone,
{
    let mut expected = expected.iter();
    let mut actual = actual.iter();
    let mut index = 0..;

    let mut found_beginning = !partial_search;
    let mut expected_item = expected.next();
    let mut actual_item = actual.next();
    loop {
        let i = index.next().expect("known to exist");
        match (expected_item, actual_item) {
            (Some(e), Some(a)) => {
                let mut a = a.clone();
                if e.partition.is_none() {
                    a.partition = None;
                }
                let e_str = format!("{:#?}", e);
                let a_str = match &regex {
                    Some(regex) => regex
                        .replace_all(&format!("{:#?}", a).to_string(), regex_replacement.as_str())
                        .to_string(),
                    _ => format!("{:#?}", a),
                };

                if e_str != a_str {
                    if found_beginning {
                        bail!(
                            "record {} did not match\nexpected:\n{}\n\nactual:\n{}",
                            i,
                            e_str,
                            a_str,
                        );
                    }
                    actual_item = actual.next();
                } else {
                    found_beginning = true;
                    expected_item = expected.next();
                    actual_item = actual.next();
                }
            }
            (Some(e), None) => bail!("missing record {}: {:#?}", i, e),
            (None, Some(a)) => {
                if !partial_search {
                    bail!("extra record {}: {:#?}", i, a);
                }
                break;
            }
            (None, None) => break,
        }
    }
    let expected: Vec<_> = expected.map(|e| format!("{:#?}", e)).collect();
    let actual: Vec<_> = actual.map(|a| format!("{:#?}", a)).collect();

    if !expected.is_empty() {
        bail!("missing records:\n{}", expected.join("\n"))
    } else if !actual.is_empty() && !partial_search {
        bail!("extra records:\n{}", actual.join("\n"))
    } else {
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[mz_ore::test]
    fn verification_start_offset_uses_committed_offset() {
        assert_eq!(
            verification_start_offset(Offset::Offset(5), 2, 8).unwrap(),
            5
        );
    }

    #[mz_ore::test]
    fn verification_start_offset_resets_out_of_range_offsets() {
        assert_eq!(
            verification_start_offset(Offset::Offset(1), 2, 8).unwrap(),
            2
        );
        assert_eq!(
            verification_start_offset(Offset::Offset(9), 2, 8).unwrap(),
            2
        );
        assert_eq!(verification_start_offset(Offset::Invalid, 2, 8).unwrap(), 2);
    }
}