diff --git a/Cargo.lock b/Cargo.lock index 729c5742a6..8d6e6f0b8d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7383,6 +7383,7 @@ dependencies = [ "serde", "serde_arrow", "serde_json", + "tempfile", "thiserror 2.0.19", "tokio", ] @@ -7444,6 +7445,7 @@ dependencies = [ "serde", "serde_json", "serde_with", + "tempfile", "testcontainers", "thiserror 2.0.19", "tokio", @@ -7851,6 +7853,7 @@ dependencies = [ "serde", "serde_arrow", "serde_json", + "tokio", "url", "uuid", ] @@ -8088,6 +8091,7 @@ dependencies = [ "sail-sql-analyzer", "sail-telemetry", "secrecy", + "tempfile", "thiserror 2.0.19", "tokio", "tonic", 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..770e4cc910 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,54 +956,131 @@ 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)?; + 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. - 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; - 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; + // 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. + 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? + { + 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. 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}" + ))); } } - let views_result = client - .list_views(prefix.clone(), ns_string.clone(), None, None) - .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; + 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? + { + 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}" + ))); } } } - match client - .drop_namespace( - catalog_config.prefix().map(ToOwned::to_owned), - catalog_config.namespace_string(database)?, - ) - .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( @@ -997,7 +1111,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 +1174,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 +1206,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 +1250,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 +1304,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 +1329,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 +1457,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 +1566,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 +1584,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 +1624,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 +1678,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 +2052,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 +2063,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::*; @@ -2651,6 +2789,386 @@ 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_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(); + + 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 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)) + .mount(&ctx.server) + .await; + Mock::given(method("GET")) + .and(path(ctx.path("/namespaces/ns1/views").as_str())) + .respond_with(ResponseTemplate::new(500)) + .expect(0) + .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(); @@ -3471,4 +3989,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/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..80991b548c 100644 --- a/crates/sail-catalog/src/credentials.rs +++ b/crates/sail-catalog/src/credentials.rs @@ -1,6 +1,7 @@ use std::fmt::Debug; +use std::path::PathBuf; -use crate::error::CatalogResult; +use crate::error::{CatalogError, CatalogResult}; #[async_trait::async_trait] pub trait CatalogCredentials: Debug + Send + Sync + 'static { @@ -34,3 +35,142 @@ impl CatalogCredentials for StaticCatalogCredentials { Ok(Some(self.credential.clone())) } } + +/// Credentials backed by a token file on disk, such as a kubelet-projected +/// 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, +} + +impl FileCatalogCredentials { + pub fn new(path: impl Into) -> Self { + Self { path: path.into() } + } +} + +#[async_trait::async_trait] +impl CatalogCredentials for FileCatalogCredentials { + async fn retrieve(&self) -> CatalogResult> { + 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(); + if credential.is_empty() { + return Err(CatalogError::External(format!( + "token file {} is empty", + self.path.display() + ))); + } + Ok(Some(credential)) + } +} + +#[cfg(test)] +mod tests { + #![expect(clippy::unwrap_used)] + + use std::fs::File; + use std::io::Write; + use std::path::Path; + + 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(); + } + + #[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_rotated_token() { + let dir = TempDir::new().unwrap(); + let path = dir.path().join("token"); + write_token(&path, "first-token"); + + let credentials = FileCatalogCredentials::new(&path); + assert_eq!( + credentials.retrieve().await.unwrap(), + Some("first-token".to_string()) + ); + + // 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"); + 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:?}" + ); + } + + #[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:?}" + ); + + // 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-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-common/src/config/application.rs b/crates/sail-common/src/config/application.rs index ae9af1a14b..aa508e6b7c 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 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)] cache: CatalogCacheConfig, }, diff --git a/crates/sail-data-source/src/formats/python/error.rs b/crates/sail-data-source/src/formats/python/error.rs index ecff918436..4cc9b0d234 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,64 @@ pub fn import_cloudpickle( )) }) } + +#[cfg(test)] +mod tests { + use pyo3::ffi::c_str; + use pyo3::prelude::*; + use pyo3::types::{PyDict, PyDictMethods}; + + use super::*; + + #[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-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()); + } +} diff --git a/crates/sail-python-udf/src/error.rs b/crates/sail-python-udf/src/error.rs index 8bbf3a049a..8e6e6b2c33 100644 --- a/crates/sail-python-udf/src/error.rs +++ b/crates/sail-python-udf/src/error.rs @@ -59,9 +59,35 @@ impl PythonErrorCauseExtractor for PyErrExtractor { Some(PythonErrorCause { summary: e.to_string(), traceback: traceback.ok(), + failure_kind: None, }) } else { None } } } + +#[cfg(test)] +mod tests { + use pyo3::exceptions::PyRuntimeError; + + use super::*; + + #[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-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 46832190fa..b147b4c68f 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}; @@ -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, @@ -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,11 @@ 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)); + 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( || { @@ -317,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 { @@ -412,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); + } } 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/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="::"}]' ``` 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..5a152a4072 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,81 @@ 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 + from pyspark import errors + from pyspark.sql import SparkSession + from pyspark.sql.datasource import DataSource, DataSourceReader, InputPartition + + import pysail.spark + + class DeclaredDataSourceError(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) + + private_detail = "detail that must not cross the boundary" + + def direct_read(_reader, _partition): + raise DeclaredDataSourceError(private_detail) + + def generator_read(_reader, _partition): + raise DeclaredDataSourceError(private_detail) + 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 private_detail not in str(caught.value) + assert caught.value.__cause__ is None + assert private_detail 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.