mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
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 <yujong@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
adb59a7af8
commit
fdc0e7dc2b
27 changed files with 1173 additions and 348 deletions
26
litellm-rust/Cargo.lock
generated
26
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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" }
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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 `<digits>_<description>.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 },
|
||||
}
|
||||
|
|
@ -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<Vec<Entry>, 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::<u64>().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<Vec<Entry>, 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<u64> = 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 { .. })));
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
@ -1,5 +0,0 @@
|
|||
# Migrations
|
||||
|
||||
`litellm-migrate` exports the `Migration` struct and the `migrate!` macro that embeds a directory of `<digits>_<description>.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
|
||||
|
|
@ -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,
|
||||
}
|
||||
|
|
@ -1 +0,0 @@
|
|||
SELECT 10;
|
||||
|
|
@ -1 +0,0 @@
|
|||
SELECT 1;
|
||||
|
|
@ -1 +0,0 @@
|
|||
SELECT 2;
|
||||
|
|
@ -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);
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
#[derive(Debug, thiserror::Error)]
|
||||
#[derive(Clone, Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("invalid ClickHouse insert row")]
|
||||
InvalidRow,
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
||||
|
|
|
|||
308
litellm-rust/crates/storage-clickhouse/src/migrate.rs
Normal file
308
litellm-rust/crates/storage-clickhouse/src/migrate.rs
Normal file
|
|
@ -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<String, Error> {
|
||||
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<Self, Error> {
|
||||
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::<Vec<_>>()
|
||||
.join("")
|
||||
}
|
||||
|
||||
fn decode_hex(value: &str) -> Result<Vec<u8>, 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<u8> {
|
||||
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<Box<dyn Future<Output = T> + Send + 'e>>;
|
||||
|
||||
impl<R> 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<Option<i64>, MigrateError>> {
|
||||
Box::pin(async { Ok(None) })
|
||||
}
|
||||
|
||||
fn list_applied_migrations<'e>(
|
||||
&'e mut self,
|
||||
table_name: &'e str,
|
||||
) -> MigrateFuture<'e, Result<Vec<AppliedMigration>, 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::<Applied>(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<Duration, MigrateError>> {
|
||||
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<Duration, MigrateError>> {
|
||||
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,
|
||||
}
|
||||
}
|
||||
476
litellm-rust/crates/storage-clickhouse/tests/migrations.rs
Normal file
476
litellm-rust/crates/storage-clickhouse/tests/migrations.rs
Normal file
|
|
@ -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<T = ()> = Result<T, Box<dyn std::error::Error>>;
|
||||
|
||||
struct ClickHouseDatabase {
|
||||
_container: ContainerAsync<ClickHouse>,
|
||||
url: String,
|
||||
client: Client,
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
async fn database() -> TestResult<ClickHouseDatabase> {
|
||||
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<Migration>) -> 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<R>(
|
||||
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<serde_json::Value> {
|
||||
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<Vec<i64>> {
|
||||
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::<Vec<_>>()
|
||||
.join("")
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn only_pending_migrations_execute_on_the_second_run(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
) -> 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<ClickHouseDatabase>,
|
||||
) -> 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<ClickHouseDatabase>,
|
||||
) -> 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<ClickHouseDatabase>,
|
||||
) -> 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<ClickHouseDatabase>,
|
||||
) -> 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
|
||||
);
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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<Error>),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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<Vec<String>, 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<Vec<Stri
|
|||
{
|
||||
return Err(Error::InvalidSchema);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn render(sql: &str, database: &str, retention_days: u32) -> String {
|
||||
sql.replace("{database}", database)
|
||||
.replace("{retention_days}", &retention_days.to_string())
|
||||
}
|
||||
|
||||
pub fn schema_statements(database: &str, retention_days: u32) -> Result<Vec<String>, 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)]
|
||||
|
|
|
|||
|
|
@ -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<u64> {
|
|||
.expect("ClickHouse returns mutation counts as unsigned integers"))
|
||||
}
|
||||
|
||||
fn migration_versions() -> Vec<u64> {
|
||||
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::<u64>().ok())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
versions.sort_unstable();
|
||||
versions
|
||||
}
|
||||
|
||||
async fn migration_ledger_versions(database: &ClickHouseDatabase) -> TestResult<Vec<u64>> {
|
||||
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<ClickHouseDatabase>,
|
||||
#[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<ClickHouseDatabase>,
|
||||
) -> 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<ClickHouseDatabase>,
|
||||
) -> 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<ClickHouseDatabase>,
|
||||
) -> 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<ClickHouseDatabase>,
|
||||
) -> 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::<Vec<_>>(),
|
||||
["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(())
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue