From 97ca76e9bce79ef817426dfa6f500f9a3cba6dec Mon Sep 17 00:00:00 2001 From: Robin Everaars Date: Fri, 24 Jul 2026 12:23:21 +0200 Subject: [PATCH 1/8] feat: support a bearer token file for the Iceberg REST catalog Add an optional bearer_access_token_file to the Iceberg REST catalog config. When the path is set, the token is read from that file and re-read when the file changes or a short refresh interval elapses, so a rotated projected service account token is picked up without restarting the server. The file takes precedence over the static bearer_access_token. main already resolves credentials per request through CatalogCredentials::retrieve(), so this adds a FileCatalogCredentials implementation plus the config field and the wiring in create_catalog_manager. No provider changes are needed. Closes #2287. --- Cargo.lock | 1 + crates/sail-catalog/Cargo.toml | 3 + crates/sail-catalog/src/credentials.rs | 171 ++++++++++++++++++- crates/sail-common/src/config/application.rs | 8 + crates/sail-session/src/catalog.rs | 29 ++-- 5 files changed, 201 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 729c5742a6..eb6e613cdd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7383,6 +7383,7 @@ dependencies = [ "serde", "serde_arrow", "serde_json", + "tempfile", "thiserror 2.0.19", "tokio", ] diff --git a/crates/sail-catalog/Cargo.toml b/crates/sail-catalog/Cargo.toml index 1e16a7528a..0345f61abe 100644 --- a/crates/sail-catalog/Cargo.toml +++ b/crates/sail-catalog/Cargo.toml @@ -23,3 +23,6 @@ log = { workspace = true } sail-common-datafusion = { path = "../sail-common-datafusion" } sail-common = { path = "../sail-common" } + +[dev-dependencies] +tempfile = { workspace = true } diff --git a/crates/sail-catalog/src/credentials.rs b/crates/sail-catalog/src/credentials.rs index 4f70dcf15c..113ae12a36 100644 --- a/crates/sail-catalog/src/credentials.rs +++ b/crates/sail-catalog/src/credentials.rs @@ -1,6 +1,9 @@ use std::fmt::Debug; +use std::path::PathBuf; +use std::sync::Mutex; +use std::time::{Duration, Instant, SystemTime}; -use crate::error::CatalogResult; +use crate::error::{CatalogError, CatalogResult}; #[async_trait::async_trait] pub trait CatalogCredentials: Debug + Send + Sync + 'static { @@ -34,3 +37,169 @@ impl CatalogCredentials for StaticCatalogCredentials { Ok(Some(self.credential.clone())) } } + +/// How long a token read from a [`FileCatalogCredentials`] file is reused before +/// the file is checked again. Kubelet refreshes projected service account tokens +/// well before they expire (by default once 80% of their lifetime has elapsed), +/// so a short interval keeps the in-memory token fresh without reading the file +/// on every catalog request. +const FILE_CREDENTIALS_REFRESH_INTERVAL: Duration = Duration::from_secs(60); + +/// Credentials backed by a token file on disk, such as a kubelet-projected +/// service account token. The token is cached in memory and re-read when the +/// file's modification time changes or the refresh interval elapses, so a +/// rotated token is picked up without restarting the server. +#[derive(Debug)] +pub struct FileCatalogCredentials { + path: PathBuf, + cached: Mutex>, +} + +#[derive(Debug, Clone)] +struct CachedCredential { + credential: String, + modified: Option, + read_at: Instant, +} + +impl FileCatalogCredentials { + pub fn new(path: impl Into) -> Self { + Self { + path: path.into(), + cached: Mutex::new(None), + } + } +} + +#[async_trait::async_trait] +impl CatalogCredentials for FileCatalogCredentials { + async fn retrieve(&self) -> CatalogResult> { + let modified = tokio::fs::metadata(&self.path) + .await + .and_then(|metadata| metadata.modified()) + .ok(); + { + let cached = self + .cached + .lock() + .map_err(|e| CatalogError::Internal(format!("token file cache poisoned: {e}")))?; + if let Some(cached) = cached.as_ref() { + let fresh = cached.read_at.elapsed() < FILE_CREDENTIALS_REFRESH_INTERVAL; + let changed = modified.is_some() && modified != cached.modified; + if fresh && !changed { + return Ok(Some(cached.credential.clone())); + } + } + } + let credential = tokio::fs::read_to_string(&self.path) + .await + .map_err(|e| { + CatalogError::External(format!( + "failed to read token file {}: {e}", + self.path.display() + )) + })? + .trim() + .to_string(); + let mut cached = self + .cached + .lock() + .map_err(|e| CatalogError::Internal(format!("token file cache poisoned: {e}")))?; + *cached = Some(CachedCredential { + credential: credential.clone(), + modified, + read_at: Instant::now(), + }); + Ok(Some(credential)) + } +} + +#[cfg(test)] +mod tests { + #![expect(clippy::unwrap_used)] + + use std::fs::{File, FileTimes}; + use std::io::Write; + use std::path::Path; + use std::time::{Duration, SystemTime}; + + use tempfile::TempDir; + + use super::*; + + fn write_token(path: &Path, contents: &str) { + let mut file = File::create(path).unwrap(); + file.write_all(contents.as_bytes()).unwrap(); + } + + fn set_modified(path: &Path, modified: SystemTime) { + let file = File::options().write(true).open(path).unwrap(); + file.set_times(FileTimes::new().set_modified(modified)) + .unwrap(); + } + + #[tokio::test] + async fn retrieve_returns_token_from_file() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("token"); + write_token(&path, "s3cr3t-token"); + + let credentials = FileCatalogCredentials::new(&path); + assert_eq!( + credentials.retrieve().await.unwrap(), + Some("s3cr3t-token".to_string()) + ); + } + + #[tokio::test] + async fn retrieve_trims_surrounding_whitespace() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("token"); + write_token(&path, " s3cr3t-token\n\n"); + + let credentials = FileCatalogCredentials::new(&path); + assert_eq!( + credentials.retrieve().await.unwrap(), + Some("s3cr3t-token".to_string()) + ); + } + + #[tokio::test] + async fn retrieve_rereads_when_modification_time_changes() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("token"); + write_token(&path, "first-token"); + let earlier = SystemTime::UNIX_EPOCH + Duration::from_secs(1_600_000_000); + set_modified(&path, earlier); + + let credentials = FileCatalogCredentials::new(&path); + assert_eq!( + credentials.retrieve().await.unwrap(), + Some("first-token".to_string()) + ); + + // Rotate the token and advance the file's modification time. The refresh + // interval has not elapsed, so the mtime change is the only thing that can + // trigger a re-read. This mirrors kubelet swapping a projected token. + write_token(&path, "second-token"); + set_modified(&path, earlier + Duration::from_secs(60)); + + assert_eq!( + credentials.retrieve().await.unwrap(), + Some("second-token".to_string()) + ); + } + + #[tokio::test] + async fn retrieve_reports_error_for_missing_file() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("does-not-exist"); + + let credentials = FileCatalogCredentials::new(&path); + let error = credentials.retrieve().await.unwrap_err(); + assert!( + matches!(error, CatalogError::External(_)), + "unexpected error variant: {error:?}" + ); + } +} diff --git a/crates/sail-common/src/config/application.rs b/crates/sail-common/src/config/application.rs index ae9af1a14b..cf220f6fed 100644 --- a/crates/sail-common/src/config/application.rs +++ b/crates/sail-common/src/config/application.rs @@ -551,6 +551,14 @@ pub enum CatalogType { serialize_with = "serialize_optional_secret" )] bearer_access_token: Option, + /// Path to a file holding the bearer token. When set, the token is + /// re-read from this file per request (with a short refresh interval) + /// so a rotated token (for example a kubelet-projected service account + /// token) is picked up without restarting the server. Takes precedence + /// over `bearer_access_token`. The path is not a secret, so it is kept + /// as a plain string. + #[serde(skip_serializing_if = "Option::is_none")] + bearer_access_token_file: Option, #[serde(flatten)] cache: CatalogCacheConfig, }, diff --git a/crates/sail-session/src/catalog.rs b/crates/sail-session/src/catalog.rs index 46832190fa..5c73e592c2 100644 --- a/crates/sail-session/src/catalog.rs +++ b/crates/sail-session/src/catalog.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use datafusion::common::{Result, plan_datafusion_err}; use datafusion_common::plan_err; use sail_catalog::credentials::{ - CatalogCredentials, EmptyCatalogCredentials, StaticCatalogCredentials, + CatalogCredentials, EmptyCatalogCredentials, FileCatalogCredentials, StaticCatalogCredentials, }; use sail_catalog::error::CatalogResult; use sail_catalog::manager::{CatalogManager, CatalogManagerOptions}; @@ -53,6 +53,7 @@ pub fn create_catalog_manager( namespace_separator, oauth_access_token, bearer_access_token, + bearer_access_token_file, cache, } => { let mut properties = HashMap::new(); @@ -69,15 +70,23 @@ pub fn create_catalog_manager( namespace_separator.to_string(), ); } - let credentials = bearer_access_token - .as_ref() - .or(oauth_access_token.as_ref()) - .map(|token| { - Arc::new(StaticCatalogCredentials::new( - token.expose_secret().to_string(), - )) as Arc - }) - .unwrap_or_else(|| Arc::new(EmptyCatalogCredentials)); + // A token file takes precedence over a static token: its + // contents are re-read per request so a rotated projected + // service account token is picked up without a restart. + let credentials = if let Some(path) = bearer_access_token_file { + Arc::new(FileCatalogCredentials::new(path.clone())) + as Arc + } else { + bearer_access_token + .as_ref() + .or(oauth_access_token.as_ref()) + .map(|token| { + Arc::new(StaticCatalogCredentials::new( + token.expose_secret().to_string(), + )) as Arc + }) + .unwrap_or_else(|| Arc::new(EmptyCatalogCredentials)) + }; let runtime_aware = RuntimeAwareCatalogProvider::try_new( || { From f4df7068d14f69928b179163963e2819a5bbe163 Mon Sep 17 00:00:00 2001 From: Robin Everaars Date: Fri, 24 Jul 2026 13:47:00 +0200 Subject: [PATCH 2/8] fix: reject an empty bearer token file and cover the credentials wiring An empty or whitespace-only token file now yields an error instead of caching an empty credential that would be sent as a bare bearer header. The empty read is never cached, so the next retrieve picks up the token as soon as the file holds one again. The Iceberg REST credentials selection moved into iceberg_rest_credentials with behavioral tests for the file precedence, the static and OAuth fallbacks and the no-token case, covering the previously untested wiring. The token file path no longer takes a needless clone. --- Cargo.lock | 1 + crates/sail-catalog/src/credentials.rs | 28 +++++++++ crates/sail-session/Cargo.toml | 3 + crates/sail-session/src/catalog.rs | 87 ++++++++++++++++++++------ 4 files changed, 101 insertions(+), 18 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index eb6e613cdd..cdfa78fe7b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8089,6 +8089,7 @@ dependencies = [ "sail-sql-analyzer", "sail-telemetry", "secrecy", + "tempfile", "thiserror 2.0.19", "tokio", "tonic", diff --git a/crates/sail-catalog/src/credentials.rs b/crates/sail-catalog/src/credentials.rs index 113ae12a36..c9ef480f88 100644 --- a/crates/sail-catalog/src/credentials.rs +++ b/crates/sail-catalog/src/credentials.rs @@ -101,6 +101,12 @@ impl CatalogCredentials for FileCatalogCredentials { })? .trim() .to_string(); + if credential.is_empty() { + return Err(CatalogError::External(format!( + "token file {} is empty", + self.path.display() + ))); + } let mut cached = self .cached .lock() @@ -202,4 +208,26 @@ mod tests { "unexpected error variant: {error:?}" ); } + + #[tokio::test] + async fn retrieve_reports_error_for_empty_file() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("token"); + write_token(&path, "\n \n"); + + let credentials = FileCatalogCredentials::new(&path); + let error = credentials.retrieve().await.unwrap_err(); + assert!( + matches!(&error, CatalogError::External(message) if message.contains("empty")), + "unexpected error: {error:?}" + ); + + // An empty read must not be cached: once the file holds a token again, + // the next retrieve returns it. + write_token(&path, "recovered-token"); + assert_eq!( + credentials.retrieve().await.unwrap(), + Some("recovered-token".to_string()) + ); + } } diff --git a/crates/sail-session/Cargo.toml b/crates/sail-session/Cargo.toml index 5c8d495263..d6f7b3b6bf 100644 --- a/crates/sail-session/Cargo.toml +++ b/crates/sail-session/Cargo.toml @@ -47,3 +47,6 @@ tokio = { workspace = true } chrono = { workspace = true } indexmap = { workspace = true } futures = { workspace = true } + +[dev-dependencies] +tempfile = { workspace = true } diff --git a/crates/sail-session/src/catalog.rs b/crates/sail-session/src/catalog.rs index 5c73e592c2..b147b4c68f 100644 --- a/crates/sail-session/src/catalog.rs +++ b/crates/sail-session/src/catalog.rs @@ -20,7 +20,7 @@ use sail_catalog_system::{SYSTEM_CATALOG_NAME, SystemCatalogProvider}; use sail_catalog_unity::{UnityCatalogConfig, UnityCatalogOptions, UnityCatalogProvider}; use sail_common::config::{AppConfig, CacheType, CatalogCacheConfig, CatalogType, OneLakeApi}; use sail_common::runtime::RuntimeHandle; -use secrecy::ExposeSecret; +use secrecy::{ExposeSecret, SecretString}; pub fn create_catalog_manager( config: &AppConfig, @@ -70,23 +70,11 @@ pub fn create_catalog_manager( namespace_separator.to_string(), ); } - // A token file takes precedence over a static token: its - // contents are re-read per request so a rotated projected - // service account token is picked up without a restart. - let credentials = if let Some(path) = bearer_access_token_file { - Arc::new(FileCatalogCredentials::new(path.clone())) - as Arc - } else { - bearer_access_token - .as_ref() - .or(oauth_access_token.as_ref()) - .map(|token| { - Arc::new(StaticCatalogCredentials::new( - token.expose_secret().to_string(), - )) as Arc - }) - .unwrap_or_else(|| Arc::new(EmptyCatalogCredentials)) - }; + let credentials = iceberg_rest_credentials( + bearer_access_token_file.as_ref(), + bearer_access_token.as_ref(), + oauth_access_token.as_ref(), + ); let runtime_aware = RuntimeAwareCatalogProvider::try_new( || { @@ -326,6 +314,28 @@ fn wrap_catalog_provider( Ok(Arc::new(provider)) } +/// Credentials for an Iceberg REST catalog. A token file takes precedence over +/// a static token: its contents are re-read per request, so a rotated projected +/// service account token is picked up without a restart. +fn iceberg_rest_credentials( + bearer_access_token_file: Option<&String>, + bearer_access_token: Option<&SecretString>, + oauth_access_token: Option<&SecretString>, +) -> Arc { + if let Some(path) = bearer_access_token_file { + Arc::new(FileCatalogCredentials::new(path)) as Arc + } else { + bearer_access_token + .or(oauth_access_token) + .map(|token| { + Arc::new(StaticCatalogCredentials::new( + token.expose_secret().to_string(), + )) as Arc + }) + .unwrap_or_else(|| Arc::new(EmptyCatalogCredentials)) + } +} + #[cfg(test)] #[expect(clippy::unwrap_used, clippy::expect_used)] mod tests { @@ -421,4 +431,45 @@ mod tests { "session cache should not be stored in global manager" ); } + + #[tokio::test] + async fn iceberg_credentials_prefer_the_token_file() { + let dir = tempfile::TempDir::new().unwrap(); + let path = dir.path().join("token"); + std::fs::write(&path, "file-token\n").unwrap(); + let path = path.to_string_lossy().to_string(); + let static_token = SecretString::from("static-token".to_string()); + + let credentials = iceberg_rest_credentials(Some(&path), Some(&static_token), None); + assert_eq!( + credentials.retrieve().await.unwrap(), + Some("file-token".to_string()) + ); + } + + #[tokio::test] + async fn iceberg_credentials_use_the_static_bearer_token() { + let token = SecretString::from("static-token".to_string()); + let credentials = iceberg_rest_credentials(None, Some(&token), None); + assert_eq!( + credentials.retrieve().await.unwrap(), + Some("static-token".to_string()) + ); + } + + #[tokio::test] + async fn iceberg_credentials_fall_back_to_the_oauth_token() { + let token = SecretString::from("oauth-token".to_string()); + let credentials = iceberg_rest_credentials(None, None, Some(&token)); + assert_eq!( + credentials.retrieve().await.unwrap(), + Some("oauth-token".to_string()) + ); + } + + #[tokio::test] + async fn iceberg_credentials_are_empty_without_any_token() { + let credentials = iceberg_rest_credentials(None, None, None); + assert_eq!(credentials.retrieve().await.unwrap(), None); + } } From b27792d577c01c2aad35970c6e764516248bc1b4 Mon Sep 17 00:00:00 2001 From: Robin Everaars Date: Mon, 27 Jul 2026 10:34:07 +0200 Subject: [PATCH 3/8] fix: retry Iceberg REST requests once on a 401 and read the token file per request The token file credential kept an in-memory cache guarded by an mtime check, a TTL and a mutex, yet still stat'd the file on every call, so the cache added state without a filesystem-free fast path. Drop the cache: retrieve() now reads the file, trims it and errors on an empty token, so every read is fresh. Because IcebergRestCatalogProvider resolved the credential once and baked the Authorization header into a single ApiClient, drop_database(cascade) reused that client across every list and drop request. A projected service account token that rotated partway through the cascade kept sending the stale header, and because the per-object drop errors are ignored the cascade could partially apply. Add a with_auth_retry helper that builds the client from a freshly read credential, runs the request and, on a 401, rebuilds the client and retries the request once. Every catalog operation goes through it, including each request in the cascade loop, so a mid-cascade rotation is recovered per request. The one-time startup config fetch is the exception: it caches its result and a 401 there is a hard startup failure that a same-token retry cannot recover. The shared reqwest::Client and its connection pool are reused across attempts. Cover the helper with unit tests (retry once on 401, no second retry, no retry on other errors) and add a mock server regression test where the server rejects the old token mid-cascade while the token file rotates, and the cascade drop completes with every request eventually authorized. --- Cargo.lock | 1 + crates/sail-catalog-iceberg/Cargo.toml | 1 + crates/sail-catalog-iceberg/src/provider.rs | 529 ++++++++++++++----- crates/sail-catalog/src/credentials.rs | 79 +-- crates/sail-common/src/config/application.rs | 10 +- 5 files changed, 407 insertions(+), 213 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index cdfa78fe7b..c27e5917f9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7445,6 +7445,7 @@ dependencies = [ "serde", "serde_json", "serde_with", + "tempfile", "testcontainers", "thiserror 2.0.19", "tokio", diff --git a/crates/sail-catalog-iceberg/Cargo.toml b/crates/sail-catalog-iceberg/Cargo.toml index b10133caa1..a60737dc9a 100644 --- a/crates/sail-catalog-iceberg/Cargo.toml +++ b/crates/sail-catalog-iceberg/Cargo.toml @@ -41,6 +41,7 @@ log = { workspace = true } tokio = { workspace = true } wiremock = { workspace = true } testcontainers = { workspace = true } +tempfile = { workspace = true } [build-dependencies] sail-build-scripts = { path = "../sail-build-scripts" } diff --git a/crates/sail-catalog-iceberg/src/provider.rs b/crates/sail-catalog-iceberg/src/provider.rs index 95397ee06e..7a72f25067 100644 --- a/crates/sail-catalog-iceberg/src/provider.rs +++ b/crates/sail-catalog-iceberg/src/provider.rs @@ -41,7 +41,7 @@ use sail_iceberg::{ }; use tokio::sync::OnceCell; -use crate::r#gen::ApiClient; +use crate::r#gen::{ApiClient, ApiError}; pub const REST_CATALOG_PROP_URI: &str = "uri"; @@ -188,6 +188,27 @@ impl IcebergRestCatalogProvider { .await } + /// Run a single outbound REST request, retrying it once if the server + /// answers `401 Unauthorized`. Each attempt builds an [`ApiClient`] from a + /// freshly resolved credential, so a projected service account token that + /// rotated mid-operation is picked up on the retry. The credential is + /// re-read per request, so every request in a `drop_database` cascade sees + /// the current token. The shared `reqwest::Client` and its connection pool + /// are reused across attempts. + async fn with_auth_retry(&self, call: F) -> CatalogResult>> + where + F: Fn(ApiClient) -> Fut, + Fut: std::future::Future>>, + { + let client = self.client().await?; + let result = call(client).await; + if matches!(&result, Err(e) if e.status() == Some(reqwest::StatusCode::UNAUTHORIZED)) { + let client = self.client().await?; + return Ok(call(client).await); + } + Ok(result) + } + // Merge the local catalog config with the [`crate::r#gen::CatalogConfig`] fetched from the REST server. // This only happens once, then the result is cached. async fn resolved_catalog_config(&self) -> CatalogResult<&CatalogConfig<'static>> { @@ -225,34 +246,39 @@ impl IcebergRestCatalogProvider { table: &str, access_delegation: Option<&str>, ) -> CatalogResult { - let client = self.client().await?; let catalog_config = self.resolved_catalog_config().await?; - client - .load_table( - catalog_config.prefix().map(ToOwned::to_owned), - catalog_config.namespace_string(database)?, - table.to_string(), - access_delegation.map(ToOwned::to_owned), - None, - None, - ) - .await - .map(|response| response.inner) - .map_err(|e| match e { - e if e.status() == Some(reqwest::StatusCode::NOT_FOUND) => CatalogError::NotFound( - CatalogObject::Table, - format!( - "{}.{}", - quote_namespace_if_needed(database), - quote_name_if_needed(table) - ), - ), - _ => CatalogError::External(format!( - "Failed to load table {}.{}: {e}", + let prefix = catalog_config.prefix().map(ToOwned::to_owned); + let namespace = catalog_config.namespace_string(database)?; + let table_name = table.to_string(); + let access_delegation = access_delegation.map(ToOwned::to_owned); + self.with_auth_retry(|client| { + let prefix = prefix.clone(); + let namespace = namespace.clone(); + let table_name = table_name.clone(); + let access_delegation = access_delegation.clone(); + async move { + client + .load_table(prefix, namespace, table_name, access_delegation, None, None) + .await + } + }) + .await? + .map(|response| response.inner) + .map_err(|e| match e { + e if e.status() == Some(reqwest::StatusCode::NOT_FOUND) => CatalogError::NotFound( + CatalogObject::Table, + format!( + "{}.{}", quote_namespace_if_needed(database), quote_name_if_needed(table) - )), - }) + ), + ), + _ => CatalogError::External(format!( + "Failed to load table {}.{}: {e}", + quote_namespace_if_needed(database), + quote_name_if_needed(table) + )), + }) } fn normalize_scan_planning_mode(value: &str) -> CatalogResult { @@ -782,7 +808,6 @@ impl CatalogProvider for IcebergRestCatalogProvider { options: CreateDatabaseOptions, ) -> CatalogResult { let catalog_config = self.resolved_catalog_config().await?; - let client = self.client().await?; let CreateDatabaseOptions { if_not_exists, @@ -803,10 +828,15 @@ impl CatalogProvider for IcebergRestCatalogProvider { namespace: Box::new(database.clone().into()), properties: if props.is_empty() { None } else { Some(props) }, }; + let prefix = catalog_config.prefix().map(ToOwned::to_owned); - let result = client - .create_namespace(catalog_config.prefix().map(ToOwned::to_owned), request) - .await + let result = self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let request = request.clone(); + async move { client.create_namespace(prefix, request).await } + }) + .await? .map(|response| response.inner); match result { @@ -841,12 +871,16 @@ impl CatalogProvider for IcebergRestCatalogProvider { async fn get_database(&self, database: &Namespace) -> CatalogResult { let catalog_config = self.resolved_catalog_config().await?; - let client = self.client().await?; let namespace = catalog_config.namespace_string(database)?; + let prefix = catalog_config.prefix().map(ToOwned::to_owned); - let result = client - .load_namespace_metadata(catalog_config.prefix().map(ToOwned::to_owned), namespace) - .await + let result = self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let namespace = namespace.clone(); + async move { client.load_namespace_metadata(prefix, namespace).await } + }) + .await? .map(|response| response.inner) .map_err(|e| match e { e if e.status() == Some(reqwest::StatusCode::NOT_FOUND) => CatalogError::NotFound( @@ -883,19 +917,22 @@ impl CatalogProvider for IcebergRestCatalogProvider { prefix: Option<&Namespace>, ) -> CatalogResult> { let catalog_config = self.resolved_catalog_config().await?; - let client = self.client().await?; let parent = prefix .map(|namespace| catalog_config.namespace_string(namespace)) .transpose()?; + let request_prefix = catalog_config.prefix().map(ToOwned::to_owned); - let result = client - .list_namespaces( - catalog_config.prefix().map(ToOwned::to_owned), - None, - None, - parent, - ) - .await + let result = self + .with_auth_retry(|client| { + let request_prefix = request_prefix.clone(); + let parent = parent.clone(); + async move { + client + .list_namespaces(request_prefix, None, None, parent) + .await + } + }) + .await? .map(|response| response.inner) .map_err(|e| CatalogError::External(format!("Failed to list namespaces: {e}")))?; @@ -919,47 +956,63 @@ impl CatalogProvider for IcebergRestCatalogProvider { options: DropDatabaseOptions, ) -> CatalogResult<()> { let catalog_config = self.resolved_catalog_config().await?; - let client = self.client().await?; let DropDatabaseOptions { if_exists, cascade } = options; + let prefix = catalog_config.prefix().map(ToOwned::to_owned); + let ns_string = catalog_config.namespace_string(database)?; + if cascade { // For CASCADE, first drop all tables and views in the namespace before dropping the namespace. - let prefix = catalog_config.prefix().map(ToOwned::to_owned); - let ns_string = catalog_config.namespace_string(database)?; - let tables_result = client - .list_tables(prefix.clone(), ns_string.clone(), None, None) - .await; + // Each request re-reads the credential and retries once on a 401, so a token that rotates + // partway through the cascade is recovered per request instead of leaving a partial drop. + let tables_result = self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let ns_string = ns_string.clone(); + async move { client.list_tables(prefix, ns_string, None, None).await } + }) + .await?; if let Ok(tables) = tables_result { for identifier in tables.inner.identifiers.unwrap_or_default() { - let _ = client - .drop_table( - prefix.clone(), - ns_string.clone(), - identifier.name, - Some(true), - ) - .await; + let _ = self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let ns_string = ns_string.clone(); + let name = identifier.name.clone(); + async move { client.drop_table(prefix, ns_string, name, Some(true)).await } + }) + .await?; } } - let views_result = client - .list_views(prefix.clone(), ns_string.clone(), None, None) - .await; + let views_result = self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let ns_string = ns_string.clone(); + async move { client.list_views(prefix, ns_string, None, None).await } + }) + .await?; if let Ok(views) = views_result { for identifier in views.inner.identifiers.unwrap_or_default() { - let _ = client - .drop_view(prefix.clone(), ns_string.clone(), identifier.name) - .await; + let _ = self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let ns_string = ns_string.clone(); + let name = identifier.name.clone(); + async move { client.drop_view(prefix, ns_string, name).await } + }) + .await?; } } } - match client - .drop_namespace( - catalog_config.prefix().map(ToOwned::to_owned), - catalog_config.namespace_string(database)?, - ) - .await + match self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let ns_string = ns_string.clone(); + async move { client.drop_namespace(prefix, ns_string).await } + }) + .await? { Ok(_) => Ok(()), Err(e) if e.status() == Some(reqwest::StatusCode::NOT_FOUND) && if_exists => Ok(()), @@ -997,7 +1050,6 @@ impl CatalogProvider for IcebergRestCatalogProvider { } let catalog_config = self.resolved_catalog_config().await?; - let client = self.client().await?; if mode.ignore_if_exists() && let Ok(existing) = self.get_table(database, table).await @@ -1061,14 +1113,16 @@ impl CatalogProvider for IcebergRestCatalogProvider { properties: if props.is_empty() { None } else { Some(props) }, }; - let result = client - .create_table( - catalog_config.prefix().map(ToOwned::to_owned), - catalog_config.namespace_string(database)?, - None, - request, - ) - .await + let prefix = catalog_config.prefix().map(ToOwned::to_owned); + let namespace = catalog_config.namespace_string(database)?; + let result = self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let namespace = namespace.clone(); + let request = request.clone(); + async move { client.create_table(prefix, namespace, None, request).await } + }) + .await? .map(|response| response.inner) .map_err(|e| CatalogError::External(format!("Failed to create table: {e}")))?; @@ -1091,16 +1145,16 @@ impl CatalogProvider for IcebergRestCatalogProvider { async fn list_tables(&self, database: &Namespace) -> CatalogResult> { let catalog_config = self.resolved_catalog_config().await?; - let client = self.client().await?; + let prefix = catalog_config.prefix().map(ToOwned::to_owned); + let namespace = catalog_config.namespace_string(database)?; - let result = client - .list_tables( - catalog_config.prefix().map(ToOwned::to_owned), - catalog_config.namespace_string(database)?, - None, - None, - ) - .await + let result = self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let namespace = namespace.clone(); + async move { client.list_tables(prefix, namespace, None, None).await } + }) + .await? .map(|response| response.inner) .map_err(|e| CatalogError::External(format!("Failed to list tables: {e}")))?; @@ -1135,16 +1189,22 @@ impl CatalogProvider for IcebergRestCatalogProvider { options: DropTableOptions, ) -> CatalogResult<()> { let catalog_config = self.resolved_catalog_config().await?; - let client = self.client().await?; let DropTableOptions { if_exists, purge } = options; - match client - .drop_table( - catalog_config.prefix().map(ToOwned::to_owned), - catalog_config.namespace_string(database)?, - table.to_string(), - Some(purge), - ) - .await + let prefix = catalog_config.prefix().map(ToOwned::to_owned); + let namespace = catalog_config.namespace_string(database)?; + let table_name = table.to_string(); + match self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let namespace = namespace.clone(); + let table_name = table_name.clone(); + async move { + client + .drop_table(prefix, namespace, table_name, Some(purge)) + .await + } + }) + .await? { Ok(_) => Ok(()), Err(e) if e.status() == Some(reqwest::StatusCode::NOT_FOUND) && if_exists => Ok(()), @@ -1183,7 +1243,6 @@ impl CatalogProvider for IcebergRestCatalogProvider { } let catalog_config = self.resolved_catalog_config().await?; - let client = self.client().await?; let namespace = catalog_config.namespace_string(database)?; let requirements = requirements .into_iter() @@ -1209,14 +1268,21 @@ impl CatalogProvider for IcebergRestCatalogProvider { requirements, updates, }; - let response = client - .update_table( - catalog_config.prefix().map(ToOwned::to_owned), - namespace, - table.to_string(), - request, - ) - .await + let prefix = catalog_config.prefix().map(ToOwned::to_owned); + let table_name = table.to_string(); + let response = self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let namespace = namespace.clone(); + let table_name = table_name.clone(); + let request = request.clone(); + async move { + client + .update_table(prefix, namespace, table_name, request) + .await + } + }) + .await? .map(|response| response.inner) .map_err(|e| match e { e if e.status() == Some(reqwest::StatusCode::NOT_FOUND) => CatalogError::NotFound( @@ -1330,7 +1396,6 @@ impl CatalogProvider for IcebergRestCatalogProvider { options: CreateViewOptions, ) -> CatalogResult { let catalog_config = self.resolved_catalog_config().await?; - let client = self.client().await?; let CreateViewOptions { columns, @@ -1440,13 +1505,16 @@ impl CatalogProvider for IcebergRestCatalogProvider { properties: props, }; - let result = client - .create_view( - catalog_config.prefix().map(ToOwned::to_owned), - catalog_config.namespace_string(database)?, - request, - ) - .await + let prefix = catalog_config.prefix().map(ToOwned::to_owned); + let namespace = catalog_config.namespace_string(database)?; + let result = self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let namespace = namespace.clone(); + let request = request.clone(); + async move { client.create_view(prefix, namespace, request).await } + }) + .await? .map(|response| response.inner) .map_err(|e| CatalogError::External(format!("Failed to create view: {e}")))?; @@ -1455,14 +1523,17 @@ impl CatalogProvider for IcebergRestCatalogProvider { async fn get_view(&self, database: &Namespace, view: &str) -> CatalogResult { let catalog_config = self.resolved_catalog_config().await?; - let client = self.client().await?; - let result = client - .load_view( - catalog_config.prefix().map(ToOwned::to_owned), - catalog_config.namespace_string(database)?, - view.to_string(), - ) - .await + let prefix = catalog_config.prefix().map(ToOwned::to_owned); + let namespace = catalog_config.namespace_string(database)?; + let view_name = view.to_string(); + let result = self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let namespace = namespace.clone(); + let view_name = view_name.clone(); + async move { client.load_view(prefix, namespace, view_name).await } + }) + .await? .map(|response| response.inner) .map_err(|e| match e { e if matches!( @@ -1492,16 +1563,16 @@ impl CatalogProvider for IcebergRestCatalogProvider { async fn list_views(&self, database: &Namespace) -> CatalogResult> { let catalog_config = self.resolved_catalog_config().await?; - let client = self.client().await?; + let prefix = catalog_config.prefix().map(ToOwned::to_owned); + let namespace = catalog_config.namespace_string(database)?; - let result = client - .list_views( - catalog_config.prefix().map(ToOwned::to_owned), - catalog_config.namespace_string(database)?, - None, - None, - ) - .await + let result = self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let namespace = namespace.clone(); + async move { client.list_views(prefix, namespace, None, None).await } + }) + .await? .map(|response| response.inner) .map_err(|e| match e { e if matches!(e.status(), Some(reqwest::StatusCode::NOT_FOUND)) => { @@ -1546,15 +1617,18 @@ impl CatalogProvider for IcebergRestCatalogProvider { options: DropViewOptions, ) -> CatalogResult<()> { let catalog_config = self.resolved_catalog_config().await?; - let client = self.client().await?; let DropViewOptions { if_exists } = options; - match client - .drop_view( - catalog_config.prefix().map(ToOwned::to_owned), - catalog_config.namespace_string(database)?, - view.to_string(), - ) - .await + let prefix = catalog_config.prefix().map(ToOwned::to_owned); + let namespace = catalog_config.namespace_string(database)?; + let view_name = view.to_string(); + match self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let namespace = namespace.clone(); + let view_name = view_name.clone(); + async move { client.drop_view(prefix, namespace, view_name).await } + }) + .await? { Ok(_) => Ok(()), Err(e) if e.status() == Some(reqwest::StatusCode::NOT_FOUND) && if_exists => Ok(()), @@ -1917,8 +1991,10 @@ fn parse_unary_sort_transform( #[expect(clippy::unwrap_used, clippy::panic)] #[cfg(test)] mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use arrow::datatypes::DataType; - use sail_catalog::credentials::EmptyCatalogCredentials; + use sail_catalog::credentials::{EmptyCatalogCredentials, FileCatalogCredentials}; use sail_catalog::lakehouse::TableAccessPurpose; use sail_common::spec; use sail_common_datafusion::catalog::{ @@ -1926,8 +2002,9 @@ mod tests { LakehouseExecutionContext, LakehouseFormat, LakehouseOperation, MetadataPointerAuthority, TableLifecycle, }; + use tempfile::TempDir; use wiremock::matchers::{header, method, path, query_param, query_param_is_missing}; - use wiremock::{Mock, MockServer, ResponseTemplate}; + use wiremock::{Mock, MockServer, Request, ResponseTemplate}; use super::*; @@ -3471,4 +3548,176 @@ mod tests { test_get_database_impl(None).await; test_get_database_impl(Some("test")).await; } + + fn error_with_status(status: reqwest::StatusCode) -> ApiError<()> { + ApiError::Unknown(crate::r#gen::Response { + inner: (), + status, + headers: reqwest::header::HeaderMap::new(), + }) + } + + #[tokio::test] + async fn with_auth_retry_retries_once_on_unauthorized() { + let ctx = TestContext::new(None).await; + let calls = AtomicUsize::new(0); + let outcome: Result<(), ApiError<()>> = ctx + .catalog + .with_auth_retry(|_client| { + let attempt = calls.fetch_add(1, Ordering::SeqCst); + async move { + if attempt == 0 { + Err(error_with_status(reqwest::StatusCode::UNAUTHORIZED)) + } else { + Ok(()) + } + } + }) + .await + .unwrap(); + assert!(outcome.is_ok()); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn with_auth_retry_does_not_retry_more_than_once() { + let ctx = TestContext::new(None).await; + let calls = AtomicUsize::new(0); + let outcome: Result<(), ApiError<()>> = ctx + .catalog + .with_auth_retry(|_client| { + calls.fetch_add(1, Ordering::SeqCst); + async move { Err(error_with_status(reqwest::StatusCode::UNAUTHORIZED)) } + }) + .await + .unwrap(); + assert!(outcome.is_err()); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn with_auth_retry_does_not_retry_non_unauthorized() { + let ctx = TestContext::new(None).await; + let calls = AtomicUsize::new(0); + let outcome: Result<(), ApiError<()>> = ctx + .catalog + .with_auth_retry(|_client| { + calls.fetch_add(1, Ordering::SeqCst); + async move { + Err(error_with_status( + reqwest::StatusCode::INTERNAL_SERVER_ERROR, + )) + } + }) + .await + .unwrap(); + assert!(outcome.is_err()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn drop_database_cascade_recovers_when_token_rotates_midway() { + let dir = TempDir::new().unwrap(); + let token_path = dir.path().join("token"); + std::fs::write(&token_path, "token-a").unwrap(); + + let server = MockServer::start().await; + + // Bootstrap config. Reachable with the original token. + Mock::given(method("GET")) + .and(path("/v1/config")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "overrides": { "warehouse": "s3://iceberg-catalog" }, + "defaults": {} + }))) + .mount(&server) + .await; + + // The namespace still holds one table when the cascade begins. Listing + // is authorized with the original token. + Mock::given(method("GET")) + .and(path("/v1/namespaces/dbc/tables")) + .and(header("authorization", "Bearer token-a")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "identifiers": [ { "namespace": ["dbc"], "name": "t1" } ] + }))) + .mount(&server) + .await; + + // The first drop of the table arrives with the old token. The server + // rejects it with 401 and, at that moment, the projected token file + // rotates to a new value (as kubelet would swap it). + let rotate_path = token_path.clone(); + Mock::given(method("DELETE")) + .and(path("/v1/namespaces/dbc/tables/t1")) + .and(header("authorization", "Bearer token-a")) + .respond_with(move |_req: &Request| { + std::fs::write(&rotate_path, "token-b").unwrap(); + ResponseTemplate::new(401).set_body_json(serde_json::json!({ + "error": { + "message": "token expired", + "type": "NotAuthorizedException", + "code": 401 + } + })) + }) + .expect(1) + .mount(&server) + .await; + + // The retry re-reads the rotated token and is authorized. + Mock::given(method("DELETE")) + .and(path("/v1/namespaces/dbc/tables/t1")) + .and(header("authorization", "Bearer token-b")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + + // The remaining cascade requests all use the rotated token. + Mock::given(method("GET")) + .and(path("/v1/namespaces/dbc/views")) + .and(header("authorization", "Bearer token-b")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "identifiers": [] + }))) + .expect(1) + .mount(&server) + .await; + + Mock::given(method("DELETE")) + .and(path("/v1/namespaces/dbc")) + .and(header("authorization", "Bearer token-b")) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&server) + .await; + + let props = HashMap::from([(REST_CATALOG_PROP_URI.to_string(), server.uri())]); + let options = IcebergRestCatalogOptions { + credentials: Arc::new(FileCatalogCredentials::new(&token_path)), + properties: props, + }; + let catalog = IcebergRestCatalogProvider::new(String::new(), options); + + let namespace = Namespace::try_from(vec!["dbc".to_string()]).unwrap(); + let result = catalog + .drop_database( + &namespace, + DropDatabaseOptions { + if_exists: false, + cascade: true, + }, + ) + .await; + + assert!(result.is_ok(), "cascade drop should succeed: {result:?}"); + // The rotated token is the one the file ends up holding, and every + // mounted request expectation (including the retried table drop) is + // verified when the server is dropped. + assert_eq!( + std::fs::read_to_string(&token_path).unwrap(), + "token-b".to_string() + ); + } } diff --git a/crates/sail-catalog/src/credentials.rs b/crates/sail-catalog/src/credentials.rs index c9ef480f88..80991b548c 100644 --- a/crates/sail-catalog/src/credentials.rs +++ b/crates/sail-catalog/src/credentials.rs @@ -1,7 +1,5 @@ use std::fmt::Debug; use std::path::PathBuf; -use std::sync::Mutex; -use std::time::{Duration, Instant, SystemTime}; use crate::error::{CatalogError, CatalogResult}; @@ -38,59 +36,25 @@ impl CatalogCredentials for StaticCatalogCredentials { } } -/// How long a token read from a [`FileCatalogCredentials`] file is reused before -/// the file is checked again. Kubelet refreshes projected service account tokens -/// well before they expire (by default once 80% of their lifetime has elapsed), -/// so a short interval keeps the in-memory token fresh without reading the file -/// on every catalog request. -const FILE_CREDENTIALS_REFRESH_INTERVAL: Duration = Duration::from_secs(60); - /// Credentials backed by a token file on disk, such as a kubelet-projected -/// service account token. The token is cached in memory and re-read when the -/// file's modification time changes or the refresh interval elapses, so a -/// rotated token is picked up without restarting the server. +/// service account token. The file is read on every call, so a rotated token +/// is picked up without restarting the server. The Iceberg REST provider reads +/// the credential fresh for each request and retries once on a `401`, so a +/// token that rotates mid-operation is recovered without an in-memory cache. #[derive(Debug)] pub struct FileCatalogCredentials { path: PathBuf, - cached: Mutex>, -} - -#[derive(Debug, Clone)] -struct CachedCredential { - credential: String, - modified: Option, - read_at: Instant, } impl FileCatalogCredentials { pub fn new(path: impl Into) -> Self { - Self { - path: path.into(), - cached: Mutex::new(None), - } + Self { path: path.into() } } } #[async_trait::async_trait] impl CatalogCredentials for FileCatalogCredentials { async fn retrieve(&self) -> CatalogResult> { - let modified = tokio::fs::metadata(&self.path) - .await - .and_then(|metadata| metadata.modified()) - .ok(); - { - let cached = self - .cached - .lock() - .map_err(|e| CatalogError::Internal(format!("token file cache poisoned: {e}")))?; - if let Some(cached) = cached.as_ref() { - let fresh = cached.read_at.elapsed() < FILE_CREDENTIALS_REFRESH_INTERVAL; - let changed = modified.is_some() && modified != cached.modified; - if fresh && !changed { - return Ok(Some(cached.credential.clone())); - } - } - } let credential = tokio::fs::read_to_string(&self.path) .await .map_err(|e| { @@ -107,15 +71,6 @@ impl CatalogCredentials for FileCatalogCredentials { self.path.display() ))); } - let mut cached = self - .cached - .lock() - .map_err(|e| CatalogError::Internal(format!("token file cache poisoned: {e}")))?; - *cached = Some(CachedCredential { - credential: credential.clone(), - modified, - read_at: Instant::now(), - }); Ok(Some(credential)) } } @@ -124,10 +79,9 @@ impl CatalogCredentials for FileCatalogCredentials { mod tests { #![expect(clippy::unwrap_used)] - use std::fs::{File, FileTimes}; + use std::fs::File; use std::io::Write; use std::path::Path; - use std::time::{Duration, SystemTime}; use tempfile::TempDir; @@ -138,12 +92,6 @@ mod tests { file.write_all(contents.as_bytes()).unwrap(); } - fn set_modified(path: &Path, modified: SystemTime) { - let file = File::options().write(true).open(path).unwrap(); - file.set_times(FileTimes::new().set_modified(modified)) - .unwrap(); - } - #[tokio::test] async fn retrieve_returns_token_from_file() { let dir = TempDir::new().unwrap(); @@ -171,12 +119,10 @@ mod tests { } #[tokio::test] - async fn retrieve_rereads_when_modification_time_changes() { + async fn retrieve_rereads_rotated_token() { let dir = TempDir::new().unwrap(); let path = dir.path().join("token"); write_token(&path, "first-token"); - let earlier = SystemTime::UNIX_EPOCH + Duration::from_secs(1_600_000_000); - set_modified(&path, earlier); let credentials = FileCatalogCredentials::new(&path); assert_eq!( @@ -184,12 +130,10 @@ mod tests { Some("first-token".to_string()) ); - // Rotate the token and advance the file's modification time. The refresh - // interval has not elapsed, so the mtime change is the only thing that can - // trigger a re-read. This mirrors kubelet swapping a projected token. + // Every call reads the file, so a rotated token (for example kubelet + // swapping a projected service account token) is picked up on the next + // retrieve without restarting the server. write_token(&path, "second-token"); - set_modified(&path, earlier + Duration::from_secs(60)); - assert_eq!( credentials.retrieve().await.unwrap(), Some("second-token".to_string()) @@ -222,8 +166,7 @@ mod tests { "unexpected error: {error:?}" ); - // An empty read must not be cached: once the file holds a token again, - // the next retrieve returns it. + // Once the file holds a token again, the next retrieve returns it. write_token(&path, "recovered-token"); assert_eq!( credentials.retrieve().await.unwrap(), diff --git a/crates/sail-common/src/config/application.rs b/crates/sail-common/src/config/application.rs index cf220f6fed..aa508e6b7c 100644 --- a/crates/sail-common/src/config/application.rs +++ b/crates/sail-common/src/config/application.rs @@ -552,11 +552,11 @@ pub enum CatalogType { )] bearer_access_token: Option, /// Path to a file holding the bearer token. When set, the token is - /// re-read from this file per request (with a short refresh interval) - /// so a rotated token (for example a kubelet-projected service account - /// token) is picked up without restarting the server. Takes precedence - /// over `bearer_access_token`. The path is not a secret, so it is kept - /// as a plain string. + /// re-read from this file for every request, so a rotated token (for + /// example a kubelet-projected service account token) is picked up + /// without restarting the server. Takes precedence over + /// `bearer_access_token`. The path is not a secret, so it is kept as a + /// plain string. #[serde(skip_serializing_if = "Option::is_none")] bearer_access_token_file: Option, #[serde(flatten)] From 084d158ce5380f83e7b0aca344ff0856bc499a47 Mon Sep 17 00:00:00 2001 From: Robin Everaars Date: Mon, 27 Jul 2026 10:34:12 +0200 Subject: [PATCH 4/8] docs: document the Iceberg REST bearer token file option Add bearer_access_token_file to the Iceberg REST catalog options and an example, noting that the token is re-read per request and that a request rejected with a 401 is retried once after reloading the file. --- docs/guide/catalog/iceberg-rest.md | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/docs/guide/catalog/iceberg-rest.md b/docs/guide/catalog/iceberg-rest.md index 06b7dd5d3d..a99cc261f4 100644 --- a/docs/guide/catalog/iceberg-rest.md +++ b/docs/guide/catalog/iceberg-rest.md @@ -20,6 +20,10 @@ An Iceberg REST catalog can be configured using the following options: If unset or empty, Sail uses the Iceberg REST default unit separator. - `oauth_access_token` (optional): The OAuth 2.0 access token. - `bearer_access_token` (optional): The bearer token for authentication. +- `bearer_access_token_file` (optional): Path to a file that holds the bearer token. + Sail reads the token from this file for every request, so a rotated token (for example a kubelet-projected service account token) is picked up without restarting the server. + If a request is rejected with `401 Unauthorized`, Sail reloads the file and retries the request once, so a token that rotates midway through a multi-step operation is recovered. + This option takes precedence over `bearer_access_token`. See [Common Options](./index.md#common-options) for caching configuration. @@ -44,6 +48,9 @@ export SAIL_CATALOG__LIST='[{type="iceberg-rest", name="sail", uri="https://cata # Bearer token authentication export SAIL_CATALOG__LIST='[{type="iceberg-rest", name="sail", uri="https://catalog.example.com", warehouse="s3://data/warehouse", bearer_access_token="..."}]' +# Bearer token read from a file (for example a kubelet-projected service account token) +export SAIL_CATALOG__LIST='[{type="iceberg-rest", name="sail", uri="https://catalog.example.com", warehouse="s3://data/warehouse", bearer_access_token_file="/var/run/secrets/tokens/catalog-token"}]' + # Client-side namespace separator fallback export SAIL_CATALOG__LIST='[{type="iceberg-rest", name="sail", uri="https://catalog.example.com", namespace_separator="::"}]' ``` From 1ddf806058976e098f07763c248d5323bac4f3b1 Mon Sep 17 00:00:00 2001 From: Robin Everaars Date: Mon, 27 Jul 2026 11:49:57 +0200 Subject: [PATCH 5/8] fix: propagate swallowed cascade errors in Iceberg drop_database The CASCADE path wrapped the table and view listings in `if let Ok(...)` and each per-object drop in `let _ =`, so any failure was silently discarded. A transient error on a table drop or a listing produced a half applied cascade that still dropped the namespace and reported success. Every list_tables, list_views, drop_table and drop_view failure now propagates as an error. Only benign statuses are tolerated: NOT_FOUND on a per-object drop (a concurrent removal), NOT_FOUND on a listing (the namespace is already gone, handled by the trailing if_exists drop) and 405 or 501 on list_views (a catalog with no views endpoint, which the old code also tolerated). Signed-off-by: HOIST IT B.V. --- crates/sail-catalog-iceberg/src/provider.rs | 428 ++++++++++++++++++-- 1 file changed, 404 insertions(+), 24 deletions(-) diff --git a/crates/sail-catalog-iceberg/src/provider.rs b/crates/sail-catalog-iceberg/src/provider.rs index 7a72f25067..464df91e79 100644 --- a/crates/sail-catalog-iceberg/src/provider.rs +++ b/crates/sail-catalog-iceberg/src/provider.rs @@ -966,42 +966,99 @@ impl CatalogProvider for IcebergRestCatalogProvider { // For CASCADE, first drop all tables and views in the namespace before dropping the namespace. // Each request re-reads the credential and retries once on a 401, so a token that rotates // partway through the cascade is recovered per request instead of leaving a partial drop. - let tables_result = self + match self .with_auth_retry(|client| { let prefix = prefix.clone(); let ns_string = ns_string.clone(); async move { client.list_tables(prefix, ns_string, None, None).await } }) - .await?; - if let Ok(tables) = tables_result { - for identifier in tables.inner.identifiers.unwrap_or_default() { - let _ = self - .with_auth_retry(|client| { - let prefix = prefix.clone(); - let ns_string = ns_string.clone(); - let name = identifier.name.clone(); - async move { client.drop_table(prefix, ns_string, name, Some(true)).await } - }) - .await?; + .await? + { + Ok(tables) => { + for identifier in tables.inner.identifiers.unwrap_or_default() { + match self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let ns_string = ns_string.clone(); + let name = identifier.name.clone(); + async move { + client.drop_table(prefix, ns_string, name, Some(true)).await + } + }) + .await? + { + Ok(_) => {} + // The table was already removed (a concurrent drop), which is + // an acceptable outcome for a cascade, so keep going. + Err(e) if e.status() == Some(reqwest::StatusCode::NOT_FOUND) => {} + Err(e) => { + return Err(CatalogError::External(format!( + "Failed to drop table '{}' while cascading namespace drop: {e}", + identifier.name + ))); + } + } + } + } + // The namespace itself is already gone; fall through to drop_namespace, + // which applies the canonical if_exists handling below. + Err(e) if e.status() == Some(reqwest::StatusCode::NOT_FOUND) => {} + Err(e) => { + return Err(CatalogError::External(format!( + "Failed to list tables while cascading namespace drop: {e}" + ))); } } - let views_result = self + match self .with_auth_retry(|client| { let prefix = prefix.clone(); let ns_string = ns_string.clone(); async move { client.list_views(prefix, ns_string, None, None).await } }) - .await?; - if let Ok(views) = views_result { - for identifier in views.inner.identifiers.unwrap_or_default() { - let _ = self - .with_auth_retry(|client| { - let prefix = prefix.clone(); - let ns_string = ns_string.clone(); - let name = identifier.name.clone(); - async move { client.drop_view(prefix, ns_string, name).await } - }) - .await?; + .await? + { + Ok(views) => { + for identifier in views.inner.identifiers.unwrap_or_default() { + match self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let ns_string = ns_string.clone(); + let name = identifier.name.clone(); + async move { client.drop_view(prefix, ns_string, name).await } + }) + .await? + { + Ok(_) => {} + // The view was already removed (a concurrent drop), which is + // an acceptable outcome for a cascade, so keep going. + Err(e) if e.status() == Some(reqwest::StatusCode::NOT_FOUND) => {} + Err(e) => { + return Err(CatalogError::External(format!( + "Failed to drop view '{}' while cascading namespace drop: {e}", + identifier.name + ))); + } + } + } + } + // The namespace itself is already gone; fall through to drop_namespace, + // which applies the canonical if_exists handling below. + Err(e) if e.status() == Some(reqwest::StatusCode::NOT_FOUND) => {} + // The views endpoint is optional in the Iceberg REST spec, so a catalog + // that does not implement it answers 405 or 501. There are then no views + // to cascade, so tolerate it and continue to the namespace drop. The tables + // endpoint is mandatory, so the list_tables arm above does not tolerate these + // statuses and a 405 or 501 there is surfaced as a genuine failure. + Err(e) + if matches!( + e.status(), + Some(reqwest::StatusCode::METHOD_NOT_ALLOWED) + | Some(reqwest::StatusCode::NOT_IMPLEMENTED) + ) => {} + Err(e) => { + return Err(CatalogError::External(format!( + "Failed to list views while cascading namespace drop: {e}" + ))); } } } @@ -2728,6 +2785,329 @@ mod tests { test_drop_database_impl(Some("test")).await; } + async fn test_drop_database_cascade_propagates_table_drop_failure_impl(name: Option<&str>) { + let ctx = TestContext::new(name).await; + let namespace = Namespace::try_from(vec!["ns1".to_string()]).unwrap(); + + ctx.mock_get_json( + &ctx.path("/namespaces/ns1/tables"), + serde_json::json!({ + "identifiers": [ + { + "namespace": ["ns1"], + "name": "table1" + } + ] + }), + ) + .await; + + // The per-object table drop hits a real server error, so the cascade must abort. + Mock::given(method("DELETE")) + .and(path(ctx.path("/namespaces/ns1/tables/table1").as_str())) + .respond_with(ResponseTemplate::new(500)) + .mount(&ctx.server) + .await; + + // The namespace drop must never be attempted once a table drop fails. + Mock::given(method("DELETE")) + .and(path(ctx.path("/namespaces/ns1").as_str())) + .respond_with(ResponseTemplate::new(204)) + .expect(0) + .mount(&ctx.server) + .await; + + let result = ctx + .catalog + .drop_database( + &namespace, + DropDatabaseOptions { + if_exists: false, + cascade: true, + }, + ) + .await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_drop_database_cascade_propagates_table_drop_failure() { + test_drop_database_cascade_propagates_table_drop_failure_impl(None).await; + test_drop_database_cascade_propagates_table_drop_failure_impl(Some("test")).await; + } + + async fn test_drop_database_cascade_tolerates_missing_table_impl(name: Option<&str>) { + let ctx = TestContext::new(name).await; + let namespace = Namespace::try_from(vec!["ns1".to_string()]).unwrap(); + + ctx.mock_get_json( + &ctx.path("/namespaces/ns1/tables"), + serde_json::json!({ + "identifiers": [ + { + "namespace": ["ns1"], + "name": "table1" + } + ] + }), + ) + .await; + + // A concurrent removal leaves the table already gone; the cascade tolerates that. + ctx.mock_delete_404( + &ctx.path("/namespaces/ns1/tables/table1"), + "NoSuchTableException", + "The given table does not exist", + ) + .await; + + ctx.mock_get_json( + &ctx.path("/namespaces/ns1/views"), + serde_json::json!({ "identifiers": [] }), + ) + .await; + + // The cascade proceeds and drops the namespace itself. + Mock::given(method("DELETE")) + .and(path(ctx.path("/namespaces/ns1").as_str())) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&ctx.server) + .await; + + let result = ctx + .catalog + .drop_database( + &namespace, + DropDatabaseOptions { + if_exists: false, + cascade: true, + }, + ) + .await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_drop_database_cascade_tolerates_missing_table() { + test_drop_database_cascade_tolerates_missing_table_impl(None).await; + test_drop_database_cascade_tolerates_missing_table_impl(Some("test")).await; + } + + async fn test_drop_database_cascade_propagates_list_tables_failure_impl(name: Option<&str>) { + let ctx = TestContext::new(name).await; + let namespace = Namespace::try_from(vec!["ns1".to_string()]).unwrap(); + + // Listing the tables fails outright, which the cascade must surface. + Mock::given(method("GET")) + .and(path(ctx.path("/namespaces/ns1/tables").as_str())) + .respond_with(ResponseTemplate::new(500)) + .mount(&ctx.server) + .await; + + // The namespace drop must never be attempted once the listing fails. + Mock::given(method("DELETE")) + .and(path(ctx.path("/namespaces/ns1").as_str())) + .respond_with(ResponseTemplate::new(204)) + .expect(0) + .mount(&ctx.server) + .await; + + let result = ctx + .catalog + .drop_database( + &namespace, + DropDatabaseOptions { + if_exists: false, + cascade: true, + }, + ) + .await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_drop_database_cascade_propagates_list_tables_failure() { + test_drop_database_cascade_propagates_list_tables_failure_impl(None).await; + test_drop_database_cascade_propagates_list_tables_failure_impl(Some("test")).await; + } + + async fn test_drop_database_cascade_tolerates_missing_views_endpoint_impl(name: Option<&str>) { + let ctx = TestContext::new(name).await; + let namespace = Namespace::try_from(vec!["ns1".to_string()]).unwrap(); + + ctx.mock_get_json( + &ctx.path("/namespaces/ns1/tables"), + serde_json::json!({ "identifiers": [] }), + ) + .await; + + // A catalog without a views endpoint answers 405, which must not abort the cascade. + Mock::given(method("GET")) + .and(path(ctx.path("/namespaces/ns1/views").as_str())) + .respond_with(ResponseTemplate::new(405)) + .mount(&ctx.server) + .await; + + Mock::given(method("DELETE")) + .and(path(ctx.path("/namespaces/ns1").as_str())) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&ctx.server) + .await; + + let result = ctx + .catalog + .drop_database( + &namespace, + DropDatabaseOptions { + if_exists: false, + cascade: true, + }, + ) + .await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_drop_database_cascade_tolerates_missing_views_endpoint() { + test_drop_database_cascade_tolerates_missing_views_endpoint_impl(None).await; + test_drop_database_cascade_tolerates_missing_views_endpoint_impl(Some("test")).await; + } + + async fn test_drop_database_cascade_tolerates_unimplemented_views_endpoint_impl( + name: Option<&str>, + ) { + let ctx = TestContext::new(name).await; + let namespace = Namespace::try_from(vec!["ns1".to_string()]).unwrap(); + + ctx.mock_get_json( + &ctx.path("/namespaces/ns1/tables"), + serde_json::json!({ "identifiers": [] }), + ) + .await; + + // A catalog without a views endpoint may answer 501 instead of 405, which + // must also be tolerated so the cascade still drops the namespace. + Mock::given(method("GET")) + .and(path(ctx.path("/namespaces/ns1/views").as_str())) + .respond_with(ResponseTemplate::new(501)) + .mount(&ctx.server) + .await; + + Mock::given(method("DELETE")) + .and(path(ctx.path("/namespaces/ns1").as_str())) + .respond_with(ResponseTemplate::new(204)) + .expect(1) + .mount(&ctx.server) + .await; + + let result = ctx + .catalog + .drop_database( + &namespace, + DropDatabaseOptions { + if_exists: false, + cascade: true, + }, + ) + .await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_drop_database_cascade_tolerates_unimplemented_views_endpoint() { + test_drop_database_cascade_tolerates_unimplemented_views_endpoint_impl(None).await; + test_drop_database_cascade_tolerates_unimplemented_views_endpoint_impl(Some("test")).await; + } + + async fn test_drop_database_cascade_propagates_list_views_failure_impl(name: Option<&str>) { + let ctx = TestContext::new(name).await; + let namespace = Namespace::try_from(vec!["ns1".to_string()]).unwrap(); + + ctx.mock_get_json( + &ctx.path("/namespaces/ns1/tables"), + serde_json::json!({ "identifiers": [] }), + ) + .await; + + // A genuine views listing failure (not a missing endpoint) must abort the cascade. + Mock::given(method("GET")) + .and(path(ctx.path("/namespaces/ns1/views").as_str())) + .respond_with(ResponseTemplate::new(500)) + .mount(&ctx.server) + .await; + + // The namespace drop must never be attempted once the listing fails. + Mock::given(method("DELETE")) + .and(path(ctx.path("/namespaces/ns1").as_str())) + .respond_with(ResponseTemplate::new(204)) + .expect(0) + .mount(&ctx.server) + .await; + + let result = ctx + .catalog + .drop_database( + &namespace, + DropDatabaseOptions { + if_exists: false, + cascade: true, + }, + ) + .await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_drop_database_cascade_propagates_list_views_failure() { + test_drop_database_cascade_propagates_list_views_failure_impl(None).await; + test_drop_database_cascade_propagates_list_views_failure_impl(Some("test")).await; + } + + async fn test_drop_database_cascade_tolerates_missing_namespace_impl(name: Option<&str>) { + let ctx = TestContext::new(name).await; + let namespace = Namespace::try_from(vec!["ns1".to_string()]).unwrap(); + + // The namespace disappeared before the cascade started; both listings answer 404. + Mock::given(method("GET")) + .and(path(ctx.path("/namespaces/ns1/tables").as_str())) + .respond_with(ResponseTemplate::new(404)) + .mount(&ctx.server) + .await; + Mock::given(method("GET")) + .and(path(ctx.path("/namespaces/ns1/views").as_str())) + .respond_with(ResponseTemplate::new(404)) + .mount(&ctx.server) + .await; + + // With if_exists set, the trailing namespace 404 is the success path. + ctx.mock_delete_404( + &ctx.path("/namespaces/ns1"), + "NoSuchNamespaceException", + "The given namespace does not exist", + ) + .await; + + let result = ctx + .catalog + .drop_database( + &namespace, + DropDatabaseOptions { + if_exists: true, + cascade: true, + }, + ) + .await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_drop_database_cascade_tolerates_missing_namespace() { + test_drop_database_cascade_tolerates_missing_namespace_impl(None).await; + test_drop_database_cascade_tolerates_missing_namespace_impl(Some("test")).await; + } + async fn test_drop_table_impl(name: Option<&str>) { let ctx = TestContext::new(name).await; let namespace = Namespace::try_from(vec!["ns1".to_string()]).unwrap(); From 186efa635aae4ff2555d9c4b93efc8e07d7e81b5 Mon Sep 17 00:00:00 2001 From: Robin Everaars Date: Mon, 27 Jul 2026 18:39:51 +0200 Subject: [PATCH 6/8] fix: skip view listing after missing namespace Signed-off-by: HOIST IT B.V. --- crates/sail-catalog-iceberg/src/provider.rs | 101 ++++++++++++++++---- 1 file changed, 81 insertions(+), 20 deletions(-) diff --git a/crates/sail-catalog-iceberg/src/provider.rs b/crates/sail-catalog-iceberg/src/provider.rs index 464df91e79..770e4cc910 100644 --- a/crates/sail-catalog-iceberg/src/provider.rs +++ b/crates/sail-catalog-iceberg/src/provider.rs @@ -958,9 +958,24 @@ impl CatalogProvider for IcebergRestCatalogProvider { let catalog_config = self.resolved_catalog_config().await?; let DropDatabaseOptions { if_exists, cascade } = options; - let prefix = catalog_config.prefix().map(ToOwned::to_owned); let ns_string = catalog_config.namespace_string(database)?; + let drop_namespace = || async { + match self + .with_auth_retry(|client| { + let prefix = prefix.clone(); + let ns_string = ns_string.clone(); + async move { client.drop_namespace(prefix, ns_string).await } + }) + .await? + { + Ok(_) => Ok(()), + Err(e) if e.status() == Some(reqwest::StatusCode::NOT_FOUND) && if_exists => Ok(()), + Err(e) => Err(CatalogError::External(format!( + "Failed to drop namespace: {e}" + ))), + } + }; if cascade { // For CASCADE, first drop all tables and views in the namespace before dropping the namespace. @@ -1000,9 +1015,11 @@ impl CatalogProvider for IcebergRestCatalogProvider { } } } - // The namespace itself is already gone; fall through to drop_namespace, - // which applies the canonical if_exists handling below. - Err(e) if e.status() == Some(reqwest::StatusCode::NOT_FOUND) => {} + // The namespace itself is already gone. Skip the optional views endpoint and + // fall through to drop_namespace, which applies canonical if_exists handling. + Err(e) if e.status() == Some(reqwest::StatusCode::NOT_FOUND) => { + return drop_namespace().await; + } Err(e) => { return Err(CatalogError::External(format!( "Failed to list tables while cascading namespace drop: {e}" @@ -1063,20 +1080,7 @@ impl CatalogProvider for IcebergRestCatalogProvider { } } - match self - .with_auth_retry(|client| { - let prefix = prefix.clone(); - let ns_string = ns_string.clone(); - async move { client.drop_namespace(prefix, ns_string).await } - }) - .await? - { - Ok(_) => Ok(()), - Err(e) if e.status() == Some(reqwest::StatusCode::NOT_FOUND) && if_exists => Ok(()), - Err(e) => Err(CatalogError::External(format!( - "Failed to drop namespace: {e}" - ))), - } + drop_namespace().await } async fn create_table( @@ -3021,6 +3025,61 @@ mod tests { test_drop_database_cascade_tolerates_unimplemented_views_endpoint_impl(Some("test")).await; } + async fn test_drop_database_cascade_propagates_view_drop_failure_impl(name: Option<&str>) { + let ctx = TestContext::new(name).await; + let namespace = Namespace::try_from(vec!["ns1".to_string()]).unwrap(); + + ctx.mock_get_json( + &ctx.path("/namespaces/ns1/tables"), + serde_json::json!({ "identifiers": [] }), + ) + .await; + ctx.mock_get_json( + &ctx.path("/namespaces/ns1/views"), + serde_json::json!({ + "identifiers": [ + { + "namespace": ["ns1"], + "name": "view1" + } + ] + }), + ) + .await; + + Mock::given(method("DELETE")) + .and(path(ctx.path("/namespaces/ns1/views/view1").as_str())) + .respond_with(ResponseTemplate::new(500)) + .mount(&ctx.server) + .await; + + // The namespace drop must never be attempted once a view drop fails. + Mock::given(method("DELETE")) + .and(path(ctx.path("/namespaces/ns1").as_str())) + .respond_with(ResponseTemplate::new(204)) + .expect(0) + .mount(&ctx.server) + .await; + + let result = ctx + .catalog + .drop_database( + &namespace, + DropDatabaseOptions { + if_exists: false, + cascade: true, + }, + ) + .await; + assert!(result.is_err()); + } + + #[tokio::test] + async fn test_drop_database_cascade_propagates_view_drop_failure() { + test_drop_database_cascade_propagates_view_drop_failure_impl(None).await; + test_drop_database_cascade_propagates_view_drop_failure_impl(Some("test")).await; + } + async fn test_drop_database_cascade_propagates_list_views_failure_impl(name: Option<&str>) { let ctx = TestContext::new(name).await; let namespace = Namespace::try_from(vec!["ns1".to_string()]).unwrap(); @@ -3069,7 +3128,8 @@ mod tests { let ctx = TestContext::new(name).await; let namespace = Namespace::try_from(vec!["ns1".to_string()]).unwrap(); - // The namespace disappeared before the cascade started; both listings answer 404. + // The mandatory tables endpoint reports that the namespace is gone. No + // optional views request should run after that definitive result. Mock::given(method("GET")) .and(path(ctx.path("/namespaces/ns1/tables").as_str())) .respond_with(ResponseTemplate::new(404)) @@ -3077,7 +3137,8 @@ mod tests { .await; Mock::given(method("GET")) .and(path(ctx.path("/namespaces/ns1/views").as_str())) - .respond_with(ResponseTemplate::new(404)) + .respond_with(ResponseTemplate::new(500)) + .expect(0) .mount(&ctx.server) .await; From 860ad594bcf445e388891bd416314d33fdf53059 Mon Sep 17 00:00:00 2001 From: Robin Everaars Date: Mon, 27 Jul 2026 15:32:22 +0200 Subject: [PATCH 7/8] feat: support selective Iceberg overwrite snapshots Signed-off-by: Robin Everaars --- Cargo.lock | 1 + crates/sail-iceberg/Cargo.toml | 3 + crates/sail-iceberg/src/operations/mod.rs | 2 + crates/sail-iceberg/src/operations/rewrite.rs | 41 + .../sail-iceberg/src/operations/snapshot.rs | 978 +++++++++++++++++- .../sail-iceberg/src/spec/manifest/writer.rs | 138 ++- 6 files changed, 1109 insertions(+), 54 deletions(-) create mode 100644 crates/sail-iceberg/src/operations/rewrite.rs diff --git a/Cargo.lock b/Cargo.lock index c27e5917f9..8d6e6f0b8d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7853,6 +7853,7 @@ dependencies = [ "serde", "serde_arrow", "serde_json", + "tokio", "url", "uuid", ] diff --git a/crates/sail-iceberg/Cargo.toml b/crates/sail-iceberg/Cargo.toml index 6fa615ac67..01d24a3449 100644 --- a/crates/sail-iceberg/Cargo.toml +++ b/crates/sail-iceberg/Cargo.toml @@ -57,6 +57,9 @@ rust_decimal = { workspace = true } murmur3 = { workspace = true } educe = { workspace = true } +[dev-dependencies] +tokio = { workspace = true } + [build-dependencies] sail-build-scripts = { path = "../sail-build-scripts" } diff --git a/crates/sail-iceberg/src/operations/mod.rs b/crates/sail-iceberg/src/operations/mod.rs index 52219ce618..848f52824c 100644 --- a/crates/sail-iceberg/src/operations/mod.rs +++ b/crates/sail-iceberg/src/operations/mod.rs @@ -15,6 +15,7 @@ pub mod append; pub mod bootstrap; pub mod helpers; pub mod overwrite; +pub mod rewrite; pub mod snapshot; pub mod write; @@ -22,6 +23,7 @@ pub use action::*; pub use append::*; pub use bootstrap::*; pub use overwrite::*; +pub use rewrite::*; pub use snapshot::*; use crate::spec::Snapshot; diff --git a/crates/sail-iceberg/src/operations/rewrite.rs b/crates/sail-iceberg/src/operations/rewrite.rs new file mode 100644 index 0000000000..78a3eb130d --- /dev/null +++ b/crates/sail-iceberg/src/operations/rewrite.rs @@ -0,0 +1,41 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::SnapshotProduceOperation; + +/// A snapshot operation that replaces a specific set of live data files. +#[derive(Debug, Clone)] +pub struct RewriteFilesOperation { + deleted_data_file_paths: Vec, +} + +impl RewriteFilesOperation { + pub fn new(deleted_data_file_paths: impl IntoIterator) -> Self { + Self { + deleted_data_file_paths: deleted_data_file_paths.into_iter().collect(), + } + } + + pub fn deleted_data_file_paths(&self) -> &[String] { + &self.deleted_data_file_paths + } +} + +impl SnapshotProduceOperation for RewriteFilesOperation { + fn operation(&self) -> &'static str { + "overwrite" + } + + fn deleted_data_file_paths_for_rewrite(&self) -> Option<&[String]> { + Some(&self.deleted_data_file_paths) + } +} diff --git a/crates/sail-iceberg/src/operations/snapshot.rs b/crates/sail-iceberg/src/operations/snapshot.rs index 76b80355ff..08eb72892f 100644 --- a/crates/sail-iceberg/src/operations/snapshot.rs +++ b/crates/sail-iceberg/src/operations/snapshot.rs @@ -10,21 +10,28 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashSet; + use bytes::Bytes; use object_store::ObjectStoreExt; use super::{ActionCommit, Transaction}; use crate::io::StoreContext; -use crate::spec::manifest::ManifestWriterBuilder; +use crate::spec::manifest::{ManifestEntry, ManifestStatus, ManifestWriter, ManifestWriterBuilder}; use crate::spec::manifest_list::ManifestListWriter; use crate::spec::{ - DataFile, FormatVersion, MAIN_BRANCH, ManifestContentType, Operation, PartitionSpec, Schema, - SnapshotBuilder, SnapshotReference, SnapshotRetention, TableRequirement, TableUpdate, + DataFile, FormatVersion, MAIN_BRANCH, ManifestContentType, ManifestFile, Operation, + PartitionSpec, Schema, SnapshotBuilder, SnapshotReference, SnapshotRetention, TableRequirement, + TableUpdate, }; use crate::utils::join_table_uri; pub trait SnapshotProduceOperation: Send + Sync { fn operation(&self) -> &'static str; + + fn deleted_data_file_paths_for_rewrite(&self) -> Option<&[String]> { + None + } } pub struct SnapshotProducer<'a> { @@ -78,14 +85,211 @@ impl<'a> SnapshotProducer<'a> { Ok(()) } + async fn write_manifest( + &self, + store_ctx: &StoreContext, + writer: ManifestWriter, + sequence_number: i64, + snapshot_id: i64, + first_row_id: Option, + ) -> Result { + let manifest_bytes = writer.to_avro_bytes_v2()?; + let manifest_len = i64::try_from(manifest_bytes.len()) + .map_err(|_| "manifest length exceeds i64".to_string())?; + let manifest_rel = format!("metadata/manifest-{}.avro", uuid::Uuid::new_v4()); + let mut manifest_file = writer.into_manifest_file( + join_table_uri(self.tx.table_uri(), &manifest_rel, &self.write_path_mode), + sequence_number, + snapshot_id, + ); + manifest_file.manifest_length = manifest_len; + manifest_file.first_row_id = first_row_id; + + store_ctx + .prefixed + .put( + &object_store::path::Path::from(manifest_rel.as_str()), + object_store::PutPayload::from(Bytes::from(manifest_bytes)), + ) + .await + .map_err(|e| e.to_string())?; + Ok(manifest_file) + } + + fn materialize_inherited_entry( + mut entry: ManifestEntry, + manifest_file: &ManifestFile, + inherited_next_row_id: &mut Option, + ) -> Result { + entry.snapshot_id = entry.snapshot_id.or(Some(manifest_file.added_snapshot_id)); + // V1 entries have no sequence columns and default to 0. In V2 and later, + // only Added entries may inherit sequence numbers from manifest metadata. + if matches!(entry.status, ManifestStatus::Added) || manifest_file.sequence_number == 0 { + entry.sequence_number = entry + .sequence_number + .or(Some(manifest_file.sequence_number)); + entry.file_sequence_number = entry + .file_sequence_number + .or(Some(manifest_file.sequence_number)); + } else if entry.sequence_number.is_none() || entry.file_sequence_number.is_none() { + return Err( + "existing and deleted manifest entries require explicit data and file sequence numbers" + .to_string(), + ); + } + + if entry.data_file.first_row_id.is_none() { + entry.data_file.first_row_id = *inherited_next_row_id; + if let Some(next_row_id) = inherited_next_row_id { + let record_count = i64::try_from(entry.data_file.record_count) + .map_err(|_| "data file record count exceeds i64".to_string())?; + *next_row_id = next_row_id + .checked_add(record_count) + .ok_or_else(|| "row lineage id overflow".to_string())?; + } + } + Ok(entry) + } + + async fn rewrite_parent_manifests( + &self, + store_ctx: &StoreContext, + parent_manifests: Vec, + deleted_data_file_paths: &HashSet, + sequence_number: i64, + snapshot_id: i64, + ) -> Result<(Vec, i64, i64, i64, i64), String> { + enum PlannedManifest { + Reuse(ManifestFile), + Rewrite(ManifestWriter), + } + + let mut planned_manifests = Vec::with_capacity(parent_manifests.len()); + let mut found_paths = HashSet::new(); + let mut parent_live_files = 0i64; + let mut parent_live_rows = 0i64; + + for parent_manifest_file in parent_manifests { + if !matches!(parent_manifest_file.content, ManifestContentType::Data) { + planned_manifests.push((parent_manifest_file, None)); + continue; + } + + let manifest = crate::io::load_manifest(store_ctx, &parent_manifest_file.manifest_path) + .await + .map_err(|e| format!("failed to load parent manifest: {e}"))?; + let mut contains_deleted_path = false; + for entry in manifest.entries().iter().filter(|entry| { + matches!( + entry.status, + ManifestStatus::Added | ManifestStatus::Existing + ) + }) { + parent_live_files = parent_live_files + .checked_add(1) + .ok_or_else(|| "parent data file count overflow".to_string())?; + let record_count = i64::try_from(entry.data_file.record_count) + .map_err(|_| "data file record count exceeds i64".to_string())?; + parent_live_rows = parent_live_rows + .checked_add(record_count) + .ok_or_else(|| "parent record count overflow".to_string())?; + if deleted_data_file_paths.contains(&entry.data_file.file_path) { + contains_deleted_path = true; + found_paths.insert(entry.data_file.file_path.clone()); + } + } + planned_manifests.push(( + parent_manifest_file, + contains_deleted_path.then_some(manifest), + )); + } + + let mut missing_paths = deleted_data_file_paths + .difference(&found_paths) + .cloned() + .collect::>(); + if !missing_paths.is_empty() { + missing_paths.sort(); + return Err(format!( + "rewrite data files are not live in the parent snapshot: {}", + missing_paths.join(", ") + )); + } + + let mut output_plan = Vec::with_capacity(planned_manifests.len()); + let mut deleted_files = 0i64; + let mut deleted_rows = 0i64; + for (parent_manifest_file, manifest) in planned_manifests { + let Some(manifest) = manifest else { + output_plan.push(PlannedManifest::Reuse(parent_manifest_file)); + continue; + }; + + let (entries, metadata) = manifest.into_parts(); + let mut writer = ManifestWriterBuilder::new(Some(snapshot_id), None, metadata).build(); + let mut inherited_next_row_id = parent_manifest_file.first_row_id; + for entry in entries { + if !matches!( + entry.status, + ManifestStatus::Added | ManifestStatus::Existing + ) { + continue; + } + let entry = Self::materialize_inherited_entry( + entry.as_ref().clone(), + &parent_manifest_file, + &mut inherited_next_row_id, + )?; + if deleted_data_file_paths.contains(&entry.data_file.file_path) { + deleted_files = deleted_files + .checked_add(1) + .ok_or_else(|| "deleted data file count overflow".to_string())?; + let record_count = i64::try_from(entry.data_file.record_count) + .map_err(|_| "data file record count exceeds i64".to_string())?; + deleted_rows = deleted_rows + .checked_add(record_count) + .ok_or_else(|| "deleted record count overflow".to_string())?; + writer.add_deleted_entry(entry)?; + } else { + writer.add_existing_entry(entry)?; + } + } + output_plan.push(PlannedManifest::Rewrite(writer)); + } + + let mut output_manifests = Vec::with_capacity(output_plan.len()); + for manifest in output_plan { + match manifest { + PlannedManifest::Reuse(manifest_file) => output_manifests.push(manifest_file), + PlannedManifest::Rewrite(writer) => output_manifests.push( + self.write_manifest(store_ctx, writer, sequence_number, snapshot_id, None) + .await?, + ), + } + } + + Ok(( + output_manifests, + parent_live_files, + parent_live_rows, + deleted_files, + deleted_rows, + )) + } + pub async fn commit(self, op: impl SnapshotProduceOperation) -> Result { let timestamp_ms = crate::utils::timestamp::monotonic_timestamp_ms(); let is_overwrite = op.operation() == Operation::Overwrite.as_str(); - let summary = if is_overwrite { - crate::spec::snapshots::Summary::new(Operation::Overwrite) - } else { - crate::spec::snapshots::Summary::new(Operation::Append) - }; + let deleted_data_file_paths = op + .deleted_data_file_paths_for_rewrite() + .map(|paths| paths.iter().cloned().collect::>()); + if deleted_data_file_paths + .as_ref() + .is_some_and(HashSet::is_empty) + { + return Err("rewrite requires at least one data file path to delete".to_string()); + } + let is_rewrite = deleted_data_file_paths.is_some(); // Build manifest metadata: prefer caller-provided metadata derived from table schema/spec // Fall back to deriving from the current transaction snapshot if not provided @@ -127,7 +331,10 @@ impl<'a> SnapshotProducer<'a> { let parent_manifest_list_path_str = parent_snapshot.manifest_list(); let mut parent_manifest_entries = Vec::new(); - if !self.is_bootstrap && !is_overwrite && !parent_manifest_list_path_str.is_empty() { + if !self.is_bootstrap + && (!is_overwrite || is_rewrite) + && !parent_manifest_list_path_str.is_empty() + { let (store_ref, manifest_list_path) = store_ctx .resolve(parent_manifest_list_path_str) .map_err(|e| format!("{}", e))?; @@ -152,11 +359,14 @@ impl<'a> SnapshotProducer<'a> { parent_manifest_entries.extend(parent_manifest_list.entries().iter().cloned()); } - let new_added_rows: i64 = self - .added_data_files - .iter() - .map(|df| df.record_count as i64) - .sum(); + let mut new_added_rows = 0i64; + for data_file in &self.added_data_files { + let record_count = i64::try_from(data_file.record_count) + .map_err(|_| "data file record count exceeds i64".to_string())?; + new_added_rows = new_added_rows + .checked_add(record_count) + .ok_or_else(|| "added record count overflow".to_string())?; + } let mut row_lineage_next_row_id = self.row_lineage_start_row_id; let mut snapshot_added_rows = 0; @@ -179,48 +389,53 @@ impl<'a> SnapshotProducer<'a> { snapshot_added_rows += new_added_rows; } - let mut writer = ManifestWriterBuilder::new(None, None, metadata.clone()).build(); let added_data_files = self.added_data_files.clone(); - for df in &added_data_files { - writer.add(df.clone()); - } - let manifest = writer.finish(); - let manifest_bytes = manifest.to_avro_bytes_v2()?; - - let manifest_len = manifest_bytes.len() as i64; - let manifest_rel = format!("metadata/manifest-{}.avro", uuid::Uuid::new_v4()); - let manifest_path = object_store::path::Path::from(manifest_rel.as_str()); - store_ctx - .prefixed - .put( - &manifest_path, - object_store::PutPayload::from(Bytes::from(manifest_bytes)), + let added_files = i64::try_from(added_data_files.len()) + .map_err(|_| "added data file count exceeds i64".to_string())?; + let ( + parent_manifest_entries, + parent_live_files, + parent_live_rows, + deleted_files, + deleted_rows, + ) = if let Some(paths) = deleted_data_file_paths.as_ref() { + self.rewrite_parent_manifests( + store_ctx, + parent_manifest_entries, + paths, + new_sequence_number, + new_snapshot_id, ) - .await - .map_err(|e| format!("{}", e))?; + .await? + } else { + (parent_manifest_entries, 0, 0, 0, 0) + }; - let mut manifest_file_builder = crate::spec::manifest_list::ManifestFile::builder() - .with_manifest_path(join_table_uri( - self.tx.table_uri(), - &manifest_rel, - &self.write_path_mode, - )) - .with_manifest_length(manifest_len) - .with_partition_spec_id(metadata.partition_spec.spec_id()) - .with_content(ManifestContentType::Data) - .with_sequence_number(new_sequence_number) - .with_min_sequence_number(new_sequence_number) - .with_added_snapshot_id(new_snapshot_id) - .with_file_counts(added_data_files.len() as i32, 0, 0) - .with_row_counts(new_added_rows, 0, 0); - if let Some(first_row_id) = new_manifest_first_row_id { - manifest_file_builder = manifest_file_builder.with_first_row_id(first_row_id); + let mut summary = if is_overwrite { + crate::spec::snapshots::Summary::new(Operation::Overwrite) + } else { + crate::spec::snapshots::Summary::new(Operation::Append) + }; + if is_rewrite { + let total_data_files = parent_live_files + .checked_sub(deleted_files) + .and_then(|count| count.checked_add(added_files)) + .ok_or_else(|| "total data file count overflow".to_string())?; + let total_records = parent_live_rows + .checked_sub(deleted_rows) + .and_then(|count| count.checked_add(new_added_rows)) + .ok_or_else(|| "total record count overflow".to_string())?; + summary = summary + .with_property("added-data-files", added_files) + .with_property("deleted-data-files", deleted_files) + .with_property("added-records", new_added_rows) + .with_property("deleted-records", deleted_rows) + .with_property("total-data-files", total_data_files) + .with_property("total-records", total_records); } - let manifest_file = manifest_file_builder.build()?; let mut list_writer = ManifestListWriter::new(); - let mut total_manifest_count = 0; - + let mut total_manifest_count = 0usize; for entry in parent_manifest_entries { list_writer.append(entry); total_manifest_count += 1; @@ -233,8 +448,23 @@ impl<'a> SnapshotProducer<'a> { self.tx.snapshot().snapshot_id() ); - list_writer.append(manifest_file); - total_manifest_count += 1; + if !is_rewrite || !added_data_files.is_empty() { + let mut writer = ManifestWriterBuilder::new(None, None, metadata.clone()).build(); + for data_file in added_data_files { + writer.add(data_file); + } + list_writer.append( + self.write_manifest( + store_ctx, + writer, + new_sequence_number, + new_snapshot_id, + new_manifest_first_row_id, + ) + .await?, + ); + total_manifest_count += 1; + } log::trace!( "snapshot producer: new manifest list will have files: {}", total_manifest_count @@ -315,3 +545,647 @@ impl<'a> SnapshotProducer<'a> { Ok(ActionCommit::new(updates, requirements)) } } + +#[cfg(test)] +mod tests { + #![expect(clippy::unwrap_used)] + + use std::collections::HashMap; + use std::sync::Arc; + + use bytes::Bytes; + use datafusion::arrow::array::Int64Array; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::datasource::TableProvider; + use datafusion::prelude::SessionContext; + use futures::TryStreamExt; + use object_store::memory::InMemory; + use parquet::arrow::ArrowWriter; + use url::Url; + + use super::*; + use crate::datasource::IcebergTableProvider; + use crate::datasource::type_converter::iceberg_schema_to_arrow; + use crate::io::{load_manifest, load_manifest_list}; + use crate::operations::RewriteFilesOperation; + use crate::spec::manifest::{ManifestEntry, ManifestStatus}; + use crate::spec::types::{NestedField, PrimitiveType, Type}; + use crate::spec::{DataContentType, DataFileFormat, ManifestListWriter}; + use crate::utils::WritePathMode; + + fn test_schema() -> Schema { + Schema::builder() + .with_schema_id(0) + .with_fields(vec![Arc::new(NestedField::required( + 1, + "id", + Type::Primitive(PrimitiveType::Long), + ))]) + .build() + .unwrap() + } + + fn data_file(path: &str, record_count: u64, file_size_in_bytes: u64) -> DataFile { + DataFile { + content: DataContentType::Data, + file_path: path.to_string(), + file_format: DataFileFormat::Parquet, + partition: vec![], + record_count, + file_size_in_bytes, + column_sizes: HashMap::new(), + value_counts: HashMap::new(), + null_value_counts: HashMap::new(), + nan_value_counts: HashMap::new(), + lower_bounds: HashMap::new(), + upper_bounds: HashMap::new(), + block_size_in_bytes: None, + key_metadata: None, + split_offsets: vec![], + equality_ids: vec![], + sort_order_id: None, + first_row_id: None, + partition_spec_id: 0, + referenced_data_file: None, + content_offset: None, + content_size_in_bytes: None, + } + } + + fn parquet_bytes(schema: Arc, ids: Vec) -> Vec { + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(ids))]).unwrap(); + let mut bytes = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut bytes, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + bytes + } + + struct SingleFileParent { + store_ctx: StoreContext, + manifest_metadata: crate::spec::manifest::ManifestMetadata, + tx: Transaction, + live_file: DataFile, + } + + async fn single_file_parent(table_url: &str) -> SingleFileParent { + let table_url = Url::parse(table_url).unwrap(); + let object_store = Arc::new(InMemory::new()); + let store_ctx = StoreContext::new(object_store, &table_url).unwrap(); + let schema = test_schema(); + let partition_spec = PartitionSpec::builder().with_spec_id(0).build(); + let manifest_metadata = crate::spec::manifest::ManifestMetadata::new( + Arc::new(schema.clone()), + schema.schema_id(), + partition_spec, + FormatVersion::V2, + ManifestContentType::Data, + ); + let parent_snapshot_id = 10; + let parent_sequence_number = 1; + let live_file = data_file("data/live.parquet", 1, 100); + let mut parent_writer = + ManifestWriterBuilder::new(Some(parent_snapshot_id), None, manifest_metadata.clone()) + .build(); + parent_writer.add(live_file.clone()); + let parent_manifest_bytes = parent_writer.to_avro_bytes_v2().unwrap(); + let mut parent_manifest_file = parent_writer.into_manifest_file( + "metadata/manifest-parent.avro".to_string(), + parent_sequence_number, + parent_snapshot_id, + ); + parent_manifest_file.manifest_length = parent_manifest_bytes.len() as i64; + store_ctx + .prefixed + .put( + &object_store::path::Path::from("metadata/manifest-parent.avro"), + object_store::PutPayload::from(Bytes::from(parent_manifest_bytes)), + ) + .await + .unwrap(); + + let mut parent_list_writer = ManifestListWriter::new(); + parent_list_writer.append(parent_manifest_file); + let parent_list_bytes = parent_list_writer.to_bytes(FormatVersion::V2).unwrap(); + store_ctx + .prefixed + .put( + &object_store::path::Path::from("metadata/snap-parent.avro"), + object_store::PutPayload::from(Bytes::from(parent_list_bytes)), + ) + .await + .unwrap(); + + let parent_snapshot = SnapshotBuilder::new() + .with_snapshot_id(parent_snapshot_id) + .with_sequence_number(parent_sequence_number) + .with_manifest_list("metadata/snap-parent.avro") + .with_summary(crate::spec::snapshots::Summary::new(Operation::Append)) + .with_schema_id(schema.schema_id()) + .build() + .unwrap(); + let tx = Transaction::new(table_url.to_string(), parent_snapshot); + SingleFileParent { + store_ctx, + manifest_metadata, + tx, + live_file, + } + } + + #[test] + fn inherited_row_ids_do_not_advance_for_preassigned_files() { + let parent_manifest = ManifestFile::builder() + .with_manifest_path("metadata/parent.avro") + .with_sequence_number(1) + .with_min_sequence_number(1) + .with_added_snapshot_id(10) + .with_file_counts(2, 0, 0) + .with_row_counts(15, 0, 0) + .with_first_row_id(100) + .build() + .unwrap(); + let mut assigned_file = data_file("data/assigned.parquet", 10, 100); + assigned_file.first_row_id = Some(50); + let assigned_entry = ManifestEntry::new( + ManifestStatus::Added, + Some(10), + Some(1), + Some(1), + assigned_file, + ); + let unassigned_entry = ManifestEntry::new( + ManifestStatus::Added, + Some(10), + Some(1), + Some(1), + data_file("data/unassigned.parquet", 5, 100), + ); + let mut inherited_next_row_id = parent_manifest.first_row_id; + + let assigned_entry = SnapshotProducer::<'static>::materialize_inherited_entry( + assigned_entry, + &parent_manifest, + &mut inherited_next_row_id, + ) + .unwrap(); + let unassigned_entry = SnapshotProducer::<'static>::materialize_inherited_entry( + unassigned_entry, + &parent_manifest, + &mut inherited_next_row_id, + ) + .unwrap(); + + assert_eq!(assigned_entry.data_file.first_row_id, Some(50)); + assert_eq!(unassigned_entry.data_file.first_row_id, Some(100)); + assert_eq!(inherited_next_row_id, Some(105)); + } + + #[test] + fn existing_entries_only_inherit_sequence_numbers_for_v1() { + let parent_manifest = ManifestFile::builder() + .with_manifest_path("metadata/parent.avro") + .with_sequence_number(2) + .with_min_sequence_number(1) + .with_added_snapshot_id(20) + .with_file_counts(0, 1, 0) + .with_row_counts(0, 1, 0) + .build() + .unwrap(); + let entry = ManifestEntry::new( + ManifestStatus::Existing, + Some(10), + None, + None, + data_file("data/existing.parquet", 1, 100), + ); + + let result = SnapshotProducer::<'static>::materialize_inherited_entry( + entry, + &parent_manifest, + &mut None, + ); + + assert!(matches!( + result, + Err(message) if message.contains("sequence numbers") + )); + + let mut v1_manifest = parent_manifest; + v1_manifest.sequence_number = 0; + let v1_entry = ManifestEntry::new( + ManifestStatus::Existing, + Some(10), + None, + None, + data_file("data/v1-existing.parquet", 1, 100), + ); + let v1_entry = SnapshotProducer::<'static>::materialize_inherited_entry( + v1_entry, + &v1_manifest, + &mut None, + ) + .unwrap(); + assert_eq!(v1_entry.sequence_number, Some(0)); + assert_eq!(v1_entry.file_sequence_number, Some(0)); + } + + #[tokio::test] + async fn rewrite_files_marks_removed_files_and_preserves_survivor_rows() { + let table_url = Url::parse("memory://rewrite-test/table/").unwrap(); + let object_store = Arc::new(InMemory::new()); + let store_ctx = StoreContext::new(object_store.clone(), &table_url).unwrap(); + let schema = test_schema(); + let arrow_schema = Arc::new(iceberg_schema_to_arrow(&schema).unwrap()); + let partition_spec = PartitionSpec::builder().with_spec_id(0).build(); + let manifest_metadata = crate::spec::manifest::ManifestMetadata::new( + Arc::new(schema.clone()), + schema.schema_id(), + partition_spec.clone(), + FormatVersion::V2, + ManifestContentType::Data, + ); + + let old_bytes = parquet_bytes(arrow_schema.clone(), vec![1]); + let survivor_bytes = parquet_bytes(arrow_schema.clone(), vec![2]); + let replacement_bytes = parquet_bytes(arrow_schema.clone(), vec![3]); + let unaffected_bytes = parquet_bytes(arrow_schema, vec![4]); + store_ctx + .prefixed + .put( + &object_store::path::Path::from("data/old.parquet"), + object_store::PutPayload::from(Bytes::from(old_bytes.clone())), + ) + .await + .unwrap(); + store_ctx + .prefixed + .put( + &object_store::path::Path::from("data/survivor.parquet"), + object_store::PutPayload::from(Bytes::from(survivor_bytes.clone())), + ) + .await + .unwrap(); + store_ctx + .prefixed + .put( + &object_store::path::Path::from("data/replacement.parquet"), + object_store::PutPayload::from(Bytes::from(replacement_bytes.clone())), + ) + .await + .unwrap(); + store_ctx + .prefixed + .put( + &object_store::path::Path::from("data/unaffected.parquet"), + object_store::PutPayload::from(Bytes::from(unaffected_bytes.clone())), + ) + .await + .unwrap(); + + let old_file = data_file("data/old.parquet", 1, old_bytes.len() as u64); + let survivor_file = data_file("data/survivor.parquet", 1, survivor_bytes.len() as u64); + let replacement_file = data_file( + "data/replacement.parquet", + 1, + replacement_bytes.len() as u64, + ); + let unaffected_file = + data_file("data/unaffected.parquet", 1, unaffected_bytes.len() as u64); + + let parent_snapshot_id = 10; + let parent_sequence_number = 1; + let mut parent_writer = + ManifestWriterBuilder::new(Some(parent_snapshot_id), None, manifest_metadata.clone()) + .build(); + parent_writer.add(old_file.clone()); + parent_writer.add(survivor_file.clone()); + let parent_manifest_bytes = parent_writer.to_avro_bytes_v2().unwrap(); + let mut parent_manifest_file = parent_writer.into_manifest_file( + "metadata/manifest-parent.avro".to_string(), + parent_sequence_number, + parent_snapshot_id, + ); + parent_manifest_file.manifest_length = parent_manifest_bytes.len() as i64; + store_ctx + .prefixed + .put( + &object_store::path::Path::from("metadata/manifest-parent.avro"), + object_store::PutPayload::from(Bytes::from(parent_manifest_bytes)), + ) + .await + .unwrap(); + + let mut unaffected_writer = + ManifestWriterBuilder::new(Some(parent_snapshot_id), None, manifest_metadata.clone()) + .build(); + unaffected_writer.add(unaffected_file.clone()); + let unaffected_manifest_bytes = unaffected_writer.to_avro_bytes_v2().unwrap(); + let mut unaffected_manifest_file = unaffected_writer.into_manifest_file( + "metadata/manifest-unaffected.avro".to_string(), + parent_sequence_number, + parent_snapshot_id, + ); + unaffected_manifest_file.manifest_length = unaffected_manifest_bytes.len() as i64; + store_ctx + .prefixed + .put( + &object_store::path::Path::from("metadata/manifest-unaffected.avro"), + object_store::PutPayload::from(Bytes::from(unaffected_manifest_bytes)), + ) + .await + .unwrap(); + + let delete_manifest_metadata = crate::spec::manifest::ManifestMetadata::new( + Arc::new(schema.clone()), + schema.schema_id(), + partition_spec.clone(), + FormatVersion::V2, + ManifestContentType::Deletes, + ); + let delete_writer = + ManifestWriterBuilder::new(Some(parent_snapshot_id), None, delete_manifest_metadata) + .build(); + let delete_manifest_bytes = delete_writer.to_avro_bytes_v2().unwrap(); + let mut delete_manifest_file = delete_writer.into_manifest_file( + "metadata/manifest-deletes.avro".to_string(), + parent_sequence_number, + parent_snapshot_id, + ); + delete_manifest_file.manifest_length = delete_manifest_bytes.len() as i64; + store_ctx + .prefixed + .put( + &object_store::path::Path::from("metadata/manifest-deletes.avro"), + object_store::PutPayload::from(Bytes::from(delete_manifest_bytes)), + ) + .await + .unwrap(); + + let mut parent_list_writer = ManifestListWriter::new(); + parent_list_writer.append(parent_manifest_file); + parent_list_writer.append(unaffected_manifest_file); + parent_list_writer.append(delete_manifest_file); + let parent_list_bytes = parent_list_writer.to_bytes(FormatVersion::V2).unwrap(); + store_ctx + .prefixed + .put( + &object_store::path::Path::from("metadata/snap-parent.avro"), + object_store::PutPayload::from(Bytes::from(parent_list_bytes)), + ) + .await + .unwrap(); + + let parent_snapshot = SnapshotBuilder::new() + .with_snapshot_id(parent_snapshot_id) + .with_sequence_number(parent_sequence_number) + .with_manifest_list("metadata/snap-parent.avro") + .with_summary( + crate::spec::snapshots::Summary::new(Operation::Append) + .with_property("total-data-files", 3) + .with_property("total-records", 3), + ) + .with_schema_id(schema.schema_id()) + .build() + .unwrap(); + let tx = Transaction::new(table_url.to_string(), parent_snapshot); + + let action_commit = SnapshotProducer::new( + &tx, + vec![replacement_file.clone()], + Some(store_ctx.clone()), + Some(manifest_metadata), + ) + .with_write_path_mode(WritePathMode::Relative) + .commit(RewriteFilesOperation::new(vec![old_file.file_path.clone()])) + .await + .unwrap(); + + let new_snapshot = action_commit + .updates() + .iter() + .find_map(|update| match update { + TableUpdate::AddSnapshot { snapshot } => Some(snapshot.clone()), + _ => None, + }) + .unwrap(); + assert_eq!(new_snapshot.parent_snapshot_id(), Some(parent_snapshot_id)); + assert_eq!(new_snapshot.summary.operation, Operation::Overwrite); + assert_eq!( + new_snapshot + .summary + .additional_properties + .get("added-data-files"), + Some(&"1".to_string()) + ); + assert_eq!( + new_snapshot + .summary + .additional_properties + .get("deleted-data-files"), + Some(&"1".to_string()) + ); + assert_eq!( + new_snapshot + .summary + .additional_properties + .get("total-data-files"), + Some(&"3".to_string()) + ); + assert_eq!( + new_snapshot + .summary + .additional_properties + .get("total-records"), + Some(&"3".to_string()) + ); + + let manifest_list = load_manifest_list(&store_ctx, new_snapshot.manifest_list()) + .await + .unwrap(); + assert!( + manifest_list + .entries() + .iter() + .any(|manifest| { manifest.manifest_path == "metadata/manifest-unaffected.avro" }) + ); + assert!( + manifest_list + .entries() + .iter() + .any(|manifest| { manifest.manifest_path == "metadata/manifest-deletes.avro" }) + ); + let mut entries_by_path = HashMap::::new(); + for manifest_file in manifest_list.entries() { + let manifest = load_manifest(&store_ctx, &manifest_file.manifest_path) + .await + .unwrap(); + for entry in manifest.entries() { + entries_by_path.insert(entry.data_file.file_path.clone(), (**entry).clone()); + } + } + assert_eq!( + entries_by_path["data/old.parquet"].status, + ManifestStatus::Deleted + ); + assert_eq!( + entries_by_path["data/survivor.parquet"].status, + ManifestStatus::Existing + ); + assert_eq!( + entries_by_path["data/replacement.parquet"].status, + ManifestStatus::Added + ); + assert_eq!( + entries_by_path["data/replacement.parquet"].snapshot_id, + None + ); + assert_eq!( + entries_by_path["data/unaffected.parquet"].status, + ManifestStatus::Added + ); + assert_eq!( + entries_by_path["data/old.parquet"].snapshot_id, + Some(new_snapshot.snapshot_id()) + ); + assert_eq!( + entries_by_path["data/old.parquet"].sequence_number, + Some(parent_sequence_number) + ); + assert_eq!( + entries_by_path["data/old.parquet"].file_sequence_number, + Some(parent_sequence_number) + ); + assert_eq!( + entries_by_path["data/survivor.parquet"].snapshot_id, + Some(parent_snapshot_id) + ); + assert_eq!( + entries_by_path["data/survivor.parquet"].sequence_number, + Some(parent_sequence_number) + ); + assert_eq!( + entries_by_path["data/survivor.parquet"].file_sequence_number, + Some(parent_sequence_number) + ); + + let context = SessionContext::new(); + context + .runtime_env() + .register_object_store(&Url::parse("memory://rewrite-test/").unwrap(), object_store); + let provider = + IcebergTableProvider::new(table_url, schema, new_snapshot, vec![partition_spec], 0) + .unwrap(); + let plan = provider + .scan(&context.state(), None, &[], None) + .await + .unwrap(); + let batches = datafusion::physical_plan::collect(plan, context.task_ctx()) + .await + .unwrap(); + let mut ids = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect::>(); + ids.sort_unstable(); + assert_eq!(ids, vec![2, 3, 4]); + } + + #[tokio::test] + async fn rewrite_files_rejects_paths_not_live_without_writing_orphan_manifests() { + let fixture = single_file_parent("memory://rewrite-missing-test/table/").await; + + let result = SnapshotProducer::new( + &fixture.tx, + vec![], + Some(fixture.store_ctx.clone()), + Some(fixture.manifest_metadata), + ) + .with_write_path_mode(WritePathMode::Relative) + .commit(RewriteFilesOperation::new(vec![ + fixture.live_file.file_path, + "data/missing.parquet".to_string(), + ])) + .await; + + assert!(matches!( + result, + Err(message) if message.contains("not live in the parent snapshot") + )); + let metadata_objects = fixture + .store_ctx + .prefixed + .list(Some(&object_store::path::Path::from("metadata"))) + .try_collect::>() + .await + .unwrap(); + assert_eq!(metadata_objects.len(), 2); + } + + #[tokio::test] + async fn rewrite_files_supports_pure_deletes_without_an_empty_added_manifest() { + let fixture = single_file_parent("memory://rewrite-pure-delete-test/table/").await; + + let action_commit = SnapshotProducer::new( + &fixture.tx, + vec![], + Some(fixture.store_ctx.clone()), + Some(fixture.manifest_metadata), + ) + .with_write_path_mode(WritePathMode::Relative) + .commit(RewriteFilesOperation::new(vec![ + fixture.live_file.file_path, + ])) + .await + .unwrap(); + + let new_snapshot = action_commit + .updates() + .iter() + .find_map(|update| match update { + TableUpdate::AddSnapshot { snapshot } => Some(snapshot), + _ => None, + }) + .unwrap(); + for (key, expected) in [ + ("added-data-files", "0"), + ("deleted-data-files", "1"), + ("added-records", "0"), + ("deleted-records", "1"), + ("total-data-files", "0"), + ("total-records", "0"), + ] { + assert_eq!( + new_snapshot + .summary + .additional_properties + .get(key) + .map(String::as_str), + Some(expected) + ); + } + + let manifest_list = load_manifest_list(&fixture.store_ctx, new_snapshot.manifest_list()) + .await + .unwrap(); + assert_eq!(manifest_list.entries().len(), 1); + let manifest = load_manifest( + &fixture.store_ctx, + &manifest_list.entries()[0].manifest_path, + ) + .await + .unwrap(); + assert_eq!(manifest.entries().len(), 1); + assert_eq!(manifest.entries()[0].status, ManifestStatus::Deleted); + } +} diff --git a/crates/sail-iceberg/src/spec/manifest/writer.rs b/crates/sail-iceberg/src/spec/manifest/writer.rs index 22a027b83b..9299051607 100644 --- a/crates/sail-iceberg/src/spec/manifest/writer.rs +++ b/crates/sail-iceberg/src/spec/manifest/writer.rs @@ -79,6 +79,31 @@ impl ManifestWriter { self.entries.push(Arc::new(entry)); } + pub fn add_existing_entry(&mut self, mut entry: ManifestEntry) -> Result<(), String> { + entry.status = ManifestStatus::Existing; + self.add_rewritten_entry(entry) + } + + pub fn add_deleted_entry(&mut self, mut entry: ManifestEntry) -> Result<(), String> { + entry.status = ManifestStatus::Deleted; + entry.snapshot_id = self.snapshot_id; + self.add_rewritten_entry(entry) + } + + fn add_rewritten_entry(&mut self, entry: ManifestEntry) -> Result<(), String> { + if entry.snapshot_id.is_none() + || entry.sequence_number.is_none() + || entry.file_sequence_number.is_none() + { + return Err( + "existing and deleted manifest entries require snapshot, data, and file sequence metadata" + .to_string(), + ); + } + self.entries.push(Arc::new(entry)); + Ok(()) + } + pub fn finish(self) -> Manifest { Manifest::new( self.metadata, @@ -125,13 +150,20 @@ impl ManifestWriter { .filter(|e| matches!(e.status, ManifestStatus::Deleted)) .map(|e| e.data_file.record_count as i64) .sum(); + let min_sequence_number = self + .entries + .iter() + .filter(|entry| !matches!(entry.status, ManifestStatus::Deleted)) + .filter_map(|entry| entry.sequence_number) + .min() + .unwrap_or(sequence_number); ManifestFile { manifest_path, manifest_length: 0, partition_spec_id: self.metadata.partition_spec.spec_id(), - content: ManifestContentType::Data, + content: self.metadata.content, sequence_number, - min_sequence_number: sequence_number, + min_sequence_number, added_snapshot_id: snapshot_id, added_files_count: Some(added), existing_files_count: Some(existing), @@ -208,3 +240,105 @@ impl ManifestWriter { .map_err(|e| format!("Avro writer finalize error: {e}")) } } + +#[cfg(test)] +mod tests { + #![expect(clippy::unwrap_used)] + + use std::collections::HashMap; + + use super::*; + use crate::spec::{ + DataContentType, DataFileFormat, ManifestContentType, PartitionSpec, Schema, + }; + + fn data_file(path: &str, record_count: u64) -> DataFile { + DataFile { + content: DataContentType::Data, + file_path: path.to_string(), + file_format: DataFileFormat::Parquet, + partition: vec![], + record_count, + file_size_in_bytes: 100, + column_sizes: HashMap::new(), + value_counts: HashMap::new(), + null_value_counts: HashMap::new(), + nan_value_counts: HashMap::new(), + lower_bounds: HashMap::new(), + upper_bounds: HashMap::new(), + block_size_in_bytes: None, + key_metadata: None, + split_offsets: vec![], + equality_ids: vec![], + sort_order_id: None, + first_row_id: None, + partition_spec_id: 0, + referenced_data_file: None, + content_offset: None, + content_size_in_bytes: None, + } + } + + fn manifest_metadata() -> ManifestMetadata { + let schema = Schema::builder().with_fields(vec![]).build().unwrap(); + ManifestMetadata::new( + Arc::new(schema), + 0, + PartitionSpec::builder().with_spec_id(0).build(), + FormatVersion::V2, + ManifestContentType::Data, + ) + } + + #[test] + fn existing_and_deleted_entries_preserve_sequence_metadata_and_counts() { + let original = ManifestEntry::new( + ManifestStatus::Added, + Some(10), + Some(1), + Some(1), + data_file("data/original.parquet", 3), + ); + let mut writer = ManifestWriterBuilder::new(Some(20), None, manifest_metadata()).build(); + + writer.add_existing_entry(original.clone()).unwrap(); + writer.add_deleted_entry(original).unwrap(); + + let bytes = writer.to_avro_bytes_v2().unwrap(); + let manifest_file = writer.into_manifest_file("metadata/rewrite.avro".to_string(), 2, 20); + let manifest = Manifest::parse_avro(&bytes).unwrap(); + + assert_eq!(manifest.entries.len(), 2); + assert_eq!(manifest.entries[0].status, ManifestStatus::Existing); + assert_eq!(manifest.entries[0].snapshot_id, Some(10)); + assert_eq!(manifest.entries[0].sequence_number, Some(1)); + assert_eq!(manifest.entries[0].file_sequence_number, Some(1)); + assert_eq!(manifest.entries[1].status, ManifestStatus::Deleted); + assert_eq!(manifest.entries[1].snapshot_id, Some(20)); + assert_eq!(manifest.entries[1].sequence_number, Some(1)); + assert_eq!(manifest.entries[1].file_sequence_number, Some(1)); + + assert_eq!(manifest_file.content, ManifestContentType::Data); + assert_eq!(manifest_file.added_files_count, Some(0)); + assert_eq!(manifest_file.existing_files_count, Some(1)); + assert_eq!(manifest_file.deleted_files_count, Some(1)); + assert_eq!(manifest_file.added_rows_count, Some(0)); + assert_eq!(manifest_file.existing_rows_count, Some(3)); + assert_eq!(manifest_file.deleted_rows_count, Some(3)); + assert_eq!(manifest_file.min_sequence_number, 1); + } + + #[test] + fn existing_entries_require_their_original_snapshot_id() { + let entry = ManifestEntry::new( + ManifestStatus::Added, + None, + Some(1), + Some(1), + data_file("data/original.parquet", 3), + ); + let mut writer = ManifestWriterBuilder::new(Some(20), None, manifest_metadata()).build(); + + assert!(writer.add_existing_entry(entry).is_err()); + } +} From eafe25f5f066b5a043396d18e68f15a67c52b619 Mon Sep 17 00:00:00 2001 From: Robin Everaars Date: Wed, 29 Jul 2026 12:48:33 +0200 Subject: [PATCH 8/8] feat: preserve typed Python data source failures Signed-off-by: Robin Everaars --- crates/sail-common-datafusion/src/error.rs | 69 +++++++++++ .../src/formats/python/error.rs | 117 +++++++++++++++++- crates/sail-python-udf/src/error.rs | 25 ++++ crates/sail-spark-connect/src/error.rs | 40 +++--- docs/guide/sources/python/index.md | 28 +++++ .../tests/spark/datasource/test_python.py | 73 +++++++++++ 6 files changed, 332 insertions(+), 20 deletions(-) diff --git a/crates/sail-common-datafusion/src/error.rs b/crates/sail-common-datafusion/src/error.rs index a9c1a81740..f9df981a2a 100644 --- a/crates/sail-common-datafusion/src/error.rs +++ b/crates/sail-common-datafusion/src/error.rs @@ -14,10 +14,37 @@ pub struct RemoteError { pub cause: CommonErrorCause, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum PythonFailureKind { + Terminal, + Transient, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)] +pub enum PythonDataSourceFailure { + #[error("Python data source reported a terminal failure")] + Terminal, + #[error("Python data source reported a transient failure")] + Transient, +} + +impl PythonDataSourceFailure { + pub fn kind(self) -> PythonFailureKind { + match self { + Self::Terminal => PythonFailureKind::Terminal, + Self::Transient => PythonFailureKind::Transient, + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] pub struct PythonErrorCause { pub summary: String, pub traceback: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub failure_kind: Option, } /// A trait to extract Python error cause from a generic error. @@ -166,6 +193,14 @@ impl CommonErrorCause { }; } + if let Some(failure) = error.downcast_ref::() { + return Self::Python(PythonErrorCause { + summary: failure.to_string(), + traceback: None, + failure_kind: Some(failure.kind()), + }); + } + if let Some(cause) = Py::extract(error) { return Self::Python(cause); } @@ -187,3 +222,37 @@ impl CommonErrorCause { Self::build::(error, &mut HashSet::new()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_python_failure_kind_round_trips() -> Result<(), Box> { + let cause = CommonErrorCause::Python(PythonErrorCause { + summary: "Python data source reported a terminal failure".to_string(), + traceback: None, + failure_kind: Some(PythonFailureKind::Terminal), + }); + + let encoded = serde_json::to_string(&cause)?; + assert!(encoded.contains(r#""failureKind":"terminal""#)); + let decoded: CommonErrorCause = serde_json::from_str(&encoded)?; + let CommonErrorCause::Python(decoded) = decoded else { + return Err(std::io::Error::other("expected Python error cause").into()); + }; + assert_eq!(decoded.failure_kind, Some(PythonFailureKind::Terminal)); + Ok(()) + } + + #[test] + fn test_legacy_python_cause_defaults_failure_kind() -> Result<(), Box> { + let encoded = r#"{"python":{"summary":"legacy","traceback":null}}"#; + let decoded: CommonErrorCause = serde_json::from_str(encoded)?; + let CommonErrorCause::Python(decoded) = decoded else { + return Err(std::io::Error::other("expected Python error cause").into()); + }; + assert_eq!(decoded.failure_kind, None); + Ok(()) + } +} diff --git a/crates/sail-data-source/src/formats/python/error.rs b/crates/sail-data-source/src/formats/python/error.rs index ecff918436..5828867944 100644 --- a/crates/sail-data-source/src/formats/python/error.rs +++ b/crates/sail-data-source/src/formats/python/error.rs @@ -3,8 +3,43 @@ //! Provides structured error types with context for debugging Python datasource issues. use datafusion_common::DataFusionError; +use sail_common_datafusion::error::PythonDataSourceFailure; use thiserror::Error; +const FAILURE_KIND_ATTRIBUTE: &str = "__sail_data_source_failure_kind__"; + +fn declared_failure_kind(error: &pyo3::PyErr) -> Option { + use pyo3::prelude::PyAnyMethods; + use pyo3::types::{PyTuple, PyTupleMethods, PyType}; + + pyo3::Python::attach(|py| { + let type_type = py.get_type::(); + let getattribute = type_type.getattr("__getattribute__").ok()?; + let exception_type = error.get_type(py); + let mro = getattribute + .call1((&exception_type, "__mro__")) + .ok()? + .cast_into::() + .ok()?; + + for base in mro.iter() { + let namespace = getattribute.call1((&base, "__dict__")).ok()?; + let Ok(value) = namespace.get_item(FAILURE_KIND_ATTRIBUTE) else { + continue; + }; + let Ok(value) = value.extract::() else { + return None; + }; + return match value.as_str() { + "terminal" => Some(PythonDataSourceFailure::Terminal), + "transient" => Some(PythonDataSourceFailure::Transient), + _ => None, + }; + } + None + }) +} + /// Result type alias for Python data source operations. #[expect(dead_code)] pub type PythonDataSourceResult = Result; @@ -30,6 +65,9 @@ pub enum PythonDataSourceError { /// Resource exhaustion (e.g., partition too large) #[error("Resource exhausted: {0}")] ResourceExhausted(String), + /// Application-declared failure with private Python details discarded. + #[error("{0}")] + DeclaredFailure(#[from] PythonDataSourceFailure), } impl PythonDataSourceError { @@ -92,7 +130,10 @@ impl PythonDataSourceContext { /// Wrap a Python error with context information, preserving traceback. pub fn wrap_py_error(&self, e: pyo3::PyErr) -> PythonDataSourceError { - self.wrap_error(format_py_error_with_traceback(e)) + match declared_failure_kind(&e) { + Some(failure) => failure.into(), + None => self.wrap_error(format_py_error_with_traceback(e)), + } } } @@ -125,7 +166,10 @@ pub fn format_py_error_with_traceback(e: pyo3::PyErr) -> String { impl From for PythonDataSourceError { fn from(e: pyo3::PyErr) -> Self { - Self::python(format_py_error_with_traceback(e)) + match declared_failure_kind(&e) { + Some(failure) => failure.into(), + None => Self::python(format_py_error_with_traceback(e)), + } } } @@ -134,9 +178,12 @@ impl From for PythonDataSourceError { /// This is a shared helper to avoid duplicating this conversion pattern /// across multiple modules (stream.rs, executor.rs, arrow_utils.rs, etc.). pub fn py_err(e: pyo3::PyErr) -> DataFusionError { - DataFusionError::External(Box::new(std::io::Error::other( - format_py_error_with_traceback(e), - ))) + match declared_failure_kind(&e) { + Some(failure) => PythonDataSourceError::from(failure).into(), + None => DataFusionError::External(Box::new(std::io::Error::other( + format_py_error_with_traceback(e), + ))), + } } /// Import cloudpickle from PySpark. @@ -154,3 +201,63 @@ pub fn import_cloudpickle( )) }) } + +#[cfg(test)] +mod tests { + use super::*; + use pyo3::ffi::c_str; + use pyo3::prelude::*; + use pyo3::types::{PyDict, PyDictMethods}; + + #[expect(clippy::unwrap_used)] + fn declared_error(kind: &str) -> PyErr { + Python::initialize(); + Python::attach(|py| { + let namespace = PyDict::new(py); + namespace.set_item("failure_kind", kind).unwrap(); + py.run( + c_str!( + "class DeclaredError(RuntimeError):\n __sail_data_source_failure_kind__ = failure_kind\n" + ), + Some(&namespace), + None, + ) + .unwrap(); + let exception_type = namespace.get_item("DeclaredError").unwrap().unwrap(); + let value = exception_type.call1(("private Python detail",)).unwrap(); + PyErr::from_value(value) + }) + } + + fn assert_declared_failure_is_constant(kind: &str, message: &str) -> Result<(), String> { + match py_err(declared_error(kind)) { + DataFusionError::External(error) => { + assert_eq!(error.to_string(), message); + assert!(!error.to_string().contains("private Python detail")); + let Some(source) = error.source() else { + return Err("classified marker source was not preserved".to_string()); + }; + assert_eq!(source.to_string(), message); + assert!(source.source().is_none()); + Ok(()) + } + other => Err(format!("expected external error, got {other:?}")), + } + } + + #[test] + fn test_terminal_declared_failure_preserves_only_finite_marker() -> Result<(), String> { + assert_declared_failure_is_constant( + "terminal", + "Python data source reported a terminal failure", + ) + } + + #[test] + fn test_transient_declared_failure_preserves_only_finite_marker() -> Result<(), String> { + assert_declared_failure_is_constant( + "transient", + "Python data source reported a transient failure", + ) + } +} diff --git a/crates/sail-python-udf/src/error.rs b/crates/sail-python-udf/src/error.rs index 8bbf3a049a..caddccfe99 100644 --- a/crates/sail-python-udf/src/error.rs +++ b/crates/sail-python-udf/src/error.rs @@ -59,9 +59,34 @@ impl PythonErrorCauseExtractor for PyErrExtractor { Some(PythonErrorCause { summary: e.to_string(), traceback: traceback.ok(), + failure_kind: None, }) } else { None } } } + +#[cfg(test)] +mod tests { + use super::*; + use pyo3::exceptions::PyRuntimeError; + + #[test] + #[expect(clippy::unwrap_used)] + fn test_generic_python_error_ignores_data_source_attribute() { + Python::initialize(); + let error = Python::attach(|py| { + let error = PyRuntimeError::new_err("generic Python detail"); + error + .value(py) + .setattr("__sail_data_source_failure_kind__", "terminal") + .unwrap(); + error + }); + + let cause = PyErrExtractor::extract(&error).unwrap(); + assert_eq!(cause.failure_kind, None); + assert!(cause.summary.contains("generic Python detail")); + } +} diff --git a/crates/sail-spark-connect/src/error.rs b/crates/sail-spark-connect/src/error.rs index 450a8ea86a..cab5547e36 100644 --- a/crates/sail-spark-connect/src/error.rs +++ b/crates/sail-spark-connect/src/error.rs @@ -8,7 +8,7 @@ use datafusion::common::DataFusionError; use prost::{DecodeError, UnknownEnumValue}; use sail_cache::error::CacheError; use sail_common::error::CommonError; -use sail_common_datafusion::error::{CommonErrorCause, PythonErrorCause}; +use sail_common_datafusion::error::{CommonErrorCause, PythonErrorCause, PythonFailureKind}; use sail_execution::error::ExecutionError; use sail_plan::error::PlanError; use sail_python_udf::error::PyErrExtractor; @@ -369,21 +369,31 @@ impl From for SparkThrowable { SparkThrowable::ArithmeticException(x) } CommonErrorCause::ArrowParse(x) => SparkThrowable::ParseException(x), - CommonErrorCause::Python(PythonErrorCause { summary, traceback }) => { - // The message must end with a newline character - // since the PySpark unit tests expect it. - let message = if let Some(traceback) = traceback { - // Each line string already ends with a newline character. - traceback.join("") - } else { - format!("{summary}\n") - }; - if message.contains("net.razorvine.pickle.PickleException") { - SparkThrowable::SparkException(message) - } else { - SparkThrowable::PythonException(message) + CommonErrorCause::Python(PythonErrorCause { + summary, + traceback, + failure_kind, + }) => match failure_kind { + Some(PythonFailureKind::Terminal) => SparkThrowable::AnalysisException(summary), + Some(PythonFailureKind::Transient) => { + SparkThrowable::SparkRuntimeException(summary) } - } + None => { + // The message must end with a newline character + // since the PySpark unit tests expect it. + let message = if let Some(traceback) = traceback { + // Each line string already ends with a newline character. + traceback.join("") + } else { + format!("{summary}\n") + }; + if message.contains("net.razorvine.pickle.PickleException") { + SparkThrowable::SparkException(message) + } else { + SparkThrowable::PythonException(message) + } + } + }, CommonErrorCause::ArrowCast(x) => cast_error_to_throwable(x), CommonErrorCause::Schema(x) | CommonErrorCause::Plan(x) diff --git a/docs/guide/sources/python/index.md b/docs/guide/sources/python/index.md index 23d344f82c..6fa8fe9380 100644 --- a/docs/guide/sources/python/index.md +++ b/docs/guide/sources/python/index.md @@ -12,6 +12,34 @@ You can define a Python class that inherits from the `pyspark.sql.datasource.Dat Currently, Sail supports Python data sources for batch reading and writing. +## Classified Reader Failures + +By default, Sail propagates a Python data source exception with its Python +traceback. A data source can instead declare a finite failure category when its +caller must distinguish a deterministic failure from a retryable one without +inspecting exception text. Define `__sail_data_source_failure_kind__` on the +exception class with one of these exact values: + +- `"terminal"` becomes a Spark `AnalysisException`. +- `"transient"` becomes a Spark `SparkRuntimeException`. + +For a declared failure, Sail replaces the Python message and traceback with a +constant category message before the error crosses Spark Connect. Unknown +values retain the default exception behavior. + +```python +class RetryableReadError(TimeoutError): + __sail_data_source_failure_kind__ = "transient" + + +class ContractReadError(ValueError): + __sail_data_source_failure_kind__ = "terminal" +``` + +Use this protocol only for data-source-controlled exception classes. Do not put +record values, credentials, endpoints, or other runtime details in the category +attribute. + ## Examples diff --git a/python/pysail/tests/spark/datasource/test_python.py b/python/pysail/tests/spark/datasource/test_python.py index ead765716e..047fd03629 100644 --- a/python/pysail/tests/spark/datasource/test_python.py +++ b/python/pysail/tests/spark/datasource/test_python.py @@ -5,6 +5,7 @@ including both Arrow RecordBatch and tuple-based paths. """ +import contextlib import json from collections.abc import Iterator from pathlib import Path @@ -580,6 +581,78 @@ def read(self, partition): # noqa: ARG002 df.collect() +@pytest.mark.parametrize( + ("failure_kind", "expected_type", "expected_message"), + [ + ("terminal", "AnalysisException", "Python data source reported a terminal failure"), + ("transient", "SparkRuntimeException", "Python data source reported a transient failure"), + ], +) +def test_python_declared_failure_kind_is_structured_and_message_free( + failure_kind: str, expected_type: str, expected_message: str +): + """A declared failure crosses Spark Connect by finite class, not Python text.""" + import pyarrow as pa + import pysail.spark + from pyspark import errors + from pyspark.sql import SparkSession + from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition + + class DeclaredDataSourceFailure(Exception): + __sail_data_source_failure_kind__ = failure_kind + + def __getattribute__(self, name: str): + if name == "__sail_data_source_failure_kind__": + return "masked-on-instance" + return super().__getattribute__(name) + + def direct_read(_reader, _partition): + raise DeclaredDataSourceFailure("detail that must not cross the boundary") + + def generator_read(_reader, _partition): + raise DeclaredDataSourceFailure("detail that must not cross the boundary") + yield (0,) # pragma: no cover - makes this a generator like streaming readers + + class DeclaredFailureReader(DataSourceReader): + read = direct_read if failure_kind == "terminal" else generator_read + + def partitions(self): + return [InputPartition(0)] + + class DeclaredFailureDataSource(DataSource): + @classmethod + def name(cls) -> str: + return f"declared_failure_{failure_kind}" + + def schema(self): + return pa.schema([("id", pa.int32())]) + + def reader(self, schema): # noqa: ARG002 + return DeclaredFailureReader() + + server = pysail.spark.SparkConnectServer("127.0.0.1", 0) + server.start() + host, port = server.listening_address + spark = SparkSession.builder.remote(f"sc://{host}:{port}").create() + try: + spark.conf.set("spark.sql.session.localRelationSizeLimit", "3g") + spark.dataSource.register(DeclaredFailureDataSource) + frame = spark.read.format(DeclaredFailureDataSource.name()).load() + exception_type = getattr(errors, expected_type) + + with pytest.raises(exception_type) as caught: + frame.collect() + assert expected_message in str(caught.value) + assert "detail that must not cross the boundary" not in str(caught.value) + assert caught.value.__cause__ is None + assert "detail that must not cross the boundary" not in repr(caught.value.__context__) + finally: + with contextlib.suppress(Exception): + spark.stop() + with contextlib.suppress(Exception): + server.stop() + + def test_python_session_isolation(remote: str): """Test that datasources registered in one session are not visible in another.