laminar_connectors/kafka/

mod.rs

//! Kafka source and sink connectors.

// Source modules
pub mod avro;
pub mod config;
mod metadata_error;
pub mod metrics;
mod offsets;
pub mod rebalance;
pub mod source;
mod vnode_routing;

// Sink modules
pub mod avro_serializer;
pub mod partitioner;
pub mod sink;
pub mod sink_config;
pub mod sink_metrics;

// Shared modules
pub mod schema_registry;

// Source re-exports
pub use avro::AvroDeserializer;
pub use config::{
    AssignmentStrategy, CompatibilityLevel, IsolationLevel, KafkaSourceConfig, OffsetReset,
    SaslMechanism, SchemaEvolutionStrategy, SecurityProtocol, SrAuth, StartupMode,
    TopicSubscription,
};
pub use metrics::KafkaSourceMetrics;
pub use source::KafkaSource;

// Sink re-exports
pub use avro_serializer::AvroSerializer;
pub use partitioner::{
    KafkaPartitioner, KeyHashPartitioner, RoundRobinPartitioner, StickyPartitioner,
};
pub use sink::KafkaSink;
pub use sink_config::{CompressionType, KafkaSinkConfig, PartitionStrategy};
pub use sink_metrics::KafkaSinkMetrics;

/// Test-only access to Kafka-specific routing probes.
#[cfg(feature = "testing")]
pub mod testing {
    use crate::error::ConnectorError;

    /// Map a complete Kafka topic inventory to engine vnodes.
    ///
    /// # Errors
    /// Returns a configuration error when the Kafka route identity, partition
    /// inventory, or vnode count is invalid.
    pub fn partition_vnodes(
        source_identity: &str,
        topic: &str,
        total_partitions: i32,
        vnode_count: u32,
    ) -> Result<Vec<u32>, ConnectorError> {
        super::vnode_routing::partition_vnodes(
            source_identity,
            topic,
            total_partitions,
            vnode_count,
        )
    }
}

// Shared re-exports
pub use schema_registry::{CachedSchema, CompatibilityResult, SchemaRegistryClient, SchemaType};

use std::sync::Arc;

use crate::config::{ConfigKeySpec, ConnectorInfo};
use crate::registry::ConnectorRegistry;

/// Registers the Kafka source connector with the given registry.
///
/// # Errors
///
/// Returns an error if a Kafka source factory is already registered or the
/// registry has been frozen.
pub fn register_kafka_source(
    registry: &ConnectorRegistry,
) -> Result<(), crate::error::ConnectorError> {
    let info = ConnectorInfo {
        name: "kafka".to_string(),
        display_name: "Apache Kafka Source".to_string(),
        version: env!("CARGO_PKG_VERSION").to_string(),
        is_source: true,
        is_sink: false,
        config_keys: kafka_source_config_keys(),
    };

    registry.register_source(
        "kafka",
        info,
        Arc::new(|registry: Option<&Arc<prometheus::Registry>>| {
            // Empty schema — filled in by discover_schema / open() / SQL DDL columns.
            let empty = Arc::new(arrow_schema::Schema::empty());
            Ok(Box::new(KafkaSource::new(
                empty,
                KafkaSourceConfig::default(),
                registry.map(Arc::as_ref),
            )))
        }),
    )
}

/// Registers the Kafka sink connector with the given registry.
///
/// # Errors
///
/// Returns an error if a Kafka sink factory is already registered or the
/// registry has been frozen.
pub fn register_kafka_sink(
    registry: &ConnectorRegistry,
) -> Result<(), crate::error::ConnectorError> {
    let info = ConnectorInfo {
        name: "kafka".to_string(),
        display_name: "Apache Kafka Sink".to_string(),
        version: env!("CARGO_PKG_VERSION").to_string(),
        is_source: false,
        is_sink: true,
        config_keys: kafka_sink_config_keys(),
    };

    registry.register_sink(
        "kafka",
        info,
        Arc::new(|_config, registry: Option<&Arc<prometheus::Registry>>| {
            // Empty schema — the sink's schema is bound from the upstream query at build time.
            let empty = Arc::new(arrow_schema::Schema::empty());
            Ok(Box::new(KafkaSink::new(
                empty,
                KafkaSinkConfig::default(),
                registry.map(Arc::as_ref),
            )))
        }),
    )
}

/// Returns the configuration key specifications for the Kafka source.
#[allow(clippy::too_many_lines)]
fn kafka_source_config_keys() -> Vec<ConfigKeySpec> {
    vec![
        // Required
        ConfigKeySpec::required("bootstrap.servers", "Kafka broker addresses"),
        ConfigKeySpec::required("group.id", "Consumer group identifier"),
        // The key schema cannot express alternatives; parsing requires one of these two.
        ConfigKeySpec::optional(
            "topic",
            "Comma-separated topics (required unless topic.pattern is set)",
            "",
        ),
        ConfigKeySpec::optional(
            "topic.pattern",
            "Topic regex (required unless topic is set)",
            "",
        ),
        // Format
        ConfigKeySpec::optional("format", "Data format (json/csv/avro/raw/debezium)", "json"),
        // Security
        ConfigKeySpec::optional(
            "security.protocol",
            "Security protocol (plaintext/ssl/sasl_plaintext/sasl_ssl)",
            "plaintext",
        ),
        ConfigKeySpec::optional(
            "sasl.mechanism",
            "SASL mechanism (PLAIN/SCRAM-SHA-256/SCRAM-SHA-512/GSSAPI/OAUTHBEARER)",
            "",
        ),
        ConfigKeySpec::optional("sasl.username", "SASL username for PLAIN/SCRAM", ""),
        ConfigKeySpec::optional("sasl.password", "SASL password for PLAIN/SCRAM", ""),
        ConfigKeySpec::optional("ssl.ca.location", "SSL CA certificate file path", ""),
        ConfigKeySpec::optional(
            "ssl.certificate.location",
            "Client SSL certificate file path",
            "",
        ),
        ConfigKeySpec::optional("ssl.key.location", "Client SSL private key file path", ""),
        ConfigKeySpec::optional("ssl.key.password", "Password for encrypted SSL key", ""),
        // Consumer tuning
        ConfigKeySpec::optional(
            "startup.mode",
            "Startup mode (group-offsets/earliest/latest)",
            "group-offsets",
        ),
        ConfigKeySpec::optional(
            "startup.specific.offsets",
            "Start from specific offsets (format: 'partition:offset,...')",
            "",
        ),
        ConfigKeySpec::optional(
            "startup.timestamp.ms",
            "Start from timestamp (milliseconds since epoch)",
            "",
        ),
        ConfigKeySpec::optional(
            "auto.offset.reset",
            "Fallback when no committed offset (earliest/latest/none)",
            "earliest",
        ),
        ConfigKeySpec::optional(
            "isolation.level",
            "Transaction isolation (read_uncommitted/read_committed)",
            "read_committed",
        ),
        ConfigKeySpec::optional("max.poll.records", "Max records per poll", "1000"),
        ConfigKeySpec::optional(
            "partition.assignment.strategy",
            "Partition assignment (range/roundrobin/cooperative-sticky)",
            "range",
        ),
        // Consumer group timing
        ConfigKeySpec::optional(
            "session.timeout.ms",
            "Consumer session timeout in milliseconds (production-safe default)",
            "45000",
        ),
        ConfigKeySpec::optional(
            "heartbeat.interval.ms",
            "Consumer heartbeat interval in milliseconds",
            "10000",
        ),
        ConfigKeySpec::optional(
            "queued.max.messages.kbytes",
            "Max per-partition pre-fetch queue size in kbytes",
            "16384",
        ),
        // Fetch tuning
        ConfigKeySpec::optional("fetch.min.bytes", "Minimum bytes per fetch request", "1"),
        ConfigKeySpec::optional(
            "fetch.max.bytes",
            "Maximum bytes per fetch request",
            "52428800",
        ),
        ConfigKeySpec::optional(
            "fetch.max.wait.ms",
            "Max wait time for fetch.min.bytes",
            "500",
        ),
        ConfigKeySpec::optional(
            "max.partition.fetch.bytes",
            "Max bytes per partition per fetch",
            "1048576",
        ),
        // Metadata
        ConfigKeySpec::optional(
            "include.metadata",
            "Include _partition/_offset/_timestamp columns",
            "false",
        ),
        ConfigKeySpec::optional("include.headers", "Include _headers column", "false"),
        // Backpressure
        ConfigKeySpec::optional(
            "backpressure.high.watermark",
            "Channel fill ratio to pause",
            "0.8",
        ),
        ConfigKeySpec::optional(
            "backpressure.low.watermark",
            "Channel fill ratio to resume",
            "0.25",
        ),
        // Error handling
        ConfigKeySpec::optional(
            "max.deser.error.rate",
            "Max tolerated deserialization error rate per batch (0.0-1.0)",
            "0.5",
        ),
        // Schema Registry
        ConfigKeySpec::optional(
            "schema.registry.url",
            "Confluent Schema Registry URL (required for Avro)",
            "",
        ),
        ConfigKeySpec::optional("schema.registry.username", "Schema Registry username", ""),
        ConfigKeySpec::optional("schema.registry.password", "Schema Registry password", ""),
        ConfigKeySpec::optional(
            "schema.registry.ssl.ca.location",
            "Schema Registry SSL CA cert path",
            "",
        ),
        ConfigKeySpec::optional(
            "schema.registry.ssl.certificate.location",
            "Schema Registry SSL client cert path",
            "",
        ),
        ConfigKeySpec::optional(
            "schema.registry.ssl.key.location",
            "Schema Registry SSL client key path",
            "",
        ),
        ConfigKeySpec::optional(
            "schema.compatibility",
            "Schema compatibility level override",
            "",
        ),
        ConfigKeySpec::optional(
            "schema.evolution.strategy",
            "Runtime schema evolution handling (log/reject/ignore)",
            "log",
        ),
        ConfigKeySpec::optional(
            "schema.registry.subject.name.strategy",
            "Schema Registry subject naming (topic-name/record-name/topic-record-name)",
            "topic-name",
        ),
        ConfigKeySpec::optional(
            "schema.registry.record.name",
            "Avro record name for record-based subject naming",
            "",
        ),
        ConfigKeySpec::optional(
            "schema.registry.discovery.timeout.ms",
            "Schema discovery timeout in milliseconds",
            "10000",
        ),
        ConfigKeySpec::optional(
            "max.poll.interval.ms",
            "Maximum interval between consumer polls in milliseconds",
            "600000",
        ),
        ConfigKeySpec::optional(
            "broker.commit.on.checkpoint",
            "Commit broker offsets after checkpoint completion",
            "true",
        ),
        ConfigKeySpec::optional(
            "reader.channel.capacity",
            "Bounded reader channel capacity in records",
            "8192",
        ),
    ]
}

