mz_testdrive/action/kafka/

ingest.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::cmp;
use std::io::{BufRead, Read};
use std::time::Duration;

use anyhow::{Context, anyhow, bail};
use byteorder::{NetworkEndian, WriteBytesExt};
use futures::stream::{FuturesUnordered, StreamExt};
use maplit::btreemap;
use prost::Message;
use prost_reflect::{DescriptorPool, DynamicMessage, MessageDescriptor};
use rdkafka::message::{Header, OwnedHeaders};
use rdkafka::producer::FutureRecord;
use serde::de::DeserializeOwned;
use tokio::fs;

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

const INGEST_BATCH_SIZE: isize = 10000;

/// Extracts ALL type names defined in an Avro schema (including nested types).
/// Returns a set of fully qualified type names.
#[allow(clippy::disallowed_types)]
fn extract_all_defined_types(
    schema_json: &str,
) -> anyhow::Result<std::collections::HashSet<String>> {
    let value: serde_json::Value = serde_json::from_str(schema_json)
        .context("parsing schema JSON to extract defined types")?;

    let mut types = std::collections::HashSet::new();
    collect_defined_types(&value, None, &mut types);
    Ok(types)
}

/// Recursively collects all named type definitions from an Avro schema.
#[allow(clippy::disallowed_types)]
fn collect_defined_types(
    value: &serde_json::Value,
    parent_namespace: Option<&str>,
    types: &mut std::collections::HashSet<String>,
) {
    match value {
        serde_json::Value::Object(map) => {
            // Get this schema's namespace (falls back to parent's namespace)
            let namespace = map
                .get("namespace")
                .and_then(|v| v.as_str())
                .or(parent_namespace);

            // Check if this is a named type definition (record, enum, or fixed)
            if let Some(type_val) = map.get("type")
                && type_val
                    .as_str()
                    .is_some_and(|typ| ["record", "enum", "fixed"].contains(&typ))
            {
                if let Some(name) = map.get("name").and_then(|v| v.as_str()) {
                    // Construct fully qualified name
                    let fullname = if name.contains('.') {
                        name.to_string()
                    } else if let Some(ns) = namespace {
                        format!("{}.{}", ns, name)
                    } else {
                        name.to_string()
                    };
                    types.insert(fullname);
                }
            }

            // The following types may have references:
            // type field, items (array types), values (map types), and fields (e.g. unions)
            for entity_type in &["type", "items", "values", "fields"] {
                if let Some(val) = map.get(*entity_type) {
                    collect_defined_types(val, namespace, types);
                }
            }
        }
        serde_json::Value::Array(arr) => {
            for item in arr {
                collect_defined_types(item, parent_namespace, types);
            }
        }
        _ => {}
    }
}

/// Extracts all type references from an Avro schema JSON string.
/// This finds all fully qualified type names that are referenced but not defined in the schema.
#[allow(clippy::disallowed_types)]
fn extract_type_references(schema_json: &str) -> anyhow::Result<std::collections::HashSet<String>> {
    let value: serde_json::Value = serde_json::from_str(schema_json)
        .context("parsing schema JSON to extract type references")?;

    let mut references = std::collections::HashSet::new();
    collect_type_references(&value, &mut references);
    Ok(references)
}

/// Recursively collects type references from an Avro schema JSON value.
#[allow(clippy::disallowed_types)]
fn collect_type_references(
    value: &serde_json::Value,
    references: &mut std::collections::HashSet<String>,
) {
    match value {
        serde_json::Value::String(s) => {
            // A string type that contains a dot is likely a fully qualified type reference
            if s.contains('.')
                && ![
                    "null", "boolean", "int", "long", "float", "double", "bytes", "string",
                ]
                .contains(&s.as_str())
            {
                references.insert(s.clone());
            }
        }
        serde_json::Value::Object(map) => {
            // For named types, we want to recurse into the fields, but the named type doesn't
            // get added to references.
            if let Some(type_val) = map.get("type")
                && type_val
                    .as_str()
                    .is_some_and(|typ| ["record", "enum", "fixed"].contains(&typ))
            {
                if let Some(fields) = map.get("fields") {
                    collect_type_references(fields, references);
                }
                return;
            }

            // The following types may have references:
            // type field, items (array types), values (map types), and fields (e.g. unions)
            for entity_type in &["type", "items", "values", "fields"] {
                if let Some(val) = map.get(*entity_type) {
                    collect_type_references(val, references);
                }
            }
        }
        serde_json::Value::Array(arr) => {
            for item in arr {
                collect_type_references(item, references);
            }
        }
        _ => {}
    }
}

#[derive(Clone)]
enum Format {
    Avro {
        schema: String,
        confluent_wire_format: bool,
        /// Schema references (subject names) for Confluent Schema Registry
        references: Vec<String>,
    },
    Protobuf {
        descriptor_file: String,
        message: String,
        confluent_wire_format: bool,
        schema_id_subject: Option<String>,
        schema_message_id: u8,
    },
    Bytes {
        terminator: Option<u8>,
    },
}

enum Transcoder {
    PlainAvro {
        schema: Schema,
    },
    ConfluentAvro {
        schema: Schema,
        schema_id: i32,
    },
    Protobuf {
        message: MessageDescriptor,
        confluent_wire_format: bool,
        schema_id: i32,
        schema_message_id: u8,
    },
    Bytes {
        terminator: Option<u8>,
    },
}

impl Transcoder {
    fn decode_json<R, T>(row: R) -> Result<Option<T>, anyhow::Error>
    where
        R: Read,
        T: DeserializeOwned,
    {
        let deserializer = serde_json::Deserializer::from_reader(row);
        deserializer
            .into_iter()
            .next()
            .transpose()
            .context("parsing json")
    }

    fn transcode<R>(&self, mut row: R) -> Result<Option<Vec<u8>>, anyhow::Error>
    where
        R: BufRead,
    {
        match self {
            Transcoder::ConfluentAvro { schema, schema_id } => {
                if let Some(val) = Self::decode_json(row)? {
                    let val = avro::from_json(&val, schema.top_node())?;
                    let mut out = vec![];
                    // The first byte is a magic byte (0) that indicates the Confluent
                    // serialization format version, and the next four bytes are a
                    // 32-bit schema ID.
                    //
                    // https://docs.confluent.io/3.3.0/schema-registry/docs/serializer-formatter.html#wire-format
                    out.write_u8(0).unwrap();
                    out.write_i32::<NetworkEndian>(*schema_id).unwrap();
                    out.extend(avro::to_avro_datum(schema, val)?);
                    Ok(Some(out))
                } else {
                    Ok(None)
                }
            }
            Transcoder::PlainAvro { schema } => {
                if let Some(val) = Self::decode_json(row)? {
                    let val = avro::from_json(&val, schema.top_node())?;
                    let mut out = vec![];
                    out.extend(avro::to_avro_datum(schema, val)?);
                    Ok(Some(out))
                } else {
                    Ok(None)
                }
            }
            Transcoder::Protobuf {
                message,
                confluent_wire_format,
                schema_id,
                schema_message_id,
            } => {
                if let Some(val) = Self::decode_json::<_, serde_json::Value>(row)? {
                    let message = DynamicMessage::deserialize(message.clone(), val)
                        .context("parsing protobuf JSON")?;
                    let mut out = vec![];
                    if *confluent_wire_format {
                        // See: https://github.com/MaterializeInc/database-issues/issues/2837
                        // The first byte is a magic byte (0) that indicates the Confluent
                        // serialization format version, and the next four bytes are a
                        // 32-bit schema ID, which we default to something fun.
                        // And, as we only support single-message proto files for now,
                        // we also set the following message id to 0.
                        out.write_u8(0).unwrap();
                        out.write_i32::<NetworkEndian>(*schema_id).unwrap();
                        out.write_u8(*schema_message_id).unwrap();
                    }
                    message.encode(&mut out)?;
                    Ok(Some(out))
                } else {
                    Ok(None)
                }
            }
            Transcoder::Bytes { terminator } => {
                let mut out = vec![];
                match terminator {
                    Some(t) => {
                        row.read_until(*t, &mut out)?;
                        if out.last() == Some(t) {
                            out.pop();
                        }
                    }
                    None => {
                        row.read_to_end(&mut out)?;
                    }
                }
                if out.is_empty() {
                    Ok(None)
                } else {
                    Ok(Some(bytes::unescape(&out)?))
                }
            }
        }
    }
}

