diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 163f60bca4a..40552b19e43 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -97,6 +97,53 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d" +[[package]] +name = "askama" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6024d73179f43f15ccd2b881bfea6fee7f3a46ec53f33b52210dea749ebebaa4" +dependencies = [ + "askama_macros", + "itoa", + "percent-encoding", + "serde", + "serde_json", +] + +[[package]] +name = "askama_derive" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071ee5ebf2138e3ad180e0aacf6940c2cab5e6d8333741d9925c7bee2b153f39" +dependencies = [ + "askama_parser", + "memchr", + "proc-macro2", + "quote", + "rustc-hash", + "syn 3.0.6", +] + +[[package]] +name = "askama_macros" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "643e1c7cbb6aec1d920332fe51a7c0d8219e273dcb8602db03f5263e4d16487b" +dependencies = [ + "askama_derive", +] + +[[package]] +name = "askama_parser" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c5ae75772275d268b03ab8bdccdd12117b6169ee23256942b34e46c9f476583" +dependencies = [ + "rustc-hash", + "unicode-ident", + "winnow 1.0.4", +] + [[package]] name = "asn1-rs" version = "0.7.2" @@ -4404,13 +4451,17 @@ dependencies = [ name = "litellm-traces" version = "0.1.0" dependencies = [ + "askama", "base64 0.22.1", "criterion", "flate2", + "futures-util", + "hmac 0.12.1", "indexmap 2.14.0", "litellm-http", "litellm-migrate", "litellm-storage-clickhouse", + "moka", "opentelemetry-proto", "prost", "rstest", @@ -4422,6 +4473,7 @@ dependencies = [ "thiserror 2.0.19", "time", "tokio", + "url", "wiremock", ] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 84b758da538..f4cb2ecbe59 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -65,6 +65,7 @@ litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" } litellm-host-python = { path = "crates/host-python" } litellm-python-compat = { path = "crates/python-compat" } +askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] } tracing = "0.1" axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] } axum-login = "0.18.0" diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs index 06d3f987f27..0ff5fe98f55 100644 --- a/litellm-rust/crates/python-bridge/src/routes/traces.rs +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -3,7 +3,9 @@ use std::collections::BTreeMap; use litellm_host_python::{FromPythonCache, ToPythonCache}; use litellm_http::ClientVariant; use litellm_storage_clickhouse::Storage; -use litellm_traces::{Error, InsertTable, Parameter, ReadQuery, Shared}; +use litellm_traces::{ + Error, InsertTable, Parameter, QueryAccessError, QueryReaders, QueryScope, ReadQuery, Shared, +}; use prost::Message; use pyo3::{ exceptions::{PyOverflowError, PyRuntimeError, PyValueError}, @@ -46,9 +48,25 @@ fn map_error(error: Error) -> PyErr { } } +fn map_sql_error(error: Error) -> PyErr { + match error { + Error::QueryFailed(400 | 404) => PyValueError::new_err(error.to_string()), + error => map_error(error), + } +} + +fn map_query_access_error(error: QueryAccessError) -> PyErr { + match error { + QueryAccessError::Storage(error) => map_sql_error(error), + QueryAccessError::InvalidScope => PyValueError::new_err(error.to_string()), + error => PyRuntimeError::new_err(error.to_string()), + } +} + #[pyclass] pub struct NativeTraceStorage { storage: Storage, + query_readers: QueryReaders, } #[pymethods] @@ -57,8 +75,13 @@ impl NativeTraceStorage { #[pyo3(signature = (database, url, reader_url = None))] fn new(database: String, url: &str, reader_url: Option<&str>) -> PyResult { litellm_traces::schema_statements(&database, 1, 1).map_err(map_error)?; + let storage = Storage::new(database, url, reader_url).map_err(map_error)?; Ok(Self { - storage: Storage::new(database, url, reader_url).map_err(map_error)?, + query_readers: QueryReaders::new( + storage.writer().clone(), + storage.database().to_owned(), + ), + storage, }) } @@ -107,6 +130,52 @@ impl NativeTraceStorage { ) } + fn query_sql<'py>( + &self, + py: Python<'py>, + sql: String, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope, + secret: String, + ) -> PyResult> { + if sql.trim().is_empty() { + return Err(map_error(Error::EmptySql)); + } + let readers = self.query_readers.clone(); + let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; + crate::execution::run_async( + py, + async move { + let _permit = readers.acquire()?; + let connection = readers.connection(&client, &scope, &secret).await?; + litellm_traces::query_sql(&client, &connection, &sql) + .await + .map_err(QueryAccessError::Storage) + }, + map_query_access_error, + ) + } + + fn query_help<'py>( + &self, + py: Python<'py>, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope, + secret: String, + ) -> PyResult> { + let readers = self.query_readers.clone(); + let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; + crate::execution::run_async( + py, + async move { + let _permit = readers.acquire()?; + let connection = readers.connection(&client, &scope, &secret).await?; + litellm_traces::query_help(&client, &connection) + .await + .map_err(QueryAccessError::Storage) + }, + map_query_access_error, + ) + } + fn lens_query<'py>( &self, py: Python<'py>, diff --git a/litellm-rust/crates/traces/Cargo.toml b/litellm-rust/crates/traces/Cargo.toml index b21bc886e6a..5e3b41e719a 100644 --- a/litellm-rust/crates/traces/Cargo.toml +++ b/litellm-rust/crates/traces/Cargo.toml @@ -6,9 +6,13 @@ license.workspace = true repository.workspace = true [dependencies] +askama.workspace = true base64.workspace = true flate2.workspace = true +futures-util.workspace = true +hmac = "0.12.1" indexmap = { version = "2", features = ["serde"] } +moka.workspace = true opentelemetry-proto = { workspace = true, features = ["gen-tonic-messages", "trace", "with-serde"] } prost.workspace = true time = { workspace = true, features = ["formatting"] } @@ -20,6 +24,8 @@ serde = { workspace = true, features = ["rc"] } serde_json.workspace = true strum.workspace = true thiserror.workspace = true +tokio.workspace = true +url.workspace = true [dev-dependencies] criterion.workspace = true diff --git a/litellm-rust/crates/traces/src/error.rs b/litellm-rust/crates/traces/src/error.rs index 9738333b113..c677e73e96f 100644 --- a/litellm-rust/crates/traces/src/error.rs +++ b/litellm-rust/crates/traces/src/error.rs @@ -7,3 +7,23 @@ pub enum DecodeError { #[error("OTLP token count is outside the storage range")] TokenCountOutOfRange, } + +#[derive(Debug, thiserror::Error)] +pub enum QueryAccessError { + #[error("trace SQL queries require a configured proxy master key")] + MissingSecret, + #[error("invalid trace query scope")] + InvalidScope, + #[error("trace SQL query concurrency limit exceeded")] + Busy, + #[error( + "ClickHouse reader provisioning failed with HTTP status {0}; the configured connection must be allowed to manage users, row policies, and SELECT grants on the trace tables" + )] + ProvisionFailed(u16), + #[error("ClickHouse reader provisioning transport failed")] + ProvisionTransport, + #[error(transparent)] + Storage(#[from] litellm_storage_clickhouse::Error), + #[error(transparent)] + Cached(#[from] std::sync::Arc), +} diff --git a/litellm-rust/crates/traces/src/lib.rs b/litellm-rust/crates/traces/src/lib.rs index 2ad8916731c..9e01d10f73e 100644 --- a/litellm-rust/crates/traces/src/lib.rs +++ b/litellm-rust/crates/traces/src/lib.rs @@ -2,17 +2,22 @@ mod error; mod insert; mod normalize; mod otlp; +mod query; +mod query_access; mod schema; mod shared; mod sql; -pub use error::DecodeError; +pub use error::{DecodeError, QueryAccessError}; pub use insert::{InsertRow, InsertTable, encode_rows, insert_rows, insert_shared_rows}; pub use litellm_storage_clickhouse::{Connection, Error, Parameter, execute_read}; pub use normalize::{ NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, NormalizedSpan, ObservationType, }; pub use otlp::{DecodedSpan, decode_otlp}; +pub use query_access::{QueryReaders, QueryScope}; pub use schema::{ensure_schema, schema_statements}; pub use shared::{Shared, SharedIdentity}; pub use sql::{LensQuery, ReadQuery, execute_named_read}; + +pub use query::{query_help, query_sql}; diff --git a/litellm-rust/crates/traces/src/query.rs b/litellm-rust/crates/traces/src/query.rs new file mode 100644 index 00000000000..b2e8ddd8702 --- /dev/null +++ b/litellm-rust/crates/traces/src/query.rs @@ -0,0 +1,346 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use futures_util::{ + StreamExt, + stream::{self, TryStreamExt}, +}; +use litellm_http::Client; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; + +use crate::{Connection, Error, NORMALIZED_FIELD_DEFINITIONS, execute_read}; + +mod guide; + +const SAMPLE_ROWS: usize = 200; +const MAX_FIELDS: usize = 200; +const MAX_DEPTH: usize = 16; +const METADATA_SQL: &str = "SELECT metadata FROM spend_logs FINAL \ + WHERE start_time >= now() - INTERVAL 7 DAY AND length(metadata) <= 8192 \ + LIMIT 201"; +const METADATA_SCOPE: &str = "Up to 200 unordered rows from the last 7 days, excluding metadata larger than 8192 bytes; up to 200 paths and 16 levels. Missing paths may exist outside this sample. Array indexes are 1-based and describe sampled positions, not a fixed schema"; +const ATTRIBUTE_SCOPE: &str = "Distinct keys from up to 200 unordered spans in the last 7 days; up to 200 keys per map. Missing keys may exist outside this sample"; + +#[derive(Deserialize)] +struct Rows { + data: Vec, +} + +#[derive(Deserialize)] +struct MetadataRow { + metadata: String, +} + +#[derive(Deserialize)] +struct AttributeRow { + key: String, +} + +#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize)] +#[serde(untagged)] +enum PathPart { + Key(String), + Index(usize), +} + +#[derive(Serialize)] +struct MetadataField { + path: Vec, + types: BTreeSet<&'static str>, + expression: String, +} + +#[derive(Deserialize, Serialize)] +struct ColumnSchema { + name: String, + #[serde(rename = "type")] + kind: String, + #[serde(flatten)] + details: BTreeMap, +} + +#[derive(Serialize)] +struct TableSchema { + name: &'static str, + columns: Vec, +} + +#[derive(Serialize)] +struct MetadataCatalog { + table: &'static str, + column: &'static str, + fields: Vec, + sampled_rows: usize, + invalid_json_rows: usize, + truncated: bool, + sample_sql: &'static str, + scope: &'static str, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, +} + +#[derive(Serialize)] +struct AttributeField { + key: String, + #[serde(rename = "type")] + kind: &'static str, + expression: String, +} + +#[derive(Serialize)] +struct AttributeCatalog { + table: &'static str, + column: &'static str, + fields: Vec, + truncated: bool, + discovery_sql: String, + scope: &'static str, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, +} + +pub async fn query_sql( + client: &Client, + connection: &Connection, + sql: &str, +) -> Result { + execute_read(client, connection, sql, &BTreeMap::new()).await +} + +async fn rows( + client: &Client, + connection: &Connection, + sql: &str, +) -> Result, Error> { + let body = query_sql(client, connection, sql).await?; + serde_json::from_str::>(&body) + .map(|result| result.data) + .map_err(|_| Error::InvalidResponse) +} + +fn literal(value: &str) -> String { + format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'")) +} + +fn metadata_expression(path: &[PathPart]) -> String { + let arguments = path + .iter() + .map(|part| match part { + PathPart::Key(key) => literal(key), + PathPart::Index(index) => index.to_string(), + }) + .collect::>() + .join(", "); + format!("JSONExtractRaw(metadata, {arguments})") +} + +fn discover( + value: &Value, + path: Vec, + fields: &mut BTreeMap, BTreeSet<&'static str>>, +) -> bool { + if path.len() > MAX_DEPTH || (fields.len() >= MAX_FIELDS && !fields.contains_key(&path)) { + return true; + } + if !path.is_empty() { + let kind = match value { + Value::Null => "null", + Value::Bool(_) => "boolean", + Value::Number(number) if number.is_i64() || number.is_u64() => "integer", + Value::Number(_) => "number", + Value::String(_) => "string", + Value::Array(_) => "array", + Value::Object(_) => "object", + }; + fields.entry(path.clone()).or_default().insert(kind); + } + match value { + Value::Object(object) => object.iter().fold(false, |limited, (key, value)| { + let child = path + .iter() + .cloned() + .chain([PathPart::Key(key.clone())]) + .collect(); + discover(value, child, fields) | limited + }), + Value::Array(array) => array + .iter() + .enumerate() + .fold(false, |limited, (index, value)| { + let child = path + .iter() + .cloned() + .chain([PathPart::Index(index + 1)]) + .collect(); + discover(value, child, fields) | limited + }), + _ => false, + } +} + +fn metadata_catalog(sample: &[MetadataRow]) -> MetadataCatalog { + let (fields, limited, invalid_rows) = sample.iter().take(SAMPLE_ROWS).fold( + (BTreeMap::new(), sample.len() > SAMPLE_ROWS, 0), + |(fields, limited, invalid_rows), row| match serde_json::from_str::(&row.metadata) { + Ok(value) => { + let mut fields = fields; + let limited = limited | discover(&value, Vec::new(), &mut fields); + (fields, limited, invalid_rows) + } + Err(_) => (fields, limited, invalid_rows + 1), + }, + ); + let fields: Vec<_> = fields + .into_iter() + .map(|(path, types)| MetadataField { + expression: metadata_expression(&path), + path, + types, + }) + .collect(); + MetadataCatalog { + table: "spend_logs", + column: "metadata", + fields, + sampled_rows: sample.len().min(SAMPLE_ROWS), + invalid_json_rows: invalid_rows, + truncated: limited, + sample_sql: METADATA_SQL, + error: None, + scope: METADATA_SCOPE, + } +} + +pub async fn query_help(client: &Client, connection: &Connection) -> Result { + let tables = stream::iter(["otel_traces", "agent_traces_by_key", "spend_logs"]) + .then(|table| async move { + Ok::<_, Error>(TableSchema { + name: table, + columns: rows::( + client, + connection, + &format!("DESCRIBE TABLE {table}"), + ) + .await?, + }) + }) + .try_collect::>() + .await?; + let metadata = match rows::(client, connection, METADATA_SQL).await { + Ok(sample) => metadata_catalog(&sample), + Err(error) => MetadataCatalog { + error: Some(error.to_string()), + truncated: true, + ..metadata_catalog(&[]) + }, + }; + let attributes = stream::iter(["SpanAttributes", "ResourceAttributes"]) + .then(|column| async move { + let sql = format!( + "SELECT DISTINCT arrayJoin(mapKeys({column})) AS key FROM \ + (SELECT {column} FROM otel_traces WHERE Timestamp >= now() - INTERVAL 7 DAY \ + LIMIT 200) ORDER BY key LIMIT 201" + ); + let (keys, error) = match rows::(client, connection, &sql).await { + Ok(keys) => (keys, None), + Err(error) => (Vec::new(), Some(error.to_string())), + }; + let fields = keys + .iter() + .take(MAX_FIELDS) + .map(|row| AttributeField { + key: row.key.clone(), + kind: "String", + expression: format!("{column}[{}]", literal(&row.key)), + }) + .collect(); + AttributeCatalog { + table: "otel_traces", + column, + fields, + truncated: error.is_some() || keys.len() > MAX_FIELDS, + discovery_sql: sql, + scope: ATTRIBUTE_SCOPE, + error, + } + }) + .collect::>() + .await; + let guide = guide::QueryGuide { + tables: &tables, + normalized_fields: &NORMALIZED_FIELD_DEFINITIONS, + metadata: &metadata, + attributes: &attributes, + }; + Ok(json!({ + "dialect": "ClickHouse SQL", + "access": "Authenticated team scope enforced by ClickHouse row policies; proxy admins can read all teams, while project-bound and teamless keys can read only their own rows", + "response": "ClickHouse JSON envelope: meta, data, rows, statistics; 64-bit integers may be strings", + "tables": tables, + "normalized_fields": NORMALIZED_FIELD_DEFINITIONS.iter().map(|field| json!({ + "table": "otel_traces", "name": field.name, "column": field.clickhouse_column, + "type": field.clickhouse_type, "meaning": field.meaning + })).collect::>(), + "metadata": metadata, + "attributes": attributes, + "relationships": [{ + "left": "otel_traces.LiteLLMRequestId", "right": "spend_logs.response_id", + "additional_predicates": "otel_traces.TeamId = spend_logs.team_id AND otel_traces.ApiKeyHash = spend_logs.api_key", + "meaning": "The normalized ID is the response ID, not request_id. Cached requests can share response_id; joins may return multiple spend rows" + }], + "examples": guide.examples()?, + "gotchas": guide.gotchas()?, + "guide": guide::render(&guide)?, + }).to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + fn metadata_discovery_preserves_mixed_types_and_reports_invalid_rows() { + let sample = [ + MetadataRow { + metadata: r#"{"x": 1}"#.into(), + }, + MetadataRow { + metadata: r#"{"x": "one"}"#.into(), + }, + MetadataRow { + metadata: "invalid".into(), + }, + ]; + let catalog = json!(metadata_catalog(&sample)); + assert_eq!( + catalog["fields"], + json!([{ + "path": ["x"], "types": ["integer", "string"], "expression": "JSONExtractRaw(metadata, 'x')" + }]) + ); + assert_eq!(catalog["invalid_json_rows"], 1); + assert_eq!(catalog["sampled_rows"], sample.len()); + } + + #[rstest] + #[case::rows(SAMPLE_ROWS + 1, 1)] + #[case::paths(1, MAX_FIELDS + 1)] + fn metadata_discovery_reports_truncation(#[case] row_count: usize, #[case] field_count: usize) { + let metadata: BTreeMap<_, _> = (0..field_count) + .map(|index| (format!("field{index}"), index)) + .collect(); + let sample: Vec<_> = (0..row_count) + .map(|_| MetadataRow { + metadata: json!(metadata).to_string(), + }) + .collect(); + let catalog = json!(metadata_catalog(&sample)); + assert_eq!(catalog["truncated"], true); + assert_eq!(catalog["sampled_rows"], row_count.min(SAMPLE_ROWS)); + assert_eq!( + catalog["fields"].as_array().unwrap().len(), + field_count.min(MAX_FIELDS) + ); + } +} diff --git a/litellm-rust/crates/traces/src/query/guide.rs b/litellm-rust/crates/traces/src/query/guide.rs new file mode 100644 index 00000000000..3bf7336648d --- /dev/null +++ b/litellm-rust/crates/traces/src/query/guide.rs @@ -0,0 +1,89 @@ +use askama::Template; +use serde::Serialize; + +use super::{AttributeCatalog, MetadataCatalog, TableSchema}; +use crate::{Error, NormalizedFieldDefinition}; + +#[derive(Template)] +#[template(path = "query_help.jinja", escape = "none", blocks = [ + "recent_spans_name", + "recent_spans_sql", + "custom_metadata_name", + "custom_metadata_sql", + "nested_metadata_name", + "nested_metadata_sql", + "correlated_calls_name", + "correlated_calls_sql", + "discover_keys_name", + "discover_keys_sql", + "time_window", + "reader_limits", + "reader_profile", + "output_format", + "json_values", + "map_values", + "literal_keys", + "time_units", + "spend_totals", + "trace_rollups", + "sampling", +])] +pub(super) struct QueryGuide<'a> { + pub tables: &'a [TableSchema], + pub normalized_fields: &'a [NormalizedFieldDefinition], + pub metadata: &'a MetadataCatalog, + pub attributes: &'a [AttributeCatalog], +} + +#[derive(Serialize)] +pub(super) struct Example { + name: String, + sql: String, +} + +impl QueryGuide<'_> { + pub fn examples(&self) -> Result<[Example; 5], Error> { + Ok([ + Example { + name: render(&self.as_recent_spans_name())?, + sql: render(&self.as_recent_spans_sql())?, + }, + Example { + name: render(&self.as_custom_metadata_name())?, + sql: render(&self.as_custom_metadata_sql())?, + }, + Example { + name: render(&self.as_nested_metadata_name())?, + sql: render(&self.as_nested_metadata_sql())?, + }, + Example { + name: render(&self.as_correlated_calls_name())?, + sql: render(&self.as_correlated_calls_sql())?, + }, + Example { + name: render(&self.as_discover_keys_name())?, + sql: render(&self.as_discover_keys_sql())?, + }, + ]) + } + + pub fn gotchas(&self) -> Result<[String; 11], Error> { + Ok([ + render(&self.as_time_window())?, + render(&self.as_reader_limits())?, + render(&self.as_reader_profile())?, + render(&self.as_output_format())?, + render(&self.as_json_values())?, + render(&self.as_map_values())?, + render(&self.as_literal_keys())?, + render(&self.as_time_units())?, + render(&self.as_spend_totals())?, + render(&self.as_trace_rollups())?, + render(&self.as_sampling())?, + ]) + } +} + +pub(super) fn render(template: &impl Template) -> Result { + template.render().map_err(|_| Error::InvalidResponse) +} diff --git a/litellm-rust/crates/traces/src/query_access.rs b/litellm-rust/crates/traces/src/query_access.rs new file mode 100644 index 00000000000..881b609177d --- /dev/null +++ b/litellm-rust/crates/traces/src/query_access.rs @@ -0,0 +1,200 @@ +use std::{sync::Arc, time::Duration}; + +use hmac::{Hmac, Mac}; +use litellm_http::Client; +use moka::future::Cache; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; + +use crate::{Connection, QueryAccessError}; + +const TABLES: [&str; 3] = ["otel_traces", "agent_traces_by_key", "spend_logs"]; + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] +pub enum QueryScope { + Admin, + Team { + team_id: String, + }, + Key { + team_id: String, + api_key_hash: String, + }, +} + +impl QueryScope { + fn validate(&self) -> Result<(), QueryAccessError> { + match self { + Self::Admin => Ok(()), + Self::Team { team_id } if !team_id.is_empty() => Ok(()), + Self::Key { api_key_hash, .. } if !api_key_hash.is_empty() => Ok(()), + _ => Err(QueryAccessError::InvalidScope), + } + } + + fn predicate(&self, table: &str) -> String { + let (team, key) = if table == "spend_logs" { + ("team_id", "api_key") + } else { + ("TeamId", "ApiKeyHash") + }; + match self { + Self::Admin => "1".to_owned(), + Self::Team { team_id } => format!("{team} = {}", literal(team_id)), + Self::Key { + team_id, + api_key_hash, + } => format!( + "{team} = {} AND {key} = {}", + literal(team_id), + literal(api_key_hash) + ), + } + } +} + +#[derive(Clone)] +pub struct QueryReaders { + writer: Connection, + database: String, + readers: Cache, + slots: Arc, +} + +impl QueryReaders { + pub fn new(writer: Connection, database: String) -> Self { + Self { + writer, + database, + readers: Cache::builder().max_capacity(1024).build(), + slots: Arc::new(Semaphore::new(8)), + } + } + + pub fn acquire(&self) -> Result { + self.slots + .clone() + .try_acquire_owned() + .map_err(|_| QueryAccessError::Busy) + } + + pub async fn connection( + &self, + client: &Client, + scope: &QueryScope, + secret: &str, + ) -> Result { + scope.validate()?; + if secret.is_empty() { + return Err(QueryAccessError::MissingSecret); + } + let identity = serde_json::to_vec(&("litellm_trace_reader_v1", &self.database, scope)) + .map_err(|_| QueryAccessError::InvalidScope)?; + let user = format!("litellm_traces_{:x}", Sha256::digest(&identity)); + let password = credential(secret, b"password", &identity)?; + self.readers + .try_get_with( + user.clone(), + self.provision(client, scope, &user, &password), + ) + .await + .map_err(QueryAccessError::Cached) + } + + async fn provision( + &self, + client: &Client, + scope: &QueryScope, + user: &str, + password: &str, + ) -> Result { + let database = &self.database; + if database.is_empty() + || !database + .bytes() + .all(|c| c.is_ascii_alphanumeric() || c == b'_') + { + return Err(QueryAccessError::InvalidScope); + } + let password_hash = format!("{:x}", Sha256::digest(password)); + self.execute( + client, + format!( + "CREATE USER IF NOT EXISTS {user} IDENTIFIED WITH sha256_hash BY '{password_hash}' \ + SETTINGS readonly = 1 CONST, max_execution_time = 10 CONST, \ + max_result_rows = 1000 CONST, max_result_bytes = 4194304 CONST, \ + result_overflow_mode = 'throw' CONST, max_memory_usage = 268435456 CONST, \ + max_threads = 2 CONST, max_concurrent_queries_for_user = 8 CONST" + ), + ) + .await?; + self.execute( + client, + format!("ALTER USER {user} IDENTIFIED WITH sha256_hash BY '{password_hash}'"), + ) + .await?; + for table in TABLES { + let predicate = scope.predicate(table); + self.execute( + client, + format!( + "CREATE ROW POLICY IF NOT EXISTS {user}_allow ON `{database}`.{table} \ + USING 1 TO {user}" + ), + ) + .await?; + self.execute( + client, + format!( + "CREATE ROW POLICY IF NOT EXISTS {user}_scope ON `{database}`.{table} \ + AS RESTRICTIVE USING {predicate} TO {user}" + ), + ) + .await?; + } + for table in TABLES { + self.execute( + client, + format!("GRANT SELECT ON `{database}`.{table} TO {user}"), + ) + .await?; + } + Connection::configured( + &self.writer.url()[..url::Position::AfterPath], + database, + user, + password, + ) + .map_err(QueryAccessError::Storage) + } + + async fn execute(&self, client: &Client, sql: String) -> Result<(), QueryAccessError> { + let response = client + .post(self.writer.url().clone()) + .timeout(Duration::from_secs(15)) + .body(sql) + .send() + .await + .map_err(|_| QueryAccessError::ProvisionTransport)?; + if !response.status().is_success() { + return Err(QueryAccessError::ProvisionFailed( + response.status().as_u16(), + )); + } + Ok(()) + } +} + +fn credential(secret: &str, purpose: &[u8], identity: &[u8]) -> Result { + let mut mac = Hmac::::new_from_slice(secret.as_bytes()) + .map_err(|_| QueryAccessError::MissingSecret)?; + mac.update(purpose); + mac.update(identity); + Ok(format!("{:x}", mac.finalize().into_bytes())) +} + +fn literal(value: &str) -> String { + format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'")) +} diff --git a/litellm-rust/crates/traces/templates/query_help.jinja b/litellm-rust/crates/traces/templates/query_help.jinja new file mode 100644 index 00000000000..4bded74588c --- /dev/null +++ b/litellm-rust/crates/traces/templates/query_help.jinja @@ -0,0 +1,65 @@ +Trace SQL query guide + +Live ClickHouse schema +{% for table in tables %} +{{ table.name }} +{% for column in table.columns %}{{ column.name }}: {{ column.kind }} +{% endfor %}{% endfor %} +Normalized span fields +{% for field in normalized_fields %}{{ field.name }}: otel_traces.{{ field.clickhouse_column }} ({{ field.clickhouse_type }}) +{{ field.meaning }} +{% endfor %} +Observed LLM call metadata +{{ metadata.scope }} +Sampled rows: {{ metadata.sampled_rows }}; invalid JSON rows: {{ metadata.invalid_json_rows }}; truncated: {{ metadata.truncated }} +{% if let Some(error) = metadata.error %}Metadata discovery unavailable: {{ error }} +{% else if metadata.fields.is_empty() %}No metadata paths found in the sampled rows +{% else %}{% for field in metadata.fields %}{{ field.expression }}: {% for kind in field.types %}{{ kind }} {% endfor %} +{% endfor %}{% endif %} +Observed span and resource attributes +{% for catalog in attributes %}{{ catalog.table }}.{{ catalog.column }} +{{ catalog.scope }} +{% if let Some(error) = catalog.error %}Attribute discovery unavailable: {{ error }} +{% else if catalog.fields.is_empty() %}No attribute keys found in the sampled spans +{% else %}{% for field in catalog.fields %}{{ field.expression }}: {{ field.kind }} +{% endfor %}{% endif %}{% endfor %} +Examples + +{% block recent_spans_name %}Recent normalized LLM spans{% endblock %} +{% block recent_spans_sql %}SELECT TraceId, SpanId, Model, InputTokens, OutputTokens, Duration / 1000000 AS duration_ms FROM otel_traces WHERE Timestamp >= now() - INTERVAL 1 DAY AND ObservationType = 'llm' ORDER BY Timestamp DESC LIMIT 100{% endblock %} + +{% block custom_metadata_name %}Find calls by custom metadata{% endblock %} +{% block custom_metadata_sql %}SELECT request_id, response_id, model, spend, JSONExtractString(metadata, 'project') AS project FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY AND JSONHas(metadata, 'project') AND JSONExtractString(metadata, 'project') = 'example' ORDER BY start_time DESC LIMIT 100{% endblock %} + +{% block nested_metadata_name %}Nested metadata with unknown types{% endblock %} +{% block nested_metadata_sql %}SELECT request_id, JSONType(metadata, 'labels', 'priority') AS type, JSONExtractRaw(metadata, 'labels', 'priority') AS value FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY AND JSONHas(metadata, 'labels', 'priority') LIMIT 100{% endblock %} + +{% block correlated_calls_name %}Traces correlated with LLM call metadata{% endblock %} +{% block correlated_calls_sql %}SELECT t.TraceId, t.SpanId, s.request_id, s.spend, s.metadata FROM otel_traces AS t INNER JOIN (SELECT * FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY) AS s ON t.LiteLLMRequestId = s.response_id AND t.TeamId = s.team_id AND t.ApiKeyHash = s.api_key WHERE t.Timestamp >= now() - INTERVAL 1 DAY AND t.LiteLLMRequestId != '' AND JSONExtractString(s.metadata, 'project') = 'example' LIMIT 100{% endblock %} + +{% block discover_keys_name %}Discover metadata keys over a different window{% endblock %} +{% block discover_keys_sql %}SELECT DISTINCT arrayJoin(JSONExtractKeys(metadata)) AS key FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 30 DAY ORDER BY key LIMIT 200{% endblock %} + +Gotchas + +{% block time_window %}Always bound Timestamp or start_time and use LIMIT; add TeamId/ApiKeyHash or team_id/api_key filters when investigating one tenant{% endblock %} + +{% block reader_limits %}The reader enforces 1000 result rows, 4 MiB response bytes, 256 MiB memory and a 10 second query limit; exceeding limits fails instead of returning partial results{% endblock %} + +{% block reader_profile %}LiteLLM provisions SELECT-only readers from the configured ClickHouse connection and enforces authenticated team scope through row policies. Project-bound and teamless keys see only their own rows. Provisioning requires CREATE USER, ALTER USER, CREATE ROW POLICY, and GRANT SELECT permissions{% endblock %} + +{% block output_format %}Do not add FORMAT clauses; the endpoint requires ClickHouse JSON output{% endblock %} + +{% block json_values %}metadata is a JSON-encoded String; use JSONHas before typed extraction to distinguish missing values from empty strings, zero and false{% endblock %} + +{% block map_values %}SpanAttributes and ResourceAttributes are Map(String, String); missing map keys return an empty string, so use mapContains for existence checks{% endblock %} + +{% block literal_keys %}Use the discovered path components as separate JSONExtract arguments; a dot inside a key is literal, not a path separator{% endblock %} + +{% block time_units %}Duration is nanoseconds; Timestamp has nanosecond precision, spend start_time has millisecond precision{% endblock %} + +{% block spend_totals %}Use spend_logs FINAL to collapse replacement rows before totals. Shared response IDs and multiple spans can multiply costs in joins; aggregate spend separately{% endblock %} + +{% block trace_rollups %}agent_traces_by_key uses SimpleAggregateFunction columns; group by TeamId, ApiKeyHash and TraceId, using min(StartTs), max(EndTs), sum(SpanCount) and groupUniqArrayArray(Models). Do not use Merge combinators{% endblock %} + +{% block sampling %}Discovery is sampled, contains no metadata values, and is not an exhaustive schema. Edit the supplied discovery SQL for older data or nested JSONExtractKeys(metadata, 'parent'){% endblock %} diff --git a/litellm-rust/crates/traces/tests/migrations.rs b/litellm-rust/crates/traces/tests/migrations.rs index ce4899681e8..d93083b052c 100644 --- a/litellm-rust/crates/traces/tests/migrations.rs +++ b/litellm-rust/crates/traces/tests/migrations.rs @@ -1149,3 +1149,289 @@ async fn lens_agent_discovery_and_selection_preserve_scope( assert_eq!(available["data"][0]["requests"], 0); Ok(()) } + +#[rstest] +#[case::empty(false)] +#[case::custom_metadata(true)] +#[tokio::test] +async fn query_help_discovers_live_schema_and_runs_its_examples( + #[future(awt)] database: TestResult, + #[case] populated: bool, +) -> TestResult { + let database = database?; + let writer = Connection::writer(&database.url)?; + ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?; + execute_write(&database, "CREATE USER help_reader").await?; + for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] { + execute_write( + &database, + &format!("GRANT SELECT ON trace_test.{table} TO help_reader"), + ) + .await?; + } + let reader = Connection::configured(&database.url, "trace_test", "help_reader", "")?; + let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64; + if populated { + execute_write(&database, "SYSTEM STOP MERGES trace_test.spend_logs").await?; + insert_rows( + &database, + "spend_logs", + vec![serde_json::from_value(serde_json::json!({ + "request_id": "request-1", "response_id": "response-1", "team_id": "team-1", + "api_key": "key-1", "metadata": r#"{"obsolete":true,"labels":{"priority":"old"}}"#, + "start_time": timestamp / 1_000_000, "end_time": timestamp / 1_000_000 + }))?], + ) + .await?; + let metadata = serde_json::json!({ + "project": "example", "labels": {"priority": 3, "enabled": true}, + "dotted.key": "literal", "quote'\\key": null, "items": [{"name": "first"}], + "&{{key}}": {"nested.key": true} + }); + insert_rows( + &database, + "spend_logs", + vec![serde_json::from_value(serde_json::json!({ + "request_id": "request-1", "response_id": "response-1", "team_id": "team-1", + "api_key": "key-1", "metadata": metadata.to_string(), "spend": 0.25, + "start_time": timestamp / 1_000_000, "end_time": timestamp / 1_000_000 + 100 + }))?], + ) + .await?; + insert_rows( + &database, + "otel_traces", + vec![serde_json::from_value(serde_json::json!({ + "Timestamp": timestamp, "TraceId": "trace-1", "SpanId": "span-1", + "TeamId": "team-1", "ApiKeyHash": "key-1", "ObservationType": "llm", + "LiteLLMRequestId": "response-1", "SpanAttributes": {"custom.tag": "value"}, + "ResourceAttributes": {"custom.resource": "value"} + }))?], + ) + .await?; + execute_write( + &database, + "ALTER TABLE trace_test.otel_traces ADD COLUMN CustomColumn String", + ) + .await?; + } + let help: serde_json::Value = + serde_json::from_str(&litellm_traces::query_help(&database.client, &reader).await?)?; + let keys: std::collections::BTreeSet<_> = help + .as_object() + .ok_or("missing help object")? + .keys() + .map(String::as_str) + .collect(); + assert_eq!( + keys, + std::collections::BTreeSet::from([ + "access", + "attributes", + "dialect", + "examples", + "gotchas", + "guide", + "metadata", + "normalized_fields", + "relationships", + "response", + "tables", + ]) + ); + let guide = help["guide"].as_str().ok_or("missing rendered guide")?; + assert!(guide.starts_with("Trace SQL query guide")); + for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] { + let described = read_json(&database, &format!("DESCRIBE TABLE {table}")).await?; + let schema = help["tables"] + .as_array() + .ok_or("missing tables")? + .iter() + .find(|schema| schema["name"] == table) + .ok_or("missing table")?; + assert_eq!(schema["columns"], described["data"]); + for column in described["data"].as_array().ok_or("missing live columns")? { + assert!(guide.contains(&format!( + "{}: {}", + column["name"].as_str().ok_or("column name")?, + column["type"].as_str().ok_or("column type")? + ))); + } + } + for gotcha in help["gotchas"].as_array().ok_or("missing gotchas")? { + assert!(guide.contains(gotcha.as_str().ok_or("gotcha text")?)); + } + let tables = help["tables"].as_array().ok_or("missing tables")?; + assert_eq!(tables.len(), 3); + let columns = tables[0]["columns"].as_array().ok_or("missing columns")?; + for field in NORMALIZED_FIELD_DEFINITIONS { + assert!( + columns + .iter() + .any(|column| column["name"] == field.clickhouse_column + && column["type"] == field.clickhouse_type) + ); + assert!( + help["normalized_fields"] + .as_array() + .ok_or("missing mappings")? + .iter() + .any(|mapped| { + mapped["name"] == field.name && mapped["column"] == field.clickhouse_column + }) + ); + } + let fields = help["metadata"]["fields"] + .as_array() + .ok_or("missing metadata fields")?; + assert_eq!(fields.is_empty(), !populated); + assert_eq!(help["metadata"]["truncated"], false); + assert!(guide.contains(help["metadata"]["scope"].as_str().ok_or("missing scope")?)); + assert_eq!( + guide.contains("No metadata paths found in the sampled rows"), + !populated + ); + if populated { + let versions = read_json(&database, "SELECT count() AS count FROM spend_logs").await?; + assert_eq!(versions["data"][0]["count"], 2); + assert_eq!(help["metadata"]["sampled_rows"], 1); + assert!( + !fields + .iter() + .any(|field| field["path"] == serde_json::json!(["obsolete"])) + ); + assert!( + columns + .iter() + .any(|column| column["name"] == "CustomColumn") + ); + assert!(fields.iter().any(|field| field["path"] + == serde_json::json!(["labels", "priority"]) + && field["types"] == serde_json::json!(["integer"]))); + assert!( + fields + .iter() + .any(|field| field["path"] == serde_json::json!(["items", 1, "name"])) + ); + assert!(guide.contains("CustomColumn: String")); + assert!(guide.contains("JSONExtractRaw(metadata, '&{{key}}', 'nested.key')")); + assert!(guide.contains("SpanAttributes['custom.tag']")); + assert!(guide.contains("ResourceAttributes['custom.resource']")); + assert_eq!(help["attributes"][0]["fields"][0]["key"], "custom.tag"); + assert_eq!(help["attributes"][1]["fields"][0]["key"], "custom.resource"); + for field in fields { + let expression = field["expression"].as_str().ok_or("missing expression")?; + assert!( + guide.contains(expression), + "missing plain-text expression: {expression}" + ); + let sql = format!("SELECT {expression} AS value FROM spend_logs FINAL"); + let body = litellm_traces::query_sql(&database.client, &reader, &sql).await?; + let values: serde_json::Value = serde_json::from_str(&body)?; + assert_ne!(values["data"][0]["value"], ""); + } + } + for example in help["examples"].as_array().ok_or("missing examples")? { + let sql = example["sql"].as_str().ok_or("missing example SQL")?; + assert!(guide.contains(example["name"].as_str().ok_or("missing example name")?)); + assert!(guide.contains(sql)); + assert_eq!( + example + .as_object() + .ok_or("example object")? + .keys() + .map(String::as_str) + .collect::>(), + std::collections::BTreeSet::from(["name", "sql"]) + ); + let body = litellm_traces::query_sql(&database.client, &reader, sql).await?; + let values: serde_json::Value = serde_json::from_str(&body)?; + assert_eq!( + values["data"].as_array().ok_or("missing data")?.is_empty(), + !populated, + "{sql}" + ); + if populated && example["name"] == "Traces correlated with LLM call metadata" { + assert_eq!(values["data"][0]["TraceId"], "trace-1"); + assert_eq!(values["data"][0]["spend"], 0.25); + } + } + Ok(()) +} + +#[rstest] +#[case::metadata(2, 1)] +#[case::attributes(1, 2)] +#[case::all(2, 2)] +#[tokio::test] +async fn query_help_preserves_schema_and_guide_when_discovery_hits_reader_limits( + #[future(awt)] database: TestResult, + #[case] spend_rows: usize, + #[case] span_rows: usize, +) -> TestResult { + let database = database?; + let writer = Connection::writer(&database.url)?; + ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?; + execute_write( + &database, + "CREATE USER help_reader SETTINGS max_rows_to_read = 1", + ) + .await?; + for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] { + execute_write( + &database, + &format!("GRANT SELECT ON trace_test.{table} TO help_reader"), + ) + .await?; + } + let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64; + let spend = (0..spend_rows) + .map(|index| { + serde_json::from_value(serde_json::json!({ + "request_id": format!("request-{index}"), "start_time": timestamp / 1_000_000, + "end_time": timestamp / 1_000_000, "metadata": r#"{"custom":{"enabled":true}}"# + })) + }) + .collect::, _>>()?; + insert_rows(&database, "spend_logs", spend).await?; + let spans = (0..span_rows).map(|index| serde_json::from_value(serde_json::json!({ + "Timestamp": timestamp, "TraceId": "trace", "SpanId": format!("span-{index}"), + "SpanAttributes": {"custom.span": "value"}, "ResourceAttributes": {"custom.resource": "value"} + }))).collect::, _>>()?; + insert_rows(&database, "otel_traces", spans).await?; + let reader = Connection::configured(&database.url, "trace_test", "help_reader", "")?; + let help: serde_json::Value = + serde_json::from_str(&litellm_traces::query_help(&database.client, &reader).await?)?; + assert_eq!(help["tables"].as_array().ok_or("tables")?.len(), 3); + assert!(!help["examples"].as_array().ok_or("examples")?.is_empty()); + assert_eq!( + help["normalized_fields"] + .as_array() + .ok_or("normalized fields")? + .len(), + NORMALIZED_FIELD_DEFINITIONS.len() + ); + let guide = help["guide"].as_str().ok_or("guide")?; + assert!(guide.contains("TraceId: String")); + assert_eq!( + guide.contains("Metadata discovery unavailable:"), + spend_rows > 1 + ); + assert_eq!( + guide.contains("Attribute discovery unavailable:"), + span_rows > 1 + ); + for (catalog, unavailable) in [ + (&help["metadata"], spend_rows > 1), + (&help["attributes"][0], span_rows > 1), + (&help["attributes"][1], span_rows > 1), + ] { + assert_eq!(catalog.get("error").is_some(), unavailable); + assert_eq!(catalog["truncated"], unavailable); + assert_eq!( + catalog["fields"].as_array().ok_or("fields")?.is_empty(), + unavailable + ); + } + Ok(()) +} diff --git a/litellm-rust/crates/traces/tests/query_access.rs b/litellm-rust/crates/traces/tests/query_access.rs new file mode 100644 index 00000000000..d6f44075e3b --- /dev/null +++ b/litellm-rust/crates/traces/tests/query_access.rs @@ -0,0 +1,268 @@ +use std::collections::BTreeMap; + +use litellm_http::Client; +use litellm_traces::{ + Connection, Error, QueryReaders, QueryScope, ensure_schema, query_help, query_sql, +}; +use rstest::{fixture, rstest}; +use serde_json::{Value, json}; +use testcontainers_modules::{ + clickhouse::ClickHouse, + testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner}, +}; + +struct Database { + _container: ContainerAsync, + client: Client, + writer: Connection, + readers: QueryReaders, +} + +#[fixture] +async fn database() -> Result> { + let container = ClickHouse::default() + .with_tag( + "26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e", + ) + .with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1") + .start() + .await?; + let writer = Connection::parse(&format!( + "http://{}:{}", + container.get_host().await?, + container.get_host_port_ipv4(8123).await? + ))?; + let client = Client::no_redirect_for_test(); + ensure_schema(&client, &writer, "trace_test", 7, 7).await?; + for sql in [ + "INSERT INTO trace_test.otel_traces (TeamId, ApiKeyHash, TraceId, SpanId, Timestamp, SpanAttributes) VALUES ('team-a', 'key-a1', 'shared-trace', 'a1', now(), map('visible', 'a')), ('team-a', 'key-a2', 'shared-trace', 'a2', now(), map('visible', 'a')), ('team-b', 'key-b', 'shared-trace', 'b', now(), map('secret-b', 'b'))", + "INSERT INTO trace_test.spend_logs (team_id, api_key, request_id, start_time, end_time, metadata) VALUES ('team-a', 'key-a1', 'a1', now(), now(), '{\"visible\":1}'), ('team-a', 'key-a2', 'a2', now(), now(), '{\"visible\":1}'), ('team-b', 'key-b', 'b', now(), now(), '{\"secret_b\":1}')", + "CREATE TABLE trace_test.private_data (secret String) ENGINE = Memory", + "INSERT INTO trace_test.private_data VALUES ('hidden')", + ] { + let response = client.post(writer.url().clone()).body(sql).send().await?; + assert!(response.status().is_success(), "{}", response.text().await?); + } + let readers = QueryReaders::new(writer.clone(), "trace_test".to_owned()); + Ok(Database { + _container: container, + client, + writer, + readers, + }) +} + +#[rstest] +#[case::team(QueryScope::Team { team_id: "team-a".to_owned() }, vec!["a1", "a2"])] +#[case::project_key(QueryScope::Key { team_id: "team-a".to_owned(), api_key_hash: "key-a1".to_owned() }, vec!["a1"])] +#[case::admin(QueryScope::Admin, vec!["a1", "a2", "b"])] +#[case::quoted_team(QueryScope::Team { team_id: "team-a' OR 1=1 --\\".to_owned() }, vec![])] +#[tokio::test] +async fn queries_and_help_are_scoped_by_the_database( + #[future(awt)] database: Result>, + #[case] scope: QueryScope, + #[case] expected: Vec<&str>, +) -> Result<(), Box> { + let database = database?; + let reader = database + .readers + .connection(&database.client, &scope, "test-master-secret") + .await?; + let queries = [ + "SELECT SpanId AS id FROM otel_traces ORDER BY id", + "SELECT SpanId AS id FROM trace_test.otel_traces WHERE 1 = 1 ORDER BY id", + "SELECT SpanId AS id FROM merge('trace_test', '^otel_traces$') ORDER BY id", + "WITH source AS (SELECT * FROM trace_test.otel_traces) SELECT SpanId AS id FROM source ORDER BY id", + "SELECT SpanId AS id FROM otel_traces UNION DISTINCT SELECT SpanId AS id FROM trace_test.otel_traces ORDER BY id", + "SELECT t.SpanId AS id FROM otel_traces t INNER JOIN spend_logs s ON t.SpanId = s.request_id ORDER BY id", + "SELECT request_id AS id FROM spend_logs FINAL ORDER BY id", + ]; + for sql in queries { + let body: Value = serde_json::from_str(&query_sql(&database.client, &reader, sql).await?)?; + assert_eq!( + body["data"], + json!( + expected + .iter() + .map(|id| json!({"id": id})) + .collect::>() + ), + "{sql}" + ); + } + let summary: Value = serde_json::from_str( + &query_sql( + &database.client, + &reader, + "SELECT sum(SpanCount) AS count FROM agent_traces_by_key", + ) + .await?, + )?; + assert_eq!(summary["data"][0]["count"], json!(expected.len())); + let help = query_help(&database.client, &reader).await?; + assert_eq!(help.contains("secret_b"), expected.contains(&"b")); + assert_eq!(help.contains("secret-b"), expected.contains(&"b")); + let recreated = QueryReaders::new(database.writer.clone(), "trace_test".to_owned()); + let repeated = recreated + .connection(&database.client, &scope, "test-master-secret") + .await?; + assert_eq!(reader.url(), repeated.url()); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn rotating_master_secret_revokes_previous_reader_credentials( + #[future(awt)] database: Result>, +) -> Result<(), Box> { + let database = database?; + let scope = QueryScope::Team { + team_id: "team-a".to_owned(), + }; + let old_reader = database + .readers + .connection(&database.client, &scope, "old-master-secret") + .await?; + let old_result = query_sql( + &database.client, + &old_reader, + "SELECT SpanId AS id FROM otel_traces ORDER BY id", + ) + .await?; + let old_rows: Value = serde_json::from_str(&old_result)?; + assert_eq!(old_rows["data"], json!([{ "id": "a1" }, { "id": "a2" }])); + + let rotated_readers = QueryReaders::new(database.writer.clone(), "trace_test".into()); + let new_reader = rotated_readers + .connection(&database.client, &scope, "new-master-secret") + .await?; + assert!( + query_sql( + &database.client, + &old_reader, + "SELECT SpanId AS id FROM otel_traces ORDER BY id", + ) + .await + .is_err() + ); + let new_result = query_sql( + &database.client, + &new_reader, + "SELECT SpanId AS id FROM otel_traces ORDER BY id", + ) + .await?; + let new_rows: Value = serde_json::from_str(&new_result)?; + assert_eq!(new_rows["data"], json!([{ "id": "a1" }, { "id": "a2" }])); + assert_eq!(old_reader.url().username(), new_reader.url().username()); + assert_ne!(old_reader.url().password(), new_reader.url().password()); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn managed_reader_rejects_privilege_and_scope_bypasses( + #[future(awt)] database: Result>, +) -> Result<(), Box> { + let database = database?; + let scope = QueryScope::Team { + team_id: "team-a".to_owned(), + }; + let reader = database + .readers + .connection(&database.client, &scope, "test-master-secret") + .await?; + for sql in [ + "INSERT INTO otel_traces (TraceId) VALUES ('injected')", + "DROP TABLE otel_traces", + "SELECT * FROM private_data", + "SELECT * FROM otel_traces SETTINGS readonly = 0", + "SELECT * FROM otel_traces SETTINGS max_memory_usage = 0", + "SELECT * FROM otel_traces SETTINGS max_execution_time = 0", + "CREATE USER scope_bypass", + "CREATE NAMED COLLECTION scope_bypass AS host = 'localhost'", + "BACKUP TABLE otel_traces TO Disk('default', 'scope-bypass')", + "SELECT * FROM url('http://127.0.0.1:1/', 'LineAsString', 'line String')", + "SELECT * FROM remote('127.0.0.1', 'trace_test', 'otel_traces')", + ] { + assert!( + matches!( + query_sql(&database.client, &reader, sql).await, + Err(Error::QueryFailed(_)) + ), + "{sql}" + ); + } + let roles: Value = serde_json::from_str( + &query_sql(&database.client, &reader, "SELECT enabledRoles() AS roles").await?, + )?; + assert_eq!(roles["data"], json!([{ "roles": [] }])); + let rows: Value = serde_json::from_str( + &query_sql( + &database.client, + &reader, + "SELECT DISTINCT TeamId FROM otel_traces", + ) + .await?, + )?; + assert_eq!(rows["data"], json!([{ "TeamId": "team-a" }])); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn provisioning_failure_never_returns_a_writer_connection( + #[future(awt)] database: Result>, +) -> Result<(), Box> { + let database = database?; + let reader = database + .readers + .connection(&database.client, &QueryScope::Admin, "test-master-secret") + .await?; + let no_provision_privileges = QueryReaders::new(reader, "trace_test".to_owned()); + let result = no_provision_privileges + .connection( + &database.client, + &QueryScope::Team { + team_id: "team-a".to_owned(), + }, + "other-secret", + ) + .await; + assert!(result.is_err()); + assert!( + database + .readers + .connection(&database.client, &QueryScope::Admin, "") + .await + .is_err() + ); + assert!( + database + .readers + .connection( + &database.client, + &QueryScope::Team { + team_id: String::new() + }, + "test-master-secret" + ) + .await + .is_err() + ); + let permits = (0..8) + .map(|_| database.readers.acquire()) + .collect::, _>>()?; + assert!(database.readers.acquire().is_err()); + drop(permits); + assert!(database.readers.acquire().is_ok()); + let rows = litellm_traces::execute_read( + &database.client, + &database.writer, + "SELECT count() AS count FROM trace_test.otel_traces", + &BTreeMap::new(), + ) + .await?; + let rows: Value = serde_json::from_str(&rows)?; + assert_eq!(rows["data"][0]["count"], 3); + Ok(()) +} diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a745213f9e6..bb113715d2c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -537,6 +537,8 @@ class LiteLLMRoutes(enum.Enum): "/lens/workers/register", "/lens/workers/{worker_id}", "/v1/traces", + "/v1/traces/query", + "/v1/traces/query/help", "/v1/traces/{trace_id}", "/v1/traces/{trace_id}/spans/{span_id}", ] diff --git a/litellm/proxy/tracing_endpoints.py b/litellm/proxy/tracing_endpoints.py index b44a3605659..e4d46560f08 100644 --- a/litellm/proxy/tracing_endpoints.py +++ b/litellm/proxy/tracing_endpoints.py @@ -15,12 +15,15 @@ from types import MappingProxyType from typing import Annotated, Final from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response +from pydantic import BaseModel, ConfigDict +from litellm._logging import verbose_proxy_logger from litellm.constants import OTLP_RETRY_AFTER_SECONDS from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.http_parsing_utils import is_otlp_trace_request from litellm.proxy.tracing_runtime import provide_receiver, require_receiver +from litellm.rust_bridge.traces import ClickHouseStorage, QueryScope from litellm.tracing import ( Tenant, TraceReceiver, @@ -137,6 +140,77 @@ async def list_agent_traces( raise HTTPException(status_code=400, detail=str(error)) from error +class TraceQueryRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + sql: str + + +@dataclass(frozen=True, slots=True) +class TraceQueryAccess: + storage: ClickHouseStorage + scope: QueryScope + secret: str + + +def provide_trace_query_secret() -> str: + from litellm.proxy.proxy_server import master_key + + if not master_key: + raise HTTPException(status_code=503, detail="Trace SQL queries require a configured proxy master key") + return master_key + + +def trace_query_scope(auth: UserAPIKeyAuth) -> QueryScope: + if auth.user_role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY): + return {"kind": "admin"} + if auth.project_id and auth.token: + return {"kind": "key", "team_id": auth.team_id or "", "api_key_hash": auth.token} + if auth.project_id: + raise HTTPException(status_code=403, detail="Project trace SQL queries require a project key") + if auth.team_id: + return {"kind": "team", "team_id": auth.team_id} + if auth.token: + return {"kind": "key", "team_id": "", "api_key_hash": auth.token} + raise HTTPException(status_code=403, detail="Trace SQL queries require an authenticated trace scope") + + +async def provide_trace_query_access( + auth: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + tracing: Annotated[TraceReceiver | None, Depends(provide_receiver)], + secret: Annotated[str, Depends(provide_trace_query_secret)], +) -> TraceQueryAccess: + return TraceQueryAccess(require_receiver(tracing).store.storage, trace_query_scope(auth), secret) + + +@router.post("/v1/traces/query") +async def query_agent_traces( + body: TraceQueryRequest, + access: Annotated[TraceQueryAccess, Depends(provide_trace_query_access)], +) -> Response: + try: + return Response( + content=await access.storage.query_sql(body.sql, access.scope, access.secret), media_type="application/json" + ) + except ValueError as error: + raise HTTPException(status_code=400, detail=str(error)) from error + except RuntimeError as error: + verbose_proxy_logger.warning("Trace SQL query unavailable: %s", error) + raise HTTPException(status_code=503, detail="Trace SQL query failed or exceeded reader limits") from error + + +@router.get("/v1/traces/query/help") +async def help_agent_trace_queries( + access: Annotated[TraceQueryAccess, Depends(provide_trace_query_access)], +) -> Response: + try: + return Response( + content=await access.storage.query_help(access.scope, access.secret), media_type="application/json" + ) + except RuntimeError as error: + verbose_proxy_logger.warning("Trace query help unavailable: %s", error) + raise HTTPException(status_code=503, detail="Trace query help is temporarily unavailable") from error + + @router.get("/v1/traces/{trace_id}", response_model=Trace) async def get_agent_trace( trace_id: str, diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 921badafb3b..5d23b048e86 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -11,7 +11,7 @@ from litellm.rust_bridge.embeddings.entrypoints import LiteLLMEmbeddingRequest from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest -from litellm.rust_bridge.traces import DecodedSpan +from litellm.rust_bridge.traces import DecodedSpan, QueryScope from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import EmbeddingResponse, ModelResponse @@ -30,6 +30,8 @@ class NativeTraceStorage: def __new__(cls, database: str, url: str, reader_url: str | None = None) -> NativeTraceStorage: ... def ensure_schema(self, trace_retention_days: int, spend_log_retention_days: int) -> Future[None]: ... def insert_rows(self, table: str, rows: Sequence[Mapping[str, object]]) -> Future[None]: ... + def query_sql(self, sql: str, scope: QueryScope, secret: str) -> Future[str]: ... + def query_help(self, scope: QueryScope, secret: str) -> Future[str]: ... def lens_query(self, name: str, parameters: Mapping[str, str | int | Sequence[str]]) -> Future[str]: ... def query(self, query: str, parameters: Mapping[str, str | int | Sequence[str]]) -> Future[str]: ... diff --git a/litellm/rust_bridge/traces.py b/litellm/rust_bridge/traces.py index d779166bade..afc1620a480 100644 --- a/litellm/rust_bridge/traces.py +++ b/litellm/rust_bridge/traces.py @@ -58,6 +58,24 @@ class DecodedSpan(TypedDict): ReadQueryName = Literal["list_traces", "trace_spans", "span_detail", "span_error", "spend_by_response_ids"] +class AdminQueryScope(TypedDict): + kind: ReadOnly[Literal["admin"]] + + +class TeamQueryScope(TypedDict): + kind: ReadOnly[Literal["team"]] + team_id: ReadOnly[str] + + +class KeyQueryScope(TypedDict): + kind: ReadOnly[Literal["key"]] + team_id: ReadOnly[str] + api_key_hash: ReadOnly[str] + + +QueryScope = AdminQueryScope | TeamQueryScope | KeyQueryScope + + class NativeStore(Protocol): def __init__(self, database: str, url: str, reader_url: str | None = None) -> None: ... @@ -65,6 +83,10 @@ class NativeStore(Protocol): def insert_rows(self, table: str, rows: Sequence[Mapping[str, object]]) -> Awaitable[None]: ... + def query_sql(self, sql: str, scope: QueryScope, secret: str) -> Awaitable[str]: ... + + def query_help(self, scope: QueryScope, secret: str) -> Awaitable[str]: ... + def lens_query(self, name: str, parameters: Mapping[str, str | int | Sequence[str]]) -> Awaitable[str]: ... def query(self, name: ReadQueryName, parameters: Mapping[str, str | int | Sequence[str]]) -> Awaitable[str]: ... @@ -138,6 +160,12 @@ class ClickHouseStorage: ) return QueryResponse.model_validate_json(result).data + async def query_sql(self, sql: str, scope: QueryScope, secret: str) -> str: + return await self._native.query_sql(sql, scope, secret) + + async def query_help(self, scope: QueryScope, secret: str) -> str: + return await self._native.query_help(scope, secret) + async def _lens_query(self, name: str, parameters: Mapping[str, object]) -> list[dict[str, JsonValue]]: result: Final = await self._native.lens_query(name, QUERY_PARAMETERS.validate_python(parameters)) return QueryResponse.model_validate_json(result).data diff --git a/litellm/tracing/receiver.py b/litellm/tracing/receiver.py index 6cef84ec6d0..1fae2361572 100644 --- a/litellm/tracing/receiver.py +++ b/litellm/tracing/receiver.py @@ -108,7 +108,7 @@ class TraceReceiver: ClickHouseStorage( database=os.getenv("CLICKHOUSE_DATABASE", "litellm"), url=os.environ["CLICKHOUSE_URL"], - reader_url=os.environ["CLICKHOUSE_READER_URL"], + reader_url=os.getenv("CLICKHOUSE_READER_URL", os.environ["CLICKHOUSE_URL"]), ) ) ) diff --git a/tests/test_litellm_rust/test_traces.py b/tests/test_litellm_rust/test_traces.py index 9b4d8c4fac1..5a09578661e 100644 --- a/tests/test_litellm_rust/test_traces.py +++ b/tests/test_litellm_rust/test_traces.py @@ -207,3 +207,124 @@ async def test_insert_validates_values_without_pydantic_copy(recording_server: R assert stored["Timestamp"] == "1970-01-01T00:00:00.000000001Z" assert stored["ResourceAttributes"] == attributes assert stored["SpanAttributes"] == attributes + + +@pytest.mark.parametrize("role", ["proxy_admin", "proxy_admin_viewer", "internal_user"]) +def test_trace_sql_endpoint_executes_for_admin_and_preserves_clickhouse_envelope( + recording_server: RecordingServer, role: str +) -> None: + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.tracing_endpoints import provide_receiver, provide_trace_query_secret, router + + envelope: Final = {"meta": [{"name": "answer", "type": "UInt8"}], "data": [{"answer": 42}], "rows": 1} + recording_server.expected_requests = 12 + for _ in range(11): + recording_server.enqueue(ResponseSpec(body="")) + recording_server.enqueue(ResponseSpec(body=envelope)) + storage: Final = ClickHouseStorage("trace_test", recording_server.base_url, recording_server.base_url) + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[provide_trace_query_secret] = lambda: "test-master-secret" + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=role, token="test") + app.dependency_overrides[provide_receiver] = lambda: TraceReceiver(TraceStore(storage)) + with TestClient(app) as client: + result: Final = client.post("/v1/traces/query", json={"sql": "SELECT 42 AS answer"}) + assert result.status_code == 200, result.text + assert result.json() == envelope + assert recording_server.requests[-1].raw_body == b"SELECT 42 AS answer" + assert client.post("/v1/traces/query", json={"sql": " "}).status_code == 400 + assert client.post("/v1/traces/query", json={}).status_code == 422 + + +def test_trace_help_endpoint_runs_native_schema_and_metadata_discovery(recording_server: RecordingServer) -> None: + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.tracing_endpoints import provide_receiver, provide_trace_query_secret, router + + recording_server.expected_requests = 17 + for _ in range(11): + recording_server.enqueue(ResponseSpec(body="")) + for response in ( + {"data": [{"name": "Model", "type": "String"}]}, + {"data": []}, + {"data": []}, + {"data": [{"metadata": '{"custom": {"label": "hello"}}'}]}, + {"data": [{"key": "custom.span"}]}, + {"data": [{"key": "custom.resource"}]}, + ): + recording_server.enqueue(ResponseSpec(body=response)) + storage: Final = ClickHouseStorage("trace_test", recording_server.base_url, recording_server.base_url) + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[provide_trace_query_secret] = lambda: "test-master-secret" + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role="proxy_admin", token="test") + app.dependency_overrides[provide_receiver] = lambda: TraceReceiver(TraceStore(storage)) + with TestClient(app) as client: + result: Final = client.get("/v1/traces/query/help") + assert result.status_code == 200, result.text + body: Final = result.json() + assert body["guide"].startswith("Trace SQL query guide") + assert "JSONExtractRaw(metadata, 'custom', 'label')" in body["guide"] + assert body["tables"][0]["columns"] == [{"name": "Model", "type": "String"}] + assert body["metadata"]["fields"][1] == { + "path": ["custom", "label"], + "types": ["string"], + "expression": "JSONExtractRaw(metadata, 'custom', 'label')", + } + assert body["attributes"][0]["fields"][0]["expression"] == "SpanAttributes['custom.span']" + assert body["attributes"][1]["fields"][0]["expression"] == "ResourceAttributes['custom.resource']" + + +@pytest.mark.parametrize("clickhouse_status, expected_status", [(400, 400), (404, 400), (500, 503), (503, 503)]) +def test_trace_sql_endpoint_distinguishes_query_errors_from_reader_failures( + recording_server: RecordingServer, clickhouse_status: int, expected_status: int +) -> None: + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.tracing_endpoints import provide_receiver, provide_trace_query_secret, router + + recording_server.expected_requests = 13 + for _ in range(11): + recording_server.enqueue(ResponseSpec(body="")) + recording_server.enqueue(ResponseSpec(status=clickhouse_status, body=b"ClickHouse rejected the query")) + envelope: Final = {"meta": [{"name": "answer", "type": "UInt8"}], "data": [{"answer": 42}], "rows": 1} + recording_server.enqueue(ResponseSpec(body=envelope)) + storage: Final = ClickHouseStorage("trace_test", recording_server.base_url, recording_server.base_url) + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[provide_trace_query_secret] = lambda: "test-master-secret" + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role="proxy_admin", token="test") + app.dependency_overrides[provide_receiver] = lambda: TraceReceiver(TraceStore(storage)) + with TestClient(app) as client: + failed: Final = client.post("/v1/traces/query", json={"sql": "SELEC 42"}) + assert failed.status_code == expected_status, failed.text + recovered: Final = client.post("/v1/traces/query", json={"sql": "SELECT 42 AS answer"}) + assert recovered.status_code == 200, recovered.text + assert recovered.json() == envelope + assert recording_server.requests[-2].raw_body == b"SELEC 42" + + +@pytest.mark.asyncio +async def test_trace_receiver_reads_with_only_one_clickhouse_url( + recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("CLICKHOUSE_URL", recording_server.base_url) + monkeypatch.setenv("CLICKHOUSE_DATABASE", "trace_test") + monkeypatch.delenv("CLICKHOUSE_READER_URL", raising=False) + recording_server.enqueue(ResponseSpec(body={"data": [{"trace_id": "trace-1"}]})) + receiver: Final = TraceReceiver.from_env() + rows: Final = await receiver.store.storage.query("trace_spans", {"trace_id": "trace-1"}) + assert rows == [{"trace_id": "trace-1"}] + parameters: Final = parse_qs(urlsplit(recording_server.requests[0].path).query) + assert parameters["database"] == ["trace_test"] + assert parameters["readonly"] == ["1"] diff --git a/tests/unit/proxy/test_tracing_endpoints.py b/tests/unit/proxy/test_tracing_endpoints.py index 287092c9f91..2c587897938 100644 --- a/tests/unit/proxy/test_tracing_endpoints.py +++ b/tests/unit/proxy/test_tracing_endpoints.py @@ -141,7 +141,9 @@ def test_501_when_tracing_not_enabled( assert response.status_code == 501 assert response.headers["content-type"] == "application/x-protobuf" assert Status.FromString(response.content).message == ( - "Agent tracing is not enabled. Set `tracing:` in general_settings and CLICKHOUSE_URL." if native_available else "" + "Agent tracing is not enabled. Set `tracing:` in general_settings and CLICKHOUSE_URL." + if native_available + else "" ) assert client.get("/v1/traces").status_code == 501 @@ -525,3 +527,96 @@ def test_lens_reads_from_injected_storage_without_receiver() -> None: assert response.status_code == 200, response.text assert response.json()["executions"] == [] storage.lens_sample.assert_awaited_once() + + +@pytest.mark.parametrize( + ("auth", "expected_scope"), + ( + (UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), {"kind": "admin"}), + (UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY), {"kind": "admin"}), + (TEAM_KEY, {"kind": "team", "team_id": "team-research"}), + ( + UserAPIKeyAuth(token="project-key", team_id="team-a", project_id="project-a"), + {"kind": "key", "team_id": "team-a", "api_key_hash": "project-key"}, + ), + (UserAPIKeyAuth(token="solo-key"), {"kind": "key", "team_id": "", "api_key_hash": "solo-key"}), + ), +) +def test_sql_and_help_use_authenticated_scope( + client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth, expected_scope: dict[str, str] +) -> None: + client.app.dependency_overrides[user_api_key_auth] = lambda: auth + client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret" + receiver.store.storage.query_sql = AsyncMock(return_value='{"data":[{"value":1}]}') + receiver.store.storage.query_help = AsyncMock(return_value='{"guide":"scoped"}') + result: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"}) + assert result.status_code == 200, result.text + assert result.json() == {"data": [{"value": 1}]} + receiver.store.storage.query_sql.assert_awaited_once_with( + "SELECT * FROM otel_traces", expected_scope, "test-secret" + ) + help_result: Final = client.get("/v1/traces/query/help") + assert help_result.status_code == 200, help_result.text + assert help_result.json() == {"guide": "scoped"} + receiver.store.storage.query_help.assert_awaited_once_with(expected_scope, "test-secret") + forged: Final = client.post("/v1/traces/query", json={"sql": "SELECT 1", "scope": {"kind": "admin"}}) + assert forged.status_code == 422, forged.text + assert receiver.store.storage.query_sql.await_count == 1 + + +@pytest.mark.parametrize("auth", (UserAPIKeyAuth(), UserAPIKeyAuth(team_id="a", project_id="p"))) +def test_sql_rejects_missing_identity_without_querying( + client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth +) -> None: + client.app.dependency_overrides[user_api_key_auth] = lambda: auth + client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret" + result: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"}) + assert result.status_code == 403, result.text + assert client.get("/v1/traces/query/help").status_code == 403 + receiver.store.storage.query_sql.assert_not_called() + receiver.store.storage.query_help.assert_not_called() + + +@pytest.mark.parametrize( + ("error", "status"), ((ValueError("invalid SQL"), 400), (RuntimeError("reader unavailable"), 503)) +) +def test_sql_reports_rejected_queries_and_unavailable_readers( + client: TestClient, receiver: MagicMock, error: Exception, status: int +) -> None: + client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret" + receiver.store.storage.query_sql = AsyncMock(side_effect=error) + result: Final = client.post("/v1/traces/query", json={"sql": "SELECT 1"}) + assert result.status_code == status, result.text + receiver.store.storage.query_sql.assert_awaited_once_with( + "SELECT 1", {"kind": "team", "team_id": "team-research"}, "test-secret" + ) + + +def test_query_help_does_not_fall_back_when_reader_provisioning_fails(client: TestClient, receiver: MagicMock) -> None: + client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret" + receiver.store.storage.query_help = AsyncMock(side_effect=RuntimeError("reader provisioning failed")) + result: Final = client.get("/v1/traces/query/help") + assert result.status_code == 503, result.text + receiver.store.storage.query_help.assert_awaited_once_with( + {"kind": "team", "team_id": "team-research"}, "test-secret" + ) + + +@pytest.mark.parametrize("secret", (None, "configured-master-key")) +def test_queries_require_a_proxy_secret( + client: TestClient, receiver: MagicMock, monkeypatch: pytest.MonkeyPatch, secret: str | None +) -> None: + from litellm.proxy import proxy_server + + monkeypatch.setattr(proxy_server, "master_key", secret) + receiver.store.storage.query_sql = AsyncMock(return_value='{"data":[]}') + result: Final = client.post("/v1/traces/query", json={"sql": "SELECT 1"}) + if secret is None: + assert result.status_code == 503, result.text + assert "master key" in result.json()["detail"] + receiver.store.storage.query_sql.assert_not_awaited() + return + assert result.status_code == 200, result.text + receiver.store.storage.query_sql.assert_awaited_once_with( + "SELECT 1", {"kind": "team", "team_id": "team-research"}, secret + ) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index dab0d6c5048..b72f2503e5d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22357,6 +22357,40 @@ export interface paths { patch?: never; trace?: never; }; + "/v1/traces/query": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** Query Agent Traces */ + post: operations["query_agent_traces_v1_traces_query_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/v1/traces/query/help": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** Help Agent Trace Queries */ + get: operations["help_agent_trace_queries_v1_traces_query_help_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1/traces/{trace_id}": { parameters: { query?: never; @@ -46804,6 +46838,11 @@ export interface components { */ truncated: boolean; }; + /** TraceQueryRequest */ + TraceQueryRequest: { + /** Sql */ + sql: string; + }; /** TraceSummary */ TraceSummary: { /** Agent Count */ @@ -79639,6 +79678,59 @@ export interface operations { }; }; }; + query_agent_traces_v1_traces_query_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["TraceQueryRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + help_agent_trace_queries_v1_traces_query_help_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; get_agent_trace_v1_traces__trace_id__get: { parameters: { query?: {