/// Returns the configuration key specifications for the Kafka sink.
fn kafka_sink_config_keys() -> Vec<ConfigKeySpec> {
    vec![
        // Required
        ConfigKeySpec::required("bootstrap.servers", "Kafka broker addresses"),
        ConfigKeySpec::required("topic", "Target Kafka topic"),
        // Format
        ConfigKeySpec::optional("format", "Serialization format (json/csv/avro/raw)", "json"),
        // Security
        ConfigKeySpec::optional(
            "security.protocol",
            "Security protocol (plaintext/ssl/sasl_plaintext/sasl_ssl)",
            "plaintext",
        ),
        ConfigKeySpec::optional(
            "sasl.mechanism",
            "SASL mechanism (PLAIN/SCRAM-SHA-256/SCRAM-SHA-512/GSSAPI/OAUTHBEARER)",
            "",
        ),
        ConfigKeySpec::optional("sasl.username", "SASL username for PLAIN/SCRAM", ""),
        ConfigKeySpec::optional("sasl.password", "SASL password for PLAIN/SCRAM", ""),
        ConfigKeySpec::optional("ssl.ca.location", "SSL CA certificate file path", ""),
        ConfigKeySpec::optional(
            "ssl.certificate.location",
            "Client SSL certificate file path",
            "",
        ),
        ConfigKeySpec::optional("ssl.key.location", "Client SSL private key file path", ""),
        ConfigKeySpec::optional("ssl.key.password", "Password for encrypted SSL key", ""),
        ConfigKeySpec::optional(
            "max.in.flight.requests",
            "Maximum in-flight producer requests per connection (1-5)",
            "5",
        ),
        ConfigKeySpec::optional(
            "delivery.timeout.ms",
            "Delivery timeout in milliseconds",
            "120000",
        ),
        // Partitioning
        ConfigKeySpec::optional("key.column", "Column name to use as Kafka message key", ""),
        ConfigKeySpec::optional("envelope", "Output envelope (append/upsert)", "append"),
        ConfigKeySpec::optional(
            "partitioner",
            "Partitioning strategy (key-hash/round-robin/sticky)",
            "key-hash",
        ),
        // Batching & Compression
        ConfigKeySpec::optional("linger.ms", "Producer linger time in milliseconds", "5"),
        ConfigKeySpec::optional("batch.size", "Producer batch size in bytes", "16384"),
        ConfigKeySpec::optional("batch.num.messages", "Max messages per batch", "10000"),
        ConfigKeySpec::optional(
            "compression.type",
            "Compression (none/gzip/snappy/lz4/zstd)",
            "none",
        ),
        // Error Handling
        ConfigKeySpec::optional(
            "dlq.topic",
            "Dead letter queue topic for failed records",
            "",
        ),
        // Schema Registry
        ConfigKeySpec::optional(
            "schema.registry.url",
            "Confluent Schema Registry URL (required for Avro)",
            "",
        ),
        ConfigKeySpec::optional("schema.registry.username", "Schema Registry username", ""),
        ConfigKeySpec::optional("schema.registry.password", "Schema Registry password", ""),
        ConfigKeySpec::optional(
            "schema.registry.ssl.ca.location",
            "Schema Registry SSL CA cert path",
            "",
        ),
        ConfigKeySpec::optional(
            "schema.compatibility",
            "Schema compatibility level override",
            "",
        ),
    ]
}

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

    #[test]
    fn test_register_kafka_source() {
        let registry = ConnectorRegistry::new();
        register_kafka_source(&registry).unwrap();

        let sources = registry.list_sources();
        assert!(sources.contains(&"kafka".to_string()));

        let info = registry.source_info("kafka");
        assert!(info.is_some());
        let info = info.unwrap();
        assert_eq!(info.name, "kafka");
        assert!(info.is_source);
        assert!(!info.is_sink);
        let required: Vec<&str> = info
            .config_keys
            .iter()
            .filter(|key| key.required)
            .map(|key| key.key.as_str())
            .collect();
        assert!(required.contains(&"bootstrap.servers"));
        assert!(required.contains(&"group.id"));
        assert!(!required.contains(&"topic"));
        assert!(!required.contains(&"topic.pattern"));
        assert!(info.config_keys.iter().any(|key| key.key == "topic"));
        assert!(info
            .config_keys
            .iter()
            .any(|key| key.key == "topic.pattern"));
        for supported in [
            "schema.registry.subject.name.strategy",
            "schema.registry.record.name",
            "schema.registry.discovery.timeout.ms",
            "max.poll.interval.ms",
            "broker.commit.on.checkpoint",
            "reader.channel.capacity",
        ] {
            assert!(
                info.config_keys.iter().any(|key| key.key == supported),
                "registered Kafka source descriptor omits {supported}"
            );
        }
    }

    #[test]
    fn test_factory_creates_source() {
        let registry = ConnectorRegistry::new();
        register_kafka_source(&registry).unwrap();

        let config = crate::config::ConnectorConfig::new("kafka");
        let source = registry.create_source(&config, None);
        assert!(source.is_ok());
    }

    #[test]
    fn test_register_kafka_sink() {
        let registry = ConnectorRegistry::new();
        register_kafka_sink(&registry).unwrap();

        let sinks = registry.list_sinks();
        assert!(sinks.contains(&"kafka".to_string()));

        let info = registry.sink_info("kafka");
        assert!(info.is_some());
        let info = info.unwrap();
        assert_eq!(info.name, "kafka");
        assert!(!info.is_source);
        assert!(info.is_sink);
        assert!(info.config_keys.iter().any(|key| key.key == "envelope"));
        assert!(info.config_keys.iter().any(|key| key.key == "key.column"));
        assert!(!info
            .config_keys
            .iter()
            .any(|key| key.key == "delivery.guarantee"));
    }

    #[test]
    fn test_factory_creates_sink() {
        let registry = ConnectorRegistry::new();
        register_kafka_sink(&registry).unwrap();

        let config = crate::config::ConnectorConfig::new("kafka");
        let sink = registry.create_sink(&config, None);
        assert!(sink.is_ok());
    }
}