refactor(sink): add a generic Sink wrapper around the Kafka sink

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 | — |

src/config.rs

@@ -5,28 +5,50 @@ 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<Self, ConfigError> {
-        let kafka_auth = parse_auth(
+        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()),
+            sink,
+            database_url: required("DATABASE_URL")?,
+        })
+    }
+}
+
+impl KafkaConfig {
+    fn from_env() -> Result<Self, ConfigError> {
+        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 +57,19 @@ fn required(name: &'static str) -> Result<String, ConfigError> {
     env::var(name).map_err(|_| ConfigError::MissingRequired(name))
 }
 
+#[derive(Debug, PartialEq)]
+enum SinkKind {
+    Kafka,
+}
+
+/// Which sink `SINK` selects; unset means Kafka.
+fn parse_sink(value: Option<&str>) -> Result<SinkKind, ConfigError> {
+    match value.unwrap_or("kafka") {
+        "kafka" => Ok(SinkKind::Kafka),
+        other => Err(ConfigError::InvalidSink(other.to_string())),
+    }
+}
+
 fn parse_auth(
     mode: &str,
     username: Option<String>,
@@ -56,6 +91,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 +135,18 @@ mod tests {
             Err(ConfigError::Invalid("KAFKA_AUTH", v)) if v == "iam"
         ));
     }
+
+    #[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(Some("clickhouse")),
+            Err(ConfigError::InvalidSink(v)) if v == "clickhouse"
+        ));
+    }
 }

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,27 +33,31 @@ 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)
         .acquire_timeout(Duration::from_secs(1))
         .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);

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<String>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub metadata: Option<String>,
+    #[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<S: Serializer>(
+    value: &Option<&serde_json::Value>,
+    serializer: S,
+) -> Result<S::Ok, S::Error> {
+    match value {
+        Some(value) => serializer.collect_str(value),
+        None => serializer.serialize_none(),
+    }
+}

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<KafkaSink>) -> (StatusCode, Json<HealthResponse>) {
+pub async fn health_check(State(sink): State<Sink>) -> (StatusCode, Json<HealthResponse>) {
     match sink.fatal_error() {
         None => (
             StatusCode::OK,

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<RejectedEvent>) {
 }
 
 /// 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<KafkaSink>,
+    State(sink): State<Sink>,
     Extension(ctx): Extension<EnvironmentContext>,
     headers: HeaderMap,
     Json(batch): Json<EventBatch>,
@@ -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<Vec<u8>> = valid
+    // TODO: move this to a method.
+    let records: Vec<EventRecord> = 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(),
+            metadata: event.metadata.as_ref(),
+            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,
     };
@@ -514,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,
         };
 
@@ -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)
     }
 

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(