diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 8d189c8c515..5fdf0266e8c 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4075,6 +4075,7 @@ dependencies = [ "litellm-secrets-aws", "litellm-secrets-types", "litellm-token-counter", + "litellm-traces", "litellm-tracing", "pyo3", "pyo3-async-runtimes", @@ -4351,6 +4352,21 @@ dependencies = [ "tiktoken-rs", ] +[[package]] +name = "litellm-traces" +version = "0.1.0" +dependencies = [ + "litellm-http", + "rstest", + "serde", + "serde_json", + "testcontainers-modules", + "thiserror 2.0.19", + "time", + "tokio", + "url", +] + [[package]] name = "litellm-tracing" version = "0.1.0" @@ -5704,6 +5720,7 @@ checksum = "16a1cfa75cc186dd73d5818e510e042e40927bccc9c236b061cea97e1eb08029" dependencies = [ "base64 0.23.1", "bytes", + "encoding_rs", "futures-core", "futures-util", "h2 0.4.15", @@ -5715,6 +5732,7 @@ dependencies = [ "hyper-util", "js-sys", "log", + "mime", "percent-encoding", "pin-project-lite", "quinn", @@ -6945,6 +6963,7 @@ dependencies = [ "memchr", "parse-display", "pin-project-lite", + "reqwest 0.13.5", "serde", "serde_json", "serde_with", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 53aaf7a4d52..257a47268e4 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -12,6 +12,7 @@ repository = "https://github.com/BerriAI/litellm" litellm-config = { path = "crates/config" } litellm-router = { path = "crates/router" } litellm-tracing = { path = "crates/tracing" } +litellm-traces = { path = "crates/traces" } litellm-core = { path = "crates/core" } litellm-gateway-mcp = { path = "crates/gateway-mcp" } litellm-gateway = { path = "crates/gateway" } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index f8ed125f229..99c95632bb3 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -21,6 +21,7 @@ tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] fancy-regex.workspace = true litellm-tracing.workspace = true +litellm-traces.workspace = true litellm-host.workspace = true bytes.workspace = true futures-util.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 37c21cec2de..f65cb350cfe 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -43,6 +43,8 @@ mod _native { use crate::routes::responses::{ResponsesWebSocketConnection, aresponses, responses}; #[pymodule_export] use crate::routes::token_counter::TokenCounter; + #[pymodule_export] + use crate::routes::traces::{trace_encode_rows, trace_ensure_schema, trace_query}; #[cfg(feature = "huggingface")] #[pymodule_export] use crate::tokenizer::HuggingFaceEncoding; @@ -107,6 +109,9 @@ mod tests { "aresponses", "ResponsesWebSocketConnection", "NativeDiagnosticProcessor", + "trace_encode_rows", + "trace_ensure_schema", + "trace_query", "TokenCounter", "Tokenizer", "gil_stats", diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 0ea10c52c08..2380274001e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -6,6 +6,7 @@ pub(crate) mod messages; pub(crate) mod ocr; pub(crate) mod responses; pub(crate) mod token_counter; +pub(crate) mod traces; use litellm_callbacks_legacy_python::LoggingOperation; use litellm_callbacks_legacy_python::{LegacyLogging, PublicCall}; diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs new file mode 100644 index 00000000000..fdd2a7eb93d --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -0,0 +1,83 @@ +use std::collections::BTreeMap; + +use litellm_http::ClientVariant; +use litellm_traces::{Connection, Error, Parameter}; +use pyo3::{ + exceptions::{PyRuntimeError, PyValueError}, + prelude::*, +}; + +fn map_error(error: Error) -> PyErr { + match error { + Error::InvalidRow | Error::InvalidSchema | Error::EmptySql => { + PyValueError::new_err(error.to_string()) + } + Error::InvalidUrl + | Error::QueryFailed(_) + | Error::SchemaFailed(_) + | Error::ResponseTooLarge + | Error::InvalidResponse + | Error::Transport => PyRuntimeError::new_err(error.to_string()), + } +} + +#[pyfunction] +pub fn trace_ensure_schema<'py>( + py: Python<'py>, + url: &str, + database: String, + user: &str, + password: &str, + trace_retention_days: u32, + spend_log_retention_days: u32, +) -> PyResult> { + let connection = Connection::writer(url, user, password).map_err(map_error)?; + let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; + crate::execution::run_async( + py, + async move { + litellm_traces::ensure_schema( + &client, + &connection, + &database, + trace_retention_days, + spend_log_retention_days, + ) + .await + }, + map_error, + ) +} + +#[pyfunction] +pub fn trace_query<'py>( + py: Python<'py>, + url: &str, + database: &str, + user: &str, + password: &str, + sql: String, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] parameters: BTreeMap< + String, + Parameter, + >, +) -> PyResult> { + let connection = Connection::configured(url, database, user, password).map_err(map_error)?; + let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; + crate::execution::run_async( + py, + async move { litellm_traces::execute_read(&client, &connection, &sql, ¶meters).await }, + map_error, + ) +} + +#[pyfunction] +pub fn trace_encode_rows( + py: Python<'_>, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] rows: Vec< + BTreeMap, + >, +) -> PyResult { + py.detach(|| litellm_traces::encode_rows(rows)) + .map_err(map_error) +} diff --git a/litellm-rust/crates/traces/AGENTS.md b/litellm-rust/crates/traces/AGENTS.md new file mode 100644 index 00000000000..d330429b315 --- /dev/null +++ b/litellm-rust/crates/traces/AGENTS.md @@ -0,0 +1,6 @@ +- Rust owns OTLP wire decoding, ClickHouse schema, row encoding, named reads, connection validation and transport +- Keep this crate independent of Python; PyO3 conversion and public Python exceptions belong in `python-bridge` +- Keep the SQL migrations here as the only ClickHouse schema definition +- Use typed query parameters and a dedicated SELECT-only reader with server-side limits +- Bound insert time and encoded bytes; make retry deduplication behavior explicit for supported ClickHouse versions +- Test storage behavior through the crate's public API against ClickHouse diff --git a/litellm-rust/crates/traces/Cargo.toml b/litellm-rust/crates/traces/Cargo.toml new file mode 100644 index 00000000000..a8091cf9ec8 --- /dev/null +++ b/litellm-rust/crates/traces/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "litellm-traces" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +time = { workspace = true, features = ["formatting"] } +litellm-http.workspace = true +serde.workspace = true +serde_json.workspace = true +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 diff --git a/litellm-rust/crates/traces/config/reader.xml b/litellm-rust/crates/traces/config/reader.xml new file mode 100644 index 00000000000..4ff2e9d9d84 --- /dev/null +++ b/litellm-rust/crates/traces/config/reader.xml @@ -0,0 +1,32 @@ + + + + 1 + 10 + 1000 + 4194304 + throw + 268435456 + + + + + + + + + + + + + + ::/0 + litellm_traces_reader + + GRANT SELECT ON default.otel_traces + GRANT SELECT ON default.agent_traces + GRANT SELECT ON default.spend_logs + + + + diff --git a/litellm-rust/crates/traces/migrations/0001_otel_traces.sql b/litellm-rust/crates/traces/migrations/0001_otel_traces.sql new file mode 100644 index 00000000000..f228ee8144c --- /dev/null +++ b/litellm-rust/crates/traces/migrations/0001_otel_traces.sql @@ -0,0 +1,48 @@ +CREATE TABLE IF NOT EXISTS {database}.otel_traces +( + Timestamp DateTime64(9) CODEC(Delta, ZSTD(1)), + TraceId String CODEC(ZSTD(1)), + SpanId String CODEC(ZSTD(1)), + ParentSpanId String CODEC(ZSTD(1)), + TraceState String CODEC(ZSTD(1)), + SpanName LowCardinality(String) CODEC(ZSTD(1)), + SpanKind LowCardinality(String) CODEC(ZSTD(1)), + ServiceName LowCardinality(String) CODEC(ZSTD(1)), + ResourceAttributes Map(LowCardinality(String), String) CODEC(ZSTD(1)), + ScopeName String CODEC(ZSTD(1)), + ScopeVersion String CODEC(ZSTD(1)), + SpanAttributes Map(LowCardinality(String), String) CODEC(ZSTD(1)), + Duration UInt64 CODEC(ZSTD(1)), + StatusCode LowCardinality(String) CODEC(ZSTD(1)), + StatusMessage String CODEC(ZSTD(1)), + `Events.Timestamp` Array(DateTime64(9)) CODEC(ZSTD(1)), + `Events.Name` Array(LowCardinality(String)) CODEC(ZSTD(1)), + `Events.Attributes` Array(Map(LowCardinality(String), String)) CODEC(ZSTD(1)), + `Links.TraceId` Array(String) CODEC(ZSTD(1)), + `Links.SpanId` Array(String) CODEC(ZSTD(1)), + `Links.TraceState` Array(String) CODEC(ZSTD(1)), + `Links.Attributes` Array(Map(LowCardinality(String), String)) CODEC(ZSTD(1)), + TeamId LowCardinality(String) DEFAULT ResourceAttributes['litellm.team_id'], + ApiKeyHash String DEFAULT ResourceAttributes['litellm.api_key_hash'], + ObservationType LowCardinality(String) DEFAULT multiIf( + ParentSpanId = '', 'agent', + SpanAttributes['gen_ai.operation.name'] = 'invoke_agent', 'agent', + SpanAttributes['gen_ai.operation.name'] IN ('chat', 'text_completion', 'generate_content'), 'llm', + SpanAttributes['gen_ai.operation.name'] = 'execute_tool', 'tool', + 'chain'), + AgentName LowCardinality(String) DEFAULT SpanAttributes['gen_ai.agent.name'], + LiteLLMRequestId String DEFAULT SpanAttributes['gen_ai.response.id'], + Model LowCardinality(String) DEFAULT SpanAttributes['gen_ai.request.model'], + InputTokens UInt32 DEFAULT toUInt32OrZero(SpanAttributes['gen_ai.usage.input_tokens']), + OutputTokens UInt32 DEFAULT toUInt32OrZero(SpanAttributes['gen_ai.usage.output_tokens']), + Input String CODEC(ZSTD(3)), + Output String CODEC(ZSTD(3)), + InputPreview String DEFAULT substring(Input, 1, 240), + INDEX idx_trace_id TraceId TYPE bloom_filter(0.001) GRANULARITY 1, + INDEX idx_req_id LiteLLMRequestId TYPE bloom_filter(0.01) GRANULARITY 1 +) +ENGINE = MergeTree +PARTITION BY toDate(Timestamp) +ORDER BY (TeamId, ServiceName, toDateTime(Timestamp), TraceId) +TTL toDateTime(Timestamp) + INTERVAL {trace_retention_days} DAY +SETTINGS ttl_only_drop_parts = 1 diff --git a/litellm-rust/crates/traces/migrations/0002_agent_traces.sql b/litellm-rust/crates/traces/migrations/0002_agent_traces.sql new file mode 100644 index 00000000000..ea8177fa12f --- /dev/null +++ b/litellm-rust/crates/traces/migrations/0002_agent_traces.sql @@ -0,0 +1,25 @@ +CREATE TABLE IF NOT EXISTS {database}.agent_traces +( + TeamId LowCardinality(String), + TraceId String, + StartTs SimpleAggregateFunction(min, DateTime64(9)), + EndTs SimpleAggregateFunction(max, DateTime64(9)), + ServiceName SimpleAggregateFunction(any, LowCardinality(String)), + RootName SimpleAggregateFunction(anyLast, String), + RootInput SimpleAggregateFunction(anyLast, String), + RootStatus SimpleAggregateFunction(anyLast, String), + SpanCount SimpleAggregateFunction(sum, UInt64), + AgentCount SimpleAggregateFunction(sum, UInt64), + LlmCount SimpleAggregateFunction(sum, UInt64), + ToolCount SimpleAggregateFunction(sum, UInt64), + ErrorCount SimpleAggregateFunction(sum, UInt64), + InputTokens SimpleAggregateFunction(sum, UInt64), + OutputTokens SimpleAggregateFunction(sum, UInt64), + Models SimpleAggregateFunction(groupUniqArrayArray, Array(String)), + AgentNames SimpleAggregateFunction(groupUniqArrayArray, Array(String)), + RequestIds SimpleAggregateFunction(groupArrayArray, Array(String)) +) +ENGINE = AggregatingMergeTree +PARTITION BY toDate(StartTs) +ORDER BY (TeamId, TraceId) +TTL toDateTime(StartTs) + INTERVAL {trace_retention_days} DAY diff --git a/litellm-rust/crates/traces/migrations/0003_agent_traces_mv.sql b/litellm-rust/crates/traces/migrations/0003_agent_traces_mv.sql new file mode 100644 index 00000000000..3b3c5ed3adc --- /dev/null +++ b/litellm-rust/crates/traces/migrations/0003_agent_traces_mv.sql @@ -0,0 +1,21 @@ +CREATE MATERIALIZED VIEW IF NOT EXISTS {database}.agent_traces_mv TO {database}.agent_traces AS +SELECT + TeamId, TraceId, + min(Timestamp) AS StartTs, + max(Timestamp + toIntervalNanosecond(Duration)) AS EndTs, + any(ServiceName) AS ServiceName, + anyLastIf(SpanName, ParentSpanId = '') AS RootName, + anyLastIf(InputPreview, ParentSpanId = '') AS RootInput, + anyLastIf(StatusCode, ParentSpanId = '') AS RootStatus, + count() AS SpanCount, + countIf(ObservationType = 'agent') AS AgentCount, + countIf(ObservationType = 'llm') AS LlmCount, + countIf(ObservationType = 'tool') AS ToolCount, + countIf(StatusCode = 'STATUS_CODE_ERROR') AS ErrorCount, + sum(InputTokens) AS InputTokens, + sum(OutputTokens) AS OutputTokens, + groupUniqArrayIf(toString(Model), Model != '') AS Models, + groupUniqArrayIf(SpanName, ObservationType = 'agent') AS AgentNames, + groupArrayIf(LiteLLMRequestId, LiteLLMRequestId != '') AS RequestIds +FROM {database}.otel_traces +GROUP BY TeamId, TraceId diff --git a/litellm-rust/crates/traces/migrations/0004_spend_logs.sql b/litellm-rust/crates/traces/migrations/0004_spend_logs.sql new file mode 100644 index 00000000000..c20c182ba3c --- /dev/null +++ b/litellm-rust/crates/traces/migrations/0004_spend_logs.sql @@ -0,0 +1,43 @@ +CREATE TABLE IF NOT EXISTS {database}.spend_logs +( + request_id String, + response_id String, + call_type LowCardinality(String), + api_key String, + key_alias String, + team_id LowCardinality(String), + team_alias String, + organization_id String, + user String, + end_user String, + model LowCardinality(String), + model_group LowCardinality(String), + model_id String, + custom_llm_provider LowCardinality(String), + api_base String, + spend Float64, + prompt_tokens UInt32, + completion_tokens UInt32, + total_tokens UInt32, + cache_read_tokens UInt32, + cache_write_tokens UInt32, + start_time DateTime64(3), + end_time DateTime64(3), + completion_start_time Nullable(DateTime64(3)), + status LowCardinality(String), + error_str String, + cache_hit Bool, + session_id String, + trace_id String, + span_id String, + request_tags Array(String), + metadata String CODEC(ZSTD(3)), + messages String CODEC(ZSTD(3)), + response String CODEC(ZSTD(3)), + INDEX idx_response_id response_id TYPE bloom_filter(0.001) GRANULARITY 1, + INDEX idx_trace_id trace_id TYPE bloom_filter(0.001) GRANULARITY 1 +) +ENGINE = ReplacingMergeTree(end_time) +PARTITION BY toYYYYMM(start_time) +ORDER BY (team_id, toDateTime(start_time), request_id) +TTL toDateTime(start_time) + INTERVAL {spend_log_retention_days} DAY diff --git a/litellm-rust/crates/traces/src/error.rs b/litellm-rust/crates/traces/src/error.rs new file mode 100644 index 00000000000..a06eb27397a --- /dev/null +++ b/litellm-rust/crates/traces/src/error.rs @@ -0,0 +1,21 @@ +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid ClickHouse insert row")] + InvalidRow, + #[error("invalid ClickHouse HTTP URL")] + InvalidUrl, + #[error("database must be a nonempty SQL identifier and retention must be positive")] + InvalidSchema, + #[error("SQL query must not be empty")] + EmptySql, + #[error("ClickHouse query failed with HTTP status {0}")] + QueryFailed(u16), + #[error("ClickHouse schema setup failed with HTTP status {0}")] + SchemaFailed(u16), + #[error("ClickHouse query exceeded the response size limit")] + ResponseTooLarge, + #[error("ClickHouse returned an invalid or failed JSON query response")] + InvalidResponse, + #[error("ClickHouse query transport failed")] + Transport, +} diff --git a/litellm-rust/crates/traces/src/insert.rs b/litellm-rust/crates/traces/src/insert.rs new file mode 100644 index 00000000000..05b14651ce0 --- /dev/null +++ b/litellm-rust/crates/traces/src/insert.rs @@ -0,0 +1,37 @@ +use std::collections::BTreeMap; + +use serde_json::Value; +use time::{OffsetDateTime, format_description::well_known::Rfc3339}; + +use crate::Error; + +pub fn encode_rows(rows: Vec>) -> Result { + rows.into_iter() + .map(|row| { + let encoded = row + .into_iter() + .map(|(name, value)| insert_value(&name, value).map(|value| (name, value))) + .collect::, _>>()?; + serde_json::to_string(&encoded).map_err(|_| Error::InvalidRow) + }) + .collect::, _>>() + .map(|rows| rows.join("\n")) +} + +fn insert_value(name: &str, value: Value) -> Result { + let multiplier = match name { + "Timestamp" => 1, + "start_time" | "end_time" | "completion_start_time" => 1_000_000, + _ => return Ok(value), + }; + if name == "completion_start_time" && value.is_null() { + return Ok(value); + } + let timestamp = value.as_i64().ok_or(Error::InvalidRow)?; + let datetime = OffsetDateTime::from_unix_timestamp_nanos(i128::from(timestamp) * multiplier) + .map_err(|_| Error::InvalidRow)?; + datetime + .format(&Rfc3339) + .map(Value::String) + .map_err(|_| Error::InvalidRow) +} diff --git a/litellm-rust/crates/traces/src/lib.rs b/litellm-rust/crates/traces/src/lib.rs new file mode 100644 index 00000000000..1f56bb5b3a6 --- /dev/null +++ b/litellm-rust/crates/traces/src/lib.rs @@ -0,0 +1,73 @@ +mod error; +mod insert; +mod schema; +mod sql; + +pub use error::Error; +pub use insert::encode_rows; +pub use schema::{ensure_schema, schema_statements}; +pub use sql::{Parameter, execute_read}; +use url::Url; + +#[derive(Clone)] +pub struct Connection { + url: Url, +} + +impl Connection { + pub fn parse(value: &str) -> Result { + let url = Url::parse(value).map_err(|_| Error::InvalidUrl)?; + if !matches!(url.scheme(), "http" | "https") || url.host().is_none() { + return Err(Error::InvalidUrl); + } + Ok(Self { url }) + } + + pub fn configured( + url: &str, + database: &str, + user: &str, + password: &str, + ) -> Result { + let mut connection = Self::parse(url)?; + connection + .url + .set_username(user) + .map_err(|_| Error::InvalidUrl)?; + connection + .url + .set_password(Some(password)) + .map_err(|_| Error::InvalidUrl)?; + let pairs: Vec<_> = connection + .url + .query_pairs() + .filter(|(key, _)| !matches!(key.as_ref(), "database" | "user" | "password")) + .map(|(key, value)| (key.into_owned(), value.into_owned())) + .collect(); + connection + .url + .query_pairs_mut() + .clear() + .extend_pairs(pairs) + .append_pair("database", database); + Ok(connection) + } + + pub fn writer(url: &str, user: &str, password: &str) -> Result { + let mut connection = Self::parse(url)?; + connection + .url + .set_username(user) + .map_err(|_| Error::InvalidUrl)?; + connection + .url + .set_password(Some(password)) + .map_err(|_| Error::InvalidUrl)?; + connection.url.set_query(None); + Ok(connection) + } + + pub fn url(&self) -> &Url { + &self.url + } +} diff --git a/litellm-rust/crates/traces/src/schema.rs b/litellm-rust/crates/traces/src/schema.rs new file mode 100644 index 00000000000..3959d78b4c6 --- /dev/null +++ b/litellm-rust/crates/traces/src/schema.rs @@ -0,0 +1,61 @@ +use litellm_http::Client; + +use crate::Connection; +use crate::Error; + +const MIGRATIONS: [&str; 4] = [ + include_str!("../migrations/0001_otel_traces.sql"), + include_str!("../migrations/0002_agent_traces.sql"), + include_str!("../migrations/0003_agent_traces_mv.sql"), + include_str!("../migrations/0004_spend_logs.sql"), +]; + +pub fn schema_statements( + database: &str, + trace_retention_days: u32, + spend_log_retention_days: u32, +) -> Result, Error> { + if database.is_empty() + || !database + .bytes() + .all(|c| c.is_ascii_alphanumeric() || c == b'_') + || trace_retention_days == 0 + || spend_log_retention_days == 0 + { + return Err(Error::InvalidSchema); + } + let database = format!("`{database}`"); + Ok( + std::iter::once(format!("CREATE DATABASE IF NOT EXISTS {database}")) + .chain(MIGRATIONS.iter().map(|sql| { + sql.replace("{database}", &database) + .replace("{trace_retention_days}", &trace_retention_days.to_string()) + .replace( + "{spend_log_retention_days}", + &spend_log_retention_days.to_string(), + ) + })) + .collect(), + ) +} + +pub async fn ensure_schema( + client: &Client, + connection: &Connection, + database: &str, + trace_retention_days: u32, + spend_log_retention_days: u32, +) -> Result<(), Error> { + for statement in schema_statements(database, trace_retention_days, spend_log_retention_days)? { + let response = client + .post(connection.url().clone()) + .body(statement) + .send() + .await + .map_err(|_| Error::Transport)?; + if !response.status().is_success() { + return Err(Error::SchemaFailed(response.status().as_u16())); + } + } + Ok(()) +} diff --git a/litellm-rust/crates/traces/src/sql.rs b/litellm-rust/crates/traces/src/sql.rs new file mode 100644 index 00000000000..1aa21a59caa --- /dev/null +++ b/litellm-rust/crates/traces/src/sql.rs @@ -0,0 +1,114 @@ +use std::{collections::BTreeMap, time::Duration}; + +use serde::Deserialize; + +use litellm_http::Client; + +use crate::{Connection, Error}; + +const MAX_RESPONSE_BYTES: usize = 4 * 1024 * 1024; + +#[derive(Debug, Deserialize)] +#[serde(untagged)] +pub enum Parameter { + Text(String), + Integer(i64), + Strings(Vec), +} + +impl Parameter { + fn encoded(&self) -> String { + match self { + Self::Text(value) => escaped(value), + Self::Integer(value) => value.to_string(), + Self::Strings(values) => format!( + "[{}]", + values + .iter() + .map(|value| format!("'{}'", escaped(value).replace('\'', "\\'"))) + .collect::>() + .join(",") + ), + } + } +} + +fn escaped(value: &str) -> String { + value + .replace('\\', "\\\\") + .replace('\t', "\\t") + .replace('\n', "\\n") + .replace('\r', "\\r") + .replace('\0', "\\0") +} + +pub async fn execute_read( + client: &Client, + connection: &Connection, + sql: &str, + parameters: &BTreeMap, +) -> Result { + if sql.trim().is_empty() { + return Err(Error::EmptySql); + } + + let mut url = connection.url().clone(); + + let existing_pairs: Vec<(String, String)> = url + .query_pairs() + .filter(|(key, _)| { + !key.starts_with("param_") + && !matches!( + key.as_ref(), + "query" + | "readonly" + | "default_format" + | "max_result_rows" + | "result_overflow_mode" + | "max_execution_time" + | "wait_end_of_query" + ) + }) + .map(|(key, value)| (key.into_owned(), value.into_owned())) + .collect(); + url.query_pairs_mut() + .clear() + .extend_pairs(existing_pairs) + .append_pair("readonly", "1") + .append_pair("max_result_rows", "1000") + .append_pair("result_overflow_mode", "throw") + .append_pair("max_execution_time", "10") + .append_pair("wait_end_of_query", "1") + .append_pair("default_format", "JSON"); + + url.query_pairs_mut().extend_pairs( + parameters + .iter() + .map(|(name, value)| (format!("param_{name}"), value.encoded())), + ); + + let request = client + .post(url) + .timeout(Duration::from_secs(15)) + .body(sql.to_owned()); + let mut response = request.send().await.map_err(|_| Error::Transport)?; + if !response.status().is_success() { + return Err(Error::QueryFailed(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() > MAX_RESPONSE_BYTES { + return Err(Error::ResponseTooLarge); + } + body.extend_from_slice(&chunk); + } + + let json: serde_json::Value = + serde_json::from_slice(&body).map_err(|_| Error::InvalidResponse)?; + if json.get("exception").is_some() || !json.get("data").is_some_and(serde_json::Value::is_array) + { + return Err(Error::InvalidResponse); + } + String::from_utf8(body).map_err(|_| Error::InvalidResponse) +} diff --git a/litellm-rust/crates/traces/tests/admin_sql.rs b/litellm-rust/crates/traces/tests/admin_sql.rs new file mode 100644 index 00000000000..b2ef6113a15 --- /dev/null +++ b/litellm-rust/crates/traces/tests/admin_sql.rs @@ -0,0 +1,255 @@ +use litellm_http::Client; +use litellm_traces::{Connection, Error, Parameter, execute_read}; +use rstest::{fixture, rstest}; +use serde_json::Value; +use std::collections::BTreeMap; +use testcontainers_modules::{ + clickhouse::ClickHouse, + testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner}, +}; + +const CLICKHOUSE_TAG: &str = + "26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e"; + +struct Database { + _container: ContainerAsync, + url: String, + admin_url: String, + client: Client, +} + +#[fixture] +async fn database() -> Result> { + let container = ClickHouse::default() + .with_tag(CLICKHOUSE_TAG) + .with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1") + .with_env_var("LITELLM_TRACES_READER_PASSWORD", "test_password") + .with_copy_to( + "/etc/clickhouse-server/users.d/litellm-traces-reader.xml", + include_bytes!("../config/reader.xml").to_vec(), + ) + .start() + .await?; + let admin_url = format!( + "http://{}:{}", + container.get_host().await?, + container.get_host_port_ipv4(8123).await?, + ); + let client = Client::no_redirect_for_test(); + for sql in [ + "CREATE TABLE otel_traces (n UInt8) ENGINE = Memory", + "INSERT INTO otel_traces VALUES (1)", + "CREATE TABLE private_traces (n UInt8) ENGINE = Memory", + ] { + client + .post(&admin_url) + .body(sql) + .send() + .await? + .error_for_status()?; + } + let url = admin_url.replacen("http://", "http://litellm_traces_reader:test_password@", 1); + Ok(Database { + _container: container, + url, + admin_url, + client, + }) +} + +#[rstest] +#[tokio::test] +async fn admin_sql_reads_rows_with_enforced_settings( + #[future(awt)] database: Result>, +) -> Result<(), Box> { + let database = database?; + let connection = Connection::parse(&format!( + "{}?readonly=0&default_format=TabSeparated&query=SELECT+2", + database.url, + ))?; + + let result = read( + &database.client, + &connection, + "SELECT n AS answer FROM otel_traces", + ) + .await?; + let json: Value = serde_json::from_str(&result)?; + assert_eq!(json["data"][0]["answer"], 1); + + Ok(()) +} + +#[rstest] +#[case::table("CREATE TABLE admin_sql_test (n UInt8) ENGINE = Memory")] +#[case::insert("INSERT INTO otel_traces VALUES (2)")] +#[case::drop("DROP TABLE otel_traces")] +#[case::named_collection("CREATE NAMED COLLECTION admin_sql_test AS host = 'localhost'")] +#[case::settings("SET readonly = 0")] +#[case::inline_settings("SELECT n FROM otel_traces SETTINGS readonly = 0")] +#[case::time_limit("SELECT n FROM otel_traces SETTINGS max_execution_time = 0")] +#[case::row_limit("SELECT n FROM otel_traces SETTINGS max_result_rows = 0")] +#[case::byte_limit("SELECT n FROM otel_traces SETTINGS max_result_bytes = 0")] +#[case::memory_limit("SELECT n FROM otel_traces SETTINGS max_memory_usage = 0")] +#[case::other_table("SELECT * FROM private_traces")] +#[tokio::test] +async fn reader_rejects_writes_and_privilege_escalation( + #[future(awt)] database: Result>, + #[case] sql: &str, +) -> Result<(), Box> { + let database = database?; + let connection = Connection::parse(&format!("{}?readonly=0", database.url))?; + + let result = read(&database.client, &connection, sql).await; + + assert!(matches!(result, Err(Error::QueryFailed(_))), "{result:?}"); + let rows = read(&database.client, &connection, "SELECT n FROM otel_traces").await?; + let json: Value = serde_json::from_str(&rows)?; + assert_eq!(json["data"], serde_json::json!([{ "n": 1 }])); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn admin_sql_rejects_errors_after_output_starts( + #[future(awt)] database: Result>, +) -> Result<(), Box> { + let database = database?; + let connection = Connection::parse(&format!( + "{}?max_block_size=1&buffer_size=1&http_write_exception_in_output_format=1\ + &send_progress_in_http_headers=1&http_headers_progress_interval_ms=0", + database.admin_url, + ))?; + + let result = read( + &database.client, + &connection, + "SELECT sleepEachRow(0.2), throwIf(number = 2) FROM numbers(5)", + ) + .await; + + assert!( + matches!(result, Err(Error::InvalidResponse)), + "expected an error embedded in a successful HTTP response: {result:?}" + ); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn admin_sql_enforces_result_row_limit( + #[future(awt)] database: Result>, +) -> Result<(), Box> { + let database = database?; + let connection = Connection::parse(&format!( + "{}?max_result_rows=0&result_overflow_mode=throw&wait_end_of_query=1", + database.url, + ))?; + + let result = read( + &database.client, + &connection, + "SELECT number FROM numbers(1001)", + ) + .await; + + assert!(matches!(result, Err(Error::QueryFailed(_))), "{result:?}"); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn admin_sql_enforces_response_byte_limit( + #[future(awt)] database: Result>, +) -> Result<(), Box> { + let database = database?; + let connection = Connection::parse(&database.admin_url)?; + + let result = read( + &database.client, + &connection, + "SELECT repeat('x', 512 * 1024) AS payload FROM numbers(9)", + ) + .await; + + assert!(matches!(result, Err(Error::ResponseTooLarge)), "{result:?}"); + Ok(()) +} + +#[rstest] +#[case::plain("test_password", "test_password")] +#[case::encoded("p@ss/word%", "p%40ss%2Fword%25")] +#[tokio::test] +async fn admin_sql_authenticates_url_credentials( + #[future(awt)] database: Result>, + #[case] password: &str, + #[case] encoded_password: &str, +) -> Result<(), Box> { + let database = database?; + database + .client + .post(&database.admin_url) + .body(format!( + "CREATE USER sql_reader IDENTIFIED WITH plaintext_password BY '{password}'" + )) + .send() + .await? + .error_for_status()?; + let connection = Connection::parse(&database.admin_url.replacen( + "http://", + &format!("http://sql_reader:{encoded_password}@"), + 1, + ))?; + + let result = read( + &database.client, + &connection, + "SELECT currentUser() AS username", + ) + .await?; + let json: Value = serde_json::from_str(&result)?; + + assert_eq!(json["data"][0]["username"], "sql_reader"); + + Ok(()) +} + +async fn read(client: &Client, connection: &Connection, sql: &str) -> Result { + execute_read(client, connection, sql, &BTreeMap::new()).await +} + +#[rstest] +#[case::sql("'; DROP TABLE otel_traces; --")] +#[case::escapes("back\\slash\ttab\nline\0null")] +#[tokio::test] +async fn query_parameters_preserve_values_and_replace_url_parameters( + #[case] value: &str, + #[future(awt)] database: Result>, +) -> Result<(), Box> { + let database = database?; + let connection = Connection::parse(&format!("{}?param_value=wrong", database.url))?; + let values = vec![ + "a'b".to_owned(), + "back\\slash".to_owned(), + "line\nbreak".to_owned(), + "雪".to_owned(), + ]; + let parameters = BTreeMap::from([ + ("value".to_owned(), Parameter::Text(value.into())), + ("teams".to_owned(), Parameter::Strings(values.clone())), + ("number".to_owned(), Parameter::Integer(-42)), + ]); + let body = execute_read(&database.client, &connection, + "SELECT {value:String} AS value, {teams:Array(String)} AS teams, toInt32({number:Int64}) AS number", + ¶meters).await?; + let json: Value = serde_json::from_str(&body)?; + assert_eq!(json["data"][0]["value"], value); + assert_eq!(json["data"][0]["teams"], serde_json::json!(values)); + assert_eq!(json["data"][0]["number"], -42); + assert!( + read(&database.client, &connection, "SELECT n FROM otel_traces") + .await + .is_ok() + ); + Ok(()) +} diff --git a/litellm-rust/crates/traces/tests/insert.rs b/litellm-rust/crates/traces/tests/insert.rs new file mode 100644 index 00000000000..cba678152b9 --- /dev/null +++ b/litellm-rust/crates/traces/tests/insert.rs @@ -0,0 +1,40 @@ +use std::collections::BTreeMap; + +use litellm_traces::encode_rows; +use rstest::rstest; +use serde_json::{Value, json}; + +#[rstest] +#[case::span("Timestamp", json!(1_234_567_890), json!("1970-01-01T00:00:01.23456789Z"))] +#[case::start("start_time", json!(1_234), json!("1970-01-01T00:00:01.234Z"))] +#[case::end("end_time", json!(2_345), json!("1970-01-01T00:00:02.345Z"))] +#[case::completion("completion_start_time", json!(1_345), json!("1970-01-01T00:00:01.345Z"))] +#[case::absent_completion("completion_start_time", Value::Null, Value::Null)] +#[case::before_epoch("Timestamp", json!(-1), json!("1969-12-31T23:59:59.999999999Z"))] +fn insert_encoding_preserves_timestamp_precision_and_other_fields( + #[case] field: &str, + #[case] value: Value, + #[case] expected: Value, +) { + let rows = vec![BTreeMap::from([ + (field.to_owned(), value), + ("SpanAttributes".into(), json!({"message": "a\nb\\c\"雪"})), + ("InputTokens".into(), json!(42)), + ])]; + let encoded = encode_rows(rows).expect("valid row"); + let actual: Value = serde_json::from_str(&encoded).expect("JSONEachRow record"); + assert_eq!( + actual, + json!({ + field: expected, "SpanAttributes": {"message": "a\nb\\c\"雪"}, "InputTokens": 42 + }) + ); +} + +#[rstest] +#[case::fractional(json!(1.25))] +#[case::out_of_range(json!(u64::MAX))] +#[case::null(Value::Null)] +fn insert_encoding_rejects_invalid_span_timestamps(#[case] timestamp: Value) { + assert!(encode_rows(vec![BTreeMap::from([("Timestamp".into(), timestamp)])]).is_err()); +} diff --git a/litellm-rust/crates/traces/tests/migrations.rs b/litellm-rust/crates/traces/tests/migrations.rs new file mode 100644 index 00000000000..a77d57cfb0f --- /dev/null +++ b/litellm-rust/crates/traces/tests/migrations.rs @@ -0,0 +1,101 @@ +use std::collections::BTreeMap; + +use litellm_http::Client; +use litellm_traces::{Connection, encode_rows, ensure_schema, execute_read, schema_statements}; +use rstest::rstest; +use testcontainers_modules::{ + clickhouse::ClickHouse, + testcontainers::{ImageExt, runners::AsyncRunner}, +}; + +const CLICKHOUSE_TAG: &str = + "26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e"; + +#[rstest] +#[tokio::test] +async fn schema_supports_span_rollups_and_spend_joins() -> Result<(), Box> { + 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? + ); + let client = Client::no_redirect_for_test(); + let writer = Connection::writer(&url, "default", "")?; + ensure_schema(&client, &writer, "trace_test", 7, 14).await?; + ensure_schema(&client, &writer, "trace_test", 7, 14).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-1", "SpanId": "span-1", "ParentSpanId": "", + "ServiceName": "proxy", "SpanName": "request", "Input": "hello world", + "ResourceAttributes": {"litellm.team_id": "team-1", "litellm.api_key_hash": "hash-1"}, + "SpanAttributes": {"gen_ai.response.id": "response-1", "gen_ai.usage.input_tokens": "12"} + }))?; + let spend = serde_json::from_value(serde_json::json!({ + "request_id": "request-1", "response_id": "response-1", "team_id": "team-1", "spend": 0.125, + "start_time": timestamp / 1_000_000, "end_time": timestamp / 1_000_000 + 100, + "completion_start_time": null + }))?; + for (table, row) in [("otel_traces", span), ("spend_logs", spend)] { + client + .post(&url) + .query(&[ + ( + "query", + format!("INSERT INTO trace_test.{table} FORMAT JSONEachRow"), + ), + ("date_time_input_format", "best_effort".into()), + ]) + .body(encode_rows(vec![row])?) + .send() + .await? + .error_for_status()?; + } + let connection = Connection::configured(&url, "trace_test", "default", "")?; + let body = execute_read(&client, &connection, + "SELECT o.TeamId, o.ApiKeyHash, o.ObservationType, o.InputPreview, s.spend, \ + toString(toUnixTimestamp64Nano(o.Timestamp)) AS timestamp_ns, \ + toString(toUnixTimestamp64Milli(s.start_time)) AS start_ms \ + FROM otel_traces o JOIN spend_logs s ON o.LiteLLMRequestId = s.response_id AND o.TeamId = s.team_id", + &BTreeMap::new()).await?; + let response: serde_json::Value = serde_json::from_str(&body)?; + assert_eq!( + response["data"], + serde_json::json!([{ + "TeamId": "team-1", "ApiKeyHash": "hash-1", "ObservationType": "agent", + "InputPreview": "hello world", "spend": 0.125, + "timestamp_ns": timestamp.to_string(), "start_ms": (timestamp / 1_000_000).to_string() + }]) + ); + let body = execute_read( + &client, + &connection, + "SELECT toUInt32(sum(SpanCount)) AS spans, toUInt32(sum(InputTokens)) AS tokens \ + FROM agent_traces WHERE TeamId = 'team-1' AND TraceId = 'trace-1'", + &BTreeMap::new(), + ) + .await?; + let response: serde_json::Value = serde_json::from_str(&body)?; + assert_eq!( + response["data"], + serde_json::json!([{"spans": 1, "tokens": 12}]) + ); + Ok(()) +} + +#[rstest] +#[case::empty("", 7, 14)] +#[case::sql("db; DROP DATABASE default", 7, 14)] +#[case::trace_retention("traces", 0, 14)] +#[case::spend_retention("traces", 7, 0)] +fn schema_rejects_invalid_configuration( + #[case] database: &str, + #[case] traces: u32, + #[case] spend: u32, +) { + assert!(schema_statements(database, traces, spend).is_err()); +} diff --git a/litellm-rust/crates/traces/tests/queries.rs b/litellm-rust/crates/traces/tests/queries.rs new file mode 100644 index 00000000000..75dfe0adc19 --- /dev/null +++ b/litellm-rust/crates/traces/tests/queries.rs @@ -0,0 +1,11 @@ +use litellm_traces::Connection; +use rstest::rstest; + +#[rstest] +#[case::http("http://localhost:8123", true)] +#[case::https("https://localhost:8443", true)] +#[case::tcp("tcp://localhost:9000", false)] +#[case::missing_host("http://", false)] +fn accepts_only_clickhouse_http_urls(#[case] value: &str, #[case] expected: bool) { + assert_eq!(Connection::parse(value).is_ok(), expected); +} diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 6a579889869..5d2af93a9bd 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -20,6 +20,19 @@ class RustUpstreamError(Exception): ... class ForkedAfterNativeRuntimeStarted(RuntimeError): ... class ProcessReservedForForking(RuntimeError): ... +def trace_encode_rows(rows: Sequence[Mapping[str, JsonValue]]) -> str: ... +def trace_ensure_schema( + url: str, database: str, user: str, password: str, trace_retention_days: int, spend_log_retention_days: int +) -> Future[None]: ... +def trace_query( + url: str, + database: str, + user: str, + password: str, + sql: str, + parameters: Mapping[str, str | int | Sequence[str]], +) -> Future[str]: ... + @final class NativeDiagnosticProcessor: def __new__(cls, minimum_custom_key_length: int) -> NativeDiagnosticProcessor: ... @@ -338,6 +351,9 @@ __all__ = [ "process_state_started", "reserve_process_for_forking", "responses", + "trace_encode_rows", + "trace_ensure_schema", + "trace_query", "transcription", ] diff --git a/litellm/rust_bridge/traces.py b/litellm/rust_bridge/traces.py new file mode 100644 index 00000000000..cb4fd4a5bd3 --- /dev/null +++ b/litellm/rust_bridge/traces.py @@ -0,0 +1,69 @@ +from collections.abc import Awaitable, Mapping, Sequence +from typing import Final, Protocol, cast + +from pydantic import BaseModel, ConfigDict, JsonValue + +from litellm.rust_bridge.loader import get_native_bridge + + +class NativeTraces(Protocol): + def trace_encode_rows(self, rows: Sequence[Mapping[str, JsonValue]]) -> str: ... + + def trace_ensure_schema( + self, + url: str, + database: str, + user: str, + password: str, + trace_retention_days: int, + spend_log_retention_days: int, + ) -> Awaitable[None]: ... + + def trace_query( + self, + url: str, + database: str, + user: str, + password: str, + sql: str, + parameters: Mapping[str, str | int | Sequence[str]], + ) -> Awaitable[str]: ... + + +class QueryResponse(BaseModel): + model_config = ConfigDict(frozen=True) + data: list[dict[str, JsonValue]] + + +def _native() -> NativeTraces: + native: Final = get_native_bridge() + if native is None: + raise RuntimeError("Agent tracing requires the Rust extension") + return cast(NativeTraces, native) # cast-ok: the native extension is validated against this protocol at call sites + + +async def ensure_schema( + url: str, + database: str, + user: str, + password: str, + trace_retention_days: int, + spend_log_retention_days: int, +) -> None: + await _native().trace_ensure_schema(url, database, user, password, trace_retention_days, spend_log_retention_days) + + +async def query( + url: str, + database: str, + user: str, + password: str, + sql: str, + parameters: Mapping[str, str | int | Sequence[str]], +) -> list[dict[str, JsonValue]]: + result: Final = await _native().trace_query(url, database, user, password, sql, parameters) + return QueryResponse.model_validate_json(result).data + + +def encode_rows(rows: Sequence[Mapping[str, JsonValue]]) -> bytes: + return _native().trace_encode_rows(rows).encode("utf-8") diff --git a/tests/test_litellm_rust/test_traces.py b/tests/test_litellm_rust/test_traces.py new file mode 100644 index 00000000000..21f61ad24bc --- /dev/null +++ b/tests/test_litellm_rust/test_traces.py @@ -0,0 +1,78 @@ +import base64 +import json +from typing import Final +from urllib.parse import parse_qs, urlsplit + +import pytest + +from litellm.rust_bridge.traces import encode_rows, ensure_schema, query +from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.mark.asyncio +async def test_trace_reader_projects_connection_and_parameters(recording_server: RecordingServer) -> None: + recording_server.enqueue(ResponseSpec(body={"data": [{"trace_id": "trace-1"}]})) + rows: Final = await query( + recording_server.base_url + "?database=wrong&user=wrong&password=wrong", + "trace_test", + "reader", + "p@ss/word%", + "SELECT {trace_id:String} AS trace_id", + {"trace_id": "trace-1"}, + ) + request: Final = recording_server.requests[0] + parameters: Final = parse_qs(urlsplit(request.path).query) + assert rows == [{"trace_id": "trace-1"}] + assert request.raw_body == b"SELECT {trace_id:String} AS trace_id" + assert parameters["database"] == ["trace_test"] + assert parameters["param_trace_id"] == ["trace-1"] + assert parameters["readonly"] == ["1"] + assert "user" not in parameters + assert "password" not in parameters + assert request.headers["authorization"] == "Basic " + base64.b64encode(b"reader:p@ss/word%").decode() + + +@pytest.mark.asyncio +async def test_trace_reader_rejects_success_status_with_embedded_error(recording_server: RecordingServer) -> None: + recording_server.enqueue(ResponseSpec(body={"data": [], "exception": "query failed"})) + with pytest.raises(RuntimeError, match="invalid or failed JSON"): + await query(recording_server.base_url, "trace_test", "reader", "password", "SELECT 1", {}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("database,retention", [("db; DROP DATABASE default", 7), ("traces", 0)]) +async def test_schema_binding_preserves_configuration_validation(database: str, retention: int) -> None: + with pytest.raises(ValueError, match=r"database.*retention"): + await ensure_schema("http://localhost:8123", database, "writer", "password", retention, 14) + + +@pytest.mark.asyncio +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(status=403, body="denied")) + with pytest.raises(RuntimeError, match="schema setup failed with HTTP status 403"): + await ensure_schema( + recording_server.base_url + "?database=wrong&readonly=1", + "trace_test", + "writer", + "p@ss/word%", + 7, + 14, + ) + assert len(recording_server.requests) == 2 + assert recording_server.requests[0].raw_body.startswith(b"CREATE DATABASE IF NOT EXISTS") + assert recording_server.requests[1].raw_body.startswith(b"CREATE TABLE IF NOT EXISTS") + assert "readonly" not in parse_qs(urlsplit(recording_server.requests[0].path).query) + assert recording_server.requests[0].headers["authorization"] == "Basic " + base64.b64encode( + b"writer:p@ss/word%" + ).decode() + + +def test_insert_encoding_preserves_nanoseconds_through_bridge() -> None: + assert json.loads(encode_rows([{"Timestamp": 1_234_567_890, "Input": "hello"}])) == { + "Input": "hello", + "Timestamp": "1970-01-01T00:00:01.23456789Z", + }