diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index eab9134..8d21937 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -31,8 +31,21 @@ jobs: --health-interval 5s --health-timeout 5s --health-retries 10 + clickhouse: + image: clickhouse/clickhouse-server:25.8 + env: + # Lets the default user connect without a password. + CLICKHOUSE_SKIP_USER_SETUP: 1 + ports: + - 8123:8123 + options: >- + --health-cmd "clickhouse-client --query 'SELECT 1'" + --health-interval 5s + --health-timeout 5s + --health-retries 10 env: DATABASE_URL: postgres://postgres:password@localhost:5432/postgres + CLICKHOUSE_URL: http://localhost:8123 steps: - uses: actions/checkout@v6 diff --git a/Cargo.lock b/Cargo.lock index a8ddebf..1d08c2d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -100,6 +100,29 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +[[package]] +name = "aws-lc-rs" +version = "1.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "faac5829c2b74c28f830747e7818ccfb684261b5f48a1118b1e2a13d36dfab13" +dependencies = [ + "aws-lc-sys", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1622d8446a2d4b2ce0c7eefc73dd43a99779028d5ee5c2dd8073a658ba8a2bc" +dependencies = [ + "cc", + "cmake", + "dunce", + "fs_extra", + "pkg-config", +] + [[package]] name = "axum" version = "0.7.9" @@ -185,6 +208,21 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bnum" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "119771309b95163ec7aaf79810da82f7cd0599c19722d48b9c03894dca833966" + +[[package]] +name = "bstr" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bb31b46c14244e20ee9984b11bf5c992b91fb6939fea616e3512c8baecdbe5f" +dependencies = [ + "memchr", +] + [[package]] name = "bumpalo" version = "3.19.0" @@ -202,6 +240,9 @@ name = "bytes" version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" +dependencies = [ + "serde", +] [[package]] name = "cc" @@ -221,6 +262,17 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "chacha20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.1", + "rand_core 0.10.1", +] + [[package]] name = "chrono" version = "0.4.39" @@ -236,6 +288,66 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "cityhash-rs" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93a719913643003b84bd13022b4b7e703c09342cd03b679c4641c7d2e50dc34d" + +[[package]] +name = "clickhouse" +version = "0.15.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6ad2112cf4ad543ceed12e48b779e82e8f3a8172374cac98d59e59feda35c69" +dependencies = [ + "bnum", + "bstr", + "bytes", + "cityhash-rs", + "clickhouse-macros", + "clickhouse-types", + "futures-channel", + "futures-util", + "half", + "hashbrown 0.17.1", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "lz4_flex", + "rand 0.9.5", + "rustls", + "serde", + "serde_json", + "thiserror 2.0.21", + "tokio", + "tracing", + "twox-hash", + "url", +] + +[[package]] +name = "clickhouse-macros" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff6669899e23cb87b43daf7996f0ea3b9c07d0fb933d745bb7b815b052515ae3" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn 2.0.119", +] + +[[package]] +name = "clickhouse-types" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7ddab58fabf3bdde7b6fa401f72b9abc7566ff015264c89aaaa71cf1d173d06" +dependencies = [ + "bytes", + "thiserror 2.0.21", +] + [[package]] name = "cmake" version = "0.1.54" @@ -292,6 +404,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ca28b0ae3115b884660db4118d803791fd6756b6e88f39c0f3f7859060d7566" +dependencies = [ + "libc", +] + [[package]] name = "crc" version = "3.4.0" @@ -349,6 +470,12 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "crypto-common" version = "0.1.7" @@ -399,6 +526,12 @@ version = "0.15.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + [[package]] name = "either" version = "1.15.0" @@ -460,6 +593,7 @@ dependencies = [ "axum", "bytes", "chrono", + "clickhouse", "futures", "moka", "rdkafka", @@ -501,6 +635,12 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "form_urlencoded" version = "1.2.2" @@ -510,6 +650,12 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + [[package]] name = "futures" version = "0.3.31" @@ -639,10 +785,33 @@ checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" dependencies = [ "cfg-if", "libc", - "r-efi", + "r-efi 5.3.0", "wasip2", ] +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", + "rand_core 0.10.1", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.15.5" @@ -651,7 +820,7 @@ checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ "allocator-api2", "equivalent", - "foldhash", + "foldhash 0.1.5", ] [[package]] @@ -659,6 +828,11 @@ name = "hashbrown" version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", +] [[package]] name = "hashlink" @@ -772,6 +946,23 @@ dependencies = [ "pin-utils", "smallvec", "tokio", + "want", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfa8e654703247911e29c23fbeaa261834bd9bb74efba2f9acddc37bfb127f53" +dependencies = [ + "http", + "hyper", + "hyper-util", + "rustls", + "tokio", + "tokio-rustls", + "tower-service", + "webpki-roots 1.0.9", ] [[package]] @@ -781,13 +972,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "727805d60e7938b76b826a6ef209eb70eaa1812794f9424d4a4e2d740662df5f" dependencies = [ "bytes", + "futures-channel", "futures-core", + "futures-util", "http", "http-body", "hyper", + "libc", "pin-project-lite", + "socket2", "tokio", "tower-service", + "tracing", ] [[package]] @@ -1028,6 +1224,12 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "lz4_flex" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226" + [[package]] name = "matchers" version = "0.2.0" @@ -1126,7 +1328,7 @@ dependencies = [ "num-integer", "num-iter", "num-traits", - "rand", + "rand 0.8.6", "smallvec", "zeroize", ] @@ -1356,6 +1558,12 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + [[package]] name = "rand" version = "0.8.6" @@ -1363,8 +1571,29 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5ca0ecfa931c29007047d1bc58e623ab12e5590e8c7cc53200d5202b69266d8a" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65c9fb96cbc91e3478eaae79a69fcd3f1ae4ad052e471fe6732fff548984b4af" +dependencies = [ + "chacha20", + "getrandom 0.4.3", + "rand_core 0.10.1", ] [[package]] @@ -1374,7 +1603,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", ] [[package]] @@ -1386,6 +1625,21 @@ dependencies = [ "getrandom 0.2.16", ] +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + [[package]] name = "rdkafka" version = "0.39.0" @@ -1481,7 +1735,7 @@ dependencies = [ "num-traits", "pkcs1", "pkcs8", - "rand_core", + "rand_core 0.6.4", "signature", "spki", "subtle", @@ -1494,6 +1748,7 @@ version = "0.23.35" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "533f54bc6a7d4f647e46ad909549eda97bf5afc1585190ef692b4286b198bd8f" dependencies = [ + "aws-lc-rs", "once_cell", "ring", "rustls-pki-types", @@ -1517,6 +1772,7 @@ version = "0.103.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2ffdfa2f5286e2247234e03f680868ac2815974dc39e00ea15adc445d0aafe52" dependencies = [ + "aws-lc-rs", "ring", "rustls-pki-types", "untrusted", @@ -1570,6 +1826,17 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "serde_derive_internals" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "serde_json" version = "1.0.151" @@ -1613,7 +1880,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a978451301f4db1d02937a4ab3ccce137717b81826e79b7d49ffe3244a13c3b8" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] @@ -1624,7 +1891,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest", ] @@ -1659,7 +1926,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" dependencies = [ "digest", - "rand_core", + "rand_core 0.6.4", ] [[package]] @@ -1830,7 +2097,7 @@ dependencies = [ "memchr", "once_cell", "percent-encoding", - "rand", + "rand 0.8.6", "rsa", "serde", "sha1", @@ -1869,7 +2136,7 @@ dependencies = [ "md-5", "memchr", "once_cell", - "rand", + "rand 0.8.6", "serde", "serde_json", "sha2", @@ -2067,6 +2334,16 @@ dependencies = [ "syn 3.0.3", ] +[[package]] +name = "tokio-rustls" +version = "0.26.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9cc2678c2cdd569ef8215e2afd7954ada2ae20b4fdd2c5fe6139a3b02d105db" +dependencies = [ + "rustls", + "tokio", +] + [[package]] name = "tokio-stream" version = "0.1.18" @@ -2244,6 +2521,21 @@ dependencies = [ "tracing-serde", ] +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "twox-hash" +version = "2.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "86a801b3cea342a06d468c8710662aa29e5e05e4f5c0d62f00bbb7f2ad7941c2" +dependencies = [ + "rand 0.10.3", +] + [[package]] name = "typenum" version = "1.20.1" @@ -2330,6 +2622,15 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "want" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec4cdd0dd910afe868b7ef477227d8d538b46b3075031afee8a9f2acb0a2ed0b" +dependencies = [ + "try-lock", +] + [[package]] name = "wasi" version = "0.11.1+wasi-snapshot-preview1" diff --git a/Cargo.toml b/Cargo.toml index 20117d8..e54e7bd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,6 +15,7 @@ sqlx = { version = "0.8", default-features = false, features = ["postgres", "run rdkafka = { version = "0.39", features = ["cmake-build", "ssl-vendored", "tokio", "zstd"] } futures = "0.3" +clickhouse = { version = "0.15", features = ["lz4", "rustls-tls"] } serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/README.md b/README.md index b9a9f54..3f225e4 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. @@ -42,7 +42,7 @@ An empty `events` array adds a top-level `error` field: `{ "accepted": 0, "rejec } ``` -`event` and `timestamp` (epoch millis) are required. `$flag_exposure` is a reserved system event and requires `feature_name` (omitting it rejects that event); `identifier`, `value`, `traits`, and `metadata` are optional. +`event` and `timestamp` (epoch millis) are required; `timestamp` must be after 2020-01-01, no more than 24h in the future and no more than 10 days in the past (otherwise that event is rejected). `$flag_exposure` is a reserved system event and requires `feature_name` (omitting it rejects that event); `identifier`, `value`, `traits`, and `metadata` are optional. Field limits — an over-limit field rejects that event: `event` ≤ 256 chars; `feature_name` and `identifier` ≤ 2000 chars, trait keys ≤ 200 chars, trait string values ≤ 2000 chars (all matching Flagsmith's model limits); serialized `value` ≤ 128 KiB (fits Flagsmith's default 20,000-char flag value limit), `traits` ≤ 32 KiB, `metadata` ≤ 8 KiB. `traits` must be a JSON object. @@ -52,17 +52,24 @@ 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 | — | +| `SINK` | no | `kafka` (or `clickhouse`) | +| `KAFKA_BOOTSTRAP_SERVERS` | with `kafka` | — | +| `KAFKA_TOPIC` | with `kafka` | — | +| `KAFKA_EXTERNAL_TOPIC` | with `kafka` | — | | `KAFKA_AUTH` | no | `scram` (`none` for plaintext local brokers) | | `KAFKA_USERNAME` | with `scram` | — | | `KAFKA_PASSWORD` | with `scram` | — | +| `CLICKHOUSE_URL` | with `clickhouse` | — | +| `CLICKHOUSE_USER` | no | `default` | +| `CLICKHOUSE_PASSWORD` | no | empty | +| `CLICKHOUSE_DATABASE` | no | `flagsmith_exp` | +| `CLICKHOUSE_TABLE` | no | `events` | +| `SINK_FLUSH_INTERVAL_MS` | no | `1000` | +| `SINK_FLUSH_MAX_EVENTS` | no | `5000` | | `DATABASE_URL` | yes | — | | `LISTEN_ADDR` | no | `0.0.0.0:8080` | -Events from environments that use an external warehouse go to `KAFKA_EXTERNAL_TOPIC`; all others go to `KAFKA_TOPIC`. 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. The server looks presented keys up in the `experimentation_environment_keys` view in Flagsmith's Postgres, at `DATABASE_URL`. ### Database role @@ -79,6 +86,10 @@ The view reads its tables with its owner's privileges, so the role needs no acce With `KAFKA_AUTH=scram`, the producer authenticates with SASL/SCRAM-SHA-512 over TLS; MSK serves it on port 9096. The task definition injects both credentials from the `AmazonMSK_*` secret in Secrets Manager. +### 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 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. + ## Run locally ```bash @@ -90,7 +101,7 @@ curl localhost:8080/health The Kafka producer and Postgres pool connect lazily, so `/health` responds even without either available. -`cargo test` needs a Postgres it can create databases in, at `DATABASE_URL`: each database test gets a fresh one. +`cargo test` needs a Postgres it can create databases in, at `DATABASE_URL`, and a ClickHouse it can create databases in, at `CLICKHOUSE_URL` (with `CLICKHOUSE_USER` and `CLICKHOUSE_PASSWORD` if needed): each database test gets a fresh one. ## CI / CD diff --git a/src/auth/middleware.rs b/src/auth/middleware.rs index 507796a..1b4128e 100644 --- a/src/auth/middleware.rs +++ b/src/auth/middleware.rs @@ -22,7 +22,8 @@ pub struct EnvironmentContext { #[derive(Clone)] pub struct ContextState { pub environment_keys: EnvironmentKeys, - pub external_topic: Arc, + /// Kafka only; other sinks write every environment to one place. + pub external_topic: Option>, } pub async fn extract_environment_context( @@ -60,7 +61,8 @@ pub async fn extract_environment_context( let destination = environment_key .uses_external_warehouse - .then(|| state.external_topic.clone()); + .then(|| state.external_topic.clone()) + .flatten(); request.extensions_mut().insert(EnvironmentContext { client_api_key: environment_key.client_api_key, @@ -87,7 +89,7 @@ mod tests { fn test_state(pool: PgPool) -> ContextState { ContextState { environment_keys: EnvironmentKeys::new(pool, Duration::from_secs(300)), - external_topic: Arc::from("external-topic"), + external_topic: Some(Arc::from("external-topic")), } } diff --git a/src/config.rs b/src/config.rs index 4e0e38f..c3c3f10 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,6 +1,8 @@ use std::env; +use std::str::FromStr; +use std::time::Duration; -use crate::sink::KafkaAuth; +use crate::sink::{BatchConfig, ClickHouseConfig, KafkaAuth}; #[derive(Debug, Clone)] pub struct Config { @@ -12,6 +14,10 @@ pub struct Config { #[derive(Debug, Clone)] pub enum SinkConfig { Kafka(KafkaConfig), + ClickHouse { + clickhouse: ClickHouseConfig, + batch: BatchConfig, + }, } #[derive(Debug, Clone)] @@ -26,6 +32,10 @@ impl Config { pub fn from_env() -> Result { let sink = match parse_sink(env::var("SINK").ok().as_deref())? { SinkKind::Kafka => SinkConfig::Kafka(KafkaConfig::from_env()?), + SinkKind::ClickHouse => SinkConfig::ClickHouse { + clickhouse: clickhouse_config_from_env()?, + batch: batch_config_from_env()?, + }, }; Ok(Config { @@ -60,16 +70,49 @@ fn required(name: &'static str) -> Result { #[derive(Debug, PartialEq)] enum SinkKind { Kafka, + ClickHouse, } /// Which sink `SINK` selects; unset means Kafka. fn parse_sink(value: Option<&str>) -> Result { match value.unwrap_or("kafka") { "kafka" => Ok(SinkKind::Kafka), + "clickhouse" => Ok(SinkKind::ClickHouse), other => Err(ConfigError::InvalidSink(other.to_string())), } } +fn clickhouse_config_from_env() -> Result { + Ok(ClickHouseConfig { + url: required("CLICKHOUSE_URL")?, + user: env::var("CLICKHOUSE_USER").unwrap_or_else(|_| "default".to_string()), + password: env::var("CLICKHOUSE_PASSWORD").unwrap_or_default(), + database: env::var("CLICKHOUSE_DATABASE").unwrap_or_else(|_| "flagsmith_exp".to_string()), + table: env::var("CLICKHOUSE_TABLE").unwrap_or_else(|_| "events".to_string()), + }) +} + +/// Shared by every batching sink. +fn batch_config_from_env() -> Result { + Ok(BatchConfig { + flush_interval: Duration::from_millis(positive_env_var("SINK_FLUSH_INTERVAL_MS", 1000)?), + flush_max_events: positive_env_var("SINK_FLUSH_MAX_EVENTS", 5000)?, + }) +} + +fn positive_env_var( + name: &'static str, + default: T, +) -> Result { + match env::var(name).ok() { + None => Ok(default), + Some(v) => match v.parse() { + Ok(n) if n > T::default() => Ok(n), + _ => Err(ConfigError::InvalidNumber(name, v)), + }, + } +} + fn parse_auth( mode: &str, username: Option, @@ -91,8 +134,10 @@ 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)")] + #[error("Invalid value for SINK: {0:?} (expected kafka or clickhouse)")] InvalidSink(String), + #[error("Invalid value for {0}: {1:?} (expected a positive integer)")] + InvalidNumber(&'static str, String), } #[cfg(test)] @@ -140,13 +185,49 @@ mod tests { fn parse_sink_defaults_to_kafka() { assert_eq!(parse_sink(None).unwrap(), SinkKind::Kafka); assert_eq!(parse_sink(Some("kafka")).unwrap(), SinkKind::Kafka); + assert_eq!( + parse_sink(Some("clickhouse")).unwrap(), + SinkKind::ClickHouse + ); } #[test] fn parse_sink_rejects_unknown_value() { assert!(matches!( - parse_sink(Some("clickhouse")), - Err(ConfigError::InvalidSink(v)) if v == "clickhouse" + parse_sink(Some("postgres")), + Err(ConfigError::InvalidSink(v)) if v == "postgres" )); } + + // Each test uses its own variable, since the environment is shared + // across test threads. + + #[test] + fn test_positive_env_var_defaults_when_unset() { + assert_eq!( + positive_env_var("TEST_POSITIVE_ENV_VAR_UNSET", 1000u64).unwrap(), + 1000 + ); + } + + #[test] + fn test_positive_env_var_reads_a_positive_integer() { + env::set_var("TEST_POSITIVE_ENV_VAR_SET", "50"); + assert_eq!( + positive_env_var("TEST_POSITIVE_ENV_VAR_SET", 1000usize).unwrap(), + 50 + ); + } + + #[test] + fn test_positive_env_var_rejects_zero_and_non_numbers() { + const NAME: &str = "TEST_POSITIVE_ENV_VAR_INVALID"; + for value in ["0", "-1", "ten", ""] { + env::set_var(NAME, value); + assert!(matches!( + positive_env_var(NAME, 1000u64), + Err(ConfigError::InvalidNumber(NAME, v)) if v == value + )); + } + } } diff --git a/src/main.rs b/src/main.rs index d5185e6..bea09fc 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::{KafkaSink, KafkaSinkConfig, Sink}; +use sink::{BatchingSink, ClickHouseWriter, KafkaSink, KafkaSinkConfig, Sink}; const ENVIRONMENT_KEY_REFRESH_AFTER: Duration = Duration::from_secs(300); @@ -47,7 +48,24 @@ async fn main() -> anyhow::Result<()> { kafka.topic.clone(), kafka.auth.clone(), ))?; - (Sink::Kafka(sink), kafka.external_topic.as_str()) + (Sink::Kafka(sink), Some(kafka.external_topic.as_str())) + } + SinkConfig::ClickHouse { clickhouse, batch } => { + info!( + listen_addr = %config.listen_addr, + clickhouse = ?clickhouse, + batch = ?batch, + "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, + ) } }; @@ -57,7 +75,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(external_topic), + external_topic: external_topic.map(Arc::from), }; let app = routes::app(sink.clone(), context_state); diff --git a/src/models/event.rs b/src/models/event.rs index 2f44b77..194a2db 100644 --- a/src/models/event.rs +++ b/src/models/event.rs @@ -60,35 +60,35 @@ pub struct IngestResponse { /// only maps flat scalar fields — a nested object field is left out of the /// pipe's field mapping and the column silently defaults to ''. #[derive(Debug, Clone, Serialize)] -pub struct EventRecord<'a> { - pub environment_key: &'a str, +pub struct EventRecord { + pub environment_key: String, pub collected_at: i64, - pub sdk_language: &'a str, - pub sdk_version: &'a str, + pub sdk_language: String, + pub sdk_version: String, - pub event: &'a str, + pub event: String, #[serde(skip_serializing_if = "Option::is_none")] - pub feature_name: Option<&'a str>, + pub feature_name: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub identifier: Option<&'a str>, + pub identifier: Option, #[serde(skip_serializing_if = "serde_json::Value::is_null")] - pub value: &'a serde_json::Value, + pub value: serde_json::Value, #[serde( skip_serializing_if = "Option::is_none", serialize_with = "json_string" )] - pub traits: Option<&'a serde_json::Value>, + pub traits: Option, #[serde( skip_serializing_if = "Option::is_none", serialize_with = "json_string" )] - pub metadata: Option<&'a serde_json::Value>, + pub metadata: Option, pub timestamp: i64, } /// Serializes a JSON value as its encoded text, in a string. fn json_string( - value: &Option<&serde_json::Value>, + value: &Option, serializer: S, ) -> Result { match value { 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/routes/ingest.rs b/src/routes/ingest.rs index 4a039e0..beb9079 100644 --- a/src/routes/ingest.rs +++ b/src/routes/ingest.rs @@ -39,8 +39,13 @@ fn parse_sdk_user_agent(header: &str) -> (String, String) { /// 2020-01-01T00:00:00Z in epoch-millis; mirrors ClickHouse's `timestamp_sane` constraint. const MIN_TIMESTAMP_MS: i64 = 1_577_836_800_000; -/// Future clock-skew allowance. -const MAX_FUTURE_SKEW_MS: i64 = 24 * 60 * 60 * 1000; +const HOUR_MS: i64 = 60 * 60 * 1000; +const DAY_MS: i64 = 24 * HOUR_MS; +/// Future clock-skew allowance, in hours after now. +const MAX_FUTURE_SKEW_HOURS: i64 = 24; +/// Oldest accepted event, in days before now. Bounds how many daily +/// ClickHouse partitions one insert can touch. +const MAX_EVENT_AGE_DAYS: i64 = 10; /// Event names live in our own namespace and a LowCardinality ClickHouse column. const MAX_EVENT_NAME_CHARS: usize = 256; @@ -72,9 +77,26 @@ fn json_size(value: &serde_json::Value) -> usize { .unwrap_or(usize::MAX) } +fn validate_event_timestamps(timestamp: i64, now: i64) -> Result<(), String> { + if timestamp <= MIN_TIMESTAMP_MS { + return Err("timestamp must be an epoch-millis value after 2020-01-01".to_string()); + } + if timestamp > now + MAX_FUTURE_SKEW_HOURS * HOUR_MS { + return Err(format!( + "timestamp must not be more than {MAX_FUTURE_SKEW_HOURS}h in the future" + )); + } + if timestamp < now - MAX_EVENT_AGE_DAYS * DAY_MS { + return Err(format!( + "timestamp must not be more than {MAX_EVENT_AGE_DAYS} days in the past" + )); + } + Ok(()) +} + /// Validates one event against the ingest schema rules. Error messages name /// the offending field; the caller prefixes the batch index. -fn validate_event(event: &Event, max_timestamp: i64) -> Result<(), String> { +fn validate_event(event: &Event, now: i64) -> Result<(), String> { if event.event.trim().is_empty() { return Err("event must be non-empty".to_string()); } @@ -110,12 +132,7 @@ fn validate_event(event: &Event, max_timestamp: i64) -> Result<(), String> { )); } } - if event.timestamp <= MIN_TIMESTAMP_MS { - return Err("timestamp must be an epoch-millis value after 2020-01-01".to_string()); - } - if event.timestamp > max_timestamp { - return Err("timestamp must not be more than 24h in the future".to_string()); - } + validate_event_timestamps(event.timestamp, now)?; if json_size(&event.value) > MAX_VALUE_BYTES { return Err(format!( "value must serialize to at most {MAX_VALUE_BYTES} bytes" @@ -160,12 +177,12 @@ fn validate_event(event: &Event, max_timestamp: i64) -> Result<(), String> { /// Splits a batch into forwardable events and per-event rejections. Each /// event is validated independently: one invalid event never blocks its /// neighbours. -fn partition_events(events: &[Event]) -> (Vec<&Event>, Vec) { - let max_timestamp = Utc::now().timestamp_millis() + MAX_FUTURE_SKEW_MS; +fn partition_events(events: Vec) -> (Vec, Vec) { + let now = Utc::now().timestamp_millis(); let mut valid = Vec::with_capacity(events.len()); let mut rejected = Vec::new(); - for (i, event) in events.iter().enumerate() { - match validate_event(event, max_timestamp) { + for (i, event) in events.into_iter().enumerate() { + match validate_event(&event, now) { Ok(()) => valid.push(event), Err(error) => rejected.push(RejectedEvent { index: i, @@ -199,12 +216,13 @@ pub async fn ingest_batch( return (StatusCode::BAD_REQUEST, Json(body)).into_response(); } - let (valid, rejected) = partition_events(&batch.events); + let batch_size = batch.events.len(); + let (valid, rejected) = partition_events(batch.events); if !rejected.is_empty() { warn!( client_api_key = %mask_key(&ctx.client_api_key), rejected_count = rejected.len(), - batch_size = batch.events.len(), + batch_size, first_error = %rejected[0].error, "Rejected invalid events from batch" ); @@ -227,24 +245,24 @@ pub async fn ingest_batch( // TODO: move this to a method. let records: Vec = valid - .iter() + .into_iter() .map(|event| EventRecord { - environment_key: &ctx.client_api_key, + environment_key: ctx.client_api_key.to_string(), 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(), + sdk_language: sdk_language.clone(), + sdk_version: sdk_version.clone(), + event: event.event, + feature_name: event.feature_name, + identifier: event.identifier, + value: event.value, + traits: event.traits, + metadata: event.metadata, 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 +306,7 @@ mod tests { use crate::sink::{KafkaAuth, KafkaSink, KafkaSinkConfig}; use crate::test_helpers::{ - assert_topic_empty, mock_kafka, read_records, FLAGSMITH_WAREHOUSE_TOPIC, + assert_topic_empty, mock_kafka, now_ms, read_records, FLAGSMITH_WAREHOUSE_TOPIC, }; fn parse(json: &str) -> EventBatch { @@ -298,7 +316,7 @@ mod tests { /// Partition the batch and return (accepted count, rejections). fn partition(json: &str) -> (usize, Vec) { let batch = parse(json); - let (valid, rejected) = partition_events(&batch.events); + let (valid, rejected) = partition_events(batch.events); (valid.len(), rejected) } @@ -348,9 +366,10 @@ mod tests { #[test] fn test_validate_accepts_exposure_with_feature_name() { - assert!(accepts( - r#"{"events":[{"event":"$flag_exposure","feature_name":"dark_mode","value":"control","timestamp":1700000000000}]}"#, - )); + let ts = now_ms(); + assert!(accepts(&format!( + r#"{{"events":[{{"event":"$flag_exposure","feature_name":"dark_mode","value":"control","timestamp":{ts}}}]}}"#, + ))); } #[test] @@ -422,10 +441,31 @@ mod tests { } #[test] - fn test_validate_accepts_custom_event() { - assert!(accepts( - r#"{"events":[{"event":"purchase","value":42,"timestamp":1700000000000}]}"# + fn test_validate_accepts_timestamp_within_max_age() { + let ts = now_ms() - 10 * DAY_MS + 60_000; + assert!(accepts(&format!( + r#"{{"events":[{{"event":"purchase","timestamp":{ts}}}]}}"# + ))); + } + + #[test] + fn test_validate_rejects_timestamp_older_than_max_age() { + let ts = now_ms() - 11 * DAY_MS; + let (_, rejected) = partition(&format!( + r#"{{"events":[{{"event":"purchase","timestamp":{ts}}}]}}"# )); + assert_eq!( + rejected[0].error, + "events[0].timestamp must not be more than 10 days in the past" + ); + } + + #[test] + fn test_validate_accepts_custom_event() { + let ts = now_ms(); + assert!(accepts(&format!( + r#"{{"events":[{{"event":"purchase","value":42,"timestamp":{ts}}}]}}"# + ))); } #[test] @@ -433,13 +473,14 @@ mod tests { // One invalid event never blocks its neighbours; the rejection names // the offending index and carries the same message the old whole-batch // 400 used. - let (accepted, rejected) = partition( - r#"{"events":[ - {"event":"purchase","value":42,"timestamp":1700000000000}, - {"event":"purchase","timestamp":1}, - {"event":"$flag_exposure","feature_name":"dark_mode","timestamp":1700000000000} - ]}"#, - ); + let ts = now_ms(); + let (accepted, rejected) = partition(&format!( + r#"{{"events":[ + {{"event":"purchase","value":42,"timestamp":{ts}}}, + {{"event":"purchase","timestamp":1}}, + {{"event":"$flag_exposure","feature_name":"dark_mode","timestamp":{ts}}} + ]}}"#, + )); assert_eq!(accepted, 2); assert_eq!(rejected.len(), 1); assert_eq!(rejected[0].index, 1); @@ -498,22 +539,22 @@ mod tests { } #[test] - fn test_event_record_omits_nulls_and_renames_fields() { + fn test_event_record_omits_absent_fields() { let batch = parse( r#"{"events":[{"event":"$flag_exposure","feature_name":"dark_mode","identifier":"user_1","value":"control","timestamp":1700000000000}]}"#, ); - let event = &batch.events[0]; + let event = batch.events.into_iter().next().unwrap(); let record = EventRecord { - environment_key: "env_key", + environment_key: "env_key".to_string(), collected_at: 1700000000001, - sdk_language: "python", - sdk_version: "1.0.0", - 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(), + sdk_language: "python".to_string(), + sdk_version: "1.0.0".to_string(), + event: event.event, + feature_name: event.feature_name, + identifier: event.identifier, + value: event.value, + traits: event.traits, + metadata: event.metadata, timestamp: event.timestamp, }; @@ -557,6 +598,11 @@ mod tests { app_with_destination(sink, None) } + fn purchase_json() -> String { + let ts = now_ms(); + format!(r#"{{"events":[{{"event":"purchase","value":42,"timestamp":{ts}}}]}}"#) + } + fn post_events(json: &str) -> Request { Request::builder() .method("POST") @@ -574,14 +620,15 @@ mod tests { #[tokio::test] async fn test_handler_mixed_batch_returns_202_and_forwards_valid_only() { let kafka = mock_kafka(&[FLAGSMITH_WAREHOUSE_TOPIC]); + let ts = now_ms(); let response = app(&kafka.sink) - .oneshot(post_events( - r#"{"events":[ - {"event":"purchase","value":42,"timestamp":1700000000000}, - {"event":"purchase","timestamp":1}, - {"event":"$flag_exposure","feature_name":"dark_mode","timestamp":1700000000000} - ]}"#, - )) + .oneshot(post_events(&format!( + r#"{{"events":[ + {{"event":"purchase","value":42,"timestamp":{ts}}}, + {{"event":"purchase","timestamp":1}}, + {{"event":"$flag_exposure","feature_name":"dark_mode","timestamp":{ts}}} + ]}}"#, + ))) .await .unwrap(); @@ -601,10 +648,11 @@ mod tests { // The destination String columns (and the ClickPipes field mapping) // take flat scalar fields, so dict fields go on the wire JSON-encoded. let kafka = mock_kafka(&[FLAGSMITH_WAREHOUSE_TOPIC]); + let ts = now_ms(); let response = app(&kafka.sink) - .oneshot(post_events( - r#"{"events":[{"event":"purchase","traits":{"logins":42,"plan":"pro"},"metadata":{"source":"csv"},"timestamp":1700000000000}]}"#, - )) + .oneshot(post_events(&format!( + r#"{{"events":[{{"event":"purchase","traits":{{"logins":42,"plan":"pro"}},"metadata":{{"source":"csv"}},"timestamp":{ts}}}]}}"#, + ))) .await .unwrap(); @@ -659,9 +707,7 @@ mod tests { let sink = unreachable_sink(); let started = std::time::Instant::now(); let response = app(&sink) - .oneshot(post_events( - r#"{"events":[{"event":"purchase","value":42,"timestamp":1700000000000}]}"#, - )) + .oneshot(post_events(&purchase_json())) .await .unwrap(); @@ -688,9 +734,7 @@ mod tests { .expect("topic error"); let started = std::time::Instant::now(); let response = app(&kafka.sink) - .oneshot(post_events( - r#"{"events":[{"event":"purchase","value":42,"timestamp":1700000000000}]}"#, - )) + .oneshot(post_events(&purchase_json())) .await .unwrap(); @@ -707,9 +751,7 @@ mod tests { async fn test_handler_forwards_to_configured_destination() { let kafka = mock_kafka(&[FLAGSMITH_WAREHOUSE_TOPIC, "events-dest-x"]); let response = app_with_destination(&kafka.sink, Some("events-dest-x")) - .oneshot(post_events( - r#"{"events":[{"event":"purchase","value":42,"timestamp":1700000000000}]}"#, - )) + .oneshot(post_events(&purchase_json())) .await .unwrap(); @@ -722,8 +764,9 @@ mod tests { #[tokio::test] async fn test_handler_delivers_large_batches() { let kafka = mock_kafka(&[FLAGSMITH_WAREHOUSE_TOPIC]); + let ts = now_ms(); let events: Vec = (0..501) - .map(|i| format!(r#"{{"event":"purchase","value":{i},"timestamp":1700000000000}}"#)) + .map(|i| format!(r#"{{"event":"purchase","value":{i},"timestamp":{ts}}}"#)) .collect(); let body = format!(r#"{{"events":[{}]}}"#, events.join(",")); let response = app(&kafka.sink).oneshot(post_events(&body)).await.unwrap(); @@ -737,7 +780,8 @@ mod tests { } fn event_json(fields: &str) -> String { - format!(r#"{{"events":[{{{fields},"timestamp":1700000000000}}]}}"#) + let ts = now_ms(); + format!(r#"{{"events":[{{{fields},"timestamp":{ts}}}]}}"#) } #[test] diff --git a/src/routes/mod.rs b/src/routes/mod.rs index 7e12d2d..667f5c3 100644 --- a/src/routes/mod.rs +++ b/src/routes/mod.rs @@ -44,7 +44,7 @@ mod tests { use super::*; use crate::environment_keys::EnvironmentKeys; use crate::test_helpers::{ - assert_topic_empty, create_environment_keys, insert_environment_key, mock_kafka, + assert_topic_empty, create_environment_keys, insert_environment_key, mock_kafka, now_ms, read_records, FLAGSMITH_WAREHOUSE_TOPIC, }; @@ -53,7 +53,7 @@ mod tests { fn context_state(pool: PgPool) -> ContextState { ContextState { environment_keys: EnvironmentKeys::new(pool, Duration::from_secs(300)), - external_topic: Arc::from(EXTERNAL_WAREHOUSE_TOPIC), + external_topic: Some(Arc::from(EXTERNAL_WAREHOUSE_TOPIC)), } } @@ -63,9 +63,10 @@ mod tests { .uri("/v1/events") .header("content-type", "application/json") .header("X-Environment-Key", key) - .body(Body::from( - r#"{"events":[{"event":"purchase","value":42,"timestamp":1700000000000}]}"#, - )) + .body(Body::from(format!( + r#"{{"events":[{{"event":"purchase","value":42,"timestamp":{}}}]}}"#, + now_ms() + ))) .unwrap() } diff --git a/src/sink/batching.rs b/src/sink/batching.rs new file mode 100644 index 0000000..e2667dc --- /dev/null +++ b/src/sink/batching.rs @@ -0,0 +1,335 @@ +use std::future::Future; +use std::mem; +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::Duration; + +use tokio::sync::Notify; +use tokio::task::JoinHandle; +use tokio::time::{Instant, MissedTickBehavior}; +use tracing::{error, warn}; + +use super::{ErrorKind, SinkError}; +use crate::models::EventRecord; + +/// 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 { + fn write(&self, records: &[EventRecord]) -> impl Future> + Send; +} + +#[derive(Debug, Clone)] +pub struct BatchConfig { + pub flush_interval: Duration, + /// Flush early once this many events are buffered. Not a cap: the buffer + /// is unbounded and keeps growing while a write is slow or failing. + pub flush_max_events: usize, +} + +struct Shared { + writer: W, + config: BatchConfig, + buffer: Mutex>, + flush_now: Notify, + stop: Notify, + /// The first config error; stays set until restart. + fatal_error: OnceLock, +} + +pub struct BatchingSink { + shared: Arc>, + flusher: Arc>>>, +} + +impl Clone for BatchingSink { + fn clone(&self) -> Self { + Self { + shared: self.shared.clone(), + flusher: self.flusher.clone(), + } + } +} + +impl BatchingSink { + pub fn new(writer: W, config: BatchConfig) -> Self { + let shared = Arc::new(Shared { + writer, + config, + buffer: Mutex::new(Vec::new()), + flush_now: Notify::new(), + stop: Notify::new(), + fatal_error: OnceLock::new(), + }); + let flusher = tokio::spawn(run(shared.clone())); + Self { + shared, + flusher: Arc::new(Mutex::new(Some(flusher))), + } + } + + /// Buffer the records; never blocks or fails. + pub fn put(&self, records: Vec) { + let mut buffer = self.shared.buffer.lock().expect("buffer lock"); + buffer.extend(records); + if buffer.len() >= self.shared.config.flush_max_events { + self.shared.flush_now.notify_one(); + } + } + + /// 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(); + let Some(flusher) = self.flusher.lock().expect("flusher lock").take() else { + return; + }; + match tokio::time::timeout(timeout, flusher).await { + Ok(Ok(())) => {} + Ok(Err(e)) => warn!(error = %e, "Sink flush task failed"), + Err(_) => warn!("Sink final flush incomplete"), + } + } +} + +async fn run(shared: Arc>) { + let period = shared.config.flush_interval; + let mut interval = tokio::time::interval_at(Instant::now() + period, period); + interval.set_missed_tick_behavior(MissedTickBehavior::Delay); + loop { + let stopping = tokio::select! { + _ = interval.tick() => false, + _ = shared.flush_now.notified() => false, + _ = shared.stop.notified() => true, + }; + flush(&shared).await; + if stopping { + return; + } + } +} + +async fn flush(shared: &Shared) { + let records = mem::take(&mut *shared.buffer.lock().expect("buffer lock")); + if records.is_empty() { + return; + } + write_with_retries(shared, &records).await; +} + +/// 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; + } + 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; + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + + /// Records every write attempt; each takes `delay`, and the first + /// `failures` fail with `error`. + #[derive(Clone, Default)] + struct FakeWriter { + failures: usize, + error: Option clickhouse::error::Error>, + delay: Duration, + attempts: Arc, + written: Arc>>>, + } + + impl BatchWriter for FakeWriter { + async fn write(&self, records: &[EventRecord]) -> Result<(), SinkError> { + let attempt = self.attempts.fetch_add(1, Ordering::SeqCst); + tokio::time::sleep(self.delay).await; + 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); + Ok(()) + } + } + + fn records(timestamps: &[i64]) -> Vec { + timestamps + .iter() + .map(|×tamp| EventRecord { + environment_key: "env".to_string(), + collected_at: timestamp, + sdk_language: "rust".to_string(), + sdk_version: "1.0.0".to_string(), + event: "purchase".to_string(), + feature_name: None, + identifier: None, + value: serde_json::Value::Null, + traits: None, + metadata: None, + timestamp, + }) + .collect() + } + + fn config(flush_interval: Duration, flush_max_events: usize) -> BatchConfig { + BatchConfig { + flush_interval, + flush_max_events, + } + } + + #[tokio::test(start_paused = true)] + async fn test_flushes_on_the_interval() { + let writer = FakeWriter::default(); + let sink = BatchingSink::new(writer.clone(), config(Duration::from_secs(1), 100)); + + sink.put(records(&[1, 2])); + tokio::time::sleep(Duration::from_millis(500)).await; + assert!(writer.written.lock().unwrap().is_empty()); + + tokio::time::sleep(Duration::from_millis(600)).await; + assert_eq!(*writer.written.lock().unwrap(), vec![vec![1, 2]]); + } + + #[tokio::test(start_paused = true)] + async fn test_flushes_early_at_flush_max_events() { + let writer = FakeWriter::default(); + let sink = BatchingSink::new(writer.clone(), config(Duration::from_secs(3600), 2)); + + sink.put(records(&[1])); + tokio::time::sleep(Duration::from_millis(10)).await; + assert!(writer.written.lock().unwrap().is_empty()); + + sink.put(records(&[2])); + tokio::time::sleep(Duration::from_millis(10)).await; + assert_eq!(*writer.written.lock().unwrap(), vec![vec![1, 2]]); + } + + #[tokio::test(start_paused = true)] + async fn test_keeps_buffering_while_a_write_is_slow() { + let writer = FakeWriter { + delay: Duration::from_secs(1), + ..FakeWriter::default() + }; + let sink = BatchingSink::new(writer.clone(), config(Duration::from_secs(3600), 2)); + + sink.put(records(&[1, 2])); + tokio::time::sleep(Duration::from_millis(1)).await; + // [1, 2] is being written; the buffer grows past the threshold. + sink.put(records(&[3, 4])); + sink.put(records(&[5])); + + sink.shutdown(Duration::from_secs(5)).await; + assert_eq!( + *writer.written.lock().unwrap(), + vec![vec![1, 2], vec![3, 4, 5]] + ); + } + + #[tokio::test(start_paused = true)] + async fn test_drops_a_batch_after_the_retries() { + let writer = FakeWriter { + 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(60)).await; + + assert_eq!( + writer.attempts.load(Ordering::SeqCst), + 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)] + async fn test_shutdown_writes_what_is_buffered() { + let writer = FakeWriter::default(); + let sink = BatchingSink::new(writer.clone(), config(Duration::from_secs(3600), 100)); + + sink.put(records(&[1, 2, 3])); + sink.shutdown(Duration::from_secs(1)).await; + + assert_eq!(*writer.written.lock().unwrap(), vec![vec![1, 2, 3]]); + } +} diff --git a/src/sink/clickhouse.rs b/src/sink/clickhouse.rs new file mode 100644 index 0000000..972d067 --- /dev/null +++ b/src/sink/clickhouse.rs @@ -0,0 +1,462 @@ +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::{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, + pub user: String, + pub password: String, + pub database: String, + pub table: String, +} + +/// Keeps the password out of logs. +impl fmt::Debug for ClickHouseConfig { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ClickHouseConfig") + .field("url", &self.url) + .field("user", &self.user) + .field("database", &self.database) + .field("table", &self.table) + .finish_non_exhaustive() + } +} + +/// Field names must match the `events` columns. Every column is +/// non-Nullable, so absent optionals are written as "", as ClickPipes does. +#[derive(Debug, PartialEq, Row, Serialize)] +struct ClickHouseRow<'a> { + environment_key: &'a str, + event: &'a str, + feature_name: &'a str, + /// DateTime64(3), sent as epoch millis. + timestamp: i64, + collected_at: i64, + identifier: &'a str, + value: Cow<'a, str>, + traits: String, + metadata: String, + sdk_language: &'a str, + sdk_version: &'a str, +} + +impl<'a> From<&'a EventRecord> for ClickHouseRow<'a> { + fn from(record: &'a EventRecord) -> Self { + ClickHouseRow { + environment_key: &record.environment_key, + event: &record.event, + feature_name: record.feature_name.as_deref().unwrap_or_default(), + timestamp: record.timestamp, + collected_at: record.collected_at, + identifier: record.identifier.as_deref().unwrap_or_default(), + value: value_to_text(&record.value), + traits: record + .traits + .as_ref() + .map(ToString::to_string) + .unwrap_or_default(), + metadata: record + .metadata + .as_ref() + .map(ToString::to_string) + .unwrap_or_default(), + sdk_language: &record.sdk_language, + sdk_version: &record.sdk_version, + } + } +} + +fn value_to_text(value: &serde_json::Value) -> Cow<'_, str> { + match value { + serde_json::Value::String(s) => Cow::Borrowed(s), + serde_json::Value::Null => Cow::Borrowed(""), + other => Cow::Owned(other.to_string()), + } +} + +/// 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, +} + +impl ClickHouseWriter { + /// Connects lazily, on the first insert. + pub fn new(config: ClickHouseConfig) -> Self { + let client = Client::default() + .with_url(config.url) + .with_user(config.user) + .with_password(config.password) + .with_database(config.database) + .with_compression(Compression::Lz4); + Self { + client, + 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 { + /// One plain `INSERT` per batch; succeeds once ClickHouse has committed it. + async fn write(&self, records: &[EventRecord]) -> Result<(), SinkError> { + let mut insert = self + .client + .insert::>(&self.table) + .await? + .with_timeouts(Some(INSERT_TIMEOUT), Some(INSERT_TIMEOUT)); + for record in records { + insert.write(&ClickHouseRow::from(record)).await?; + } + insert.end().await?; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use serde::Deserialize; + use serde_json::json; + + use super::*; + use crate::sink::{BatchConfig, BatchingSink}; + use crate::test_helpers::{now_ms, ClickHouseTestDatabase}; + + 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, + } + } + + #[test] + fn test_absent_optionals_become_empty_strings() { + let record = EventRecord { + value: json!(42), + traits: Some(json!({"plan": "pro"})), + ..record("purchase", 1_700_000_000_000) + }; + + assert_eq!( + ClickHouseRow::from(&record), + ClickHouseRow { + environment_key: "env", + event: "purchase", + feature_name: "", + timestamp: 1_700_000_000_000, + collected_at: 1_700_000_000_001, + identifier: "", + value: Cow::Borrowed("42"), + traits: r#"{"plan":"pro"}"#.to_string(), + metadata: String::new(), + sdk_language: "python", + sdk_version: "1.0.0", + } + ); + } + + #[test] + fn test_value_to_text_leaves_strings_unquoted() { + assert_eq!(value_to_text(&json!("control")), "control"); + assert_eq!(value_to_text(&json!(null)), ""); + assert_eq!(value_to_text(&json!(42)), "42"); + assert_eq!(value_to_text(&json!(true)), "true"); + assert_eq!(value_to_text(&json!({"a": [1, 2]})), r#"{"a":[1,2]}"#); + } + + #[test] + fn test_debug_output_omits_the_password() { + let config = ClickHouseConfig { + url: "https://clickhouse:8443".to_string(), + user: "ingest".to_string(), + password: "hunter2".to_string(), + database: "flagsmith_exp".to_string(), + table: "events".to_string(), + }; + let shown = format!("{config:?}"); + assert!(shown.contains("ingest")); + 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, + flush_max_events: 1000, + } + } + + /// A row as read back from the table. + #[derive(Debug, PartialEq, Row, Deserialize)] + struct StoredRow { + environment_key: String, + event: String, + feature_name: String, + timestamp: i64, + collected_at: i64, + identifier: String, + value: String, + traits: String, + metadata: String, + sdk_language: String, + sdk_version: String, + } + + #[tokio::test] + async fn test_flushed_rows_land_in_the_events_table() { + let db = ClickHouseTestDatabase::create().await; + let sink = BatchingSink::new( + ClickHouseWriter::new(db.config.clone()), + batch_config(Duration::from_millis(100)), + ); + let ts = now_ms(); + let records = vec![ + EventRecord { + feature_name: Some("dark_mode".to_string()), + identifier: Some("user_1".to_string()), + value: json!("control"), + traits: Some(json!({"plan": "pro"})), + metadata: Some(json!({"source": "csv"})), + ..record("$flag_exposure", ts) + }, + record("purchase", ts + 2), + ]; + + sink.put(records); + db.wait_for_rows(2).await; + + let stored: Vec = db + .client + .query("SELECT ?fields FROM events ORDER BY timestamp") + .fetch_all() + .await + .unwrap(); + assert_eq!( + stored, + vec![ + StoredRow { + environment_key: "env".to_string(), + event: "$flag_exposure".to_string(), + feature_name: "dark_mode".to_string(), + timestamp: ts, + collected_at: ts + 1, + identifier: "user_1".to_string(), + value: "control".to_string(), + traits: r#"{"plan":"pro"}"#.to_string(), + metadata: r#"{"source":"csv"}"#.to_string(), + sdk_language: "python".to_string(), + sdk_version: "1.0.0".to_string(), + }, + StoredRow { + environment_key: "env".to_string(), + event: "purchase".to_string(), + feature_name: String::new(), + timestamp: ts + 2, + collected_at: ts + 3, + identifier: String::new(), + value: String::new(), + traits: String::new(), + metadata: String::new(), + sdk_language: "python".to_string(), + sdk_version: "1.0.0".to_string(), + }, + ] + ); + + sink.shutdown(Duration::from_secs(5)).await; + db.drop_db().await; + } + + #[tokio::test] + async fn test_shutdown_writes_buffered_rows() { + let db = ClickHouseTestDatabase::create().await; + let sink = BatchingSink::new( + ClickHouseWriter::new(db.config.clone()), + batch_config(Duration::from_secs(3600)), + ); + let record = record("purchase", now_ms()); + + sink.put(vec![record.clone(), record.clone(), record]); + assert_eq!(db.count().await, 0); + sink.shutdown(Duration::from_secs(5)).await; + + assert_eq!(db.count().await, 3); + db.drop_db().await; + } +} diff --git a/src/sink/kafka.rs b/src/sink/kafka.rs index d41fb13..6a79477 100644 --- a/src/sink/kafka.rs +++ b/src/sink/kafka.rs @@ -148,7 +148,7 @@ impl KafkaSink { pub async fn put( &self, destination: Option<&str>, - records: &[EventRecord<'_>], + records: &[EventRecord], ) -> Result<(), DeliveryError> { let topic = destination.unwrap_or(&self.default_topic); let total = records.len(); diff --git a/src/sink/mod.rs b/src/sink/mod.rs index 4369129..c37f1c3 100644 --- a/src/sink/mod.rs +++ b/src/sink/mod.rs @@ -1,7 +1,11 @@ +mod batching; +mod clickhouse; mod kafka; use std::time::Duration; +pub use batching::{BatchConfig, BatchingSink}; +pub use clickhouse::{ClickHouseConfig, ClickHouseWriter}; pub use kafka::{KafkaAuth, KafkaSink, KafkaSinkConfig}; use crate::models::EventRecord; @@ -10,6 +14,7 @@ use crate::models::EventRecord; #[derive(Clone)] pub enum Sink { Kafka(KafkaSink), + ClickHouse(BatchingSink), } impl Sink { @@ -18,10 +23,15 @@ impl Sink { pub async fn put( &self, destination: Option<&str>, - records: &[EventRecord<'_>], + records: Vec, ) -> Result<(), SinkError> { match self { - Self::Kafka(sink) => Ok(sink.put(destination, records).await?), + Self::Kafka(sink) => Ok(sink.put(destination, &records).await?), + // A single table: there is no per-environment destination. + Self::ClickHouse(sink) => { + sink.put(records); + Ok(()) + } } } @@ -29,6 +39,7 @@ impl Sink { pub fn fatal_error(&self) -> Option { match self { Self::Kafka(sink) => sink.fatal_error(), + Self::ClickHouse(sink) => sink.fatal_error(), } } @@ -36,6 +47,7 @@ impl Sink { pub async fn shutdown(self, timeout: Duration) { match self { Self::Kafka(sink) => sink.shutdown(timeout).await, + Self::ClickHouse(sink) => sink.shutdown(timeout).await, } } } @@ -44,4 +56,28 @@ impl Sink { pub enum SinkError { #[error(transparent)] Kafka(#[from] kafka::DeliveryError), + #[error("ClickHouse insert failed: {0}")] + ClickHouse(#[from] ::clickhouse::error::Error), +} + +/// 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), + } + } } diff --git a/src/test_helpers.rs b/src/test_helpers.rs index 9545c21..ed98364 100644 --- a/src/test_helpers.rs +++ b/src/test_helpers.rs @@ -1,6 +1,9 @@ +use std::env; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use chrono::{DateTime, Utc}; +use clickhouse::Client; use rdkafka::config::ClientConfig; use rdkafka::consumer::{Consumer, StreamConsumer}; use rdkafka::mocking::MockCluster; @@ -8,12 +11,17 @@ use rdkafka::producer::DefaultProducerContext; use rdkafka::Message; use sqlx::PgPool; -use crate::sink::{KafkaAuth, KafkaSink, KafkaSinkConfig, Sink}; +use crate::sink::{ClickHouseConfig, KafkaAuth, KafkaSink, KafkaSinkConfig, Sink}; pub(crate) const FLAGSMITH_WAREHOUSE_TOPIC: &str = "events-default"; pub(crate) type MockKafkaCluster = MockCluster<'static, DefaultProducerContext>; +/// Event timestamps must be recent to pass validation. +pub(crate) fn now_ms() -> i64 { + Utc::now().timestamp_millis() +} + /// The cluster is torn down on drop, so it must outlive every request in the /// test. pub(crate) struct MockKafka { @@ -107,3 +115,103 @@ pub(crate) async fn insert_environment_key( .await .unwrap(); } + +/// The production `events` table, from Flagsmith's +/// experimentation/clickhouse/events.sql. +const EVENTS_TABLE: &str = " + CREATE TABLE {database}.events + ( + environment_key LowCardinality(String), + event LowCardinality(String), + feature_name LowCardinality(String), + timestamp DateTime64(3), + collected_at DateTime64(3), + identifier String, + value String CODEC(ZSTD(3)), + traits String CODEC(ZSTD(3)), + metadata String CODEC(ZSTD(3)), + sdk_language LowCardinality(String), + sdk_version LowCardinality(String), + INDEX idx_identity identifier TYPE bloom_filter GRANULARITY 4, + CONSTRAINT environment_key_not_empty CHECK environment_key != '', + CONSTRAINT event_not_empty CHECK event != '', + CONSTRAINT timestamp_sane CHECK timestamp > toDateTime64('2020-01-01 00:00:00', 3) + ) + ENGINE = MergeTree + PARTITION BY toYYYYMMDD(timestamp) + ORDER BY (environment_key, event, feature_name, timestamp, identifier) +"; + +static NEXT_DATABASE: AtomicUsize = AtomicUsize::new(0); + +/// A fresh database holding the `events` table, on the ClickHouse at +/// `CLICKHOUSE_URL` (with `CLICKHOUSE_USER` and `CLICKHOUSE_PASSWORD`, as the +/// server reads them). Call `drop` at the end of the test. +pub(crate) struct ClickHouseTestDatabase { + pub(crate) config: ClickHouseConfig, + pub(crate) client: Client, +} + +impl ClickHouseTestDatabase { + pub(crate) async fn create() -> Self { + let url = env::var("CLICKHOUSE_URL").expect("CLICKHOUSE_URL must be set"); + let user = env::var("CLICKHOUSE_USER").unwrap_or_else(|_| "default".to_string()); + let password = env::var("CLICKHOUSE_PASSWORD").unwrap_or_default(); + let database = format!( + "test_{}_{}", + Utc::now().timestamp_millis(), + NEXT_DATABASE.fetch_add(1, Ordering::SeqCst) + ); + let server = Client::default() + .with_url(&url) + .with_user(&user) + .with_password(&password); + server + .query(&format!("CREATE DATABASE {database}")) + .execute() + .await + .expect("create database"); + server + .query(&EVENTS_TABLE.replace("{database}", &database)) + .execute() + .await + .expect("create events table"); + Self { + client: server.with_database(&database), + config: ClickHouseConfig { + url, + user, + password, + database, + table: "events".to_string(), + }, + } + } + + pub(crate) async fn count(&self) -> u64 { + self.client + .query("SELECT count() FROM events") + .fetch_one() + .await + .expect("count events") + } + + /// Polls until `n` rows are stored; flushes run in the background. + pub(crate) async fn wait_for_rows(&self, n: u64) { + for _ in 0..100 { + if self.count().await >= n { + return; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + panic!("expected {n} rows within 10s, found {}", self.count().await); + } + + pub(crate) async fn drop_db(self) { + self.client + .query(&format!("DROP DATABASE {}", self.config.database)) + .execute() + .await + .expect("drop database"); + } +}