From fdc0e7dc2b98d68daa5b83e61b7e2ff9e7c348c1 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 5 Oct 2026 16:04:44 +0000 Subject: [PATCH] refactor(rust): track ClickHouse migrations in a checksummed ledger (#44580) * refactor(rust): track ClickHouse migrations in a checksummed ledger Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(clickhouse): harden migration execution and reuse retention SQL * refactor(rust): track ClickHouse migrations in a checksummed ledger Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(rust): keep ClickHouse retention TTLs in the current policy list Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yujong Lee Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 26 +- litellm-rust/Cargo.toml | 2 - litellm-rust/crates/migrate-macros/Cargo.toml | 19 - .../crates/migrate-macros/src/error.rs | 21 - litellm-rust/crates/migrate-macros/src/lib.rs | 199 -------- litellm-rust/crates/migrate/Cargo.toml | 12 - litellm-rust/crates/migrate/README.md | 5 - litellm-rust/crates/migrate/src/lib.rs | 8 - .../tests/fixtures/migrations/10_tenth.sql | 1 - .../tests/fixtures/migrations/1_first.sql | 1 - .../tests/fixtures/migrations/2_second.sql | 1 - litellm-rust/crates/migrate/tests/migrate.rs | 21 - litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../crates/python-bridge/src/routes/traces.rs | 12 +- .../crates/storage-clickhouse/AGENTS.md | 4 + .../crates/storage-clickhouse/Cargo.toml | 3 + .../crates/storage-clickhouse/src/error.rs | 2 +- .../crates/storage-clickhouse/src/lib.rs | 2 + .../crates/storage-clickhouse/src/migrate.rs | 308 ++++++++++++ .../storage-clickhouse/tests/migrations.rs | 476 ++++++++++++++++++ .../crates/traces-clickhouse/AGENTS.md | 3 +- .../crates/traces-clickhouse/Cargo.toml | 2 +- .../crates/traces-clickhouse/src/error.rs | 6 +- .../crates/traces-clickhouse/src/lib.rs | 3 +- .../crates/traces-clickhouse/src/schema.rs | 148 +++++- .../traces-clickhouse/tests/migrations.rs | 226 ++++++++- tests/test_litellm_rust/test_traces.py | 9 +- 27 files changed, 1173 insertions(+), 348 deletions(-) delete mode 100644 litellm-rust/crates/migrate-macros/Cargo.toml delete mode 100644 litellm-rust/crates/migrate-macros/src/error.rs delete mode 100644 litellm-rust/crates/migrate-macros/src/lib.rs delete mode 100644 litellm-rust/crates/migrate/Cargo.toml delete mode 100644 litellm-rust/crates/migrate/README.md delete mode 100644 litellm-rust/crates/migrate/src/lib.rs delete mode 100644 litellm-rust/crates/migrate/tests/fixtures/migrations/10_tenth.sql delete mode 100644 litellm-rust/crates/migrate/tests/fixtures/migrations/1_first.sql delete mode 100644 litellm-rust/crates/migrate/tests/fixtures/migrations/2_second.sql delete mode 100644 litellm-rust/crates/migrate/tests/migrate.rs create mode 100644 litellm-rust/crates/storage-clickhouse/src/migrate.rs create mode 100644 litellm-rust/crates/storage-clickhouse/tests/migrations.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f7b667c8ab2..934a1f0bb6b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4085,26 +4085,6 @@ dependencies = [ "strum", ] -[[package]] -name = "litellm-migrate" -version = "0.1.0" -dependencies = [ - "litellm-migrate-macros", - "rstest", -] - -[[package]] -name = "litellm-migrate-macros" -version = "0.1.0" -dependencies = [ - "proc-macro2", - "quote", - "rstest", - "syn 2.0.119", - "tempfile", - "thiserror 2.0.19", -] - [[package]] name = "litellm-model-catalog" version = "0.1.0" @@ -4170,6 +4150,7 @@ dependencies = [ "serde_json", "serde_with", "sha2 0.10.9", + "sqlx", "strum", "thiserror 2.0.19", "tokio", @@ -4368,6 +4349,9 @@ dependencies = [ "rstest", "serde", "serde_json", + "serde_with", + "sqlx", + "testcontainers-modules", "thiserror 2.0.19", "tokio", "url", @@ -4498,7 +4482,6 @@ dependencies = [ "hmac 0.12.1", "jsonschema", "litellm-http", - "litellm-migrate", "litellm-storage-clickhouse", "litellm-traces", "litellm-traces-cache", @@ -4509,6 +4492,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.9", + "sqlx", "strum", "testcontainers-modules", "thiserror 2.0.19", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index b0766f11e87..6ef90c59d2f 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -16,8 +16,6 @@ litellm-traces = { path = "crates/traces" } litellm-traces-cache = { path = "crates/traces-cache" } litellm-traces-clickhouse = { path = "crates/traces-clickhouse" } litellm-storage-clickhouse = { path = "crates/storage-clickhouse" } -litellm-migrate = { path = "crates/migrate" } -litellm-migrate-macros = { path = "crates/migrate-macros" } litellm-core = { path = "crates/core" } litellm-gateway-mcp = { path = "crates/gateway-mcp" } litellm-gateway = { path = "crates/gateway" } diff --git a/litellm-rust/crates/migrate-macros/Cargo.toml b/litellm-rust/crates/migrate-macros/Cargo.toml deleted file mode 100644 index 5cd68415ca2..00000000000 --- a/litellm-rust/crates/migrate-macros/Cargo.toml +++ /dev/null @@ -1,19 +0,0 @@ -[package] -name = "litellm-migrate-macros" -version = "0.1.0" -edition.workspace = true -license.workspace = true -repository.workspace = true - -[lib] -proc-macro = true - -[dependencies] -proc-macro2.workspace = true -quote.workspace = true -syn = { workspace = true, features = ["parsing", "printing", "proc-macro"] } -thiserror.workspace = true - -[dev-dependencies] -rstest.workspace = true -tempfile.workspace = true diff --git a/litellm-rust/crates/migrate-macros/src/error.rs b/litellm-rust/crates/migrate-macros/src/error.rs deleted file mode 100644 index 9833009517b..00000000000 --- a/litellm-rust/crates/migrate-macros/src/error.rs +++ /dev/null @@ -1,21 +0,0 @@ -use std::io; - -#[derive(Debug, thiserror::Error)] -pub enum Error { - #[error("could not read migrations directory `{path}`")] - ReadDirectory { - path: String, - #[source] - source: io::Error, - }, - #[error( - "migration name `{name}` must be `_.sql` with a `[a-z0-9_]` description" - )] - InvalidName { name: String }, - #[error("migration version `{version}` is declared more than once")] - DuplicateVersion { version: u64 }, - #[error("migrations directory `{path}` contains no migrations")] - Empty { path: String }, - #[error("migration path `{path}` is not valid UTF-8")] - NonUtf8Path { path: String }, -} diff --git a/litellm-rust/crates/migrate-macros/src/lib.rs b/litellm-rust/crates/migrate-macros/src/lib.rs deleted file mode 100644 index 501f59e6fc2..00000000000 --- a/litellm-rust/crates/migrate-macros/src/lib.rs +++ /dev/null @@ -1,199 +0,0 @@ -mod error; - -use std::path::{Path, PathBuf}; - -use error::Error; -use proc_macro::TokenStream; -use quote::quote; -use syn::LitStr; - -struct Entry { - version: u64, - description: String, - path: PathBuf, -} - -fn resolve(dir: &Path) -> Result, Error> { - let mut entries = Vec::new(); - let files = std::fs::read_dir(dir).map_err(|source| Error::ReadDirectory { - path: dir.display().to_string(), - source, - })?; - for file in files { - let file = file.map_err(|source| Error::ReadDirectory { - path: dir.display().to_string(), - source, - })?; - let path = file.path(); - let name = path - .file_name() - .and_then(|name| name.to_str()) - .ok_or_else(|| Error::NonUtf8Path { - path: path.display().to_string(), - })? - .to_owned(); - let invalid = || Error::InvalidName { name: name.clone() }; - let stem = name - .strip_suffix(".sql") - .filter(|_| file.file_type().is_ok_and(|kind| kind.is_file())) - .and_then(|stem| stem.split_once('_')) - .filter(|(version, description)| { - !version.is_empty() - && version.bytes().all(|b| b.is_ascii_digit()) - && !description.is_empty() - && description - .bytes() - .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'_') - }) - .ok_or_else(invalid)?; - let version = stem.0.parse::().map_err(|_| invalid())?; - entries.push(Entry { - version, - description: stem.1.to_owned(), - path, - }); - } - if entries.is_empty() { - return Err(Error::Empty { - path: dir.display().to_string(), - }); - } - entries.sort_by_key(|entry| entry.version); - for pair in entries.windows(2) { - if pair[0].version == pair[1].version { - return Err(Error::DuplicateVersion { - version: pair[0].version, - }); - } - } - Ok(entries) -} - -fn resolve_input(lit: &LitStr) -> Result, Error> { - let root = std::env::var("CARGO_MANIFEST_DIR") - .map(PathBuf::from) - .unwrap_or_default(); - let dir = root.join(lit.value()); - let dir = dir.canonicalize().map_err(|source| Error::ReadDirectory { - path: dir.display().to_string(), - source, - })?; - if dir.to_str().is_none() { - return Err(Error::NonUtf8Path { - path: dir.display().to_string(), - }); - } - resolve(&dir) -} - -#[proc_macro] -pub fn migrate(input: TokenStream) -> TokenStream { - let lit = syn::parse_macro_input!(input as LitStr); - match resolve_input(&lit) { - Ok(entries) => { - let migrations = entries.iter().map(|entry| { - let version = entry.version; - let description = &entry.description; - let path = entry - .path - .to_str() - .expect("canonical migration path is UTF-8"); - quote! { - ::litellm_migrate::Migration { - version: #version, - description: #description, - sql: ::core::include_str!(#path), - } - } - }); - quote! { &[#(#migrations),*] }.into() - } - Err(err) => syn::Error::new(lit.span(), err).to_compile_error().into(), - } -} - -#[cfg(test)] -mod tests { - use std::fs; - - use rstest::rstest; - use tempfile::TempDir; - - use super::{Error, resolve}; - - fn migrations_dir(files: &[&str]) -> TempDir { - let dir = TempDir::new().expect("tempdir"); - for file in files { - fs::write(dir.path().join(file), "SELECT 1").expect("write fixture"); - } - dir - } - - #[rstest] - fn orders_versions_numerically() { - let dir = migrations_dir(&["10_tenth.sql", "2_second.sql", "1_first.sql"]); - let entries = resolve(dir.path()).expect("resolves"); - let versions: Vec = entries.iter().map(|entry| entry.version).collect(); - let descriptions: Vec<&str> = entries - .iter() - .map(|entry| entry.description.as_str()) - .collect(); - assert_eq!(versions, [1, 2, 10]); - assert_eq!(descriptions, ["first", "second", "tenth"]); - } - - #[rstest] - #[case::dash_in_version(&["0001-dash.sql"])] - #[case::not_sql(&["notes.txt"])] - #[case::empty_description(&["0001_.sql"])] - #[case::non_digit_version(&["x_name.sql"])] - #[case::uppercase_description(&["0001_Upper.sql"])] - #[case::no_underscore(&["0001.sql"])] - #[case::plus_sign_version(&["+10_add.sql"])] - fn rejects_invalid_names(#[case] files: &[&str]) { - let dir = migrations_dir(files); - assert!(matches!( - resolve(dir.path()), - Err(Error::InvalidName { .. }) - )); - } - - #[rstest] - fn rejects_subdirectories() { - let dir = migrations_dir(&["0001_a.sql"]); - fs::create_dir(dir.path().join("0002_b.sql")).expect("subdir"); - assert!(matches!( - resolve(dir.path()), - Err(Error::InvalidName { .. }) - )); - } - - #[cfg(unix)] - #[rstest] - fn rejects_symlinks() { - let dir = migrations_dir(&["0001_a.sql"]); - let target = TempDir::new().expect("tempdir"); - let target_file = target.path().join("real.sql"); - fs::write(&target_file, "SELECT 2").expect("write fixture"); - std::os::unix::fs::symlink(&target_file, dir.path().join("0002_b.sql")).expect("symlink"); - assert!(matches!( - resolve(dir.path()), - Err(Error::InvalidName { .. }) - )); - } - - #[rstest] - fn rejects_duplicate_versions() { - let dir = migrations_dir(&["0001_a.sql", "1_b.sql"]); - assert!(matches!( - resolve(dir.path()), - Err(Error::DuplicateVersion { version: 1 }) - )); - } - - #[rstest] - fn rejects_empty_directory() { - let dir = migrations_dir(&[]); - assert!(matches!(resolve(dir.path()), Err(Error::Empty { .. }))); - } -} diff --git a/litellm-rust/crates/migrate/Cargo.toml b/litellm-rust/crates/migrate/Cargo.toml deleted file mode 100644 index bb1ecaa3128..00000000000 --- a/litellm-rust/crates/migrate/Cargo.toml +++ /dev/null @@ -1,12 +0,0 @@ -[package] -name = "litellm-migrate" -version = "0.1.0" -edition.workspace = true -license.workspace = true -repository.workspace = true - -[dependencies] -litellm-migrate-macros.workspace = true - -[dev-dependencies] -rstest.workspace = true diff --git a/litellm-rust/crates/migrate/README.md b/litellm-rust/crates/migrate/README.md deleted file mode 100644 index 4817029451c..00000000000 --- a/litellm-rust/crates/migrate/README.md +++ /dev/null @@ -1,5 +0,0 @@ -# Migrations - -`litellm-migrate` exports the `Migration` struct and the `migrate!` macro that embeds a directory of `_.sql` files at compile time, sorted by numeric version - -The crate does not apply or track migrations; callers decide how and when the embedded SQL runs diff --git a/litellm-rust/crates/migrate/src/lib.rs b/litellm-rust/crates/migrate/src/lib.rs deleted file mode 100644 index f4e065e1b53..00000000000 --- a/litellm-rust/crates/migrate/src/lib.rs +++ /dev/null @@ -1,8 +0,0 @@ -pub use litellm_migrate_macros::migrate; - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct Migration { - pub version: u64, - pub description: &'static str, - pub sql: &'static str, -} diff --git a/litellm-rust/crates/migrate/tests/fixtures/migrations/10_tenth.sql b/litellm-rust/crates/migrate/tests/fixtures/migrations/10_tenth.sql deleted file mode 100644 index 31807719e9c..00000000000 --- a/litellm-rust/crates/migrate/tests/fixtures/migrations/10_tenth.sql +++ /dev/null @@ -1 +0,0 @@ -SELECT 10; diff --git a/litellm-rust/crates/migrate/tests/fixtures/migrations/1_first.sql b/litellm-rust/crates/migrate/tests/fixtures/migrations/1_first.sql deleted file mode 100644 index e0ac49d1ecf..00000000000 --- a/litellm-rust/crates/migrate/tests/fixtures/migrations/1_first.sql +++ /dev/null @@ -1 +0,0 @@ -SELECT 1; diff --git a/litellm-rust/crates/migrate/tests/fixtures/migrations/2_second.sql b/litellm-rust/crates/migrate/tests/fixtures/migrations/2_second.sql deleted file mode 100644 index e7f8100648d..00000000000 --- a/litellm-rust/crates/migrate/tests/fixtures/migrations/2_second.sql +++ /dev/null @@ -1 +0,0 @@ -SELECT 2; diff --git a/litellm-rust/crates/migrate/tests/migrate.rs b/litellm-rust/crates/migrate/tests/migrate.rs deleted file mode 100644 index 61c80351cf4..00000000000 --- a/litellm-rust/crates/migrate/tests/migrate.rs +++ /dev/null @@ -1,21 +0,0 @@ -use litellm_migrate::Migration; -use rstest::rstest; - -const MIGRATIONS: &[Migration] = litellm_migrate::migrate!("tests/fixtures/migrations"); - -#[rstest] -#[case::first(0, 1, "first", include_str!("fixtures/migrations/1_first.sql"))] -#[case::second(1, 2, "second", include_str!("fixtures/migrations/2_second.sql"))] -#[case::tenth(2, 10, "tenth", include_str!("fixtures/migrations/10_tenth.sql"))] -fn embeds_every_file_sorted_by_numeric_version( - #[case] index: usize, - #[case] version: u64, - #[case] description: &str, - #[case] sql: &str, -) { - assert_eq!(MIGRATIONS.len(), 3); - let migration = &MIGRATIONS[index]; - assert_eq!(migration.version, version); - assert_eq!(migration.description, description); - assert_eq!(migration.sql, sql); -} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 226cecd5b55..7458e64d374 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -74,6 +74,7 @@ serde_with.workspace = true criterion.workspace = true futures-util.workspace = true rstest.workspace = true +sqlx = { workspace = true, features = ["migrate"] } sha2.workspace = true tokio-tungstenite.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs index f442d6c31de..65db1439d13 100644 --- a/litellm-rust/crates/python-bridge/src/routes/traces.rs +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -52,8 +52,7 @@ fn map_error_ref(error: &Error) -> PyErr { | Error::InvalidParameters | Error::InvalidScope => PyValueError::new_err(error.to_string()), Error::Task - | Error::SchemaFailed(_) - | Error::SchemaTransport + | Error::Migration(_) | Error::MissingSecret | Error::Busy | Error::ProvisionFailed(_) @@ -452,7 +451,14 @@ mod tests { )] #[case::insert_budget(Error::InsertTooLarge, "OverflowError")] #[case::scope(Error::InvalidScope, "ValueError")] - #[case::schema(Error::SchemaFailed(503), "RuntimeError")] + #[case::schema( + Error::Storage(litellm_storage_clickhouse::Error::SchemaFailed(503)), + "RuntimeError" + )] + #[case::migration( + Error::Migration(sqlx::migrate::MigrateError::VersionMismatch(1)), + "RuntimeError" + )] #[case::reader(Error::MissingSecret, "RuntimeError")] #[case::storage( Error::Storage(litellm_storage_clickhouse::Error::InvalidUrl), diff --git a/litellm-rust/crates/storage-clickhouse/AGENTS.md b/litellm-rust/crates/storage-clickhouse/AGENTS.md index 959ffdffb88..3f87e5d0759 100644 --- a/litellm-rust/crates/storage-clickhouse/AGENTS.md +++ b/litellm-rust/crates/storage-clickhouse/AGENTS.md @@ -3,3 +3,7 @@ `litellm-storage-clickhouse` exports `Storage`, a writer and bounded reader derived from one ClickHouse URL and database. It also exports bounded HTTP read and insert execution The crate has no trace tables, OTLP types, or named trace queries. `litellm-traces-clickhouse` supplies those rules and uses this storage for both trace rows and spend rows + +It also applies embedded SQLx migrations through the `_sqlx_migrations` ledger + +Migration files are append-only, and changed applied files are rejected by their checksums. Startup migrations must be replay-safe schema changes because the runner records success after execution without dirty states or locks. Backfills belong in coordinated jobs outside proxy startup. The replay policy lives in `ClickHouseMigrate::apply`, `dirty_version`, and `lock`; a Keeper-backed or deploy-time runner changes only those methods diff --git a/litellm-rust/crates/storage-clickhouse/Cargo.toml b/litellm-rust/crates/storage-clickhouse/Cargo.toml index acce941f2c9..4e72e823dc2 100644 --- a/litellm-rust/crates/storage-clickhouse/Cargo.toml +++ b/litellm-rust/crates/storage-clickhouse/Cargo.toml @@ -11,11 +11,14 @@ flate2.workspace = true litellm-http.workspace = true serde.workspace = true serde_json.workspace = true +serde_with.workspace = true +sqlx = { workspace = true, features = ["migrate"] } thiserror.workspace = true url.workspace = true [dev-dependencies] litellm-http = { workspace = true, features = ["test-support"] } rstest.workspace = true +testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] } tokio.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/storage-clickhouse/src/error.rs b/litellm-rust/crates/storage-clickhouse/src/error.rs index acec9b91675..fa0c5eabd73 100644 --- a/litellm-rust/crates/storage-clickhouse/src/error.rs +++ b/litellm-rust/crates/storage-clickhouse/src/error.rs @@ -1,4 +1,4 @@ -#[derive(Debug, thiserror::Error)] +#[derive(Clone, Debug, thiserror::Error)] pub enum Error { #[error("invalid ClickHouse insert row")] InvalidRow, diff --git a/litellm-rust/crates/storage-clickhouse/src/lib.rs b/litellm-rust/crates/storage-clickhouse/src/lib.rs index 7ab2aa9bc0a..e2ef62d1bc3 100644 --- a/litellm-rust/crates/storage-clickhouse/src/lib.rs +++ b/litellm-rust/crates/storage-clickhouse/src/lib.rs @@ -1,9 +1,11 @@ mod error; mod insert; +mod migrate; mod read; pub use error::Error; pub use insert::{insert_compressed_rows, insert_encoded_rows}; +pub use migrate::{ClickHouseMigrate, execute_statement, storage_error}; pub use read::{Parameter, Query, READ_LIMITS, ReadLimits, execute_read, fetch, fetch_json}; use url::Url; diff --git a/litellm-rust/crates/storage-clickhouse/src/migrate.rs b/litellm-rust/crates/storage-clickhouse/src/migrate.rs new file mode 100644 index 00000000000..c249020faaf --- /dev/null +++ b/litellm-rust/crates/storage-clickhouse/src/migrate.rs @@ -0,0 +1,308 @@ +use std::{ + future::Future, + pin::Pin, + time::{Duration, Instant}, +}; + +use litellm_http::Client; +use serde::Deserialize; +use serde_with::{DisplayFromStr, PickFirst, serde_as}; +use sqlx::{ + Error as SqlxError, + migrate::{AppliedMigration, Migrate, MigrateError, Migration}, +}; + +use crate::{Connection, Error, READ_LIMITS, valid_identifier}; + +pub async fn execute_statement( + client: &Client, + connection: &Connection, + sql: &str, + timeout: Duration, +) -> Result<(), Error> { + let body = execute_sql(client, connection, sql, timeout).await?; + if !body.trim().is_empty() { + return Err(Error::InvalidResponse); + } + Ok(()) +} + +async fn execute_sql( + client: &Client, + connection: &Connection, + sql: &str, + timeout: Duration, +) -> Result { + let mut url = connection.url().clone(); + let pairs: Vec<_> = url + .query_pairs() + .filter(|(key, _)| { + !matches!( + key.as_ref(), + "query" | "wait_end_of_query" | "send_progress_in_http_headers" | "async_insert" + ) + }) + .map(|(key, value)| (key.into_owned(), value.into_owned())) + .collect(); + url.query_pairs_mut() + .clear() + .extend_pairs(pairs) + .append_pair("wait_end_of_query", "1") + .append_pair("send_progress_in_http_headers", "0") + .append_pair("async_insert", "0"); + let mut response = client + .post(url) + .timeout(timeout) + .body(sql.to_owned()) + .send() + .await + .map_err(|_| Error::Transport)?; + if !response.status().is_success() { + return Err(Error::SchemaFailed(response.status().as_u16())); + } + let mut body = Vec::new(); + while let Some(chunk) = response.chunk().await.map_err(|_| Error::Transport)? { + if body.len() + chunk.len() > READ_LIMITS.response_bytes { + return Err(Error::ResponseTooLarge); + } + body.extend_from_slice(&chunk); + } + String::from_utf8(body).map_err(|_| Error::InvalidResponse) +} + +/// The startup runner records success after execution, never marks migrations dirty, and skips locks +/// A Keeper-backed or deploy-time runner changes only `apply`, `dirty_version`, and `lock` +pub struct ClickHouseMigrate<'a, R> { + client: &'a Client, + connection: &'a Connection, + database: &'a str, + render: R, + timeout: Duration, +} + +impl<'a, R> ClickHouseMigrate<'a, R> +where + R: Fn(&str) -> String + Send + Sync, +{ + pub fn new( + client: &'a Client, + connection: &'a Connection, + database: &'a str, + render: R, + timeout: Duration, + ) -> Result { + if !valid_identifier(database) { + return Err(Error::InvalidSchema); + } + Ok(Self { + client, + connection, + database, + render, + timeout, + }) + } +} + +#[serde_as] +#[derive(Deserialize)] +struct Applied { + #[serde_as(as = "PickFirst<(_, DisplayFromStr)>")] + version: i64, + checksum: String, +} + +fn migrate_error(error: Error) -> MigrateError { + MigrateError::Execute(SqlxError::AnyDriverError(Box::new(error))) +} + +fn migrate_execution_error(error: Error, version: i64) -> MigrateError { + MigrateError::ExecuteMigration(SqlxError::AnyDriverError(Box::new(error)), version) +} + +fn encode_hex(bytes: &[u8]) -> String { + bytes + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::>() + .join("") +} + +fn decode_hex(value: &str) -> Result, Error> { + let (pairs, remainder) = value.as_bytes().as_chunks::<2>(); + if !remainder.is_empty() { + return Err(Error::InvalidResponse); + } + pairs + .iter() + .map(|pair| { + let high = decode_hex_digit(pair[0]).ok_or(Error::InvalidResponse)?; + let low = decode_hex_digit(pair[1]).ok_or(Error::InvalidResponse)?; + Ok((high << 4) | low) + }) + .collect() +} + +fn decode_hex_digit(value: u8) -> Option { + match value { + b'0'..=b'9' => Some(value - b'0'), + b'a'..=b'f' => Some(value - b'a' + 10), + b'A'..=b'F' => Some(value - b'A' + 10), + _ => None, + } +} + +fn escape_sql_string(value: &str) -> String { + value.replace('\\', "\\\\").replace('\'', "\\'") +} + +type MigrateFuture<'e, T> = Pin + Send + 'e>>; + +impl Migrate for ClickHouseMigrate<'_, R> +where + R: Fn(&str) -> String + Send + Sync, +{ + fn create_schema_if_not_exists<'e>( + &'e mut self, + schema_name: &'e str, + ) -> MigrateFuture<'e, Result<(), MigrateError>> { + Box::pin(async move { + if !valid_identifier(schema_name) { + return Err(migrate_error(Error::InvalidSchema)); + } + let statement = format!("CREATE DATABASE IF NOT EXISTS `{schema_name}`"); + execute_statement(self.client, self.connection, &statement, self.timeout) + .await + .map_err(migrate_error) + }) + } + + fn ensure_migrations_table<'e>( + &'e mut self, + table_name: &'e str, + ) -> MigrateFuture<'e, Result<(), MigrateError>> { + Box::pin(async move { + let database = format!("`{}`", self.database); + execute_statement( + self.client, + self.connection, + &format!("CREATE DATABASE IF NOT EXISTS {database}"), + self.timeout, + ) + .await + .map_err(migrate_error)?; + execute_statement( + self.client, + self.connection, + &format!( + "CREATE TABLE IF NOT EXISTS {database}.{table_name} \ + (version Int64, description String, installed_on DateTime64(3) DEFAULT now64(3), \ + success Bool, checksum String, execution_time Int64) ENGINE = MergeTree ORDER BY version" + ), + self.timeout, + ) + .await + .map_err(migrate_error) + }) + } + + fn dirty_version<'e>( + &'e mut self, + _table_name: &'e str, + ) -> MigrateFuture<'e, Result, MigrateError>> { + Box::pin(async { Ok(None) }) + } + + fn list_applied_migrations<'e>( + &'e mut self, + table_name: &'e str, + ) -> MigrateFuture<'e, Result, MigrateError>> { + Box::pin(async move { + let database = format!("`{}`", self.database); + let statement = format!( + "SELECT DISTINCT version, checksum FROM {database}.{table_name} \ + WHERE success ORDER BY version FORMAT JSONEachRow" + ); + let response = execute_sql(self.client, self.connection, &statement, self.timeout) + .await + .map_err(migrate_error)?; + response + .lines() + .filter(|line| !line.trim().is_empty()) + .map(|line| { + let row = serde_json::from_str::(line) + .map_err(|_| migrate_error(Error::InvalidResponse))?; + let checksum = decode_hex(&row.checksum).map_err(migrate_error)?; + Ok(AppliedMigration { + version: row.version, + checksum: checksum.into(), + }) + }) + .collect() + }) + } + + fn lock(&mut self) -> MigrateFuture<'_, Result<(), MigrateError>> { + Box::pin(async { Ok(()) }) + } + + fn unlock(&mut self) -> MigrateFuture<'_, Result<(), MigrateError>> { + Box::pin(async { Ok(()) }) + } + + fn apply<'e>( + &'e mut self, + table_name: &'e str, + migration: &'e Migration, + ) -> MigrateFuture<'e, Result> { + Box::pin(async move { + let started_at = Instant::now(); + let statement = (self.render)(migration.sql.as_str()); + execute_statement(self.client, self.connection, &statement, self.timeout) + .await + .map_err(|error| migrate_execution_error(error, migration.version))?; + let elapsed = started_at.elapsed(); + let execution_time = elapsed.as_nanos().min(i64::MAX as u128) as i64; + let description = escape_sql_string(&migration.description); + let checksum = encode_hex(&migration.checksum); + let database = format!("`{}`", self.database); + execute_statement( + self.client, + self.connection, + &format!( + "INSERT INTO {database}.{table_name} \ + (version, description, success, checksum, execution_time) \ + VALUES ({}, '{}', true, '{}', {execution_time})", + migration.version, description, checksum + ), + self.timeout, + ) + .await + .map_err(migrate_error)?; + Ok(elapsed) + }) + } + + fn revert<'e>( + &'e mut self, + _table_name: &'e str, + _migration: &'e Migration, + ) -> MigrateFuture<'e, Result> { + Box::pin(async { + Err(MigrateError::Execute(SqlxError::AnyDriverError(Box::new( + std::io::Error::other("ClickHouse migrations are forward-only"), + )))) + }) + } +} + +pub fn storage_error(error: &MigrateError) -> Option<&Error> { + let error = match error { + MigrateError::Execute(error) | MigrateError::ExecuteMigration(error, _) => error, + _ => return None, + }; + match error { + SqlxError::AnyDriverError(error) => error.downcast_ref(), + _ => None, + } +} diff --git a/litellm-rust/crates/storage-clickhouse/tests/migrations.rs b/litellm-rust/crates/storage-clickhouse/tests/migrations.rs new file mode 100644 index 00000000000..29a85a8c4ec --- /dev/null +++ b/litellm-rust/crates/storage-clickhouse/tests/migrations.rs @@ -0,0 +1,476 @@ +use std::time::Duration; + +use litellm_http::Client; +use litellm_storage_clickhouse::{ + ClickHouseMigrate, Connection, Error, READ_LIMITS, execute_statement, storage_error, +}; +use rstest::{fixture, rstest}; +use sqlx::{ + SqlStr, + migrate::{Migrate, MigrateError, Migration, MigrationType, Migrator}, +}; +use testcontainers_modules::{ + clickhouse::ClickHouse, + testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner}, +}; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_string, method, query_param}, +}; + +const CLICKHOUSE_TAG: &str = + "26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e"; +const DATABASE: &str = "storage_migrate_test"; +const REQUEST_TIMEOUT: Duration = Duration::from_secs(10); +const SELECT_APPLIED: &str = "SELECT DISTINCT version, checksum FROM `trace_test`._sqlx_migrations \ + WHERE success ORDER BY version FORMAT JSONEachRow"; + +type TestResult = Result>; + +struct ClickHouseDatabase { + _container: ContainerAsync, + url: String, + client: Client, +} + +#[fixture] +async fn database() -> TestResult { + let container = ClickHouse::default() + .with_tag(CLICKHOUSE_TAG) + .with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1") + .start() + .await?; + let url = format!( + "http://{}:{}", + container.get_host().await?, + container.get_host_port_ipv4(8123).await? + ); + Ok(ClickHouseDatabase { + _container: container, + url, + client: Client::no_redirect_for_test(), + }) +} + +#[fixture] +async fn mock_server() -> MockServer { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .with_priority(10) + .mount(&server) + .await; + server +} + +fn migration(version: i64, sql: &'static str) -> Migration { + Migration::new( + version, + format!("migration_{version}").into(), + MigrationType::Simple, + SqlStr::from_static(sql), + false, + ) +} + +fn migrator(migrations: Vec) -> Migrator { + Migrator { + ignore_missing: true, + locking: false, + ..Migrator::with_migrations(migrations) + } +} + +fn render_database(sql: &str) -> String { + sql.replace("{database}", &format!("`{DATABASE}`")) +} + +async fn run_migrations( + database: &ClickHouseDatabase, + migrator: &Migrator, + schema: &str, + render: R, +) -> Result<(), MigrateError> +where + R: Fn(&str) -> String + Send + Sync, +{ + let connection = Connection::writer(&database.url).expect("valid ClickHouse URL"); + let mut adapter = ClickHouseMigrate::new( + &database.client, + &connection, + schema, + render, + REQUEST_TIMEOUT, + ) + .expect("valid schema"); + migrator.run_direct(None, &mut adapter, false).await +} + +async fn execute_write(database: &ClickHouseDatabase, sql: &str) -> TestResult { + database + .client + .post(&database.url) + .body(sql.to_owned()) + .send() + .await? + .error_for_status()?; + Ok(()) +} + +async fn read_json(database: &ClickHouseDatabase, sql: &str) -> TestResult { + let response = database + .client + .post(&database.url) + .body(sql.to_owned()) + .send() + .await? + .error_for_status()?; + Ok(serde_json::from_str(&response.text().await?)?) +} + +async fn ledger_versions(database: &ClickHouseDatabase) -> TestResult> { + let response = read_json( + database, + &format!( + "SELECT version FROM `{DATABASE}`._sqlx_migrations \ + GROUP BY version ORDER BY version FORMAT JSON" + ), + ) + .await?; + Ok(response["data"] + .as_array() + .expect("ClickHouse returns versions") + .iter() + .map(|row| row["version"].as_i64().expect("version is Int64")) + .collect()) +} + +fn encode_hex(bytes: &[u8]) -> String { + bytes + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::>() + .join("") +} + +#[rstest] +#[tokio::test] +async fn only_pending_migrations_execute_on_the_second_run( + #[future(awt)] database: TestResult, +) -> TestResult { + let database = database?; + let migrator = migrator(vec![ + migration( + 1, + "CREATE TABLE IF NOT EXISTS {database}.migration_one (id UInt8) ENGINE = MergeTree ORDER BY id", + ), + migration( + 2, + "CREATE TABLE IF NOT EXISTS {database}.migration_two (id UInt8) ENGINE = MergeTree ORDER BY id", + ), + ]); + run_migrations(&database, &migrator, DATABASE, render_database).await?; + run_migrations(&database, &migrator, DATABASE, render_database).await?; + execute_statement( + &database.client, + &Connection::writer(&database.url)?, + "SYSTEM FLUSH LOGS", + REQUEST_TIMEOUT, + ) + .await?; + + let queries = read_json( + &database, + "SELECT count() AS executions FROM system.query_log \ + WHERE type = 'QueryFinish' AND query LIKE \ + 'CREATE TABLE IF NOT EXISTS `storage_migrate_test`.migration_%' FORMAT JSON", + ) + .await?; + assert_eq!(queries["data"][0]["executions"].as_u64(), Some(2)); + assert_eq!(ledger_versions(&database).await?, vec![1, 2]); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn edited_migration_checksum_returns_version_mismatch( + #[future(awt)] database: TestResult, +) -> TestResult { + let database = database?; + let original = migrator(vec![migration( + 1, + "CREATE TABLE IF NOT EXISTS {database}.original (id UInt8) ENGINE = MergeTree ORDER BY id", + )]); + let changed = migrator(vec![migration( + 1, + "CREATE TABLE IF NOT EXISTS {database}.changed (id UInt8) ENGINE = MergeTree ORDER BY id", + )]); + run_migrations(&database, &original, DATABASE, render_database).await?; + + assert!(matches!( + run_migrations(&database, &changed, DATABASE, render_database).await, + Err(MigrateError::VersionMismatch(1)) + )); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn failed_migration_is_not_recorded_and_retains_storage_error( + #[future(awt)] database: TestResult, +) -> TestResult { + let database = database?; + let migrator = migrator(vec![migration(1, "THIS IS NOT VALID CLICKHOUSE SQL")]); + let error = run_migrations(&database, &migrator, DATABASE, render_database) + .await + .expect_err("invalid SQL must fail"); + assert!(matches!(&error, MigrateError::ExecuteMigration(_, 1))); + assert!(matches!( + storage_error(&error), + Some(Error::SchemaFailed(_)) + )); + + let rows = read_json( + &database, + &format!( + "SELECT count() AS rows FROM `{DATABASE}`._sqlx_migrations \ + WHERE version = 1 FORMAT JSON" + ), + ) + .await?; + assert_eq!(rows["data"][0]["rows"].as_u64(), Some(0)); + Ok(()) +} + +#[rstest] +fn invalid_database_identifier_is_rejected() { + let client = Client::no_redirect_for_test(); + let connection = Connection::writer("http://127.0.0.1:1").expect("valid URL"); + assert!(matches!( + ClickHouseMigrate::new( + &client, + &connection, + "storage_test; DROP DATABASE default", + str::to_owned, + REQUEST_TIMEOUT, + ), + Err(Error::InvalidSchema) + )); +} + +#[rstest] +#[tokio::test] +async fn unknown_source_version_is_tolerated( + #[future(awt)] database: TestResult, +) -> TestResult { + let database = database?; + run_migrations(&database, &migrator(vec![]), DATABASE, render_database).await?; + execute_write( + &database, + &format!( + "INSERT INTO `{DATABASE}`._sqlx_migrations \ + (version, description, success, checksum, execution_time) \ + VALUES (99, 'unknown', true, '{}', 0)", + "00".repeat(48) + ), + ) + .await?; + let migrator = migrator(vec![migration( + 1, + "CREATE TABLE IF NOT EXISTS {database}.known (id UInt8) ENGINE = MergeTree ORDER BY id", + )]); + + run_migrations(&database, &migrator, DATABASE, render_database).await?; + + assert_eq!(ledger_versions(&database).await?, vec![1, 99]); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn duplicate_ledger_rows_are_tolerated( + #[future(awt)] database: TestResult, +) -> TestResult { + let database = database?; + let applied = migration( + 1, + "CREATE TABLE IF NOT EXISTS {database}.duplicate_test (id UInt8) ENGINE = MergeTree ORDER BY id", + ); + let checksum = encode_hex(&applied.checksum); + let migrator = migrator(vec![applied]); + run_migrations(&database, &migrator, DATABASE, render_database).await?; + execute_write( + &database, + &format!( + "INSERT INTO `{DATABASE}`._sqlx_migrations \ + (version, description, success, checksum, execution_time) \ + VALUES (1, 'migration_1', true, '{checksum}', 0)" + ), + ) + .await?; + + run_migrations(&database, &migrator, DATABASE, render_database).await?; + + let rows = read_json( + &database, + &format!( + "SELECT count() AS rows, uniqExact(version) AS versions \ + FROM `{DATABASE}`._sqlx_migrations FORMAT JSON" + ), + ) + .await?; + assert_eq!(rows["data"][0]["rows"].as_u64(), Some(2)); + assert_eq!(rows["data"][0]["versions"].as_u64(), Some(1)); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn revert_reports_forward_only_error() { + let client = Client::no_redirect_for_test(); + let connection = Connection::writer("http://127.0.0.1:1").expect("valid URL"); + let migration = migration( + 1, + "CREATE TABLE IF NOT EXISTS {database}.revert_test (id UInt8) ENGINE = MergeTree ORDER BY id", + ); + let mut adapter = ClickHouseMigrate::new( + &client, + &connection, + DATABASE, + render_database, + REQUEST_TIMEOUT, + ) + .expect("valid schema"); + + let error = adapter + .revert("_sqlx_migrations", &migration) + .await + .expect_err("ClickHouse migrations cannot be reverted"); + assert!( + error + .to_string() + .contains("ClickHouse migrations are forward-only") + ); +} + +#[rstest] +#[case::numeric("1")] +#[case::quoted("\"1\"")] +#[tokio::test] +async fn applied_int64_versions_accept_numeric_and_quoted_json( + #[future(awt)] mock_server: MockServer, + #[case] version: &str, +) { + Mock::given(method("POST")) + .and(body_string(SELECT_APPLIED)) + .respond_with( + ResponseTemplate::new(200) + .set_body_string(format!("{{\"version\":{version},\"checksum\":\"00\"}}\n")), + ) + .mount(&mock_server) + .await; + let client = Client::no_redirect_for_test(); + let connection = Connection::writer(&mock_server.uri()).expect("valid URL"); + let migrator = migrator(vec![]); + let mut adapter = ClickHouseMigrate::new( + &client, + &connection, + "trace_test", + str::to_owned, + REQUEST_TIMEOUT, + ) + .expect("valid schema"); + + migrator + .run_direct(None, &mut adapter, false) + .await + .expect("applied version parses"); +} + +#[rstest] +#[tokio::test] +async fn schema_requests_override_unsafe_connection_settings( + #[future(awt)] mock_server: MockServer, +) { + Mock::given(method("POST")) + .and(query_param("wait_end_of_query", "1")) + .and(query_param("send_progress_in_http_headers", "0")) + .and(query_param("async_insert", "0")) + .and(query_param("custom_setting", "preserved")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&mock_server) + .await; + let url = format!( + "{}/?wait_end_of_query=0&send_progress_in_http_headers=1&async_insert=1&custom_setting=preserved", + mock_server.uri() + ); + execute_statement( + &Client::no_redirect_for_test(), + &Connection::writer(&url).expect("valid URL"), + "CREATE DATABASE IF NOT EXISTS trace_test", + REQUEST_TIMEOUT, + ) + .await + .expect("schema execution succeeds"); + let requests = mock_server + .received_requests() + .await + .expect("requests recorded"); + for name in [ + "wait_end_of_query", + "send_progress_in_http_headers", + "async_insert", + ] { + assert_eq!( + requests[0] + .url + .query_pairs() + .filter(|(key, _)| key == name) + .count(), + 1 + ); + } +} + +#[rstest] +#[tokio::test] +async fn oversized_ledger_response_is_rejected_before_migrations( + #[future(awt)] mock_server: MockServer, +) { + Mock::given(method("POST")) + .and(body_string(SELECT_APPLIED)) + .respond_with( + ResponseTemplate::new(200).set_body_string(" ".repeat(READ_LIMITS.response_bytes + 1)), + ) + .mount(&mock_server) + .await; + let client = Client::no_redirect_for_test(); + let connection = Connection::writer(&mock_server.uri()).expect("valid URL"); + let migrator = migrator(vec![]); + let mut adapter = ClickHouseMigrate::new( + &client, + &connection, + "trace_test", + str::to_owned, + REQUEST_TIMEOUT, + ) + .expect("valid schema"); + let error = migrator + .run_direct(None, &mut adapter, false) + .await + .expect_err("oversized result is rejected"); + + assert!(matches!( + storage_error(&error), + Some(Error::ResponseTooLarge) + )); + assert_eq!( + mock_server + .received_requests() + .await + .expect("requests recorded") + .len(), + 3 + ); +} diff --git a/litellm-rust/crates/traces-clickhouse/AGENTS.md b/litellm-rust/crates/traces-clickhouse/AGENTS.md index ae0c1c8eeb9..462100b10a0 100644 --- a/litellm-rust/crates/traces-clickhouse/AGENTS.md +++ b/litellm-rust/crates/traces-clickhouse/AGENTS.md @@ -1,6 +1,7 @@ - Own trace schema, row encoding, SQL query adapters and reader provisioning; consume domain types from `litellm-traces` - Keep generic ClickHouse connections and HTTP execution in `litellm-storage-clickhouse`; keep PyO3 conversion in `python-bridge` -- Keep schema definitions only in `migrations/NNNN_description.sql`, embedded by `litellm_migrate::migrate!` +- Keep schema definitions only in `migrations/NNNN_description.sql`, embedded by `sqlx::migrate!` +- Treat retention TTLs as current configuration: change them in `RETENTION` in `src/schema.rs`, which every startup reapplies, never in a new migration - Require typed query parameters and SELECT-only readers with server-side limits and tenant isolation - Bound insert time and encoded bytes; preserve shared values and explicit retry deduplication - Test storage behavior through the public API against ClickHouse diff --git a/litellm-rust/crates/traces-clickhouse/Cargo.toml b/litellm-rust/crates/traces-clickhouse/Cargo.toml index b90ad8a7bf8..13aa6d06af6 100644 --- a/litellm-rust/crates/traces-clickhouse/Cargo.toml +++ b/litellm-rust/crates/traces-clickhouse/Cargo.toml @@ -16,7 +16,6 @@ flate2.workspace = true futures-util.workspace = true hmac = "0.12.1" litellm-http.workspace = true -litellm-migrate.workspace = true litellm-storage-clickhouse.workspace = true litellm-traces.workspace = true litellm-traces-cache.workspace = true @@ -24,6 +23,7 @@ moka.workspace = true serde.workspace = true serde_json.workspace = true sha2.workspace = true +sqlx = { workspace = true, features = ["migrate", "macros"] } strum.workspace = true thiserror.workspace = true time = { workspace = true, features = ["formatting"] } diff --git a/litellm-rust/crates/traces-clickhouse/src/error.rs b/litellm-rust/crates/traces-clickhouse/src/error.rs index 1d15c556316..e0ef127a0e9 100644 --- a/litellm-rust/crates/traces-clickhouse/src/error.rs +++ b/litellm-rust/crates/traces-clickhouse/src/error.rs @@ -16,10 +16,6 @@ pub enum Error { InvalidResponse, #[error("ClickHouse insert exceeds the encoded size limit")] InsertTooLarge, - #[error("ClickHouse schema setup failed with HTTP status {0}")] - SchemaFailed(u16), - #[error("ClickHouse schema setup transport failed")] - SchemaTransport, #[error("trace SQL queries require a configured proxy master key")] MissingSecret, #[error("invalid trace query scope")] @@ -39,5 +35,7 @@ pub enum Error { #[error(transparent)] Storage(#[from] litellm_storage_clickhouse::Error), #[error(transparent)] + Migration(#[from] sqlx::migrate::MigrateError), + #[error(transparent)] Cached(#[from] std::sync::Arc), } diff --git a/litellm-rust/crates/traces-clickhouse/src/lib.rs b/litellm-rust/crates/traces-clickhouse/src/lib.rs index 43ff0b8bf33..aed87691958 100644 --- a/litellm-rust/crates/traces-clickhouse/src/lib.rs +++ b/litellm-rust/crates/traces-clickhouse/src/lib.rs @@ -33,7 +33,8 @@ pub use query::{QueryHelp, execute_read, query_help, query_sql}; pub use query_access::QueryReaders; pub use reads::ClickHouseTraces; pub use schema::{ - NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, ensure_schema, schema_statements, + NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, apply_migrations, ensure_schema, + reconcile_retention, schema_statements, }; pub use span_row::span_rows; pub use sql::execute_named_read; diff --git a/litellm-rust/crates/traces-clickhouse/src/schema.rs b/litellm-rust/crates/traces-clickhouse/src/schema.rs index 6a39bd24041..562dbb976c2 100644 --- a/litellm-rust/crates/traces-clickhouse/src/schema.rs +++ b/litellm-rust/crates/traces-clickhouse/src/schema.rs @@ -1,15 +1,26 @@ use litellm_http::Client; -use litellm_migrate::Migration; +use litellm_storage_clickhouse::{ClickHouseMigrate, execute_statement, storage_error}; use serde::Serialize; +use sqlx::migrate::Migrator; use std::time::Duration; use super::{Connection, Error}; const SCHEMA_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); -const MIGRATIONS: &[Migration] = litellm_migrate::migrate!("migrations"); +static MIGRATOR: Migrator = Migrator { + ignore_missing: true, + locking: false, + ..sqlx::migrate!("./migrations") +}; -pub fn schema_statements(database: &str, retention_days: u32) -> Result, Error> { +const RETENTION: [(&str, &str); 3] = [ + ("otel_traces", "toDateTime(Timestamp)"), + ("agent_traces_by_key", "toDateTime(StartTs)"), + ("spend_logs", "toDateTime(start_time)"), +]; + +fn validate_schema(database: &str, retention_days: u32) -> Result<(), Error> { if database.is_empty() || !database .bytes() @@ -18,19 +29,109 @@ pub fn schema_statements(database: &str, retention_days: u32) -> Result String { + sql.replace("{database}", database) + .replace("{retention_days}", &retention_days.to_string()) +} + +pub fn schema_statements(database: &str, retention_days: u32) -> Result, Error> { + validate_schema(database, retention_days)?; let database = format!("`{database}`"); Ok( std::iter::once(format!("CREATE DATABASE IF NOT EXISTS {database}")) - .chain(MIGRATIONS.iter().map(|migration| { - migration - .sql - .replace("{database}", &database) - .replace("{retention_days}", &retention_days.to_string()) - })) + .chain( + MIGRATOR + .migrations + .iter() + .map(|migration| render(migration.sql.as_str(), &database, retention_days)), + ) .collect(), ) } +pub async fn apply_migrations( + client: &Client, + connection: &Connection, + database: &str, + retention_days: u32, +) -> Result<(), Error> { + apply_migrations_with_timeout( + client, + connection, + database, + retention_days, + SCHEMA_REQUEST_TIMEOUT, + ) + .await +} + +async fn apply_migrations_with_timeout( + client: &Client, + connection: &Connection, + database: &str, + retention_days: u32, + request_timeout: Duration, +) -> Result<(), Error> { + validate_schema(database, retention_days)?; + let quoted_database = format!("`{database}`"); + let mut adapter = ClickHouseMigrate::new( + client, + connection, + database, + |sql| render(sql, "ed_database, retention_days), + request_timeout, + )?; + MIGRATOR + .run_direct(None, &mut adapter, false) + .await + .map_err(|error| match storage_error(&error) { + Some(storage_error) => Error::Storage(storage_error.clone()), + None => Error::Migration(error), + }) +} + +pub async fn reconcile_retention( + client: &Client, + connection: &Connection, + database: &str, + retention_days: u32, +) -> Result<(), Error> { + reconcile_retention_with_timeout( + client, + connection, + database, + retention_days, + SCHEMA_REQUEST_TIMEOUT, + ) + .await +} + +async fn reconcile_retention_with_timeout( + client: &Client, + connection: &Connection, + database: &str, + retention_days: u32, + request_timeout: Duration, +) -> Result<(), Error> { + validate_schema(database, retention_days)?; + let database = format!("`{database}`"); + for (table, expression) in RETENTION { + execute_statement( + client, + connection, + &format!( + "ALTER TABLE {database}.{table} MODIFY TTL {expression} + INTERVAL {retention_days} DAY" + ), + request_timeout, + ) + .await?; + } + Ok(()) +} + pub async fn ensure_schema( client: &Client, connection: &Connection, @@ -54,19 +155,22 @@ async fn ensure_schema_with_timeout( retention_days: u32, request_timeout: Duration, ) -> Result<(), Error> { - for statement in schema_statements(database, retention_days)? { - let response = client - .post(connection.url().clone()) - .timeout(request_timeout) - .body(statement) - .send() - .await - .map_err(|_| Error::SchemaTransport)?; - if !response.status().is_success() { - return Err(Error::SchemaFailed(response.status().as_u16())); - } - } - Ok(()) + apply_migrations_with_timeout( + client, + connection, + database, + retention_days, + request_timeout, + ) + .await?; + reconcile_retention_with_timeout( + client, + connection, + database, + retention_days, + request_timeout, + ) + .await } #[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)] diff --git a/litellm-rust/crates/traces-clickhouse/tests/migrations.rs b/litellm-rust/crates/traces-clickhouse/tests/migrations.rs index 09fd3e78da8..8976ef208d8 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/migrations.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/migrations.rs @@ -1,11 +1,14 @@ use std::{collections::BTreeMap, time::Duration}; use litellm_http::Client; +use litellm_storage_clickhouse::Error as StorageError; use litellm_traces_clickhouse::{ Connection, Error, InsertTable, NORMALIZED_FIELD_DEFINITIONS, Parameter, ReadQuery, - encode_rows, ensure_schema, execute_named_read, execute_read, schema_statements, + apply_migrations, encode_rows, ensure_schema, execute_named_read, execute_read, + reconcile_retention, schema_statements, }; use rstest::rstest; +use sqlx::migrate::MigrateError; mod support; use support::{ClickHouseDatabase, TestResult, database}; @@ -71,6 +74,44 @@ async fn mutation_rows(database: &ClickHouseDatabase) -> TestResult { .expect("ClickHouse returns mutation counts as unsigned integers")) } +fn migration_versions() -> Vec { + let mut versions = std::fs::read_dir(concat!(env!("CARGO_MANIFEST_DIR"), "/migrations")) + .expect("migration directory exists") + .map(|entry| { + entry + .expect("migration directory entry is readable") + .file_name() + .into_string() + .expect("migration file name is UTF-8") + }) + .filter_map(|name| { + name.strip_suffix(".sql") + .and_then(|stem| stem.split('_').next()) + .and_then(|version| version.parse::().ok()) + }) + .collect::>(); + versions.sort_unstable(); + versions +} + +async fn migration_ledger_versions(database: &ClickHouseDatabase) -> TestResult> { + let response = read_json( + database, + "SELECT version FROM trace_test._sqlx_migrations GROUP BY version ORDER BY version", + ) + .await?; + Ok(response["data"] + .as_array() + .expect("ClickHouse returns version rows") + .iter() + .map(|row| { + row["version"] + .as_u64() + .expect("ClickHouse returns versions as unsigned integers") + }) + .collect()) +} + #[rstest] #[tokio::test] async fn schema_supports_span_rollups_and_spend_joins( @@ -80,6 +121,27 @@ async fn schema_supports_span_rollups_and_spend_joins( let writer = Connection::writer(&database.url)?; ensure_schema(&database.client, &writer, "trace_test", 7).await?; ensure_schema(&database.client, &writer, "trace_test", 7).await?; + let expected_versions = migration_versions(); + let ledger = read_json( + &database, + "SELECT count() AS rows, uniqExact(version) AS versions, \ + countIf(NOT match(checksum, '^[0-9a-f]{96}$')) AS invalid_checksums \ + FROM trace_test._sqlx_migrations", + ) + .await?; + assert_eq!( + ledger["data"][0]["rows"].as_u64(), + Some(expected_versions.len() as u64) + ); + assert_eq!( + ledger["data"][0]["versions"].as_u64(), + Some(expected_versions.len() as u64) + ); + assert_eq!(ledger["data"][0]["invalid_checksums"].as_u64(), Some(0)); + assert_eq!( + migration_ledger_versions(&database).await?, + expected_versions + ); let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64; let span = serde_json::from_value(serde_json::json!({ "Timestamp": timestamp, "TraceId": "trace-1", "SpanId": "span-1", "ParentSpanId": "", @@ -202,6 +264,112 @@ async fn schema_supports_span_rollups_and_spend_joins( Ok(()) } +#[rstest] +#[case::quoted_versions("output_format_json_quote_64bit_integers=1")] +#[case::asynchronous_inserts( + "async_insert=1&wait_for_async_insert=0&async_insert_busy_timeout_ms=20000" +)] +#[tokio::test] +async fn schema_setup_records_migrations_synchronously_with_configured_settings( + #[future(awt)] database: TestResult, + #[case] settings: &str, +) -> TestResult { + let database = database?; + let writer = Connection::writer(&format!("{}?{settings}", database.url))?; + ensure_schema(&database.client, &writer, "trace_test", 7).await?; + assert_eq!( + migration_ledger_versions(&database).await?, + migration_versions() + ); + ensure_schema(&database.client, &writer, "trace_test", 7).await?; + assert_eq!( + migration_ledger_versions(&database).await?, + migration_versions() + ); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn changed_migration_is_rejected( + #[future(awt)] database: TestResult, +) -> TestResult { + let database = database?; + let writer = Connection::writer(&database.url)?; + ensure_schema(&database.client, &writer, "trace_test", 7).await?; + execute_write( + &database, + "ALTER TABLE trace_test._sqlx_migrations UPDATE checksum = '00' \ + WHERE version = 1 SETTINGS mutations_sync = 1", + ) + .await?; + + assert!(matches!( + ensure_schema(&database.client, &writer, "trace_test", 7).await, + Err(Error::Migration(MigrateError::VersionMismatch(1))) + )); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn concurrent_schema_setup_succeeds( + #[future(awt)] database: TestResult, +) -> TestResult { + let database = database?; + let writer = Connection::writer(&database.url)?; + let (first, second, third, fourth) = tokio::join!( + ensure_schema(&database.client, &writer, "trace_test", 7), + ensure_schema(&database.client, &writer, "trace_test", 7), + ensure_schema(&database.client, &writer, "trace_test", 7), + ensure_schema(&database.client, &writer, "trace_test", 7), + ); + for result in [first, second, third, fourth] { + result?; + } + let tables = read_json( + &database, + "SELECT count() AS tables FROM system.tables \ + WHERE database = 'trace_test' AND name IN \ + ('otel_traces', 'agent_traces_by_key', 'spend_logs')", + ) + .await?; + assert_eq!(tables["data"][0]["tables"].as_u64(), Some(3)); + assert_eq!( + migration_ledger_versions(&database).await?, + migration_versions() + ); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn existing_schema_without_ledger_is_adopted( + #[future(awt)] database: TestResult, +) -> TestResult { + let database = database?; + let writer = Connection::writer(&database.url)?; + for statement in schema_statements("trace_test", 7)? { + execute_write(&database, &statement).await?; + } + let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64; + let span = serde_json::from_value(serde_json::json!({ + "Timestamp": timestamp, "TraceId": "trace-adopted", "SpanId": "span-adopted", + "ParentSpanId": "", "ServiceName": "proxy", "SpanName": "request", + "Input": "existing row", "ResourceAttributes": {}, "SpanAttributes": {} + }))?; + insert_rows(&database, "otel_traces", vec![span]).await?; + + ensure_schema(&database.client, &writer, "trace_test", 7).await?; + + assert_eq!(table_rows(&database, "otel_traces").await?, 1); + assert_eq!( + migration_ledger_versions(&database).await?, + migration_versions() + ); + Ok(()) +} + #[rstest] #[tokio::test] async fn normalized_fields_match_clickhouse_catalog( @@ -785,6 +953,60 @@ async fn retention_changes_materialize_existing_rows_and_remain_idempotent( Ok(()) } +#[rstest] +#[tokio::test] +async fn retention_reconciliation_updates_each_table_ttl( + #[future(awt)] database: TestResult, +) -> TestResult { + let database = database?; + let writer = Connection::writer(&database.url)?; + apply_migrations(&database.client, &writer, "trace_test", 7).await?; + reconcile_retention(&database.client, &writer, "trace_test", 7).await?; + let ttl_queries = read_json( + &database, + "SELECT name, create_table_query FROM system.tables \ + WHERE database = 'trace_test' AND name IN \ + ('otel_traces', 'agent_traces_by_key', 'spend_logs') ORDER BY name", + ) + .await?; + let ttl_queries = ttl_queries["data"].as_array().expect("retention tables"); + assert_eq!( + ttl_queries + .iter() + .map(|row| row["name"].as_str().expect("table name")) + .collect::>(), + ["agent_traces_by_key", "otel_traces", "spend_logs"] + ); + for row in ttl_queries { + let query = row["create_table_query"] + .as_str() + .expect("table creation query"); + assert!( + query.contains("toIntervalDay(7)") || query.contains("INTERVAL 7 DAY"), + "{query}" + ); + } + + reconcile_retention(&database.client, &writer, "trace_test", 3).await?; + let ttl_queries = read_json( + &database, + "SELECT name, create_table_query FROM system.tables \ + WHERE database = 'trace_test' AND name IN \ + ('otel_traces', 'agent_traces_by_key', 'spend_logs') ORDER BY name", + ) + .await?; + for row in ttl_queries["data"].as_array().expect("retention tables") { + let query = row["create_table_query"] + .as_str() + .expect("table creation query"); + assert!( + query.contains("toIntervalDay(3)") || query.contains("INTERVAL 3 DAY"), + "{query}" + ); + } + Ok(()) +} + #[rstest] #[tokio::test] async fn schema_statement_timeout_maps_to_transport_error() -> TestResult { @@ -804,7 +1026,7 @@ async fn schema_statement_timeout_maps_to_transport_error() -> TestResult { .await; server.abort(); assert!( - matches!(result, Ok(Err(Error::SchemaTransport))), + matches!(result, Ok(Err(Error::Storage(StorageError::Transport)))), "{result:?}" ); Ok(()) diff --git a/tests/test_litellm_rust/test_traces.py b/tests/test_litellm_rust/test_traces.py index b8af0f0c374..43950a2f011 100644 --- a/tests/test_litellm_rust/test_traces.py +++ b/tests/test_litellm_rust/test_traces.py @@ -167,13 +167,18 @@ async def test_from_env_reads_with_clickhouse_url( @pytest.mark.asyncio async def test_schema_setup_uses_configured_retention(recording_server: RecordingServer) -> None: recording_server.expected_requests = None + recording_server.default_response = ResponseSpec(body=b"") storage: Final = _native_storage("trace_test", recording_server.base_url, 7) await storage.ensure_schema() ttl_statements: Final = tuple( request.raw_body for request in recording_server.requests if b"MODIFY TTL" in request.raw_body ) - assert len(ttl_statements) == 3 assert all(b"INTERVAL 7 DAY" in statement for statement in ttl_statements) + assert tuple(request.raw_body.strip() for request in recording_server.requests[-3:]) == ( + b"ALTER TABLE `trace_test`.otel_traces MODIFY TTL toDateTime(Timestamp) + INTERVAL 7 DAY", + b"ALTER TABLE `trace_test`.agent_traces_by_key MODIFY TTL toDateTime(StartTs) + INTERVAL 7 DAY", + b"ALTER TABLE `trace_test`.spend_logs MODIFY TTL toDateTime(start_time) + INTERVAL 7 DAY", + ) @pytest.mark.asyncio @@ -181,7 +186,7 @@ async def test_schema_setup_uses_writer_credentials_and_rejects_failed_statement recording_server: RecordingServer, ) -> None: recording_server.expected_requests = 2 - recording_server.enqueue(ResponseSpec(body="")) + recording_server.enqueue(ResponseSpec(body=b"")) recording_server.enqueue(ResponseSpec(status=403, body="denied")) writer_url: Final = recording_server.base_url.replace("http://", "http://writer:p%40ss%2Fword%25@") storage: Final = _native_storage("trace_test", writer_url + "?database=wrong&readonly=1", 7)