diff --git a/README.md b/README.md index 81c3e0e..b8a8687 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ Rust + axum, deployed as an ARM64 container on ECS Fargate. | Method | Path | Description | |---|---|---| -| GET | `/health` | Liveness probe → `{"status":"healthy"}`. Touches no downstream; returns `503` once the Kafka producer has hit a fatal error, so ECS replaces the task. | +| GET | `/health` | Liveness probe → `{"status":"healthy"}`. Touches no downstream; returns `503` once the sink has hit a fatal error (the Kafka producer's, or a ClickHouse config error), so ECS replaces the task. | | POST | `/v1/events` | Ingest an event batch. Requires the `X-Environment-Key` header. | `POST /v1/events` validates each event independently — invalid events are rejected individually and never block valid ones in the same batch. @@ -52,7 +52,7 @@ On the wire to Kafka, `traits` and `metadata` are produced as JSON-encoded strin | Env var | Required | Default | |---|---|---| -| `SINK` | no | `kafka` (or `clickhouse`) | +| `SINK` | no | `kafka` (or `clickhouse`, `postgres`) | | `KAFKA_BOOTSTRAP_SERVERS` | with `kafka` | — | | `KAFKA_TOPIC` | with `kafka` | — | | `KAFKA_EXTERNAL_TOPIC` | with `kafka` | — | @@ -69,7 +69,7 @@ On the wire to Kafka, `traits` and `metadata` are produced as JSON-encoded strin | `DATABASE_URL` | yes | — | | `LISTEN_ADDR` | no | `0.0.0.0:8080` | -With `SINK=kafka`, events from environments that use an external warehouse go to `KAFKA_EXTERNAL_TOPIC`; all others go to `KAFKA_TOPIC`. With `SINK=clickhouse`, every event goes to the one ClickHouse table. The server looks presented keys up in the `experimentation_environment_keys` view in Flagsmith's Postgres, at `DATABASE_URL`. +With `SINK=kafka`, events from environments that use an external warehouse go to `KAFKA_EXTERNAL_TOPIC`; all others go to `KAFKA_TOPIC`. With `SINK=clickhouse`, every event goes to the one ClickHouse table, and with `SINK=postgres`, to `flagsmith_exp.events` at `DATABASE_URL`. The server looks presented keys up in the `experimentation_environment_keys` view in Flagsmith's Postgres, at `DATABASE_URL`. ### Database role @@ -88,7 +88,16 @@ With `KAFKA_AUTH=scram`, the producer authenticates with SASL/SCRAM-SHA-512 over ### ClickHouse sink -With `SINK=clickhouse`, `202` means the events are buffered in memory, not yet stored. A background task takes the whole buffer every `SINK_FLUSH_INTERVAL_MS`, or sooner once `SINK_FLUSH_MAX_EVENTS` events are waiting, and writes it as one `INSERT`. A failed batch is retried 3 times and then dropped with an error log. The buffer is unbounded: requests never wait or get `503`, so while a write is slow or failing (up to about 2 minutes with retries) it grows by the ingest rate times the write time, and the next flush takes it all. A normal shutdown (SIGTERM) flushes the buffer, but events still buffered are lost if the process crashes. +With `SINK=clickhouse`, `202` means the events are buffered in memory, not yet stored. A background task takes the whole buffer every `SINK_FLUSH_INTERVAL_MS`, or sooner once `SINK_FLUSH_MAX_EVENTS` events are waiting, and writes it as one `INSERT`. A batch that fails with a transient error (network, timeout, overload, or anything unrecognised) is retried 5 times over about 17 seconds, enough to ride out a ClickHouse restart, and then dropped with an error log; these never affect `/health`. Rows the table rejects (parse, type or constraint errors) drop the batch at once, with an error log, and leave `/health` alone. A config error (bad credentials or grants, or a missing database, table or column) drops the batch at once and fails `/health` until restart. At startup the server checks that ClickHouse is reachable, accepts the credentials, and has a table whose columns match, and exits with an error if not. The buffer is unbounded: requests never wait or get `503`, so while a write is slow or failing (up to about 3 minutes with retries) it grows by the ingest rate times the write time, and the next flush takes it all. A normal shutdown (SIGTERM) flushes the buffer, but events still buffered are lost if the process crashes. + +### PostgreSQL sink + +With `SINK=postgres`, events are buffered, flushed and retried as with the ClickHouse sink, and each batch is written to `flagsmith_exp.events` in Flagsmith's Postgres, at `DATABASE_URL`, as one `INSERT` in a transaction. Flagsmith's migrations create the table. The sink shares the server's connection pool. A missing table or column, or a permission or authentication error, fails `/health` until restart. If the server's role does not own the table, grant it: + +```sql +GRANT USAGE ON SCHEMA flagsmith_exp TO ingestion_server; +GRANT INSERT ON flagsmith_exp.events TO ingestion_server; +``` ## Run locally diff --git a/src/config.rs b/src/config.rs index c3c3f10..caa0c8e 100644 --- a/src/config.rs +++ b/src/config.rs @@ -18,6 +18,9 @@ pub enum SinkConfig { clickhouse: ClickHouseConfig, batch: BatchConfig, }, + Postgres { + batch: BatchConfig, + }, } #[derive(Debug, Clone)] @@ -36,6 +39,9 @@ impl Config { clickhouse: clickhouse_config_from_env()?, batch: batch_config_from_env()?, }, + SinkKind::Postgres => SinkConfig::Postgres { + batch: batch_config_from_env()?, + }, }; Ok(Config { @@ -71,6 +77,7 @@ fn required(name: &'static str) -> Result { enum SinkKind { Kafka, ClickHouse, + Postgres, } /// Which sink `SINK` selects; unset means Kafka. @@ -78,6 +85,7 @@ fn parse_sink(value: Option<&str>) -> Result { match value.unwrap_or("kafka") { "kafka" => Ok(SinkKind::Kafka), "clickhouse" => Ok(SinkKind::ClickHouse), + "postgres" => Ok(SinkKind::Postgres), other => Err(ConfigError::InvalidSink(other.to_string())), } } @@ -134,7 +142,7 @@ 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 or clickhouse)")] + #[error("Invalid value for SINK: {0:?} (expected kafka, clickhouse or postgres)")] InvalidSink(String), #[error("Invalid value for {0}: {1:?} (expected a positive integer)")] InvalidNumber(&'static str, String), @@ -189,13 +197,14 @@ mod tests { parse_sink(Some("clickhouse")).unwrap(), SinkKind::ClickHouse ); + assert_eq!(parse_sink(Some("postgres")).unwrap(), SinkKind::Postgres); } #[test] fn parse_sink_rejects_unknown_value() { assert!(matches!( - parse_sink(Some("postgres")), - Err(ConfigError::InvalidSink(v)) if v == "postgres" + parse_sink(Some("bigquery")), + Err(ConfigError::InvalidSink(v)) if v == "bigquery" )); } diff --git a/src/main.rs b/src/main.rs index 199f48c..fe180de 100644 --- a/src/main.rs +++ b/src/main.rs @@ -11,6 +11,7 @@ mod test_helpers; use std::sync::Arc; use std::time::Duration; +use anyhow::Context; use sqlx::postgres::PgPoolOptions; use tokio::signal; use tracing::info; @@ -18,7 +19,7 @@ use tracing::info; use auth::ContextState; use config::{Config, SinkConfig}; use environment_keys::EnvironmentKeys; -use sink::{BatchingSink, ClickHouseWriter, KafkaSink, KafkaSinkConfig, Sink}; +use sink::{BatchingSink, ClickHouseWriter, KafkaSink, KafkaSinkConfig, PostgresWriter, Sink}; const ENVIRONMENT_KEY_REFRESH_AFTER: Duration = Duration::from_secs(300); @@ -33,6 +34,10 @@ async fn main() -> anyhow::Result<()> { .init(); let config = Config::from_env()?; + let pool = PgPoolOptions::new() + .max_connections(8) + .acquire_timeout(Duration::from_secs(1)) + .connect_lazy(&config.database_url)?; let (sink, external_topic) = match &config.sink { SinkConfig::Kafka(kafka) => { info!( @@ -57,17 +62,30 @@ async fn main() -> anyhow::Result<()> { "Starting flagsmith-analytics server" ); let writer = ClickHouseWriter::new(clickhouse.clone()); + writer + .check() + .await + .context("ClickHouse startup check failed")?; ( Sink::ClickHouse(BatchingSink::new(writer, batch.clone())), None, ) } + SinkConfig::Postgres { batch } => { + info!( + listen_addr = %config.listen_addr, + sink = "postgres", + batch = ?batch, + "Starting flagsmith-analytics server" + ); + let writer = PostgresWriter::new(pool.clone()); + ( + Sink::Postgres(BatchingSink::new(writer, batch.clone())), + None, + ) + } }; - 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: external_topic.map(Arc::from), diff --git a/src/routes/health.rs b/src/routes/health.rs index d696d36..ae5783e 100644 --- a/src/routes/health.rs +++ b/src/routes/health.rs @@ -10,7 +10,7 @@ pub struct HealthResponse { error: Option, } -/// Liveness plus one local check: a fatal producer error is permanent, so +/// Liveness plus one local check: a fatal sink 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) { match sink.fatal_error() { diff --git a/src/sink/batching.rs b/src/sink/batching.rs index 6848a8d..e2667dc 100644 --- a/src/sink/batching.rs +++ b/src/sink/batching.rs @@ -1,7 +1,6 @@ use std::future::Future; -use std::hash::{BuildHasher, RandomState}; use std::mem; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, OnceLock}; use std::time::Duration; use tokio::sync::Notify; @@ -9,13 +8,18 @@ use tokio::task::JoinHandle; use tokio::time::{Instant, MissedTickBehavior}; use tracing::{error, warn}; -use super::SinkError; +use super::{ErrorKind, SinkError}; use crate::models::EventRecord; -/// Retries after the first failed write, before the batch is dropped. -const MAX_RETRIES: u32 = 3; -/// Doubles per retry: about 200 ms, 400 ms, 800 ms before jitter. -const RETRY_BASE_DELAY: Duration = Duration::from_millis(200); +/// Waits between attempts to write a batch that failed with a transient error, +/// before it is dropped: about 17 s in all, to ride out a ClickHouse restart. +const RETRY_DELAYS: [Duration; 5] = [ + Duration::from_secs(1), + Duration::from_secs(2), + Duration::from_secs(4), + Duration::from_secs(5), + Duration::from_secs(5), +]; /// Writes one batch of records to a store. pub trait BatchWriter: Send + Sync + 'static { @@ -36,6 +40,8 @@ struct Shared { buffer: Mutex>, flush_now: Notify, stop: Notify, + /// The first config error; stays set until restart. + fatal_error: OnceLock, } pub struct BatchingSink { @@ -60,6 +66,7 @@ impl BatchingSink { buffer: Mutex::new(Vec::new()), flush_now: Notify::new(), stop: Notify::new(), + fatal_error: OnceLock::new(), }); let flusher = tokio::spawn(run(shared.clone())); Self { @@ -77,6 +84,12 @@ impl BatchingSink { } } + /// A config error, such as bad credentials or a missing table, that + /// stops every write until the server is fixed. + pub fn fatal_error(&self) -> Option { + self.shared.fatal_error.get().cloned() + } + /// Stop the flusher and write what is buffered, bounded by `timeout`. pub async fn shutdown(self, timeout: Duration) { self.shared.stop.notify_one(); @@ -113,45 +126,51 @@ async fn flush(shared: &Shared) { if records.is_empty() { return; } - write_with_retries(&shared.writer, &records).await; + write_with_retries(shared, &records).await; } -async fn write_with_retries(writer: &W, records: &[EventRecord]) { - let mut retry = 0; - loop { - match writer.write(records).await { - Ok(()) => return, - Err(e) if retry < MAX_RETRIES => { - warn!(rows = records.len(), retry, error = %e, "Sink batch write failed, retrying"); - tokio::time::sleep(backoff(retry)).await; - retry += 1; +/// Retries transient errors per `RETRY_DELAYS`. Drops the batch at once on +/// bad data or a config error, and records a config error to fail the health +/// check. +async fn write_with_retries(shared: &Shared, records: &[EventRecord]) { + let rows = records.len(); + for retry in 0.. { + let Err(e) = shared.writer.write(records).await else { + return; + }; + match e.kind() { + ErrorKind::Transient => {} + ErrorKind::BadData => { + error!(rows, error = %e, "Sink batch dropped: rows rejected"); + return; } - Err(e) => { - error!(rows = records.len(), error = %e, "Sink batch dropped after retries"); + ErrorKind::Config => { + error!(rows, error = %e, "Sink batch dropped: sink misconfigured"); + let _ = shared.fatal_error.set(e.to_string()); return; } } + let Some(&delay) = RETRY_DELAYS.get(retry) else { + error!(rows, error = %e, "Sink batch dropped after retries"); + return; + }; + warn!(rows, retry, error = %e, "Sink batch write failed, retrying"); + tokio::time::sleep(delay).await; } } -/// Exponential backoff with ±50% jitter, so tasks that failed together do not -/// retry together. -fn backoff(retry: u32) -> Duration { - let jitter_permille = RandomState::new().hash_one(retry) % 1000; - (RETRY_BASE_DELAY * 2u32.pow(retry)).mul_f64(0.5 + jitter_permille as f64 / 1000.0) -} - #[cfg(test)] mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use super::*; - /// Records every write attempt; each takes `delay`, and all fail when - /// `fail` is set. + /// Records every write attempt; each takes `delay`, and the first + /// `failures` fail with `error`. #[derive(Clone, Default)] struct FakeWriter { - fail: bool, + failures: usize, + error: Option clickhouse::error::Error>, delay: Duration, attempts: Arc, written: Arc>>>, @@ -159,10 +178,11 @@ mod tests { impl BatchWriter for FakeWriter { async fn write(&self, records: &[EventRecord]) -> Result<(), SinkError> { - self.attempts.fetch_add(1, Ordering::SeqCst); + let attempt = self.attempts.fetch_add(1, Ordering::SeqCst); tokio::time::sleep(self.delay).await; - if self.fail { - return Err(clickhouse::error::Error::Custom("failed".to_string()).into()); + if attempt < self.failures { + let error = self.error.expect("an error to fail with"); + return Err(error().into()); } let timestamps = records.iter().map(|record| record.timestamp).collect(); self.written.lock().unwrap().push(timestamps); @@ -247,19 +267,59 @@ mod tests { #[tokio::test(start_paused = true)] async fn test_drops_a_batch_after_the_retries() { let writer = FakeWriter { - fail: true, + failures: usize::MAX, + error: Some(|| clickhouse::error::Error::TimedOut), ..FakeWriter::default() }; let sink = BatchingSink::new(writer.clone(), config(Duration::from_secs(1), 1)); sink.put(records(&[1])); - tokio::time::sleep(Duration::from_secs(10)).await; + tokio::time::sleep(Duration::from_secs(60)).await; assert_eq!( writer.attempts.load(Ordering::SeqCst), - 1 + MAX_RETRIES as usize + 1 + RETRY_DELAYS.len() ); + assert_eq!(sink.fatal_error(), None); + } + + #[tokio::test(start_paused = true)] + async fn test_drops_bad_data_without_retries_and_stays_healthy() { + let writer = FakeWriter { + failures: 1, + error: Some(|| clickhouse::error::Error::Custom("bad row".to_string())), + ..FakeWriter::default() + }; + let sink = BatchingSink::new(writer.clone(), config(Duration::from_secs(3600), 1)); + + sink.put(records(&[1])); + tokio::time::sleep(Duration::from_secs(60)).await; + assert_eq!(writer.attempts.load(Ordering::SeqCst), 1); + assert!(writer.written.lock().unwrap().is_empty()); + assert_eq!(sink.fatal_error(), None); + } + + #[tokio::test(start_paused = true)] + async fn test_drops_a_batch_without_retries_on_a_config_error() { + let writer = FakeWriter { + failures: 1, + error: Some(|| clickhouse::error::Error::SchemaMismatch("no column".to_string())), + ..FakeWriter::default() + }; + let sink = BatchingSink::new(writer.clone(), config(Duration::from_secs(3600), 1)); + + sink.put(records(&[1])); + tokio::time::sleep(Duration::from_secs(60)).await; + assert_eq!(writer.attempts.load(Ordering::SeqCst), 1); + assert!(writer.written.lock().unwrap().is_empty()); + let fatal_error = sink.fatal_error(); + assert!(fatal_error.is_some()); + + // Stays set after a later write succeeds. sink.put(records(&[2])); + tokio::time::sleep(Duration::from_millis(10)).await; + assert_eq!(*writer.written.lock().unwrap(), vec![vec![2]]); + assert_eq!(sink.fatal_error(), fatal_error); } #[tokio::test(start_paused = true)] @@ -272,13 +332,4 @@ mod tests { assert_eq!(*writer.written.lock().unwrap(), vec![vec![1, 2, 3]]); } - - #[test] - fn test_backoff_doubles_within_the_jitter_band() { - for retry in 0..MAX_RETRIES { - let base = RETRY_BASE_DELAY * 2u32.pow(retry); - let delay = backoff(retry); - assert!(delay >= base / 2 && delay < base * 3 / 2, "{delay:?}"); - } - } } diff --git a/src/sink/clickhouse.rs b/src/sink/clickhouse.rs index 0443234..972d067 100644 --- a/src/sink/clickhouse.rs +++ b/src/sink/clickhouse.rs @@ -2,16 +2,38 @@ use std::borrow::Cow; use std::fmt; use std::time::Duration; +use clickhouse::error::Error; use clickhouse::{Client, Compression, Row}; use serde::Serialize; use super::batching::BatchWriter; -use super::SinkError; +use super::{ErrorKind, SinkError}; use crate::models::EventRecord; /// Bounds sending and committing one batch of up to `SINK_FLUSH_MAX_EVENTS` rows. const INSERT_TIMEOUT: Duration = Duration::from_secs(30); +/// Server error codes for rows the table rejects. +const BAD_DATA_ERROR_CODES: [u32; 5] = [ + 6, // CANNOT_PARSE_TEXT + 27, // CANNOT_PARSE_INPUT_ASSERTION_FAILED + 53, // TYPE_MISMATCH + 117, // INCORRECT_DATA + 469, // VIOLATED_CONSTRAINT +]; + +/// Server error codes for bad credentials or grants, or a missing database, +/// table or column. +const CONFIG_ERROR_CODES: [u32; 7] = [ + 16, // NO_SUCH_COLUMN_IN_TABLE + 60, // UNKNOWN_TABLE + 81, // UNKNOWN_DATABASE + 192, // UNKNOWN_USER + 194, // REQUIRED_PASSWORD + 497, // ACCESS_DENIED + 516, // AUTHENTICATION_FAILED +]; + #[derive(Clone)] pub struct ClickHouseConfig { pub url: String, @@ -85,6 +107,29 @@ fn value_to_text(value: &serde_json::Value) -> Cow<'_, str> { } } +/// Anything unrecognised, such as a network error, timeout or overloaded +/// server, counts as transient. +pub(super) fn error_kind(error: &Error) -> ErrorKind { + match error { + // A row that fails to serialize. + Error::Custom(_) => ErrorKind::BadData, + // `ClickHouseRow` doesn't match the table. + Error::SchemaMismatch(_) => ErrorKind::Config, + Error::BadResponse(reason) => match error_code(reason) { + Some(code) if BAD_DATA_ERROR_CODES.contains(&code) => ErrorKind::BadData, + Some(code) if CONFIG_ERROR_CODES.contains(&code) => ErrorKind::Config, + _ => ErrorKind::Transient, + }, + _ => ErrorKind::Transient, + } +} + +/// The code of a server exception, e.g. 60 in "Code: 60. DB::Exception: ...". +fn error_code(reason: &str) -> Option { + let rest = reason.strip_prefix("Code: ")?; + rest[..rest.find('.')?].parse().ok() +} + pub struct ClickHouseWriter { client: Client, table: String, @@ -104,6 +149,16 @@ impl ClickHouseWriter { table: config.table, } } + + /// Fails if ClickHouse is unreachable, rejects the credentials, or the + /// table's columns don't match `ClickHouseRow`. Column types are only + /// checked on the first write. + pub async fn check(&self) -> Result<(), SinkError> { + // Opening an insert fetches the table's schema and matches it against + // the row; dropping it unused sends nothing. + let _ = self.client.insert::>(&self.table).await?; + Ok(()) + } } impl BatchWriter for ClickHouseWriter { @@ -196,6 +251,110 @@ mod tests { assert!(!shown.contains("hunter2")); } + #[test] + fn test_error_code_reads_the_server_exception_code() { + assert_eq!( + error_code("Code: 60. DB::Exception: Table x.events does not exist."), + Some(60) + ); + assert_eq!(error_code("Code: abc. DB::Exception"), None); + assert_eq!(error_code("502 Bad Gateway"), None); + } + + #[test] + fn test_error_kind_separates_bad_data_and_config_from_outages() { + let bad_response = |reason: &str| Error::BadResponse(reason.to_string()); + let cases = [ + ( + bad_response("Code: 469. DB::Exception: Constraint violated"), + ErrorKind::BadData, + ), + (Error::Custom("bad row".to_string()), ErrorKind::BadData), + ( + bad_response("Code: 516. DB::Exception: Authentication failed"), + ErrorKind::Config, + ), + ( + bad_response("Code: 81. DB::Exception: Database x does not exist"), + ErrorKind::Config, + ), + (Error::SchemaMismatch("x".to_string()), ErrorKind::Config), + ( + bad_response("Code: 202. DB::Exception: Too many simultaneous queries"), + ErrorKind::Transient, + ), + ( + bad_response("503 Service Unavailable"), + ErrorKind::Transient, + ), + (Error::TimedOut, ErrorKind::Transient), + (Error::Network("refused".into()), ErrorKind::Transient), + ]; + for (error, kind) in cases { + assert_eq!(error_kind(&error), kind, "{error}"); + } + } + + #[tokio::test] + async fn test_check_passes_for_the_events_table() { + let db = ClickHouseTestDatabase::create().await; + ClickHouseWriter::new(db.config.clone()) + .check() + .await + .unwrap(); + db.drop_db().await; + } + + #[tokio::test] + async fn test_check_fails_with_a_config_error_on_bad_config() { + let db = ClickHouseTestDatabase::create().await; + db.client + .query("ALTER TABLE events RENAME COLUMN sdk_version TO version") + .execute() + .await + .unwrap(); + let configs = [ + ClickHouseConfig { + password: "wrong".to_string(), + ..db.config.clone() + }, + ClickHouseConfig { + database: "missing".to_string(), + ..db.config.clone() + }, + ClickHouseConfig { + table: "missing".to_string(), + ..db.config.clone() + }, + // The renamed column. + db.config.clone(), + ]; + + for config in configs { + let error = ClickHouseWriter::new(config.clone()) + .check() + .await + .unwrap_err(); + assert_eq!(error.kind(), ErrorKind::Config, "{config:?}: {error}"); + } + db.drop_db().await; + } + + #[tokio::test] + async fn test_check_fails_with_a_transient_error_when_unreachable() { + let writer = ClickHouseWriter::new(ClickHouseConfig { + url: "http://127.0.0.1:1".to_string(), + user: "default".to_string(), + password: String::new(), + database: "default".to_string(), + table: "events".to_string(), + }); + assert_eq!( + writer.check().await.unwrap_err().kind(), + ErrorKind::Transient + ); + } + fn batch_config(flush_interval: Duration) -> BatchConfig { BatchConfig { flush_interval, diff --git a/src/sink/mod.rs b/src/sink/mod.rs index f22d2e2..89d4363 100644 --- a/src/sink/mod.rs +++ b/src/sink/mod.rs @@ -1,12 +1,14 @@ mod batching; mod clickhouse; mod kafka; +mod postgres; use std::time::Duration; pub use batching::{BatchConfig, BatchingSink}; pub use clickhouse::{ClickHouseConfig, ClickHouseWriter}; pub use kafka::{KafkaAuth, KafkaSink, KafkaSinkConfig}; +pub use postgres::PostgresWriter; use crate::models::EventRecord; @@ -15,6 +17,7 @@ use crate::models::EventRecord; pub enum Sink { Kafka(KafkaSink), ClickHouse(BatchingSink), + Postgres(BatchingSink), } impl Sink { @@ -32,6 +35,10 @@ impl Sink { sink.put(records); Ok(()) } + Self::Postgres(sink) => { + sink.put(records); + Ok(()) + } } } @@ -39,7 +46,8 @@ impl Sink { pub fn fatal_error(&self) -> Option { match self { Self::Kafka(sink) => sink.fatal_error(), - Self::ClickHouse(_) => None, + Self::ClickHouse(sink) => sink.fatal_error(), + Self::Postgres(sink) => sink.fatal_error(), } } @@ -48,6 +56,7 @@ impl Sink { match self { Self::Kafka(sink) => sink.shutdown(timeout).await, Self::ClickHouse(sink) => sink.shutdown(timeout).await, + Self::Postgres(sink) => sink.shutdown(timeout).await, } } } @@ -58,4 +67,32 @@ pub enum SinkError { Kafka(#[from] kafka::DeliveryError), #[error("ClickHouse insert failed: {0}")] ClickHouse(#[from] ::clickhouse::error::Error), + #[error("PostgreSQL insert failed: {0}")] + Postgres(#[from] sqlx::Error), + #[error("PostgreSQL insert timed out")] + PostgresTimeout, +} + +/// What a failed write means for the batch and the server. +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum ErrorKind { + /// May pass on retry: network, timeout, overload, or unrecognised. + Transient, + /// The store rejects the rows; retrying cannot help. + BadData, + /// The server is misconfigured, e.g. bad credentials or a missing + /// table; nothing will be written until it is fixed. + Config, +} + +impl SinkError { + /// Only batching sinks retry; the Kafka producer retries internally. + pub fn kind(&self) -> ErrorKind { + match self { + Self::Kafka(_) => ErrorKind::Transient, + Self::ClickHouse(e) => clickhouse::error_kind(e), + Self::Postgres(e) => postgres::error_kind(e), + Self::PostgresTimeout => ErrorKind::Transient, + } + } } diff --git a/src/sink/postgres.rs b/src/sink/postgres.rs new file mode 100644 index 0000000..c76ce19 --- /dev/null +++ b/src/sink/postgres.rs @@ -0,0 +1,291 @@ +use std::borrow::Cow; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use sqlx::PgPool; + +use super::batching::BatchWriter; +use super::{ErrorKind, SinkError}; +use crate::models::EventRecord; + +/// SQLSTATE classes for rows the table rejects: data exceptions and integrity +/// constraint violations. +const BAD_DATA_SQLSTATE_CLASSES: [&str; 2] = ["22", "23"]; + +/// SQLSTATEs for bad credentials or grants, or a missing database, schema, +/// table or column. +const CONFIG_SQLSTATES: [&str; 7] = [ + "28000", // invalid_authorization_specification + "28P01", // invalid_password + "3D000", // invalid_catalog_name + "3F000", // invalid_schema_name + "42501", // insufficient_privilege + "42P01", // undefined_table + "42703", // undefined_column +]; + +/// `statement_timeout` is enforced by the server, so a stalled connection is +/// caught by this instead. +const INSERT_TIMEOUT: Duration = Duration::from_secs(30); +const SET_STATEMENT_TIMEOUT: &str = "SET LOCAL statement_timeout = 30000"; + +const INSERT_EVENTS: &str = " + INSERT INTO flagsmith_exp.events ( + environment_key, event, feature_name, timestamp, collected_at, identifier, + value, traits, metadata, sdk_language, sdk_version + ) + SELECT * FROM UNNEST( + $1::text[], $2::text[], $3::text[], $4::timestamptz[], $5::timestamptz[], $6::text[], + $7::text[], $8::text[], $9::text[], $10::text[], $11::text[] + )"; + +#[derive(Debug, PartialEq)] +struct PostgresRow<'a> { + environment_key: Cow<'a, str>, + event: Cow<'a, str>, + feature_name: Option>, + timestamp: DateTime, + collected_at: DateTime, + identifier: Cow<'a, str>, + value: Option>, + traits: Option, + metadata: Option, + sdk_language: Cow<'a, str>, + sdk_version: Cow<'a, str>, +} + +impl<'a> From<&'a EventRecord> for PostgresRow<'a> { + fn from(record: &'a EventRecord) -> Self { + PostgresRow { + environment_key: storable_text(&record.environment_key), + event: storable_text(&record.event), + feature_name: record.feature_name.as_deref().map(storable_text), + timestamp: epoch_ms_to_datetime(record.timestamp), + collected_at: epoch_ms_to_datetime(record.collected_at), + identifier: storable_text(record.identifier.as_deref().unwrap_or_default()), + value: value_to_text(&record.value), + traits: record.traits.as_ref().map(ToString::to_string), + metadata: record.metadata.as_ref().map(ToString::to_string), + sdk_language: storable_text(&record.sdk_language), + sdk_version: storable_text(&record.sdk_version), + } + } +} + +fn value_to_text(value: &serde_json::Value) -> Option> { + match value { + serde_json::Value::Null => None, + serde_json::Value::String(s) => Some(storable_text(s)), + other => Some(Cow::Owned(other.to_string())), + } +} + +/// Postgres text cannot hold NUL, which JSON strings allow; one would fail the +/// whole batch. +fn storable_text(text: &str) -> Cow<'_, str> { + if text.contains('\0') { + Cow::Owned(text.replace('\0', "")) + } else { + Cow::Borrowed(text) + } +} + +/// Ingest validation keeps timestamps in range, so the epoch fallback is +/// never used; it avoids a panic, which aborts the process. +fn epoch_ms_to_datetime(epoch_ms: i64) -> DateTime { + DateTime::from_timestamp_millis(epoch_ms).unwrap_or_default() +} + +/// Anything unrecognised, such as a network error, pool timeout or +/// cancelled statement, counts as transient. +pub(super) fn error_kind(error: &sqlx::Error) -> ErrorKind { + let Some(sqlstate) = error.as_database_error().and_then(|e| e.code()) else { + return ErrorKind::Transient; + }; + if CONFIG_SQLSTATES.contains(&sqlstate.as_ref()) { + ErrorKind::Config + } else if sqlstate + .get(..2) + .is_some_and(|class| BAD_DATA_SQLSTATE_CLASSES.contains(&class)) + { + ErrorKind::BadData + } else { + ErrorKind::Transient + } +} + +fn column_array<'r, 'a, T>( + rows: &'r [PostgresRow<'a>], + field: impl Fn(&'r PostgresRow<'a>) -> T, +) -> Vec { + rows.iter().map(field).collect() +} + +pub struct PostgresWriter { + pool: PgPool, +} + +impl BatchWriter for PostgresWriter { + async fn write(&self, records: &[EventRecord]) -> Result<(), SinkError> { + tokio::time::timeout(INSERT_TIMEOUT, self.insert_in_transaction(records)) + .await + .map_err(|_| SinkError::PostgresTimeout)? + } +} + +impl PostgresWriter { + pub fn new(pool: PgPool) -> Self { + Self { pool } + } + + async fn insert_in_transaction(&self, records: &[EventRecord]) -> Result<(), SinkError> { + let rows: Vec> = records.iter().map(PostgresRow::from).collect(); + let mut transaction = self.pool.begin().await?; + sqlx::query(SET_STATEMENT_TIMEOUT) + .execute(&mut *transaction) + .await?; + sqlx::query(INSERT_EVENTS) + .bind(column_array(&rows, |row| row.environment_key.as_ref())) + .bind(column_array(&rows, |row| row.event.as_ref())) + .bind(column_array(&rows, |row| row.feature_name.as_deref())) + .bind(column_array(&rows, |row| row.timestamp)) + .bind(column_array(&rows, |row| row.collected_at)) + .bind(column_array(&rows, |row| row.identifier.as_ref())) + .bind(column_array(&rows, |row| row.value.as_deref())) + .bind(column_array(&rows, |row| row.traits.as_deref())) + .bind(column_array(&rows, |row| row.metadata.as_deref())) + .bind(column_array(&rows, |row| row.sdk_language.as_ref())) + .bind(column_array(&rows, |row| row.sdk_version.as_ref())) + .execute(&mut *transaction) + .await?; + transaction.commit().await?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::sink::{BatchConfig, BatchingSink}; + use crate::test_helpers::{now_ms, POSTGRES_EVENTS_TABLE}; + + fn record(event: &str, timestamp: i64) -> EventRecord { + EventRecord { + environment_key: "env".to_string(), + collected_at: timestamp + 1, + sdk_language: "python".to_string(), + sdk_version: "1.0.0".to_string(), + event: event.to_string(), + feature_name: None, + identifier: None, + value: serde_json::Value::Null, + traits: None, + metadata: None, + timestamp, + } + } + + fn utc_datetime(epoch_ms: i64) -> DateTime { + DateTime::from_timestamp_millis(epoch_ms).unwrap() + } + + #[derive(Debug, PartialEq, sqlx::FromRow)] + struct StoredRow { + environment_key: String, + event: String, + feature_name: Option, + timestamp: DateTime, + collected_at: Option>, + identifier: String, + value: Option, + traits: Option, + metadata: Option, + sdk_language: Option, + sdk_version: Option, + } + + #[sqlx::test] + async fn test_shutdown_writes_buffered_rows_to_the_events_table(pool: PgPool) { + sqlx::raw_sql(POSTGRES_EVENTS_TABLE) + .execute(&pool) + .await + .unwrap(); + let sink = BatchingSink::new( + PostgresWriter::new(pool.clone()), + BatchConfig { + flush_interval: Duration::from_secs(3600), + flush_max_events: 1000, + }, + ); + let timestamp = now_ms(); + + sink.put(vec![ + EventRecord { + feature_name: Some("dark_mode".to_string()), + identifier: Some("user\0_1".to_string()), + value: json!("control"), + traits: Some(json!({"plan": "pro"})), + metadata: Some(json!({"source": "csv"})), + ..record("$flag_exposure", timestamp) + }, + record("purchase", timestamp + 2), + EventRecord { + value: json!(42), + ..record("purchase", timestamp + 4) + }, + ]); + sink.shutdown(Duration::from_secs(5)).await; + + let stored_rows: Vec = + sqlx::query_as("SELECT * FROM flagsmith_exp.events ORDER BY timestamp") + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!( + stored_rows, + vec![ + StoredRow { + environment_key: "env".to_string(), + event: "$flag_exposure".to_string(), + feature_name: Some("dark_mode".to_string()), + timestamp: utc_datetime(timestamp), + collected_at: Some(utc_datetime(timestamp + 1)), + identifier: "user_1".to_string(), + value: Some("control".to_string()), + traits: Some(r#"{"plan":"pro"}"#.to_string()), + metadata: Some(r#"{"source":"csv"}"#.to_string()), + sdk_language: Some("python".to_string()), + sdk_version: Some("1.0.0".to_string()), + }, + StoredRow { + environment_key: "env".to_string(), + event: "purchase".to_string(), + feature_name: None, + timestamp: utc_datetime(timestamp + 2), + collected_at: Some(utc_datetime(timestamp + 3)), + identifier: String::new(), + value: None, + traits: None, + metadata: None, + sdk_language: Some("python".to_string()), + sdk_version: Some("1.0.0".to_string()), + }, + StoredRow { + environment_key: "env".to_string(), + event: "purchase".to_string(), + feature_name: None, + timestamp: utc_datetime(timestamp + 4), + collected_at: Some(utc_datetime(timestamp + 5)), + identifier: String::new(), + value: Some("42".to_string()), + traits: None, + metadata: None, + sdk_language: Some("python".to_string()), + sdk_version: Some("1.0.0".to_string()), + }, + ] + ); + } +} diff --git a/src/test_helpers.rs b/src/test_helpers.rs index ed98364..ab00fa2 100644 --- a/src/test_helpers.rs +++ b/src/test_helpers.rs @@ -116,6 +116,27 @@ pub(crate) async fn insert_environment_key( .unwrap(); } +pub(crate) const POSTGRES_EVENTS_TABLE: &str = " + CREATE SCHEMA IF NOT EXISTS flagsmith_exp; + + CREATE TABLE IF NOT EXISTS flagsmith_exp.events ( + environment_key TEXT NOT NULL, + event TEXT NOT NULL, + feature_name TEXT, + timestamp TIMESTAMPTZ NOT NULL, + collected_at TIMESTAMPTZ, + identifier TEXT NOT NULL, + value TEXT, + traits TEXT, + metadata TEXT, + sdk_language TEXT, + sdk_version TEXT + ); + + CREATE INDEX IF NOT EXISTS events_env_event_feature_ts ON flagsmith_exp.events (environment_key, event, feature_name, timestamp); + CREATE INDEX IF NOT EXISTS events_env_event_identifier_ts ON flagsmith_exp.events (environment_key, event, identifier, timestamp); +"; + /// The production `events` table, from Flagsmith's /// experimentation/clickhouse/events.sql. const EVENTS_TABLE: &str = "