pub async fn run_ingest(
    mut cmd: BuiltinCommand,
    state: &mut State,
) -> Result<ControlFlow, anyhow::Error> {
    let topic_prefix = format!("testdrive-{}", cmd.args.string("topic")?);
    let partition = cmd.args.opt_parse::<i32>("partition")?;
    let start_iteration = cmd.args.opt_parse::<isize>("start-iteration")?.unwrap_or(0);
    let repeat = cmd.args.opt_parse::<isize>("repeat")?.unwrap_or(1);
    let omit_key = cmd.args.opt_bool("omit-key")?.unwrap_or(false);
    let omit_value = cmd.args.opt_bool("omit-value")?.unwrap_or(false);
    let schema_id_var = cmd.args.opt_parse("set-schema-id-var")?;
    let key_schema_id_var = cmd.args.opt_parse("set-key-schema-id-var")?;
    let format = match cmd.args.string("format")?.as_str() {
        "avro" => Format::Avro {
            schema: cmd.args.string("schema")?,
            confluent_wire_format: cmd.args.opt_bool("confluent-wire-format")?.unwrap_or(true),
            // TODO (maz): update README!
            references: cmd
                .args
                .opt_string("references")
                .map(|s| s.split(',').map(|s| s.to_string()).collect())
                .unwrap_or_default(),
        },
        "protobuf" => {
            let descriptor_file = cmd.args.string("descriptor-file")?;
            let message = cmd.args.string("message")?;
            Format::Protobuf {
                descriptor_file,
                message,
                // This was introduced after the avro format's confluent-wire-format, so it defaults to
                // false
                confluent_wire_format: cmd.args.opt_bool("confluent-wire-format")?.unwrap_or(false),
                schema_id_subject: cmd.args.opt_string("schema-id-subject"),
                schema_message_id: cmd.args.opt_parse::<u8>("schema-message-id")?.unwrap_or(0),
            }
        }
        "bytes" => Format::Bytes { terminator: None },
        f => bail!("unknown format: {}", f),
    };
    let mut key_schema = cmd.args.opt_string("key-schema");
    let key_format = match cmd.args.opt_string("key-format").as_deref() {
        Some("avro") => Some(Format::Avro {
            schema: key_schema.take().ok_or_else(|| {
                anyhow!("key-schema parameter required when key-format is present")
            })?,
            confluent_wire_format: cmd.args.opt_bool("confluent-wire-format")?.unwrap_or(true),
            references: cmd
                .args
                .opt_string("key-references")
                .map(|s| s.split(',').map(|s| s.to_string()).collect())
                .unwrap_or_default(),
        }),
        Some("protobuf") => {
            let descriptor_file = cmd.args.string("key-descriptor-file")?;
            let message = cmd.args.string("key-message")?;
            Some(Format::Protobuf {
                descriptor_file,
                message,
                confluent_wire_format: cmd.args.opt_bool("confluent-wire-format")?.unwrap_or(false),
                schema_id_subject: cmd.args.opt_string("key-schema-id-subject"),
                schema_message_id: cmd
                    .args
                    .opt_parse::<u8>("key-schema-message-id")?
                    .unwrap_or(0),
            })
        }
        Some("bytes") => Some(Format::Bytes {
            terminator: match cmd.args.opt_parse::<char>("key-terminator")? {
                Some(c) => match u8::try_from(c) {
                    Ok(c) => Some(c),
                    Err(_) => bail!("key terminator must be single ASCII character"),
                },
                None => Some(b':'),
            },
        }),
        Some(f) => bail!("unknown key format: {}", f),
        None => None,
    };
    if key_schema.is_some() {
        anyhow::bail!("key-schema specified without a matching key-format");
    }

    let timestamp = cmd.args.opt_parse("timestamp")?;

    use serde_json::Value;
    let headers = if let Some(headers_val) = cmd.args.opt_parse::<serde_json::Value>("headers")? {
        let mut headers = Vec::new();
        let headers_maps = match headers_val {
            Value::Array(values) => {
                let mut headers_map = Vec::new();
                for value in values {
                    if let Value::Object(m) = value {
                        headers_map.push(m)
                    } else {
                        bail!("`headers` array values must be maps")
                    }
                }
                headers_map
            }
            Value::Object(v) => vec![v],
            _ => bail!("`headers` must be a map or an array"),
        };

        for headers_map in headers_maps {
            for (k, v) in headers_map.iter() {
                headers.push((k.clone(), match v {
                    Value::String(val) => Some(val.as_bytes().to_vec()),
                    Value::Array(val) => {
                        let mut values = Vec::new();
                        for value in val {
                            if let Value::Number(int) = value {
                                values.push(u8::try_from(int.as_i64().unwrap()).unwrap())
                            } else {
                                bail!("`headers` value arrays must only contain numbers (to represent bytes)")
                            }
                        }
                        Some(values.clone())
                    },
                    Value::Null => None,
                    _ => bail!("`headers` must have string, int array or null values")
                }));
            }
        }
        Some(headers)
    } else {
        None
    };

    cmd.args.done()?;

    if let Some(kf) = &key_format {
        fn is_confluent_format(fmt: &Format) -> Option<bool> {
            match fmt {
                Format::Avro {
                    confluent_wire_format,
                    ..
                } => Some(*confluent_wire_format),
                Format::Protobuf {
                    confluent_wire_format,
                    ..
                } => Some(*confluent_wire_format),
                Format::Bytes { .. } => None,
            }
        }
        match (is_confluent_format(kf), is_confluent_format(&format)) {
            (Some(false), Some(true)) | (Some(true), Some(false)) => {
                bail!(
                    "It does not make sense to have the key be in confluent format and not the value, or vice versa."
                );
            }
            _ => {}
        }
    }

    let topic_name = &format!("{}-{}", topic_prefix, state.seed);
    println!(
        "Ingesting data into Kafka topic {} with start_iteration = {}, repeat = {}",
        topic_name, start_iteration, repeat
    );

    let set_schema_id_var = |state: &mut State, schema_id_var, transcoder| match transcoder {
        &Transcoder::ConfluentAvro { schema_id, .. } | &Transcoder::Protobuf { schema_id, .. } => {
            state.cmd_vars.insert(schema_id_var, schema_id.to_string());
        }
        _ => (),
    };

    let value_transcoder =
        make_transcoder(state, format.clone(), format!("{}-value", topic_name)).await?;
    if let Some(var) = schema_id_var {
        set_schema_id_var(state, var, &value_transcoder);
    }

    let key_transcoder = match key_format.clone() {
        None => None,
        Some(f) => {
            let transcoder = make_transcoder(state, f, format!("{}-key", topic_name)).await?;
            if let Some(var) = key_schema_id_var {
                set_schema_id_var(state, var, &transcoder);
            }
            Some(transcoder)
        }
    };

    let mut futs = FuturesUnordered::new();

    for iteration in start_iteration..(start_iteration + repeat) {
        let iter = &mut cmd.input.iter().peekable();

        for row in iter {
            let row = action::substitute_vars(
                row,
                &btreemap! { "kafka-ingest.iteration".into() => iteration.to_string() },
                &None,
                false,
            )?;
            let mut row = row.as_bytes();
            let key = match (omit_key, &key_transcoder) {
                (true, _) => None,
                (false, None) => None,
                (false, Some(kt)) => kt.transcode(&mut row)?,
            };
            let value = if omit_value {
                None
            } else {
                value_transcoder
                    .transcode(&mut row)
                    .with_context(|| format!("parsing row: {}", String::from_utf8_lossy(row)))?
            };
            let producer = &state.kafka_producer;
            let timeout = cmp::max(state.default_timeout, Duration::from_secs(1));
            let headers = headers.clone();
            futs.push(async move {
                let mut record: FutureRecord<_, _> = FutureRecord::to(topic_name);

                if let Some(partition) = partition {
                    record = record.partition(partition);
                }
                if let Some(key) = &key {
                    record = record.key(key);
                }
                if let Some(value) = &value {
                    record = record.payload(value);
                }
                if let Some(timestamp) = timestamp {
                    record = record.timestamp(timestamp);
                }
                if let Some(headers) = headers {
                    let mut rd_meta = OwnedHeaders::new();
                    for (k, v) in &headers {
                        rd_meta = rd_meta.insert(Header {
                            key: k,
                            value: v.as_deref(),
                        });
                    }
                    record = record.headers(rd_meta);
                }
                producer.send(record, timeout).await
            });
        }

        // Reap the futures thus produced periodically or after the last iteration
        if iteration % INGEST_BATCH_SIZE == 0 || iteration == (start_iteration + repeat - 1) {
            while let Some(res) = futs.next().await {
                res.map_err(|(e, _message)| e)?;
            }
        }
    }
    Ok(ControlFlow::Continue)
}

async fn make_transcoder(
    state: &State,
    format: Format,
    ccsr_subject: String,
) -> Result<Transcoder, anyhow::Error> {
    match format {
        Format::Avro {
            schema,
            confluent_wire_format,
            references,
        } => {
            if confluent_wire_format {
                // Build references list by fetching each subject from the registry.
                // Start with immediate references and automatically resolve transitive ones.
                // We need ALL references for local parsing, but only DIRECT references for the registry.
                #[allow(clippy::disallowed_types)]
                let mut reference_subjects = vec![];
                #[allow(clippy::disallowed_types)]
                let mut seen_subjects: std::collections::HashSet<String> =
                    std::collections::HashSet::new();
                let mut queue: Vec<String> = references.clone();

                // Process queue (as a stack), adding transitive dependencies as we discover them
                while let Some(ref_name) = queue.pop() {
                    if seen_subjects.contains(&ref_name) {
                        continue;
                    }
                    seen_subjects.insert(ref_name.clone());

                    let (subject, ref_deps) = state
                        .ccsr_client
                        .get_subject_with_references(&ref_name)
                        .await
                        .with_context(|| format!("fetching reference {}", ref_name))?;

                    // Add newly discovered dependencies to the queue
                    for dep in ref_deps {
                        if !seen_subjects.contains(&dep.subject) {
                            queue.push(dep.subject);
                        }
                    }

                    // Extract ALL type names defined in this schema (including nested types)
                    let defined_types = extract_all_defined_types(&subject.schema.raw)
                        .with_context(|| {
                            format!("extracting type names from reference schema {}", ref_name)
                        })?;
                    reference_subjects.push((
                        ref_name,
                        subject.version,
                        subject.schema.raw,
                        defined_types,
                    ));
                }

                // Reverse to get dependency order: since we use a stack, dependencies are
                // discovered and added after the schemas that depend on them, so reversing
                // puts dependencies first (required for incremental schema parsing)
                reference_subjects.reverse();

                // Extract types directly referenced by the primary schema
                let direct_refs = extract_type_references(&schema)
                    .context("extracting type references from schema")?;

                // For the registry, create a reference for each type in direct_refs
                // that is defined in one of the reference subjects
                let mut schema_references = vec![];
                for type_name in &direct_refs {
                    for (subject_name, version, _, defined_types) in &reference_subjects {
                        if defined_types.contains(type_name) {
                            schema_references.push(mz_ccsr::SchemaReference {
                                name: type_name.clone(),
                                subject: subject_name.clone(),
                                version: *version,
                            });
                            break;
                        }
                    }
                }

                // For local parsing, we need all reference schemas
                let reference_raw_schemas: Vec<_> = reference_subjects
                    .into_iter()
                    .map(|(_, _, raw, _)| raw)
                    .collect();

                let schema_id = state
                    .ccsr_client
                    .publish_schema(
                        &ccsr_subject,
                        &schema,
                        mz_ccsr::SchemaType::Avro,
                        &schema_references,
                    )
                    .await
                    .context("publishing to schema registry")?;

                // Parse schema, handling references if any
                let schema = if reference_raw_schemas.is_empty() {
                    avro::parse_schema(&schema, &[])
                        .with_context(|| format!("parsing avro schema: {}", schema))?
                } else {
                    // Parse reference schemas incrementally (each may depend on previous ones).
                    // References must be specified in dependency order (dependencies first).
                    let mut parsed_refs: Vec<Schema> = vec![];
                    for raw in &reference_raw_schemas {
                        let schema_value: serde_json::Value = serde_json::from_str(raw)
                            .with_context(|| format!("parsing reference schema JSON: {}", raw))?;
                        let parsed = Schema::parse_with_references(&schema_value, &parsed_refs)
                            .with_context(|| format!("parsing reference avro schema: {}", raw))?;
                        parsed_refs.push(parsed);
                    }

                    // Parse primary schema with all reference types available
                    let schema_value: serde_json::Value = serde_json::from_str(&schema)
                        .with_context(|| format!("parsing schema JSON: {}", schema))?;
                    Schema::parse_with_references(&schema_value, &parsed_refs).with_context(
                        || format!("parsing avro schema with references: {}", schema),
                    )?
                };

                Ok::<_, anyhow::Error>(Transcoder::ConfluentAvro { schema, schema_id })
            } else {
                let schema = avro::parse_schema(&schema, &[])
                    .with_context(|| format!("parsing avro schema: {}", schema))?;
                Ok(Transcoder::PlainAvro { schema })
            }
        }
        Format::Protobuf {
            descriptor_file,
            message,
            confluent_wire_format,
            schema_id_subject,
            schema_message_id,
        } => {
            let schema_id = if confluent_wire_format {
                state
                    .ccsr_client
                    .get_schema_by_subject(schema_id_subject.as_deref().unwrap_or(&ccsr_subject))
                    .await
                    .context("fetching schema from registry")?
                    .id
            } else {
                0
            };

            let bytes = fs::read(state.temp_path.join(descriptor_file))
                .await
                .context("reading protobuf descriptor file")?;
            let fd = DescriptorPool::decode(&*bytes).context("parsing protobuf descriptor file")?;
            let message = fd
                .get_message_by_name(&message)
                .ok_or_else(|| anyhow!("unknown message name {}", message))?;
            Ok(Transcoder::Protobuf {
                message,
                confluent_wire_format,
                schema_id,
                schema_message_id,
            })
        }
        Format::Bytes { terminator } => Ok(Transcoder::Bytes { terminator }),
    }
}