diff --git a/README.md b/README.md index 7dbe3f1..17e9a7d 100644 --- a/README.md +++ b/README.md @@ -180,6 +180,9 @@ and [synchronous example](https://github.com/braintrustdata/braintrust-sdk-rust/ - `eval` separates definitions, case sources, execution, recording and reporting. - `api::Client` exposes direct endpoint operations with one HTTP attempt per call. Use it when your application owns delivery and retry policy. +- `ingest::BatchUploader` uploads caller-owned bulk import batches with SDK + splitting, retries and overflow handling. It accepts `api::Config`, returns + eager completions from `upload_rows(rows, max_items)`, and supports `shutdown`. - Usage extractors support OpenAI and Anthropic responses; streaming helpers support provider response assembly. diff --git a/src/api/client.rs b/src/api/client.rs index fe95709..d457b22 100644 --- a/src/api/client.rs +++ b/src/api/client.rs @@ -39,6 +39,15 @@ impl Client { T: Send + 'static, F: FnOnce(Transport) -> W, W: Future> + Send + 'static, + { + self.submit_operation(make) + } + pub(crate) fn submit_operation(&self, make: F) -> Completion + where + T: Send + 'static, + E: From + From + Send + 'static, + F: FnOnce(Transport) -> W, + W: Future> + Send + 'static, { self.execution .response_try(|| Ok(make(self.transport.clone()))) diff --git a/src/api/transport.rs b/src/api/transport.rs index 7dbee48..2be41fd 100644 --- a/src/api/transport.rs +++ b/src/api/transport.rs @@ -28,6 +28,9 @@ struct Inner { headers: HeaderMap, } impl Transport { + pub(crate) fn http_client(&self) -> reqwest::Client { + self.inner.http.clone() + } pub fn new(config: Config) -> Result { let api_url = base_url(&config.api_url)?; let app_url = base_url(&config.app_url)?; diff --git a/src/ingest.rs b/src/ingest.rs new file mode 100644 index 0000000..781d365 --- /dev/null +++ b/src/ingest.rs @@ -0,0 +1,137 @@ +//! Reliable ingestion of caller-owned batches, separate from resource logging. +//! +//! Use this adapter for bulk imports whose batch membership belongs to the +//! application. It shares SDK encoding, splitting, retry and overflow handling. +//! `api::Client::logs3()` remains the one-attempt endpoint surface. +//! +//! ```no_run +//! use braintrust_sdk_rust::{api::Config, ingest::BatchUploader}; +//! use serde_json::{Map, Value}; +//! use std::time::Duration; +//! # async fn example() -> anyhow::Result<()> { +//! let uploader = BatchUploader::new(Config::new("your-api-key"))?; +//! let row = Map::from_iter([ +//! ("id".into(), Value::String("row-id".into())), +//! ("dataset_id".into(), Value::String("dataset-id".into())), +//! ("input".into(), Value::String("hello".into())), +//! ]); +//! let delivery = uploader.upload_rows(vec![row], 100); +//! assert!(delivery.is_ok()); // admission has already happened +//! let result = delivery.await?; +//! assert_eq!(result.rows_uploaded, 1); +//! uploader.shutdown(Duration::from_secs(10)).await?; +//! # Ok(()) +//! # } +//! ``` +use crate::internal::{ + ingest::{EncodedBatch, EncodedRow}, + upload::{Capabilities, UploadCoordinator, UploadPolicy}, +}; +use crate::{api, BraintrustError, Completion, ShutdownError}; +use serde_json::{Map, Value}; +use std::{sync::Arc, time::Duration}; +use tokio::sync::Mutex; + +#[derive(Debug, Clone, Default)] +pub struct UploadResult { + pub rows_uploaded: usize, + pub bytes_processed: usize, + pub requests_sent: usize, +} + +/// Concrete bulk-ingestion client. Clones share a runtime, admission capacity, +/// immutable credentials, capability discovery and shutdown. +#[derive(Clone)] +pub struct BatchUploader { + api: api::Client, + limits: Arc>>, +} +impl BatchUploader { + /// Validate explicit configuration without login; runtime selection matches + /// `api::Client`: None owns a runtime, Some uses the caller's handle. + pub fn new(config: api::Config) -> Result { + Ok(Self { + api: api::Client::new(config)?, + limits: Arc::new(Mutex::new(None)), + }) + } + /// Submit this owned batch immediately. Await for full delivery; dropping + /// the observer leaves accepted work running. A terminal error belongs only + /// to this operation and does not prevent the next batch from proceeding. + pub fn upload_rows( + &self, + rows: Vec>, + batch_max_num_items: usize, + ) -> Completion { + let cache = self.limits.clone(); + self.api.submit_operation(move |transport| async move { + if rows.is_empty() { + return Ok(UploadResult::default()); + } + let coordinator = UploadCoordinator::for_transport(transport, UploadPolicy::Standalone); + let limits = { + let mut cached = cache.lock().await; + match *cached { + Some(limits) => limits, + None => { + let limits = coordinator.capabilities(6 * 1024 * 1024, None).await; + *cached = Some(limits); + limits + } + } + }; + let row_payloads = rows + .iter() + .map(EncodedRow::from_map) + .collect::, _>>() + .map_err(|error| { + BraintrustError::InvalidConfig(format!( + "failed to serialize logs3 row: {error}" + )) + })?; + let batches = EncodedBatch::partition( + row_payloads, + batch_max_num_items.max(1), + (limits.max_request_size / 2).max(1), + ); + + let mut result = UploadResult::default(); + for batch in batches { + let batch_result = coordinator.upload(batch, limits).await?; + result.bytes_processed += batch_result.bytes_processed; + result.requests_sent += batch_result.requests_sent; + result.rows_uploaded += batch_result.rows_uploaded; + } + Ok(result) + }) + } + /// Close admission and drain accepted batch operations. Individual upload + /// failures stay on their tokens; this reports execution drain/timeout. + pub fn shutdown(&self, timeout: Duration) -> Completion<(), ShutdownError> { + self.api.shutdown(timeout) + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[tokio::test] + async fn uploader_rejects_malformed_environment_ca_bundle() { + let _lock = crate::http::tests::ENV_LOCK.lock().await; + let original = std::env::var_os(crate::http::CUSTOM_CA_BUNDLE_ENV); + std::env::set_var( + crate::http::CUSTOM_CA_BUNDLE_ENV, + "-----BEGIN CERTIFICATE-----\nnot-base64\n-----END CERTIFICATE-----\n", + ); + let result = BatchUploader::new(api::Config::new("test-key")); + match original { + Some(value) => std::env::set_var(crate::http::CUSTOM_CA_BUNDLE_ENV, value), + None => std::env::remove_var(crate::http::CUSTOM_CA_BUNDLE_ENV), + } + assert!(result + .err() + .expect("malformed CA bundle should fail") + .to_string() + .contains(crate::http::CUSTOM_CA_BUNDLE_ENV)); + } +} diff --git a/src/internal/upload.rs b/src/internal/upload.rs index d03b63c..ea31f10 100644 --- a/src/internal/upload.rs +++ b/src/internal/upload.rs @@ -102,6 +102,16 @@ impl UploadCoordinator { }) } + pub(crate) fn for_transport(logs: api::Transport, policy: UploadPolicy) -> Self { + Self { + objects: ObjectStoreTransport { + http: logs.http_client(), + }, + logs, + policy, + } + } + pub async fn capabilities( &self, fallback: usize, diff --git a/src/lib.rs b/src/lib.rs index fc64fdd..ba2d796 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,11 +10,11 @@ pub mod eval; mod experiments; mod extractors; mod http; +pub mod ingest; mod internal; mod json_merge; mod log_queue; mod logger; -mod logs3; mod runtime; mod span; mod span_components; @@ -45,7 +45,6 @@ pub use logger::{ BraintrustClient, BraintrustClientBuilder, LoginState, OrgInfo, DEFAULT_API_URL, DEFAULT_APP_URL, }; -pub use logs3::{Logs3BatchUploader, Logs3UploadResult}; pub use runtime::current_runtime; pub use span::{ SpanBuilder, SpanHandle, SpanLog, SpanLogBuilder, SpanLogBuilderError, SpanOrigin, diff --git a/src/logs3.rs b/src/logs3.rs deleted file mode 100644 index 1582f31..0000000 --- a/src/logs3.rs +++ /dev/null @@ -1,158 +0,0 @@ -use crate::error::{BraintrustError, Result}; -use crate::http::{build_http_client, CaBundleConfig}; -use crate::internal::ingest::{EncodedBatch, EncodedRow}; -use crate::internal::upload::{Capabilities, UploadCoordinator, UploadPolicy}; -use reqwest::{Client, Url}; -use serde_json::{Map, Value}; -use std::time::Duration; - -const DEFAULT_MAX_REQUEST_SIZE: usize = 6 * 1024 * 1024; - -#[derive(Debug, Clone, Default)] -pub struct Logs3UploadResult { - pub rows_uploaded: usize, - pub bytes_processed: usize, - pub requests_sent: usize, -} - -#[derive(Debug, Clone)] -pub struct Logs3BatchUploader { - client: Client, - api_url: Url, - api_key: String, - org_name: Option, - max_request_size_override: Option, - payload_limits: Option, -} - -impl Logs3BatchUploader { - pub fn new( - api_url: impl AsRef, - api_key: impl Into, - org_name: Option, - ) -> Result { - Self::new_with_ca_bundle_config( - api_url, - api_key, - org_name, - CaBundleConfig::from_environment(), - ) - } - - /// Create an uploader with additional PEM certificates to trust. - /// - /// This overrides `BRAINTRUST_CUSTOM_CA_BUNDLE` for this uploader. - pub fn new_with_custom_ca_bundle_pem( - api_url: impl AsRef, - api_key: impl Into, - org_name: Option, - pem: impl AsRef<[u8]>, - ) -> Result { - Self::new_with_ca_bundle_config( - api_url, - api_key, - org_name, - Some(CaBundleConfig::from_pem(pem, "custom CA bundle")), - ) - } - - fn new_with_ca_bundle_config( - api_url: impl AsRef, - api_key: impl Into, - org_name: Option, - ca_bundle: Option, - ) -> Result { - let api_url = Url::parse(api_url.as_ref()) - .map_err(|e| BraintrustError::InvalidConfig(format!("invalid api_url: {e}")))?; - let client = build_http_client(Duration::from_secs(30), ca_bundle.as_ref(), None)?; - Ok(Self { - client, - api_url, - api_key: api_key.into(), - org_name: org_name.filter(|s| !s.trim().is_empty()), - max_request_size_override: None, - payload_limits: None, - }) - } - - pub fn with_max_request_size_override(mut self, max_bytes: usize) -> Self { - self.max_request_size_override = Some(max_bytes.max(1)); - self - } - - pub async fn upload_rows( - &mut self, - rows: &[Map], - batch_max_num_items: usize, - ) -> Result { - if rows.is_empty() { - return Ok(Logs3UploadResult::default()); - } - - let coordinator = UploadCoordinator::new( - self.client.clone(), - &self.api_url, - &self.api_key, - self.org_name.as_deref(), - UploadPolicy::Standalone, - )?; - let limits = match self.payload_limits { - Some(limits) => limits, - None => { - let limits = coordinator - .capabilities(DEFAULT_MAX_REQUEST_SIZE, self.max_request_size_override) - .await; - self.payload_limits = Some(limits); - limits - } - }; - let row_payloads = rows - .iter() - .map(EncodedRow::from_map) - .collect::, _>>() - .map_err(|error| { - BraintrustError::InvalidConfig(format!("failed to serialize logs3 row: {error}")) - })?; - let batches = EncodedBatch::partition( - row_payloads, - batch_max_num_items.max(1), - (limits.max_request_size / 2).max(1), - ); - - let mut result = Logs3UploadResult::default(); - for batch in batches { - let batch_result = coordinator.upload(batch, limits).await?; - result.bytes_processed += batch_result.bytes_processed; - result.requests_sent += batch_result.requests_sent; - result.rows_uploaded += batch_result.rows_uploaded; - } - Ok(result) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn uploader_rejects_malformed_environment_ca_bundle() { - let _lock = crate::http::tests::ENV_LOCK.lock().await; - let original = std::env::var_os(crate::http::CUSTOM_CA_BUNDLE_ENV); - std::env::set_var( - crate::http::CUSTOM_CA_BUNDLE_ENV, - "-----BEGIN CERTIFICATE-----\nnot-base64\n-----END CERTIFICATE-----\n", - ); - - let result = Logs3BatchUploader::new("https://example.com", "test-key", None); - - match original { - Some(value) => std::env::set_var(crate::http::CUSTOM_CA_BUNDLE_ENV, value), - None => std::env::remove_var(crate::http::CUSTOM_CA_BUNDLE_ENV), - } - - let error = result.expect_err("malformed environment CA bundle should fail"); - assert!(error - .to_string() - .contains(crate::http::CUSTOM_CA_BUNDLE_ENV)); - } -} diff --git a/tests/logs3_uploader.rs b/tests/logs3_uploader.rs index 54779e1..59b6a5d 100644 --- a/tests/logs3_uploader.rs +++ b/tests/logs3_uploader.rs @@ -1,20 +1,19 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; -use braintrust_sdk_rust::Logs3BatchUploader; +use braintrust_sdk_rust::{api::Config, ingest::BatchUploader}; use serde_json::{json, Map, Value}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate}; #[test] fn uploader_rejects_malformed_custom_ca_bundle() { - let error = Logs3BatchUploader::new_with_custom_ca_bundle_pem( - "https://example.com", - "test-key", - None, - "-----BEGIN CERTIFICATE-----\nnot-base64\n-----END CERTIFICATE-----\n", - ) - .expect_err("malformed CA bundle should fail"); + let mut config = Config::new("test-key"); + config.custom_ca_bundle_pem = + Some(b"-----BEGIN CERTIFICATE-----\nnot-base64\n-----END CERTIFICATE-----\n".to_vec()); + let error = BatchUploader::new(config) + .err() + .expect("malformed CA bundle should fail"); assert!(error .to_string() @@ -88,10 +87,12 @@ async fn uploader_splits_batches_on_413_without_overflow() { .mount(&server) .await; - let mut uploader = Logs3BatchUploader::new(server.uri(), "test-key", None).expect("uploader"); + let mut config = Config::new("test-key"); + config.api_url = server.uri(); + let uploader = BatchUploader::new(config).expect("uploader"); let rows = vec![make_row("1", 32), make_row("2", 32), make_row("3", 32)]; - let result = uploader.upload_rows(&rows, 100).await.expect("upload"); + let result = uploader.upload_rows(rows, 100).await.expect("upload"); assert_eq!(result.rows_uploaded, 3); assert!( calls.load(Ordering::SeqCst) >= 3, @@ -142,10 +143,12 @@ async fn uploader_uses_overflow_protocol_for_large_payloads() { .mount(&server) .await; - let mut uploader = Logs3BatchUploader::new(server.uri(), "test-key", None).expect("uploader"); + let mut config = Config::new("test-key"); + config.api_url = server.uri(); + let uploader = BatchUploader::new(config).expect("uploader"); let rows = vec![make_row("big", 8_000)]; - let result = uploader.upload_rows(&rows, 100).await.expect("upload"); + let result = uploader.upload_rows(rows, 100).await.expect("upload"); assert_eq!(result.rows_uploaded, 1); let requests = server.received_requests().await.unwrap(); @@ -170,3 +173,62 @@ async fn uploader_uses_overflow_protocol_for_large_payloads() { assert_eq!(body["rows"]["type"], "logs3_overflow"); assert_eq!(body["rows"]["key"], "overflow-key-1"); } + +#[tokio::test] +async fn accepted_bulk_work_progresses_without_polling_and_shutdown_closes_clones() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + let mut config = Config::new("test-key"); + config.api_url = server.uri(); + let uploader = BatchUploader::new(config).unwrap(); + let observer = std::thread::spawn({ + let uploader = uploader.clone(); + move || { + let observer = uploader.upload_rows(vec![make_row("eager", 8)], 100); + assert!(observer.is_ok()); + observer + } + }) + .join() + .unwrap(); + uploader + .shutdown(std::time::Duration::from_secs(5)) + .await + .unwrap(); + assert_eq!(observer.await.unwrap().rows_uploaded, 1); + assert!(uploader.clone().upload_rows(vec![], 100).is_err()); +} + +#[tokio::test] +async fn a_terminal_bulk_failure_does_not_poison_the_next_batch() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(400)) + .mount(&server) + .await; + let mut config = Config::new("test-key"); + config.api_url = server.uri(); + let uploader = BatchUploader::new(config).unwrap(); + assert!(uploader + .upload_rows(vec![make_row("failed", 8)], 100) + .await + .is_err()); + server.reset().await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + drop(uploader.upload_rows(vec![make_row("fresh", 8)], 100)); + uploader + .shutdown(std::time::Duration::from_secs(5)) + .await + .unwrap(); +}