From 3779a7f1d9de2be3e682375e7aecda1a90847e51 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Tue, 27 Oct 2026 10:51:28 +0530 Subject: [PATCH 1/3] refactor(sink): add a generic Sink wrapper around the Kafka sink The handler passes typed EventRecords to the sink, which encodes them. New SINK env var (default kafka); Kafka settings are only read when SINK=kafka. No behaviour change. Refs #3 --- README.md | 1 + src/config.rs | 61 +++++++++++++++++++++++++++++++++++--------- src/main.rs | 36 ++++++++++++++------------ src/routes/health.rs | 4 +-- src/routes/ingest.rs | 52 ++++++++++++++++++------------------- src/routes/mod.rs | 4 +-- src/sink/kafka.rs | 12 +++++++-- src/sink/mod.rs | 38 +++++++++++++++++++++++++++ src/test_helpers.rs | 6 ++--- 9 files changed, 150 insertions(+), 64 deletions(-) diff --git a/README.md b/README.md index fe941aa..b9a9f54 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,7 @@ On the wire to Kafka, `traits` and `metadata` are produced as JSON-encoded strin | Env var | Required | Default | |---|---|---| +| `SINK` | no | `kafka` (the only sink so far) | | `KAFKA_BOOTSTRAP_SERVERS` | yes | — | | `KAFKA_TOPIC` | yes | — | | `KAFKA_EXTERNAL_TOPIC` | yes | — | diff --git a/src/config.rs b/src/config.rs index d7582af..771e8d3 100644 --- a/src/config.rs +++ b/src/config.rs @@ -5,28 +5,48 @@ use crate::sink::KafkaAuth; #[derive(Debug, Clone)] pub struct Config { pub listen_addr: String, - pub kafka_bootstrap_servers: String, - pub kafka_topic: String, - pub kafka_external_topic: String, - pub kafka_auth: KafkaAuth, + pub sink: SinkConfig, pub database_url: String, } +#[derive(Debug, Clone)] +pub enum SinkConfig { + Kafka(KafkaConfig), +} + +#[derive(Debug, Clone)] +pub struct KafkaConfig { + pub bootstrap_servers: String, + pub topic: String, + pub external_topic: String, + pub auth: KafkaAuth, +} + impl Config { pub fn from_env() -> Result { - let kafka_auth = parse_auth( + let sink = parse_sink(&env::var("SINK").unwrap_or_else(|_| "kafka".to_string()))?; + + Ok(Config { + listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()), + sink, + database_url: required("DATABASE_URL")?, + }) + } +} + +impl KafkaConfig { + fn from_env() -> Result { + let auth = parse_auth( &env::var("KAFKA_AUTH").unwrap_or_else(|_| "scram".to_string()), env::var("KAFKA_USERNAME").ok(), env::var("KAFKA_PASSWORD").ok(), )?; - Ok(Config { - listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()), - kafka_bootstrap_servers: required("KAFKA_BOOTSTRAP_SERVERS")?, - kafka_topic: required("KAFKA_TOPIC")?, - kafka_external_topic: required("KAFKA_EXTERNAL_TOPIC")?, - kafka_auth, - database_url: required("DATABASE_URL")?, + Ok(KafkaConfig { + bootstrap_servers: required("KAFKA_BOOTSTRAP_SERVERS")?, + topic: required("KAFKA_TOPIC")?, + external_topic: required("KAFKA_EXTERNAL_TOPIC")?, + auth, }) } } @@ -35,6 +55,13 @@ fn required(name: &'static str) -> Result { env::var(name).map_err(|_| ConfigError::MissingRequired(name)) } +fn parse_sink(name: &str) -> Result { + match name { + "kafka" => Ok(SinkConfig::Kafka(KafkaConfig::from_env()?)), + other => Err(ConfigError::InvalidSink(other.to_string())), + } +} + fn parse_auth( mode: &str, username: Option, @@ -56,6 +83,8 @@ pub enum ConfigError { MissingRequired(&'static str), #[error("Invalid value for {0}: {1:?} (expected scram or none)")] Invalid(&'static str, String), + #[error("Invalid value for SINK: {0:?} (expected kafka)")] + InvalidSink(String), } #[cfg(test)] @@ -98,4 +127,12 @@ mod tests { Err(ConfigError::Invalid("KAFKA_AUTH", v)) if v == "iam" )); } + + #[test] + fn parse_sink_rejects_unknown_value() { + assert!(matches!( + parse_sink("clickhouse"), + Err(ConfigError::InvalidSink(v)) if v == "clickhouse" + )); + } } diff --git a/src/main.rs b/src/main.rs index e9b507f..d5185e6 100644 --- a/src/main.rs +++ b/src/main.rs @@ -16,9 +16,9 @@ use tokio::signal; use tracing::info; use auth::ContextState; -use config::Config; +use config::{Config, SinkConfig}; use environment_keys::EnvironmentKeys; -use sink::{KafkaSink, KafkaSinkConfig}; +use sink::{KafkaSink, KafkaSinkConfig, Sink}; const ENVIRONMENT_KEY_REFRESH_AFTER: Duration = Duration::from_secs(300); @@ -33,19 +33,23 @@ async fn main() -> anyhow::Result<()> { .init(); let config = Config::from_env()?; - info!( - listen_addr = %config.listen_addr, - kafka_bootstrap_servers = %config.kafka_bootstrap_servers, - kafka_topic = %config.kafka_topic, - kafka_auth = ?config.kafka_auth, - "Starting flagsmith-analytics server" - ); - - let sink = KafkaSink::new(KafkaSinkConfig::new( - config.kafka_bootstrap_servers.clone(), - config.kafka_topic.clone(), - config.kafka_auth.clone(), - ))?; + let (sink, external_topic) = match &config.sink { + SinkConfig::Kafka(kafka) => { + info!( + listen_addr = %config.listen_addr, + kafka_bootstrap_servers = %kafka.bootstrap_servers, + kafka_topic = %kafka.topic, + kafka_auth = ?kafka.auth, + "Starting flagsmith-analytics server" + ); + let sink = KafkaSink::new(KafkaSinkConfig::new( + kafka.bootstrap_servers.clone(), + kafka.topic.clone(), + kafka.auth.clone(), + ))?; + (Sink::Kafka(sink), kafka.external_topic.as_str()) + } + }; let pool = PgPoolOptions::new() .max_connections(8) @@ -53,7 +57,7 @@ async fn main() -> anyhow::Result<()> { .connect_lazy(&config.database_url)?; let context_state = ContextState { environment_keys: EnvironmentKeys::new(pool, ENVIRONMENT_KEY_REFRESH_AFTER), - external_topic: Arc::from(config.kafka_external_topic.as_str()), + external_topic: Arc::from(external_topic), }; let app = routes::app(sink.clone(), context_state); diff --git a/src/routes/health.rs b/src/routes/health.rs index 942b003..d696d36 100644 --- a/src/routes/health.rs +++ b/src/routes/health.rs @@ -1,7 +1,7 @@ use axum::{extract::State, http::StatusCode, Json}; use serde::Serialize; -use crate::sink::KafkaSink; +use crate::sink::Sink; #[derive(Serialize)] pub struct HealthResponse { @@ -12,7 +12,7 @@ pub struct HealthResponse { /// Liveness plus one local check: a fatal producer error is permanent, so /// report it and let ECS replace the task. No downstream call, so this cannot hang. -pub async fn health_check(State(sink): State) -> (StatusCode, Json) { +pub async fn health_check(State(sink): State) -> (StatusCode, Json) { match sink.fatal_error() { None => ( StatusCode::OK, diff --git a/src/routes/ingest.rs b/src/routes/ingest.rs index 1cea1ab..4678989 100644 --- a/src/routes/ingest.rs +++ b/src/routes/ingest.rs @@ -12,7 +12,7 @@ use crate::lookup::mask_key; use crate::models::{ Event, EventBatch, EventRecord, IngestResponse, RejectedEvent, FLAG_EXPOSURE_EVENT, }; -use crate::sink::KafkaSink; +use crate::sink::Sink; /// Header-derived fields are copied into every record of the batch, so an /// unbounded header would inflate every record. @@ -177,15 +177,15 @@ fn partition_events(events: &[Event]) -> (Vec<&Event>, Vec) { } /// POST /v1/events -/// Validates each event independently: valid events are serialized and -/// forwarded to the sink, invalid ones are reported back as `rejected` -/// entries without blocking the rest of the batch. +/// Validates each event independently: valid events are forwarded to the +/// sink, invalid ones are reported back as `rejected` entries without +/// blocking the rest of the batch. /// Returns 202 when at least one event was accepted and forwarded, 400 when /// none were (all invalid, or empty batch), 503 on sink failure /// (delivery not confirmed; retrying is safe but may duplicate events from /// a partially delivered batch). pub async fn ingest_batch( - State(sink): State, + State(sink): State, Extension(ctx): Extension, headers: HeaderMap, Json(batch): Json, @@ -225,28 +225,26 @@ pub async fn ingest_batch( .map(parse_sdk_user_agent) .unwrap_or_else(|| ("unknown".to_string(), "unknown".to_string())); - let records: Vec> = valid + // TODO: move this to a method. + let records: Vec = valid .iter() - .map(|event| { - let record = EventRecord { - environment_key: &ctx.client_api_key, - collected_at, - sdk_language: &sdk_language, - sdk_version: &sdk_version, - event: &event.event, - feature_name: event.feature_name.as_deref(), - identifier: event.identifier.as_deref(), - value: &event.value, - traits: event.traits.as_ref().map(ToString::to_string), - metadata: event.metadata.as_ref().map(ToString::to_string), - timestamp: event.timestamp, - }; - serde_json::to_vec(&record).expect("EventRecord serialization cannot fail") + .map(|event| EventRecord { + environment_key: &ctx.client_api_key, + collected_at, + sdk_language: &sdk_language, + sdk_version: &sdk_version, + event: &event.event, + feature_name: event.feature_name.as_deref(), + identifier: event.identifier.as_deref(), + value: &event.value, + traits: event.traits.as_ref().map(ToString::to_string), + metadata: event.metadata.as_ref().map(ToString::to_string), + timestamp: event.timestamp, }) .collect(); let record_count = records.len(); - match sink.put(ctx.destination.as_deref(), records).await { + match sink.put(ctx.destination.as_deref(), &records).await { Ok(()) => { debug!( client_api_key = %mask_key(&ctx.client_api_key), @@ -288,7 +286,7 @@ mod tests { use rdkafka::types::RDKafkaRespErr; use tower::ServiceExt; - use crate::sink::{KafkaAuth, KafkaSinkConfig}; + use crate::sink::{KafkaAuth, KafkaSink, KafkaSinkConfig}; use crate::test_helpers::{ assert_topic_empty, mock_kafka, read_records, FLAGSMITH_WAREHOUSE_TOPIC, }; @@ -534,17 +532,17 @@ mod tests { } /// A sink whose bootstrap server never answers, so every put times out. - fn unreachable_sink() -> KafkaSink { + fn unreachable_sink() -> Sink { let mut config = KafkaSinkConfig::new( "127.0.0.1:1".to_string(), FLAGSMITH_WAREHOUSE_TOPIC.to_string(), KafkaAuth::None, ); config.message_timeout = Duration::from_secs(2); - KafkaSink::new(config).expect("sink") + Sink::Kafka(KafkaSink::new(config).expect("sink")) } - fn app_with_destination(sink: &KafkaSink, destination: Option<&str>) -> Router { + fn app_with_destination(sink: &Sink, destination: Option<&str>) -> Router { let ctx = EnvironmentContext { client_api_key: Arc::from("test-client-key"), destination: destination.map(Arc::from), @@ -555,7 +553,7 @@ mod tests { .with_state(sink.clone()) } - fn app(sink: &KafkaSink) -> Router { + fn app(sink: &Sink) -> Router { app_with_destination(sink, None) } diff --git a/src/routes/mod.rs b/src/routes/mod.rs index fd2b277..7e12d2d 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -6,12 +6,12 @@ use tower_http::cors::{Any, CorsLayer}; use tower_http::decompression::RequestDecompressionLayer; use crate::auth::{extract_environment_context, ContextState}; -use crate::sink::KafkaSink; +use crate::sink::Sink; pub use health::health_check; pub use ingest::ingest_batch; -pub fn app(sink: KafkaSink, context_state: ContextState) -> Router { +pub fn app(sink: Sink, context_state: ContextState) -> Router { Router::new() .route("/health", get(health_check)) .route( diff --git a/src/sink/kafka.rs b/src/sink/kafka.rs index 33ba54b..09cb405 100644 --- a/src/sink/kafka.rs +++ b/src/sink/kafka.rs @@ -10,6 +10,8 @@ use rdkafka::statistics::Statistics; use rdkafka::util::Timeout; use tracing::{debug, error, info, warn}; +use crate::models::EventRecord; + const CLIENT_ID: &str = "events-ingestion-api"; /// How long `put` waits for room in the local queue before failing the batch. const ENQUEUE_TIMEOUT: Duration = Duration::from_millis(500); @@ -146,11 +148,17 @@ impl KafkaSink { pub async fn put( &self, destination: Option<&str>, - records: Vec>, + records: &[EventRecord<'_>], ) -> Result<(), SinkError> { let topic = destination.unwrap_or(&self.default_topic); let total = records.len(); - let sends = records.iter().map(|bytes| { + let payloads: Vec> = records + .iter() + .map(|record| { + serde_json::to_vec(record).expect("EventRecord serialization cannot fail") + }) + .collect(); + let sends = payloads.iter().map(|bytes| { self.producer.send( FutureRecord::<(), _>::to(topic).payload(bytes), Timeout::After(ENQUEUE_TIMEOUT), diff --git a/src/sink/mod.rs b/src/sink/mod.rs index 2a73a9b..07812b4 100644 --- a/src/sink/mod.rs +++ b/src/sink/mod.rs @@ -1,3 +1,41 @@ mod kafka; +use std::time::Duration; + pub use kafka::{KafkaAuth, KafkaSink, KafkaSinkConfig}; + +use crate::models::EventRecord; + +/// Where accepted events are written. +#[derive(Clone)] +pub enum Sink { + Kafka(KafkaSink), +} + +impl Sink { + /// Write every record to the destination (the sink's default when `None`). + /// Returns once the sink has accepted all of them. + pub async fn put( + &self, + destination: Option<&str>, + records: &[EventRecord<'_>], + ) -> Result<(), kafka::SinkError> { + match self { + Self::Kafka(sink) => sink.put(destination, records).await, + } + } + + /// A permanent sink failure that should fail the health check. + pub fn fatal_error(&self) -> Option { + match self { + Self::Kafka(sink) => sink.fatal_error(), + } + } + + /// Drain anything still buffered, bounded by `timeout`. + pub async fn shutdown(self, timeout: Duration) { + match self { + Self::Kafka(sink) => sink.shutdown(timeout).await, + } + } +} diff --git a/src/test_helpers.rs b/src/test_helpers.rs index 35c85b3..9545c21 100644 --- a/src/test_helpers.rs +++ b/src/test_helpers.rs @@ -8,7 +8,7 @@ use rdkafka::producer::DefaultProducerContext; use rdkafka::Message; use sqlx::PgPool; -use crate::sink::{KafkaAuth, KafkaSink, KafkaSinkConfig}; +use crate::sink::{KafkaAuth, KafkaSink, KafkaSinkConfig, Sink}; pub(crate) const FLAGSMITH_WAREHOUSE_TOPIC: &str = "events-default"; @@ -18,7 +18,7 @@ pub(crate) type MockKafkaCluster = MockCluster<'static, DefaultProducerContext>; /// test. pub(crate) struct MockKafka { pub(crate) cluster: MockKafkaCluster, - pub(crate) sink: KafkaSink, + pub(crate) sink: Sink, } pub(crate) fn mock_kafka(topics: &[&str]) -> MockKafka { @@ -32,7 +32,7 @@ pub(crate) fn mock_kafka(topics: &[&str]) -> MockKafka { KafkaAuth::None, ); config.message_timeout = Duration::from_secs(5); - let sink = KafkaSink::new(config).expect("sink"); + let sink = Sink::Kafka(KafkaSink::new(config).expect("sink")); MockKafka { cluster, sink } } From 631540a0b8523457027c6058e7d5165382a4bc9f Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Tue, 27 Oct 2026 12:55:35 +0530 Subject: [PATCH 2/3] refactor(sink): make SinkError sink-agnostic Move SinkError to the sink module, with a Kafka variant ready for other sinks. No behaviour change. Refs #3 --- src/sink/kafka.rs | 16 +++++++--------- src/sink/mod.rs | 10 ++++++++-- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/src/sink/kafka.rs b/src/sink/kafka.rs index 09cb405..d41fb13 100644 --- a/src/sink/kafka.rs +++ b/src/sink/kafka.rs @@ -149,7 +149,7 @@ impl KafkaSink { &self, destination: Option<&str>, records: &[EventRecord<'_>], - ) -> Result<(), SinkError> { + ) -> Result<(), DeliveryError> { let topic = destination.unwrap_or(&self.default_topic); let total = records.len(); let payloads: Vec> = records @@ -167,7 +167,7 @@ impl KafkaSink { let mut failures = join_all(sends).await.into_iter().filter_map(Result::err); match failures.next() { None => Ok(()), - Some((error, _record)) => Err(SinkError::Delivery { + Some((error, _record)) => Err(DeliveryError { failed: 1 + failures.count(), total, error, @@ -198,13 +198,11 @@ impl KafkaSink { } #[derive(Debug, thiserror::Error)] -pub enum SinkError { - #[error("Kafka delivery failed for {failed} of {total} records: {error}")] - Delivery { - failed: usize, - total: usize, - error: KafkaError, - }, +#[error("Kafka delivery failed for {failed} of {total} records: {error}")] +pub struct DeliveryError { + failed: usize, + total: usize, + error: KafkaError, } #[cfg(test)] diff --git a/src/sink/mod.rs b/src/sink/mod.rs index 07812b4..4369129 100644 --- a/src/sink/mod.rs +++ b/src/sink/mod.rs @@ -19,9 +19,9 @@ impl Sink { &self, destination: Option<&str>, records: &[EventRecord<'_>], - ) -> Result<(), kafka::SinkError> { + ) -> Result<(), SinkError> { match self { - Self::Kafka(sink) => sink.put(destination, records).await, + Self::Kafka(sink) => Ok(sink.put(destination, records).await?), } } @@ -39,3 +39,9 @@ impl Sink { } } } + +#[derive(Debug, thiserror::Error)] +pub enum SinkError { + #[error(transparent)] + Kafka(#[from] kafka::DeliveryError), +} From b00b032dbbe14b0f71c206b81b45943bb0d69510 Mon Sep 17 00:00:00 2001 From: Gagan Trivedi Date: Wed, 28 Oct 2026 12:16:32 +0530 Subject: [PATCH 3/3] refactor(ingest): encode traits and metadata in the sink EventRecord borrows traits and metadata as JSON values; serializing the record encodes them as JSON strings, so the handler does no encoding and Kafka output is unchanged. Deciding the sink kind no longer reads the environment, so an unset SINK defaulting to kafka is tested. Refs #3 --- src/config.rs | 24 +++++++++++++++++++----- src/models/event.rs | 27 ++++++++++++++++++++++----- src/routes/ingest.rs | 8 ++++---- 3 files changed, 45 insertions(+), 14 deletions(-) diff --git a/src/config.rs b/src/config.rs index 771e8d3..4e0e38f 100644 --- a/src/config.rs +++ b/src/config.rs @@ -24,7 +24,9 @@ pub struct KafkaConfig { impl Config { pub fn from_env() -> Result { - let sink = parse_sink(&env::var("SINK").unwrap_or_else(|_| "kafka".to_string()))?; + let sink = match parse_sink(env::var("SINK").ok().as_deref())? { + SinkKind::Kafka => SinkConfig::Kafka(KafkaConfig::from_env()?), + }; Ok(Config { listen_addr: env::var("LISTEN_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".to_string()), @@ -55,9 +57,15 @@ fn required(name: &'static str) -> Result { env::var(name).map_err(|_| ConfigError::MissingRequired(name)) } -fn parse_sink(name: &str) -> Result { - match name { - "kafka" => Ok(SinkConfig::Kafka(KafkaConfig::from_env()?)), +#[derive(Debug, PartialEq)] +enum SinkKind { + Kafka, +} + +/// Which sink `SINK` selects; unset means Kafka. +fn parse_sink(value: Option<&str>) -> Result { + match value.unwrap_or("kafka") { + "kafka" => Ok(SinkKind::Kafka), other => Err(ConfigError::InvalidSink(other.to_string())), } } @@ -128,10 +136,16 @@ mod tests { )); } + #[test] + fn parse_sink_defaults_to_kafka() { + assert_eq!(parse_sink(None).unwrap(), SinkKind::Kafka); + assert_eq!(parse_sink(Some("kafka")).unwrap(), SinkKind::Kafka); + } + #[test] fn parse_sink_rejects_unknown_value() { assert!(matches!( - parse_sink("clickhouse"), + parse_sink(Some("clickhouse")), Err(ConfigError::InvalidSink(v)) if v == "clickhouse" )); } diff --git a/src/models/event.rs b/src/models/event.rs index d9d2d6b..2f44b77 100644 --- a/src/models/event.rs +++ b/src/models/event.rs @@ -1,4 +1,4 @@ -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Serialize, Serializer}; /// Reserved system event for flag-exposure tracking. `feature_name` is required /// for events of this type. @@ -73,9 +73,26 @@ pub struct EventRecord<'a> { pub identifier: Option<&'a str>, #[serde(skip_serializing_if = "serde_json::Value::is_null")] pub value: &'a serde_json::Value, - #[serde(skip_serializing_if = "Option::is_none")] - pub traits: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub metadata: Option, + #[serde( + skip_serializing_if = "Option::is_none", + serialize_with = "json_string" + )] + pub traits: Option<&'a serde_json::Value>, + #[serde( + skip_serializing_if = "Option::is_none", + serialize_with = "json_string" + )] + pub metadata: Option<&'a serde_json::Value>, pub timestamp: i64, } + +/// Serializes a JSON value as its encoded text, in a string. +fn json_string( + value: &Option<&serde_json::Value>, + serializer: S, +) -> Result { + match value { + Some(value) => serializer.collect_str(value), + None => serializer.serialize_none(), + } +} diff --git a/src/routes/ingest.rs b/src/routes/ingest.rs index 4678989..4a039e0 100644 --- a/src/routes/ingest.rs +++ b/src/routes/ingest.rs @@ -237,8 +237,8 @@ pub async fn ingest_batch( feature_name: event.feature_name.as_deref(), identifier: event.identifier.as_deref(), value: &event.value, - traits: event.traits.as_ref().map(ToString::to_string), - metadata: event.metadata.as_ref().map(ToString::to_string), + traits: event.traits.as_ref(), + metadata: event.metadata.as_ref(), timestamp: event.timestamp, }) .collect(); @@ -512,8 +512,8 @@ mod tests { feature_name: event.feature_name.as_deref(), identifier: event.identifier.as_deref(), value: &event.value, - traits: event.traits.as_ref().map(ToString::to_string), - metadata: event.metadata.as_ref().map(ToString::to_string), + traits: event.traits.as_ref(), + metadata: event.metadata.as_ref(), timestamp: event.timestamp, };