From ac73567324f2f67b4fffd84e92dc42285c4298c9 Mon Sep 17 00:00:00 2001 From: David Dallakyan Date: Sun, 16 Aug 2026 16:37:48 +0400 Subject: [PATCH 1/6] feat(catalog-rest): add REST server-side scan planning client Port plan / fetch-result / cancel / fetch-tasks onto RestCatalog with a WaitForPlan poller. Task decoding and TableScan routing stay follow-ups, so supports_remote_scan_planning remains false. Part of #1690. --- Cargo.lock | 1 + crates/catalog/rest/Cargo.toml | 5 +- crates/catalog/rest/public-api.txt | 175 +++ crates/catalog/rest/src/catalog.rs | 14 +- crates/catalog/rest/src/endpoint.rs | 28 + crates/catalog/rest/src/lib.rs | 2 + crates/catalog/rest/src/scan_planning.rs | 1323 ++++++++++++++++++++++ crates/catalog/rest/src/types.rs | 6 + 8 files changed, 1545 insertions(+), 9 deletions(-) create mode 100644 crates/catalog/rest/src/scan_planning.rs diff --git a/Cargo.lock b/Cargo.lock index ae367e7d8c..36f522f311 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3899,6 +3899,7 @@ dependencies = [ "iceberg_test_utils", "itertools 0.13.0", "mockito", + "rand 0.9.5", "reqwest 0.12.28", "serde", "serde_derive", diff --git a/crates/catalog/rest/Cargo.toml b/crates/catalog/rest/Cargo.toml index 8dc9a86d7f..9c78c1b24f 100644 --- a/crates/catalog/rest/Cargo.toml +++ b/crates/catalog/rest/Cargo.toml @@ -35,14 +35,15 @@ chrono = { workspace = true } http = { workspace = true } iceberg = { workspace = true } itertools = { workspace = true } +rand = { workspace = true } reqwest = { workspace = true } serde = { workspace = true } serde_derive = { workspace = true } serde_json = { workspace = true } -tokio = { workspace = true } +tokio = { workspace = true, features = ["time"] } tracing = { workspace = true } typed-builder = { workspace = true } -uuid = { workspace = true, features = ["v4"] } +uuid = { workspace = true, features = ["v4", "v7"] } [dev-dependencies] bytes = { workspace = true } diff --git a/crates/catalog/rest/public-api.txt b/crates/catalog/rest/public-api.txt index 6f3fa58a47..8d5334b6e6 100644 --- a/crates/catalog/rest/public-api.txt +++ b/crates/catalog/rest/public-api.txt @@ -14,6 +14,24 @@ impl<'a> core::fmt::Debug for iceberg_catalog_rest::HttpRequestBody<'a> pub fn iceberg_catalog_rest::HttpRequestBody<'a>::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl<'a> core::marker::Copy for iceberg_catalog_rest::HttpRequestBody<'a> impl<'a> core::marker::StructuralPartialEq for iceberg_catalog_rest::HttpRequestBody<'a> +pub enum iceberg_catalog_rest::PlanStatus +pub iceberg_catalog_rest::PlanStatus::Cancelled +pub iceberg_catalog_rest::PlanStatus::Completed +pub iceberg_catalog_rest::PlanStatus::Failed +pub iceberg_catalog_rest::PlanStatus::Submitted +impl core::clone::Clone for iceberg_catalog_rest::PlanStatus +pub fn iceberg_catalog_rest::PlanStatus::clone(&self) -> iceberg_catalog_rest::PlanStatus +impl core::cmp::Eq for iceberg_catalog_rest::PlanStatus +impl core::cmp::PartialEq for iceberg_catalog_rest::PlanStatus +pub fn iceberg_catalog_rest::PlanStatus::eq(&self, other: &iceberg_catalog_rest::PlanStatus) -> bool +impl core::fmt::Debug for iceberg_catalog_rest::PlanStatus +pub fn iceberg_catalog_rest::PlanStatus::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::marker::Copy for iceberg_catalog_rest::PlanStatus +impl core::marker::StructuralPartialEq for iceberg_catalog_rest::PlanStatus +impl serde_core::ser::Serialize for iceberg_catalog_rest::PlanStatus +pub fn iceberg_catalog_rest::PlanStatus::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::PlanStatus +pub fn iceberg_catalog_rest::PlanStatus::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub struct iceberg_catalog_rest::CommitTableRequest pub iceberg_catalog_rest::CommitTableRequest::identifier: core::option::Option pub iceberg_catalog_rest::CommitTableRequest::requirements: alloc::vec::Vec @@ -44,6 +62,22 @@ impl serde_core::ser::Serialize for iceberg_catalog_rest::CommitTableResponse pub fn iceberg_catalog_rest::CommitTableResponse::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::CommitTableResponse pub fn iceberg_catalog_rest::CommitTableResponse::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg_catalog_rest::CompletedPlanningResult +pub iceberg_catalog_rest::CompletedPlanningResult::scan_tasks: iceberg_catalog_rest::ScanTasks +pub iceberg_catalog_rest::CompletedPlanningResult::status: iceberg_catalog_rest::PlanStatus +pub iceberg_catalog_rest::CompletedPlanningResult::storage_credentials: core::option::Option> +impl core::clone::Clone for iceberg_catalog_rest::CompletedPlanningResult +pub fn iceberg_catalog_rest::CompletedPlanningResult::clone(&self) -> iceberg_catalog_rest::CompletedPlanningResult +impl core::cmp::Eq for iceberg_catalog_rest::CompletedPlanningResult +impl core::cmp::PartialEq for iceberg_catalog_rest::CompletedPlanningResult +pub fn iceberg_catalog_rest::CompletedPlanningResult::eq(&self, other: &iceberg_catalog_rest::CompletedPlanningResult) -> bool +impl core::fmt::Debug for iceberg_catalog_rest::CompletedPlanningResult +pub fn iceberg_catalog_rest::CompletedPlanningResult::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::marker::StructuralPartialEq for iceberg_catalog_rest::CompletedPlanningResult +impl serde_core::ser::Serialize for iceberg_catalog_rest::CompletedPlanningResult +pub fn iceberg_catalog_rest::CompletedPlanningResult::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::CompletedPlanningResult +pub fn iceberg_catalog_rest::CompletedPlanningResult::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub struct iceberg_catalog_rest::CreateNamespaceRequest pub iceberg_catalog_rest::CreateNamespaceRequest::namespace: iceberg::catalog::NamespaceIdent pub iceberg_catalog_rest::CreateNamespaceRequest::properties: std::collections::hash::map::HashMap @@ -124,6 +158,52 @@ impl serde_core::ser::Serialize for iceberg_catalog_rest::ErrorResponse pub fn iceberg_catalog_rest::ErrorResponse::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::ErrorResponse pub fn iceberg_catalog_rest::ErrorResponse::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg_catalog_rest::FetchPlanningResultOptions +pub iceberg_catalog_rest::FetchPlanningResultOptions::access_delegation: core::option::Option +impl core::clone::Clone for iceberg_catalog_rest::FetchPlanningResultOptions +pub fn iceberg_catalog_rest::FetchPlanningResultOptions::clone(&self) -> iceberg_catalog_rest::FetchPlanningResultOptions +impl core::default::Default for iceberg_catalog_rest::FetchPlanningResultOptions +pub fn iceberg_catalog_rest::FetchPlanningResultOptions::default() -> iceberg_catalog_rest::FetchPlanningResultOptions +impl core::fmt::Debug for iceberg_catalog_rest::FetchPlanningResultOptions +pub fn iceberg_catalog_rest::FetchPlanningResultOptions::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +pub struct iceberg_catalog_rest::FetchPlanningResultResponse +pub iceberg_catalog_rest::FetchPlanningResultResponse::error: core::option::Option +pub iceberg_catalog_rest::FetchPlanningResultResponse::scan_tasks: iceberg_catalog_rest::ScanTasks +pub iceberg_catalog_rest::FetchPlanningResultResponse::status: iceberg_catalog_rest::PlanStatus +pub iceberg_catalog_rest::FetchPlanningResultResponse::storage_credentials: core::option::Option> +impl core::fmt::Debug for iceberg_catalog_rest::FetchPlanningResultResponse +pub fn iceberg_catalog_rest::FetchPlanningResultResponse::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl serde_core::ser::Serialize for iceberg_catalog_rest::FetchPlanningResultResponse +pub fn iceberg_catalog_rest::FetchPlanningResultResponse::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::FetchPlanningResultResponse +pub fn iceberg_catalog_rest::FetchPlanningResultResponse::deserialize>(deserializer: D) -> core::result::Result::Error> +pub struct iceberg_catalog_rest::FetchScanTasksRequest +pub iceberg_catalog_rest::FetchScanTasksRequest::idempotency_key: core::option::Option +pub iceberg_catalog_rest::FetchScanTasksRequest::plan_task: alloc::string::String +impl core::clone::Clone for iceberg_catalog_rest::FetchScanTasksRequest +pub fn iceberg_catalog_rest::FetchScanTasksRequest::clone(&self) -> iceberg_catalog_rest::FetchScanTasksRequest +impl core::fmt::Debug for iceberg_catalog_rest::FetchScanTasksRequest +pub fn iceberg_catalog_rest::FetchScanTasksRequest::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl serde_core::ser::Serialize for iceberg_catalog_rest::FetchScanTasksRequest +pub fn iceberg_catalog_rest::FetchScanTasksRequest::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::FetchScanTasksRequest +pub fn iceberg_catalog_rest::FetchScanTasksRequest::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg_catalog_rest::FetchScanTasksResponse +pub iceberg_catalog_rest::FetchScanTasksResponse::scan_tasks: iceberg_catalog_rest::ScanTasks +impl core::clone::Clone for iceberg_catalog_rest::FetchScanTasksResponse +pub fn iceberg_catalog_rest::FetchScanTasksResponse::clone(&self) -> iceberg_catalog_rest::FetchScanTasksResponse +impl core::cmp::Eq for iceberg_catalog_rest::FetchScanTasksResponse +impl core::cmp::PartialEq for iceberg_catalog_rest::FetchScanTasksResponse +pub fn iceberg_catalog_rest::FetchScanTasksResponse::eq(&self, other: &iceberg_catalog_rest::FetchScanTasksResponse) -> bool +impl core::default::Default for iceberg_catalog_rest::FetchScanTasksResponse +pub fn iceberg_catalog_rest::FetchScanTasksResponse::default() -> iceberg_catalog_rest::FetchScanTasksResponse +impl core::fmt::Debug for iceberg_catalog_rest::FetchScanTasksResponse +pub fn iceberg_catalog_rest::FetchScanTasksResponse::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::marker::StructuralPartialEq for iceberg_catalog_rest::FetchScanTasksResponse +impl serde_core::ser::Serialize for iceberg_catalog_rest::FetchScanTasksResponse +pub fn iceberg_catalog_rest::FetchScanTasksResponse::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::FetchScanTasksResponse +pub fn iceberg_catalog_rest::FetchScanTasksResponse::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub struct iceberg_catalog_rest::HttpClient impl iceberg_catalog_rest::HttpClient pub async fn iceberg_catalog_rest::HttpClient::post_form(&self, url: &str, headers: &http::header::map::HeaderMap, form: &std::collections::hash::map::HashMap<&str, &str>) -> iceberg::error::Result<(http::status::StatusCode, alloc::vec::Vec)> @@ -225,6 +305,40 @@ pub fn iceberg_catalog_rest::OAuth2Manager::fmt(&self, f: &mut core::fmt::Format impl iceberg_catalog_rest::AuthManager for iceberg_catalog_rest::OAuth2Manager pub fn iceberg_catalog_rest::OAuth2Manager::catalog_session<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, client: &'life1 iceberg_catalog_rest::HttpClient, props: &'life2 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait pub fn iceberg_catalog_rest::OAuth2Manager::init_session<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, client: &'life1 iceberg_catalog_rest::HttpClient, props: &'life2 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub struct iceberg_catalog_rest::PlanTableScanRequest +pub iceberg_catalog_rest::PlanTableScanRequest::access_delegation: core::option::Option +pub iceberg_catalog_rest::PlanTableScanRequest::case_sensitive: core::option::Option +pub iceberg_catalog_rest::PlanTableScanRequest::end_snapshot_id: core::option::Option +pub iceberg_catalog_rest::PlanTableScanRequest::filter: core::option::Option +pub iceberg_catalog_rest::PlanTableScanRequest::idempotency_key: core::option::Option +pub iceberg_catalog_rest::PlanTableScanRequest::min_rows_requested: core::option::Option +pub iceberg_catalog_rest::PlanTableScanRequest::select: alloc::vec::Vec +pub iceberg_catalog_rest::PlanTableScanRequest::snapshot_id: core::option::Option +pub iceberg_catalog_rest::PlanTableScanRequest::start_snapshot_id: core::option::Option +pub iceberg_catalog_rest::PlanTableScanRequest::stats_fields: alloc::vec::Vec +pub iceberg_catalog_rest::PlanTableScanRequest::use_snapshot_schema: core::option::Option +impl core::clone::Clone for iceberg_catalog_rest::PlanTableScanRequest +pub fn iceberg_catalog_rest::PlanTableScanRequest::clone(&self) -> iceberg_catalog_rest::PlanTableScanRequest +impl core::default::Default for iceberg_catalog_rest::PlanTableScanRequest +pub fn iceberg_catalog_rest::PlanTableScanRequest::default() -> iceberg_catalog_rest::PlanTableScanRequest +impl core::fmt::Debug for iceberg_catalog_rest::PlanTableScanRequest +pub fn iceberg_catalog_rest::PlanTableScanRequest::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl serde_core::ser::Serialize for iceberg_catalog_rest::PlanTableScanRequest +pub fn iceberg_catalog_rest::PlanTableScanRequest::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::PlanTableScanRequest +pub fn iceberg_catalog_rest::PlanTableScanRequest::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg_catalog_rest::PlanTableScanResponse +pub iceberg_catalog_rest::PlanTableScanResponse::error: core::option::Option +pub iceberg_catalog_rest::PlanTableScanResponse::plan_id: core::option::Option +pub iceberg_catalog_rest::PlanTableScanResponse::scan_tasks: iceberg_catalog_rest::ScanTasks +pub iceberg_catalog_rest::PlanTableScanResponse::status: iceberg_catalog_rest::PlanStatus +pub iceberg_catalog_rest::PlanTableScanResponse::storage_credentials: core::option::Option> +impl core::fmt::Debug for iceberg_catalog_rest::PlanTableScanResponse +pub fn iceberg_catalog_rest::PlanTableScanResponse::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl serde_core::ser::Serialize for iceberg_catalog_rest::PlanTableScanResponse +pub fn iceberg_catalog_rest::PlanTableScanResponse::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::PlanTableScanResponse +pub fn iceberg_catalog_rest::PlanTableScanResponse::deserialize>(deserializer: D) -> core::result::Result::Error> pub struct iceberg_catalog_rest::RegisterTableRequest pub iceberg_catalog_rest::RegisterTableRequest::metadata_location: alloc::string::String pub iceberg_catalog_rest::RegisterTableRequest::name: alloc::string::String @@ -257,6 +371,15 @@ pub fn iceberg_catalog_rest::RenameTableRequest::serialize<__S>(&self, __seriali impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::RenameTableRequest pub fn iceberg_catalog_rest::RenameTableRequest::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub struct iceberg_catalog_rest::RestCatalog +impl iceberg_catalog_rest::RestCatalog +pub async fn iceberg_catalog_rest::RestCatalog::cancel_planning(&self, table: &iceberg::catalog::TableIdent, plan_id: &str) -> iceberg::error::Result<()> +pub async fn iceberg_catalog_rest::RestCatalog::fetch_planning_result(&self, table: &iceberg::catalog::TableIdent, plan_id: &str, opts: iceberg_catalog_rest::FetchPlanningResultOptions) -> iceberg::error::Result +pub async fn iceberg_catalog_rest::RestCatalog::fetch_scan_tasks(&self, table: &iceberg::catalog::TableIdent, request: iceberg_catalog_rest::FetchScanTasksRequest) -> iceberg::error::Result +pub async fn iceberg_catalog_rest::RestCatalog::plan_table_scan(&self, table: &iceberg::catalog::TableIdent, request: iceberg_catalog_rest::PlanTableScanRequest) -> iceberg::error::Result +pub async fn iceberg_catalog_rest::RestCatalog::supports_full_remote_scan_planning(&self) -> iceberg::error::Result +pub async fn iceberg_catalog_rest::RestCatalog::supports_plan_table_scan(&self) -> iceberg::error::Result +pub async fn iceberg_catalog_rest::RestCatalog::supports_remote_scan_planning(&self) -> iceberg::error::Result +pub async fn iceberg_catalog_rest::RestCatalog::wait_for_plan(&self, table: &iceberg::catalog::TableIdent, plan_id: &str, opts: iceberg_catalog_rest::WaitForPlanOptions) -> iceberg::error::Result impl core::fmt::Debug for iceberg_catalog_rest::RestCatalog pub fn iceberg_catalog_rest::RestCatalog::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl iceberg::catalog::Catalog for iceberg_catalog_rest::RestCatalog @@ -289,6 +412,40 @@ pub fn iceberg_catalog_rest::RestCatalogBuilder::load(self, name: impl core::con pub fn iceberg_catalog_rest::RestCatalogBuilder::with_kms_client_factory(self, kms_client_factory: alloc::sync::Arc) -> Self pub fn iceberg_catalog_rest::RestCatalogBuilder::with_runtime(self, runtime: iceberg::runtime::Runtime) -> Self pub fn iceberg_catalog_rest::RestCatalogBuilder::with_storage_factory(self, storage_factory: alloc::sync::Arc) -> Self +pub struct iceberg_catalog_rest::RestFileScanTask +pub iceberg_catalog_rest::RestFileScanTask::data_file: serde_json::value::Value +pub iceberg_catalog_rest::RestFileScanTask::delete_file_references: core::option::Option> +pub iceberg_catalog_rest::RestFileScanTask::residual_filter: core::option::Option +impl core::clone::Clone for iceberg_catalog_rest::RestFileScanTask +pub fn iceberg_catalog_rest::RestFileScanTask::clone(&self) -> iceberg_catalog_rest::RestFileScanTask +impl core::cmp::Eq for iceberg_catalog_rest::RestFileScanTask +impl core::cmp::PartialEq for iceberg_catalog_rest::RestFileScanTask +pub fn iceberg_catalog_rest::RestFileScanTask::eq(&self, other: &iceberg_catalog_rest::RestFileScanTask) -> bool +impl core::fmt::Debug for iceberg_catalog_rest::RestFileScanTask +pub fn iceberg_catalog_rest::RestFileScanTask::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::marker::StructuralPartialEq for iceberg_catalog_rest::RestFileScanTask +impl serde_core::ser::Serialize for iceberg_catalog_rest::RestFileScanTask +pub fn iceberg_catalog_rest::RestFileScanTask::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::RestFileScanTask +pub fn iceberg_catalog_rest::RestFileScanTask::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg_catalog_rest::ScanTasks +pub iceberg_catalog_rest::ScanTasks::delete_files: alloc::vec::Vec +pub iceberg_catalog_rest::ScanTasks::file_scan_tasks: alloc::vec::Vec +pub iceberg_catalog_rest::ScanTasks::plan_tasks: alloc::vec::Vec +impl core::clone::Clone for iceberg_catalog_rest::ScanTasks +pub fn iceberg_catalog_rest::ScanTasks::clone(&self) -> iceberg_catalog_rest::ScanTasks +impl core::cmp::Eq for iceberg_catalog_rest::ScanTasks +impl core::cmp::PartialEq for iceberg_catalog_rest::ScanTasks +pub fn iceberg_catalog_rest::ScanTasks::eq(&self, other: &iceberg_catalog_rest::ScanTasks) -> bool +impl core::default::Default for iceberg_catalog_rest::ScanTasks +pub fn iceberg_catalog_rest::ScanTasks::default() -> iceberg_catalog_rest::ScanTasks +impl core::fmt::Debug for iceberg_catalog_rest::ScanTasks +pub fn iceberg_catalog_rest::ScanTasks::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::marker::StructuralPartialEq for iceberg_catalog_rest::ScanTasks +impl serde_core::ser::Serialize for iceberg_catalog_rest::ScanTasks +pub fn iceberg_catalog_rest::ScanTasks::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::ScanTasks +pub fn iceberg_catalog_rest::ScanTasks::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub struct iceberg_catalog_rest::StorageCredential pub iceberg_catalog_rest::StorageCredential::config: std::collections::hash::map::HashMap pub iceberg_catalog_rest::StorageCredential::prefix: alloc::string::String @@ -335,6 +492,19 @@ impl serde_core::ser::Serialize for iceberg_catalog_rest::UpdateNamespacePropert pub fn iceberg_catalog_rest::UpdateNamespacePropertiesResponse::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::UpdateNamespacePropertiesResponse pub fn iceberg_catalog_rest::UpdateNamespacePropertiesResponse::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg_catalog_rest::WaitForPlanOptions +pub iceberg_catalog_rest::WaitForPlanOptions::access_delegation: core::option::Option +pub iceberg_catalog_rest::WaitForPlanOptions::cancel_grace_period: core::time::Duration +pub iceberg_catalog_rest::WaitForPlanOptions::max_delay: core::time::Duration +pub iceberg_catalog_rest::WaitForPlanOptions::max_retries: u32 +pub iceberg_catalog_rest::WaitForPlanOptions::min_delay: core::time::Duration +pub iceberg_catalog_rest::WaitForPlanOptions::timeout: core::option::Option +impl core::clone::Clone for iceberg_catalog_rest::WaitForPlanOptions +pub fn iceberg_catalog_rest::WaitForPlanOptions::clone(&self) -> iceberg_catalog_rest::WaitForPlanOptions +impl core::default::Default for iceberg_catalog_rest::WaitForPlanOptions +pub fn iceberg_catalog_rest::WaitForPlanOptions::default() -> Self +impl core::fmt::Debug for iceberg_catalog_rest::WaitForPlanOptions +pub fn iceberg_catalog_rest::WaitForPlanOptions::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub const iceberg_catalog_rest::AUTH_TYPE_NONE: &str pub const iceberg_catalog_rest::AUTH_TYPE_OAUTH2: &str pub const iceberg_catalog_rest::REST_CATALOG_PROP_AUTH_TYPE: &str @@ -352,3 +522,8 @@ pub fn iceberg_catalog_rest::OAuth2Manager::catalog_session<'life0, 'life1, 'lif pub fn iceberg_catalog_rest::OAuth2Manager::init_session<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, client: &'life1 iceberg_catalog_rest::HttpClient, props: &'life2 std::collections::hash::map::HashMap) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait pub trait iceberg_catalog_rest::AuthSession: core::fmt::Debug + core::marker::Send + core::marker::Sync pub fn iceberg_catalog_rest::AuthSession::authenticate<'life0, 'life1, 'async_trait>(&'life0 self, request: &'life1 mut iceberg_catalog_rest::HttpRequest) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait +pub fn iceberg_catalog_rest::is_no_such_plan_task(err: &iceberg::error::Error) -> bool +pub fn iceberg_catalog_rest::is_plan_cancelled(err: &iceberg::error::Error) -> bool +pub fn iceberg_catalog_rest::is_plan_expired(err: &iceberg::error::Error) -> bool +pub fn iceberg_catalog_rest::is_plan_failed(err: &iceberg::error::Error) -> bool +pub fn iceberg_catalog_rest::is_plan_poll_exhausted(err: &iceberg::error::Error) -> bool diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 0c15c50663..be57fbe830 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -285,7 +285,7 @@ impl RestCatalogConfig { self.url_prefixed(&["namespaces", &ns.to_url_string(), "register"]) } - fn table_endpoint(&self, table: &TableIdent) -> String { + pub(crate) fn table_endpoint(&self, table: &TableIdent) -> String { self.url_prefixed(&[ "namespaces", &table.namespace.to_url_string(), @@ -459,14 +459,14 @@ pub(crate) fn oauth_params_from_props(props: &HashMap) -> HashMa } #[derive(Debug)] -struct RestClient { +pub(crate) struct RestClient { /// Carries the session the auth manager derived from the merged /// configuration, so every request below is authenticated. - http_client: HttpClient, + pub(crate) http_client: HttpClient, /// Runtime config is fetched from rest server and stored here. /// /// It's could be different from the user config. - config: RestCatalogConfig, + pub(crate) config: RestCatalogConfig, /// Capabilities the server advertises (see [`RestCatalog::supports_endpoint`]). endpoints: HashSet, } @@ -479,7 +479,7 @@ impl RestClient { } /// Sends `request`, authenticated by the client's session. - async fn query_catalog(&self, request: HttpRequest) -> Result { + pub(crate) async fn query_catalog(&self, request: HttpRequest) -> Result { self.http_client.query_catalog(request).await } } @@ -504,7 +504,7 @@ pub struct RestCatalog { impl RestCatalog { /// Creates a `RestCatalog` from a [`RestCatalogConfig`]. - fn new( + pub(crate) fn new( config: RestCatalogConfig, auth_manager: Option>, storage_factory: Option>, @@ -623,7 +623,7 @@ impl RestCatalog { } /// Gets the [`RestClient`] from the catalog. - async fn client(&self) -> Result<&RestClient> { + pub(crate) async fn client(&self) -> Result<&RestClient> { self.client .get_or_try_init(|| async { let http_client = HttpClient::new(&self.user_config)?; diff --git a/crates/catalog/rest/src/endpoint.rs b/crates/catalog/rest/src/endpoint.rs index 1f4e7f354a..5a5c9b8129 100644 --- a/crates/catalog/rest/src/endpoint.rs +++ b/crates/catalog/rest/src/endpoint.rs @@ -169,6 +169,10 @@ endpoints! { V1_REGISTER_TABLE => POST "/v1/{prefix}/namespaces/{namespace}/register", V1_REPORT_METRICS => POST "/v1/{prefix}/namespaces/{namespace}/tables/{table}/metrics", V1_COMMIT_TRANSACTION => POST "/v1/{prefix}/transactions/commit", + V1_PLAN_TABLE_SCAN => POST "/v1/{prefix}/namespaces/{namespace}/tables/{table}/plan", + V1_FETCH_PLAN_RESULT => GET "/v1/{prefix}/namespaces/{namespace}/tables/{table}/plan/{plan-id}", + V1_CANCEL_PLANNING => DELETE "/v1/{prefix}/namespaces/{namespace}/tables/{table}/plan/{plan-id}", + V1_FETCH_SCAN_TASKS => POST "/v1/{prefix}/namespaces/{namespace}/tables/{table}/tasks", } /// The standard v1 endpoints assumed to be supported when a server's @@ -232,4 +236,28 @@ mod tests { fn normalizes_http_method_to_uppercase() { assert_eq!("get /v1/x".parse::().unwrap().method(), "GET"); } + + #[test] + fn scan_planning_endpoints_are_optional_and_not_in_the_default_set() { + assert_eq!( + V1_PLAN_TABLE_SCAN.to_string(), + "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan" + ); + assert_eq!( + V1_FETCH_PLAN_RESULT.to_string(), + "GET /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan/{plan-id}" + ); + assert_eq!( + V1_CANCEL_PLANNING.to_string(), + "DELETE /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan/{plan-id}" + ); + assert_eq!( + V1_FETCH_SCAN_TASKS.to_string(), + "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/tasks" + ); + assert!(!DEFAULT_ENDPOINTS.contains(&V1_PLAN_TABLE_SCAN)); + assert!(!DEFAULT_ENDPOINTS.contains(&V1_FETCH_PLAN_RESULT)); + assert!(!DEFAULT_ENDPOINTS.contains(&V1_CANCEL_PLANNING)); + assert!(!DEFAULT_ENDPOINTS.contains(&V1_FETCH_SCAN_TASKS)); + } } diff --git a/crates/catalog/rest/src/lib.rs b/crates/catalog/rest/src/lib.rs index 5670ef3c9c..465a4a10e7 100644 --- a/crates/catalog/rest/src/lib.rs +++ b/crates/catalog/rest/src/lib.rs @@ -58,9 +58,11 @@ pub use client::HttpClient; mod request; pub use request::{HttpRequest, HttpRequestBody}; mod endpoint; +mod scan_planning; mod types; pub use auth::*; pub use catalog::*; pub use endpoint::Endpoint; +pub use scan_planning::*; pub use types::*; diff --git a/crates/catalog/rest/src/scan_planning.rs b/crates/catalog/rest/src/scan_planning.rs new file mode 100644 index 0000000000..3bd9216293 --- /dev/null +++ b/crates/catalog/rest/src/scan_planning.rs @@ -0,0 +1,1323 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +//! REST server-side scan planning client. +//! +//! Implements the plan / fetch-result / cancel / fetch-tasks endpoints and a +//! [`RestCatalog::wait_for_plan`] poller. Task decoding and `TableScan` +//! auto-routing are follow-ups: [`RestCatalog::supports_remote_scan_planning`] +//! stays `false` until those land. + +use std::time::Duration; + +use iceberg::{Error, ErrorKind, Result, TableIdent}; +use rand::Rng; +use reqwest::{Method, Response, StatusCode}; +use serde::de::{self, Deserializer}; +use serde::{Deserialize, Serialize}; +use uuid::{Uuid, Variant, Version}; + +use crate::catalog::RestCatalog; +use crate::client::{deserialize_catalog_response, deserialize_unexpected_catalog_error}; +use crate::endpoint::{ + Endpoint, V1_CANCEL_PLANNING, V1_FETCH_PLAN_RESULT, V1_FETCH_SCAN_TASKS, V1_PLAN_TABLE_SCAN, +}; +use crate::request::HttpRequest; +use crate::types::{ErrorModel, ErrorResponse, StorageCredential}; + +const HEADER_IDEMPOTENCY_KEY: &str = "Idempotency-Key"; +const HEADER_ACCESS_DELEGATION: &str = "X-Iceberg-Access-Delegation"; + +const MSG_PLAN_EXPIRED: &str = "scan plan expired"; +const MSG_PLAN_FAILED: &str = "scan plan failed"; +const MSG_PLAN_CANCELLED: &str = "scan plan cancelled"; +const MSG_NO_SUCH_PLAN_TASK: &str = "scan plan task not found"; +const MSG_PLAN_POLL_EXHAUSTED: &str = "scan plan polling exhausted retries"; + +const ERR_TYPE_NO_SUCH_PLAN_ID: &str = "NoSuchPlanIdException"; +const ERR_TYPE_NO_SUCH_PLAN_TASK: &str = "NoSuchPlanTaskException"; +const ERR_TYPE_NO_SUCH_TABLE: &str = "NoSuchTableException"; +const ERR_TYPE_NO_SUCH_NAMESPACE: &str = "NoSuchNamespaceException"; + +/// Status of a server-side scan plan. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum PlanStatus { + /// Planning finished and tasks (or plan-task handles) are available. + Completed, + /// Planning is still running; poll [`RestCatalog::fetch_planning_result`]. + Submitted, + /// The plan was cancelled. Valid on fetch-result, not on planTableScan. + Cancelled, + /// Planning failed. The error detail is on the failed arm. + Failed, +} + +/// Task payload shared by completed planning responses and fetchScanTasks. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "kebab-case")] +pub struct ScanTasks { + /// Opaque plan-task handles that still need [`RestCatalog::fetch_scan_tasks`]. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub plan_tasks: Vec, + /// File scan tasks. `data-file` is left as JSON until a decoder lands. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub file_scan_tasks: Vec, + /// Delete files referenced by the scan tasks, as raw REST JSON. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub delete_files: Vec, +} + +/// REST `FileScanTask` wire payload. Nested content-files stay opaque. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub struct RestFileScanTask { + /// REST ContentFile JSON for the data file. + pub data_file: serde_json::Value, + /// Indices into the sibling delete-files array. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub delete_file_references: Option>, + /// Optional residual filter in ExpressionParser JSON. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub residual_filter: Option, +} + +/// POST `.../plan` request body. Header-only fields are skipped on the wire. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub struct PlanTableScanRequest { + /// `Idempotency-Key` header. `None` generates a fresh UUIDv7 per call. + #[serde(skip)] + pub idempotency_key: Option, + /// `X-Iceberg-Access-Delegation` header. `None` sends no such header. + #[serde(skip)] + pub access_delegation: Option, + /// Snapshot to scan. Omitted for the current snapshot. + #[serde(skip_serializing_if = "Option::is_none")] + pub snapshot_id: Option, + /// Selected schema fields. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub select: Vec, + /// Row filter as ExpressionParser JSON, not `iceberg::expr::Predicate`. + #[serde(skip_serializing_if = "Option::is_none")] + pub filter: Option, + /// Hint for the minimum number of rows the server should return. + #[serde(skip_serializing_if = "Option::is_none")] + pub min_rows_requested: Option, + /// Case-sensitive field matching for filter and select. + #[serde(skip_serializing_if = "Option::is_none")] + pub case_sensitive: Option, + /// When true, use the schema at the scanned snapshot. + #[serde(skip_serializing_if = "Option::is_none")] + pub use_snapshot_schema: Option, + /// Incremental scan start (exclusive). Wire-only in this PR. + #[serde(skip_serializing_if = "Option::is_none")] + pub start_snapshot_id: Option, + /// Incremental scan end (inclusive). Wire-only in this PR. + #[serde(skip_serializing_if = "Option::is_none")] + pub end_snapshot_id: Option, + /// Fields for which the server should send column stats. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub stats_fields: Vec, +} + +/// POST `.../plan` response. `completed` and `submitted` require `plan-id`. +#[derive(Debug, Serialize)] +#[serde(rename_all = "kebab-case")] +pub struct PlanTableScanResponse { + /// Discriminator for the planning-result union. + pub status: PlanStatus, + /// Server-issued plan id. Required for completed and submitted. + #[serde(skip_serializing_if = "Option::is_none")] + pub plan_id: Option, + /// Failed-arm error detail. + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, + /// Task payload. Empty unless status is completed. + #[serde(flatten)] + pub scan_tasks: ScanTasks, + /// Optional vended credentials for reading the returned files. + #[serde(skip_serializing_if = "Option::is_none")] + pub storage_credentials: Option>, +} + +/// GET `.../plan/{plan-id}` response. +#[derive(Debug, Serialize)] +#[serde(rename_all = "kebab-case")] +pub struct FetchPlanningResultResponse { + /// Discriminator for the planning-result union. + pub status: PlanStatus, + /// Failed-arm error detail. + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, + /// Task payload. Empty unless status is completed. + #[serde(flatten)] + pub scan_tasks: ScanTasks, + /// Optional vended credentials for reading the returned files. + #[serde(skip_serializing_if = "Option::is_none")] + pub storage_credentials: Option>, +} + +/// Completed arm of a planning result, as returned by [`RestCatalog::wait_for_plan`]. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub struct CompletedPlanningResult { + /// Always [`PlanStatus::Completed`]. + pub status: PlanStatus, + /// Task payload, which may still include plan-task handles. + #[serde(flatten)] + pub scan_tasks: ScanTasks, + /// Optional vended credentials for reading the returned files. + #[serde(skip_serializing_if = "Option::is_none")] + pub storage_credentials: Option>, +} + +/// Per-call options for [`RestCatalog::fetch_planning_result`]. +#[derive(Debug, Clone, Default)] +pub struct FetchPlanningResultOptions { + /// `X-Iceberg-Access-Delegation` header. `None` sends no such header. + pub access_delegation: Option, +} + +/// POST `.../tasks` request body. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub struct FetchScanTasksRequest { + /// `Idempotency-Key` header. `None` generates a fresh UUIDv7 per call. + #[serde(skip)] + pub idempotency_key: Option, + /// Opaque plan-task handle from a completed plan. + pub plan_task: String, +} + +/// POST `.../tasks` response. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "kebab-case")] +pub struct FetchScanTasksResponse { + /// Task payload, which may itself contain further plan-task handles. + #[serde(flatten)] + pub scan_tasks: ScanTasks, +} + +/// Polling backoff and bounds for [`RestCatalog::wait_for_plan`]. +#[derive(Debug, Clone)] +pub struct WaitForPlanOptions { + /// Backoff floor. Zero uses 100ms. + pub min_delay: Duration, + /// Backoff cap. Zero uses 5s. + pub max_delay: Duration, + /// Bound on the best-effort cancel after giving up. Zero uses 5s. + pub cancel_grace_period: Duration, + /// Poll attempts after the first. Zero uses 10. + pub max_retries: u32, + /// Optional overall deadline. `None` relies on `max_retries`. + pub timeout: Option, + /// `X-Iceberg-Access-Delegation` forwarded on each poll. + pub access_delegation: Option, +} + +impl Default for WaitForPlanOptions { + fn default() -> Self { + Self { + min_delay: Duration::from_millis(100), + max_delay: Duration::from_secs(5), + cancel_grace_period: Duration::from_secs(5), + max_retries: 10, + timeout: None, + access_delegation: None, + } + } +} + +impl<'de> Deserialize<'de> for PlanTableScanResponse { + fn deserialize>(deserializer: D) -> std::result::Result { + let raw = RawPlanningResponse::deserialize(deserializer)?; + match raw.status { + PlanStatus::Completed | PlanStatus::Submitted => { + if raw.plan_id.is_none() { + return Err(de::Error::custom(format!( + "planTableScan response with status {:?} missing plan-id", + raw.status + ))); + } + } + PlanStatus::Cancelled => { + return Err(de::Error::custom( + "planTableScan response has invalid status cancelled", + )); + } + PlanStatus::Failed => {} + } + Ok(PlanTableScanResponse { + status: raw.status, + plan_id: raw.plan_id, + error: decode_planning_error(raw.error), + scan_tasks: raw.scan_tasks, + storage_credentials: raw.storage_credentials, + }) + } +} + +impl<'de> Deserialize<'de> for FetchPlanningResultResponse { + fn deserialize>(deserializer: D) -> std::result::Result { + let raw = RawPlanningResponse::deserialize(deserializer)?; + Ok(FetchPlanningResultResponse { + status: raw.status, + error: decode_planning_error(raw.error), + scan_tasks: raw.scan_tasks, + storage_credentials: raw.storage_credentials, + }) + } +} + +#[derive(Deserialize)] +#[serde(rename_all = "kebab-case")] +struct RawPlanningResponse { + status: PlanStatus, + #[serde(default)] + plan_id: Option, + #[serde(default)] + error: Option, + #[serde(flatten)] + scan_tasks: ScanTasks, + #[serde(default)] + storage_credentials: Option>, +} + +fn decode_planning_error(raw: Option) -> Option { + let value = raw?; + serde_json::from_value(value).ok() +} + +/// True when a fetch-result 404 was a forgotten plan-id. +pub fn is_plan_expired(err: &Error) -> bool { + err.message() == MSG_PLAN_EXPIRED +} + +/// True when the server reported a failed plan. +pub fn is_plan_failed(err: &Error) -> bool { + err.message() == MSG_PLAN_FAILED || err.message().starts_with(&format!("{MSG_PLAN_FAILED}: ")) +} + +/// True when a submitted plan was cancelled. +pub fn is_plan_cancelled(err: &Error) -> bool { + err.message() == MSG_PLAN_CANCELLED +} + +/// True when a plan-task handle is gone. +pub fn is_no_such_plan_task(err: &Error) -> bool { + err.message() == MSG_NO_SUCH_PLAN_TASK +} + +/// True when [`RestCatalog::wait_for_plan`] hit its retry or timeout bound. +pub fn is_plan_poll_exhausted(err: &Error) -> bool { + err.message() == MSG_PLAN_POLL_EXHAUSTED +} + +impl RestCatalog { + /// Whether the server advertised the synchronous plan endpoint. + pub async fn supports_plan_table_scan(&self) -> Result { + self.supports_endpoint(&V1_PLAN_TABLE_SCAN).await + } + + /// Whether the server advertised all four scan-planning endpoints. + pub async fn supports_full_remote_scan_planning(&self) -> Result { + Ok(self.supports_plan_table_scan().await? + && self.supports_endpoint(&V1_FETCH_PLAN_RESULT).await? + && self.supports_endpoint(&V1_CANCEL_PLANNING).await? + && self.supports_endpoint(&V1_FETCH_SCAN_TASKS).await?) + } + + /// Whether this catalog can complete a remote plan end-to-end. + /// + /// Stays `false` until task decoding is wired into `TableScan`. Routing + /// auto-mode scans on endpoint capability alone would fail with + /// [`ErrorKind::FeatureUnsupported`] instead of falling back to local + /// planning. + pub async fn supports_remote_scan_planning(&self) -> Result { + // Touch config so a down server still surfaces as an error, but never + // claim end-to-end remote planning until TableScan can decode tasks. + let _ = self.supports_plan_table_scan().await?; + Ok(false) + } + + /// Submits a server-side scan plan. + /// + /// Completed and submitted plans return `Ok`. A failed plan returns + /// [`is_plan_failed`]. + pub async fn plan_table_scan( + &self, + table: &TableIdent, + request: PlanTableScanRequest, + ) -> Result { + self.require_endpoint(&V1_PLAN_TABLE_SCAN, "planTableScan") + .await?; + let client = self.client().await?; + let url = format!("{}/plan", client.config.table_endpoint(table)); + let headers = scan_planning_headers( + request.idempotency_key.as_deref(), + request.access_delegation.as_deref(), + true, + )?; + let mut builder = client.http_client.request(Method::POST, url).json(&request); + for (name, value) in headers { + builder = builder.header(name, value); + } + let http_response = client.query_catalog(HttpRequest::build(builder)?).await?; + let resp: PlanTableScanResponse = match http_response.status() { + StatusCode::OK => deserialize_catalog_response(http_response).await?, + StatusCode::NOT_FOUND => { + return Err(map_plan_not_found(http_response, PlanNotFoundKind::Submit).await); + } + _ => { + return Err(deserialize_unexpected_catalog_error( + http_response, + client.http_client.disable_header_redaction(), + ) + .await); + } + }; + match resp.status { + PlanStatus::Completed | PlanStatus::Submitted => Ok(resp), + PlanStatus::Failed => Err(failed_plan_error(resp.error.as_ref())), + PlanStatus::Cancelled => Err(Error::new( + ErrorKind::DataInvalid, + "planTableScan response has invalid status cancelled", + )), + } + } + + /// Polls a previously submitted plan. + pub async fn fetch_planning_result( + &self, + table: &TableIdent, + plan_id: &str, + opts: FetchPlanningResultOptions, + ) -> Result { + match self.fetch_planning_result_raw(table, plan_id, opts).await { + Ok(resp) => Ok(resp), + Err(PollError::Retry { error, .. }) | Err(PollError::Terminal(error)) => Err(error), + } + } + + /// Cancels a server-side plan. Best-effort: a 404 is returned as a generic + /// unexpected error rather than [`is_plan_expired`]. + pub async fn cancel_planning(&self, table: &TableIdent, plan_id: &str) -> Result<()> { + require_plan_id(plan_id)?; + self.require_endpoint(&V1_CANCEL_PLANNING, "cancelPlanning") + .await?; + let client = self.client().await?; + let url = format!( + "{}/plan/{}", + client.config.table_endpoint(table), + escape_opaque_path_segment(plan_id) + ); + let request = HttpRequest::build(client.http_client.request(Method::DELETE, url))?; + let http_response = client.query_catalog(request).await?; + match http_response.status() { + StatusCode::NO_CONTENT | StatusCode::OK => Ok(()), + _ => Err(deserialize_unexpected_catalog_error( + http_response, + client.http_client.disable_header_redaction(), + ) + .await), + } + } + + /// Fetches scan tasks for a plan-task handle. + pub async fn fetch_scan_tasks( + &self, + table: &TableIdent, + request: FetchScanTasksRequest, + ) -> Result { + self.require_endpoint(&V1_FETCH_SCAN_TASKS, "fetchScanTasks") + .await?; + let client = self.client().await?; + let url = format!("{}/tasks", client.config.table_endpoint(table)); + let headers = scan_planning_headers(request.idempotency_key.as_deref(), None, true)?; + let mut builder = client.http_client.request(Method::POST, url).json(&request); + for (name, value) in headers { + builder = builder.header(name, value); + } + let http_response = client.query_catalog(HttpRequest::build(builder)?).await?; + match http_response.status() { + StatusCode::OK => { + let bytes = http_response.bytes().await?; + if bytes.is_empty() { + return Err(Error::new( + ErrorKind::Unexpected, + "fetchScanTasks response was empty", + )); + } + serde_json::from_slice(&bytes).map_err(|e| { + Error::new( + ErrorKind::Unexpected, + "Failed to parse response from rest catalog server", + ) + .with_context("json", String::from_utf8_lossy(&bytes)) + .with_source(e) + }) + } + StatusCode::NOT_FOUND => { + Err(map_plan_not_found(http_response, PlanNotFoundKind::Tasks).await) + } + _ => Err(deserialize_unexpected_catalog_error( + http_response, + client.http_client.disable_header_redaction(), + ) + .await), + } + } + + /// Polls a submitted plan until it completes, fails, or the retry budget + /// is spent. Does not expand plan-task handles. + pub async fn wait_for_plan( + &self, + table: &TableIdent, + plan_id: &str, + opts: WaitForPlanOptions, + ) -> Result { + require_plan_id(plan_id)?; + self.require_endpoint(&V1_FETCH_PLAN_RESULT, "fetchPlanningResult") + .await?; + + let work = self.wait_for_plan_loop(table, plan_id, opts.clone()); + if let Some(timeout) = opts.timeout { + match tokio::time::timeout(timeout, work).await { + Ok(result) => result, + Err(_) => { + self.abandon_plan(table, plan_id, resolve_wait_options(&opts).2) + .await; + Err(Error::new(ErrorKind::Unexpected, MSG_PLAN_POLL_EXHAUSTED)) + } + } + } else { + work.await + } + } + + async fn wait_for_plan_loop( + &self, + table: &TableIdent, + plan_id: &str, + opts: WaitForPlanOptions, + ) -> Result { + let (min_delay, max_delay, grace, max_retries) = resolve_wait_options(&opts); + let fetch_opts = FetchPlanningResultOptions { + access_delegation: opts.access_delegation.clone(), + }; + let mut sleep = min_delay; + let mut retries = 0u32; + loop { + let retry_after = match self + .fetch_planning_result_raw(table, plan_id, fetch_opts.clone()) + .await + { + Ok(resp) => match resp.status { + PlanStatus::Completed => { + return Ok(CompletedPlanningResult { + status: PlanStatus::Completed, + scan_tasks: resp.scan_tasks, + storage_credentials: resp.storage_credentials, + }); + } + PlanStatus::Submitted => None, + PlanStatus::Failed => return Err(failed_plan_error(resp.error.as_ref())), + PlanStatus::Cancelled => { + return Err(Error::new(ErrorKind::Unexpected, MSG_PLAN_CANCELLED)); + } + }, + Err(PollError::Retry { retry_after, .. }) => retry_after, + Err(PollError::Terminal(err)) => return Err(err), + }; + + if retries >= max_retries { + self.abandon_plan(table, plan_id, grace).await; + return Err(Error::new(ErrorKind::Unexpected, MSG_PLAN_POLL_EXHAUSTED)); + } + retries += 1; + sleep = next_scan_plan_backoff(sleep, min_delay, max_delay); + sleep = apply_retry_after(sleep, retry_after, min_delay, max_delay); + tokio::time::sleep(sleep).await; + } + } + + async fn fetch_planning_result_raw( + &self, + table: &TableIdent, + plan_id: &str, + opts: FetchPlanningResultOptions, + ) -> std::result::Result { + if let Err(err) = require_plan_id(plan_id) { + return Err(PollError::Terminal(err)); + } + if let Err(err) = self + .require_endpoint(&V1_FETCH_PLAN_RESULT, "fetchPlanningResult") + .await + { + return Err(PollError::Terminal(err)); + } + let client = match self.client().await { + Ok(client) => client, + Err(err) => return Err(PollError::Terminal(err)), + }; + let url = format!( + "{}/plan/{}", + client.config.table_endpoint(table), + escape_opaque_path_segment(plan_id) + ); + let mut builder = client.http_client.request(Method::GET, url); + if let Some(delegation) = opts.access_delegation.as_deref() { + builder = builder.header(HEADER_ACCESS_DELEGATION, delegation); + } + let http_response = match client + .query_catalog(HttpRequest::build(builder).map_err(PollError::Terminal)?) + .await + { + Ok(resp) => resp, + Err(err) => return Err(PollError::Terminal(err)), + }; + let status = http_response.status(); + let retry_after = + parse_retry_after(http_response.headers().get(reqwest::header::RETRY_AFTER)); + match status { + StatusCode::OK => { + let resp: FetchPlanningResultResponse = deserialize_catalog_response(http_response) + .await + .map_err(PollError::Terminal)?; + match resp.status { + PlanStatus::Failed => { + Err(PollError::Terminal(failed_plan_error(resp.error.as_ref()))) + } + PlanStatus::Cancelled => Err(PollError::Terminal(Error::new( + ErrorKind::Unexpected, + MSG_PLAN_CANCELLED, + ))), + PlanStatus::Completed | PlanStatus::Submitted => Ok(resp), + } + } + StatusCode::NOT_FOUND => Err(PollError::Terminal( + map_plan_not_found(http_response, PlanNotFoundKind::Fetch).await, + )), + StatusCode::REQUEST_TIMEOUT + | StatusCode::TOO_MANY_REQUESTS + | StatusCode::INTERNAL_SERVER_ERROR + | StatusCode::BAD_GATEWAY + | StatusCode::SERVICE_UNAVAILABLE + | StatusCode::GATEWAY_TIMEOUT => { + let error = deserialize_unexpected_catalog_error( + http_response, + client.http_client.disable_header_redaction(), + ) + .await + .with_retryable(true); + Err(PollError::Retry { retry_after, error }) + } + _ => Err(PollError::Terminal( + deserialize_unexpected_catalog_error( + http_response, + client.http_client.disable_header_redaction(), + ) + .await, + )), + } + } + + async fn require_endpoint(&self, endpoint: &Endpoint, name: &str) -> Result<()> { + if self.supports_endpoint(endpoint).await? { + Ok(()) + } else { + Err(Error::new( + ErrorKind::FeatureUnsupported, + format!("{name} is not advertised by the REST catalog"), + )) + } + } + + async fn abandon_plan(&self, table: &TableIdent, plan_id: &str, grace: Duration) { + let _ = tokio::time::timeout(grace, self.cancel_planning(table, plan_id)).await; + } +} + +enum PollError { + Retry { + retry_after: Option, + error: Error, + }, + Terminal(Error), +} + +enum PlanNotFoundKind { + Submit, + Fetch, + Tasks, +} + +async fn map_plan_not_found(response: Response, kind: PlanNotFoundKind) -> Error { + let bytes = match response.bytes().await { + Ok(bytes) => bytes, + Err(err) => return err.into(), + }; + let err_type = serde_json::from_slice::(&bytes) + .ok() + .map(|parsed| parsed.error_type().to_string()); + match (kind, err_type.as_deref()) { + (_, Some(ERR_TYPE_NO_SUCH_TABLE)) => Error::new( + ErrorKind::TableNotFound, + "Tried to plan a table that does not exist", + ), + (_, Some(ERR_TYPE_NO_SUCH_NAMESPACE)) => Error::new( + ErrorKind::NamespaceNotFound, + "Tried to plan a table in a namespace that does not exist", + ), + (PlanNotFoundKind::Fetch, Some(ERR_TYPE_NO_SUCH_PLAN_ID)) => { + Error::new(ErrorKind::Unexpected, MSG_PLAN_EXPIRED) + } + (PlanNotFoundKind::Tasks, Some(ERR_TYPE_NO_SUCH_PLAN_TASK)) => { + Error::new(ErrorKind::Unexpected, MSG_NO_SUCH_PLAN_TASK) + } + _ => Error::new( + ErrorKind::Unexpected, + "Received response with unexpected status code", + ) + .with_context("status", "404") + .with_context("json", String::from_utf8_lossy(&bytes)), + } +} + +fn failed_plan_error(detail: Option<&ErrorModel>) -> Error { + match detail { + Some(detail) if !detail.message.is_empty() => Error::new( + ErrorKind::Unexpected, + format!("{}: {}", MSG_PLAN_FAILED, detail.message), + ) + .with_context("type", detail.r#type.clone()), + _ => Error::new(ErrorKind::Unexpected, MSG_PLAN_FAILED), + } +} + +fn scan_planning_headers( + idempotency_key: Option<&str>, + access_delegation: Option<&str>, + include_idempotency: bool, +) -> Result> { + let mut headers = Vec::with_capacity(2); + if include_idempotency { + headers.push(( + HEADER_IDEMPOTENCY_KEY, + idempotency_header_value(idempotency_key)?, + )); + } + if let Some(delegation) = access_delegation { + headers.push((HEADER_ACCESS_DELEGATION, delegation.to_string())); + } + Ok(headers) +} + +fn idempotency_header_value(idempotency_key: Option<&str>) -> Result { + match idempotency_key { + None => Ok(Uuid::now_v7().to_string()), + Some(key) => { + let parsed = Uuid::parse_str(key).map_err(|_| { + Error::new( + ErrorKind::DataInvalid, + format!("invalid idempotency key {key:?}"), + ) + })?; + if !parsed.to_string().eq_ignore_ascii_case(key) { + return Err(Error::new( + ErrorKind::DataInvalid, + format!("idempotency key {key:?} must be a canonical hyphenated UUID"), + )); + } + if parsed.get_version() != Some(Version::SortRand) + || parsed.get_variant() != Variant::RFC4122 + { + return Err(Error::new( + ErrorKind::DataInvalid, + format!("idempotency key {key:?} must be an RFC 4122 UUIDv7"), + )); + } + Ok(key.to_string()) + } + } +} + +fn escape_opaque_path_segment(s: &str) -> String { + let mut out = String::new(); + for b in s.bytes() { + match b { + b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => { + out.push(b as char); + } + _ => out.push_str(&format!("%{b:02X}")), + } + } + match out.as_str() { + "." => "%2E".to_string(), + ".." => "%2E%2E".to_string(), + _ => out, + } +} + +fn resolve_wait_options(opts: &WaitForPlanOptions) -> (Duration, Duration, Duration, u32) { + let mut min_delay = opts.min_delay; + let mut max_delay = opts.max_delay; + let mut grace = opts.cancel_grace_period; + let mut max_retries = opts.max_retries; + if min_delay.is_zero() { + min_delay = Duration::from_millis(100); + } + if max_delay.is_zero() { + max_delay = Duration::from_secs(5); + } + if max_delay < min_delay { + max_delay = min_delay; + } + if grace.is_zero() { + grace = Duration::from_secs(5); + } + if max_retries == 0 { + max_retries = 10; + } + (min_delay, max_delay, grace, max_retries) +} + +fn next_scan_plan_backoff(prev: Duration, min_delay: Duration, max_delay: Duration) -> Duration { + if min_delay >= max_delay { + return min_delay; + } + let mut ceiling = max_delay; + if prev <= max_delay / 3 { + ceiling = prev * 3; + if ceiling < min_delay { + ceiling = min_delay; + } + if ceiling > max_delay { + ceiling = max_delay; + } + } + let lo = min_delay.as_nanos().min(u64::MAX as u128) as u64; + let hi = ceiling.as_nanos().min(u64::MAX as u128) as u64; + Duration::from_nanos(rand::rng().random_range(lo..=hi)) +} + +fn apply_retry_after( + backoff: Duration, + retry_after: Option, + min_delay: Duration, + max_delay: Duration, +) -> Duration { + let Some(retry_after) = retry_after.filter(|d| !d.is_zero()) else { + return backoff; + }; + retry_after.max(min_delay).min(max_delay) +} + +fn parse_retry_after(value: Option<&reqwest::header::HeaderValue>) -> Option { + let value = value?.to_str().ok()?.trim(); + let seconds: i64 = value.parse().ok()?; + (seconds > 0).then(|| Duration::from_secs(seconds as u64)) +} + +fn require_plan_id(plan_id: &str) -> Result<()> { + if plan_id.is_empty() { + Err(Error::new(ErrorKind::DataInvalid, "empty plan-id")) + } else { + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use iceberg::{NamespaceIdent, Runtime}; + use mockito::Server; + use serde_json::json; + use uuid::Uuid; + + use super::*; + use crate::catalog::RestCatalogConfig; + + fn table() -> TableIdent { + TableIdent::new(NamespaceIdent::new("ns".into()), "tbl".into()) + } + + fn catalog(uri: &str) -> RestCatalog { + RestCatalog::new( + RestCatalogConfig::builder().uri(uri.to_string()).build(), + None, + None, + Runtime::current(), + None, + ) + } + + async fn config_with_endpoints( + server: &mut mockito::ServerGuard, + endpoints: &[&str], + ) -> mockito::Mock { + let body = json!({ + "overrides": {}, + "defaults": {}, + "endpoints": endpoints, + }); + server + .mock("GET", "/v1/config") + .with_status(200) + .with_body(body.to_string()) + .create_async() + .await + } + + const ALL_PLAN: &[&str] = &[ + "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan", + "GET /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan/{plan-id}", + "DELETE /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan/{plan-id}", + "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/tasks", + ]; + + #[test] + fn plan_table_scan_requires_plan_id_on_completed_and_submitted() { + let completed = json!({"status": "completed"}); + assert!(serde_json::from_value::(completed).is_err()); + let submitted = json!({"status": "submitted", "plan-id": "p1"}); + let parsed: PlanTableScanResponse = serde_json::from_value(submitted).unwrap(); + assert_eq!(parsed.status, PlanStatus::Submitted); + assert_eq!(parsed.plan_id.as_deref(), Some("p1")); + } + + #[test] + fn plan_table_scan_rejects_cancelled() { + let body = json!({"status": "cancelled", "plan-id": "p1"}); + assert!(serde_json::from_value::(body).is_err()); + } + + #[test] + fn fetch_result_accepts_cancelled_and_failed_without_error() { + let cancelled: FetchPlanningResultResponse = + serde_json::from_value(json!({"status": "cancelled"})).unwrap(); + assert_eq!(cancelled.status, PlanStatus::Cancelled); + let failed: FetchPlanningResultResponse = + serde_json::from_value(json!({"status": "failed", "error": "not-an-object"})).unwrap(); + assert_eq!(failed.status, PlanStatus::Failed); + assert!(failed.error.is_none()); + } + + #[test] + fn completed_plan_round_trips_opaque_tasks_and_credentials() { + let body = json!({ + "status": "completed", + "plan-id": "p1", + "plan-tasks": ["t1"], + "file-scan-tasks": [{"data-file": {"file-path": "s3://b/f.parquet"}}], + "storage-credentials": [{"prefix": "s3://b/", "config": {"s3.access-key-id": "k"}}] + }); + let parsed: PlanTableScanResponse = serde_json::from_value(body.clone()).unwrap(); + assert_eq!(parsed.scan_tasks.plan_tasks, ["t1"]); + assert_eq!( + parsed.scan_tasks.file_scan_tasks[0].data_file["file-path"], + "s3://b/f.parquet" + ); + assert_eq!( + parsed.storage_credentials.as_ref().unwrap()[0].prefix, + "s3://b/" + ); + let encoded = serde_json::to_value(&parsed).unwrap(); + assert_eq!(encoded["plan-id"], "p1"); + assert_eq!(encoded["plan-tasks"][0], "t1"); + } + + #[test] + fn plan_request_omits_header_fields_from_json() { + let req = PlanTableScanRequest { + idempotency_key: Some("not-on-wire".into()), + access_delegation: Some("vended-credentials".into()), + snapshot_id: Some(7), + ..Default::default() + }; + let value = serde_json::to_value(&req).unwrap(); + assert_eq!(value["snapshot-id"], 7); + assert!(value.get("idempotency-key").is_none()); + assert!(value.get("access-delegation").is_none()); + } + + #[test] + fn idempotency_key_must_be_canonical_uuidv7() { + let v7 = Uuid::now_v7().to_string(); + assert!(idempotency_header_value(Some(&v7)).is_ok()); + assert!(idempotency_header_value(Some(&v7.to_uppercase())).is_ok()); + let v4 = Uuid::new_v4().to_string(); + assert!(idempotency_header_value(Some(&v4)).is_err()); + assert!(idempotency_header_value(Some("not-a-uuid")).is_err()); + let unhyphenated: String = v7.chars().filter(|c| *c != '-').collect(); + assert!(idempotency_header_value(Some(&unhyphenated)).is_err()); + } + + #[test] + fn opaque_plan_id_is_a_single_path_segment() { + assert_eq!(escape_opaque_path_segment("a/b"), "a%2Fb"); + assert_eq!(escape_opaque_path_segment("."), "%2E"); + assert_eq!(escape_opaque_path_segment(".."), "%2E%2E"); + assert_eq!(escape_opaque_path_segment("plain"), "plain"); + } + + #[tokio::test] + async fn default_config_does_not_advertise_plan() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, &[]).await; + // Empty endpoints list falls back to DEFAULT_ENDPOINTS. + let catalog = catalog(&server.url()); + assert!(!catalog.supports_plan_table_scan().await.unwrap()); + assert!(!catalog.supports_full_remote_scan_planning().await.unwrap()); + assert!(!catalog.supports_remote_scan_planning().await.unwrap()); + config.assert_async().await; + } + + #[tokio::test] + async fn remote_scan_planning_stays_false_when_all_endpoints_are_advertised() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let catalog = catalog(&server.url()); + assert!(catalog.supports_plan_table_scan().await.unwrap()); + assert!(catalog.supports_full_remote_scan_planning().await.unwrap()); + assert!(!catalog.supports_remote_scan_planning().await.unwrap()); + config.assert_async().await; + } + + #[tokio::test] + async fn plan_table_scan_is_feature_unsupported_when_not_advertised() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, &["GET /v1/{prefix}/namespaces"]).await; + let catalog = catalog(&server.url()); + let err = catalog + .plan_table_scan(&table(), PlanTableScanRequest::default()) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::FeatureUnsupported); + config.assert_async().await; + } + + #[tokio::test] + async fn plan_table_scan_completed_sends_uuidv7_idempotency_key() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let plan = server + .mock("POST", "/v1/namespaces/ns/tables/tbl/plan") + .match_header( + "idempotency-key", + mockito::Matcher::Regex( + r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-7[0-9a-fA-F]{3}-[89abAB][0-9a-fA-F]{3}-[0-9a-fA-F]{12}$" + .into(), + ), + ) + .with_status(200) + .with_body(r#"{"status":"completed","plan-id":"p1","file-scan-tasks":[]}"#) + .create_async() + .await; + let catalog = catalog(&server.url()); + let resp = catalog + .plan_table_scan(&table(), PlanTableScanRequest::default()) + .await + .unwrap(); + assert_eq!(resp.status, PlanStatus::Completed); + assert_eq!(resp.plan_id.as_deref(), Some("p1")); + config.assert_async().await; + plan.assert_async().await; + } + + #[tokio::test] + async fn plan_table_scan_failed_is_an_error() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let plan = server + .mock("POST", "/v1/namespaces/ns/tables/tbl/plan") + .with_status(200) + .with_body(r#"{"status":"failed","error":{"message":"boom","type":"IcebergException","code":500}}"#) + .create_async() + .await; + let catalog = catalog(&server.url()); + let err = catalog + .plan_table_scan(&table(), PlanTableScanRequest::default()) + .await + .unwrap_err(); + assert!(is_plan_failed(&err)); + assert!(err.message().contains("boom")); + config.assert_async().await; + plan.assert_async().await; + } + + #[tokio::test] + async fn plan_404_splits_table_from_unrecognized() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let missing = server + .mock("POST", "/v1/namespaces/ns/tables/tbl/plan") + .with_status(404) + .with_body(r#"{"error":{"message":"gone","type":"NoSuchTableException","code":404}}"#) + .create_async() + .await; + let catalog = catalog(&server.url()); + let err = catalog + .plan_table_scan(&table(), PlanTableScanRequest::default()) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::TableNotFound); + config.assert_async().await; + missing.assert_async().await; + } + + #[tokio::test] + async fn fetch_result_expired_plan_id() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let fetch = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(404) + .with_body( + r#"{"error":{"message":"expired","type":"NoSuchPlanIdException","code":404}}"#, + ) + .create_async() + .await; + let catalog = catalog(&server.url()); + let err = catalog + .fetch_planning_result(&table(), "p1", FetchPlanningResultOptions::default()) + .await + .unwrap_err(); + assert!(is_plan_expired(&err)); + config.assert_async().await; + fetch.assert_async().await; + } + + #[tokio::test] + async fn fetch_scan_tasks_and_empty_body() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let tasks = server + .mock("POST", "/v1/namespaces/ns/tables/tbl/tasks") + .match_body(r#"{"plan-task":"h1"}"#) + .with_status(200) + .with_body( + r#"{"plan-tasks":["h2"],"file-scan-tasks":[{"data-file":{"file-path":"f"}}]}"#, + ) + .create_async() + .await; + let catalog = catalog(&server.url()); + let resp = catalog + .fetch_scan_tasks(&table(), FetchScanTasksRequest { + idempotency_key: None, + plan_task: "h1".into(), + }) + .await + .unwrap(); + assert_eq!(resp.scan_tasks.plan_tasks, ["h2"]); + config.assert_async().await; + tasks.assert_async().await; + } + + #[tokio::test] + async fn fetch_scan_tasks_rejects_empty_200() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let tasks = server + .mock("POST", "/v1/namespaces/ns/tables/tbl/tasks") + .with_status(200) + .with_body("") + .create_async() + .await; + let catalog = catalog(&server.url()); + let err = catalog + .fetch_scan_tasks(&table(), FetchScanTasksRequest { + idempotency_key: None, + plan_task: "h1".into(), + }) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::Unexpected); + config.assert_async().await; + tasks.assert_async().await; + } + + #[tokio::test] + async fn fetch_scan_tasks_no_such_plan_task() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let tasks = server + .mock("POST", "/v1/namespaces/ns/tables/tbl/tasks") + .with_status(404) + .with_body( + r#"{"error":{"message":"gone","type":"NoSuchPlanTaskException","code":404}}"#, + ) + .create_async() + .await; + let catalog = catalog(&server.url()); + let err = catalog + .fetch_scan_tasks(&table(), FetchScanTasksRequest { + idempotency_key: None, + plan_task: "h1".into(), + }) + .await + .unwrap_err(); + assert!(is_no_such_plan_task(&err)); + config.assert_async().await; + tasks.assert_async().await; + } + + #[tokio::test] + async fn wait_for_plan_polls_submitted_then_completes() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let first = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"submitted"}"#) + .expect(1) + .create_async() + .await; + let second = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"completed","plan-tasks":["h1"]}"#) + .expect(1) + .create_async() + .await; + let catalog = catalog(&server.url()); + let result = catalog + .wait_for_plan(&table(), "p1", WaitForPlanOptions { + min_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(result.scan_tasks.plan_tasks, ["h1"]); + config.assert_async().await; + first.assert_async().await; + second.assert_async().await; + } + + #[tokio::test] + async fn wait_for_plan_retries_503_then_completes() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let busy = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(503) + .with_header("Retry-After", "0") + .expect(1) + .create_async() + .await; + let done = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"completed"}"#) + .expect(1) + .create_async() + .await; + let catalog = catalog(&server.url()); + let result = catalog + .wait_for_plan(&table(), "p1", WaitForPlanOptions { + min_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(result.status, PlanStatus::Completed); + config.assert_async().await; + busy.assert_async().await; + done.assert_async().await; + } + + #[tokio::test] + async fn wait_for_plan_cancels_after_max_retries() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let poll = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"submitted"}"#) + .expect(2) + .create_async() + .await; + let cancel = server + .mock("DELETE", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(204) + .expect(1) + .create_async() + .await; + let catalog = catalog(&server.url()); + let err = catalog + .wait_for_plan(&table(), "p1", WaitForPlanOptions { + min_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + max_retries: 1, + cancel_grace_period: Duration::from_secs(1), + ..Default::default() + }) + .await + .unwrap_err(); + assert!(is_plan_poll_exhausted(&err)); + config.assert_async().await; + poll.assert_async().await; + cancel.assert_async().await; + } + + #[tokio::test] + async fn wait_for_plan_rejects_empty_plan_id() { + let catalog = catalog("http://127.0.0.1:1"); + let err = catalog + .wait_for_plan(&table(), "", WaitForPlanOptions::default()) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + } + + #[tokio::test] + async fn plan_id_with_slash_stays_one_path_segment() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let fetch = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/a%2Fb") + .with_status(200) + .with_body(r#"{"status":"submitted"}"#) + .create_async() + .await; + let catalog = catalog(&server.url()); + let resp = catalog + .fetch_planning_result(&table(), "a/b", FetchPlanningResultOptions::default()) + .await + .unwrap(); + assert_eq!(resp.status, PlanStatus::Submitted); + config.assert_async().await; + fetch.assert_async().await; + } + + #[tokio::test] + async fn cancel_planning_accepts_204() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let cancel = server + .mock("DELETE", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(204) + .create_async() + .await; + let catalog = catalog(&server.url()); + catalog.cancel_planning(&table(), "p1").await.unwrap(); + config.assert_async().await; + cancel.assert_async().await; + } +} diff --git a/crates/catalog/rest/src/types.rs b/crates/catalog/rest/src/types.rs index 390521229e..50cdb68a09 100644 --- a/crates/catalog/rest/src/types.rs +++ b/crates/catalog/rest/src/types.rs @@ -42,6 +42,12 @@ pub struct ErrorResponse { error: ErrorModel, } +impl ErrorResponse { + pub(crate) fn error_type(&self) -> &str { + &self.error.r#type + } +} + impl From for Error { fn from(resp: ErrorResponse) -> Error { resp.error.into() From 2abcef8d499c1d38e1aafa4cabb484a2fc11a1f6 Mon Sep 17 00:00:00 2001 From: David Dallakyan Date: Sun, 16 Aug 2026 16:51:55 +0400 Subject: [PATCH 2/6] fix(catalog-rest): cancel wait_for_plan on drop and bound Retry-After A huge Retry-After used to panic Duration::from_secs. Ignore overflow, await cancel when the poller times out, and DELETE the plan if the wait future is dropped. --- crates/catalog/rest/src/catalog.rs | 16 ++ crates/catalog/rest/src/scan_planning.rs | 259 +++++++++++++++++++++-- 2 files changed, 261 insertions(+), 14 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index be57fbe830..9431a6c2ae 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -521,6 +521,22 @@ impl RestCatalog { } } + /// Same catalog identity with an empty HTTP client cache, for fire-and-forget + /// work (best-effort plan cancel on drop) that must not borrow `self`. + pub(crate) fn clone_uninitialized(&self) -> Self { + Self::new( + self.user_config.clone(), + self.auth_manager.clone(), + self.storage_factory.clone(), + self.runtime.clone(), + self.kms_client.clone(), + ) + } + + pub(crate) fn runtime(&self) -> &Runtime { + &self.runtime + } + /// Sends a DELETE request for the given table, optionally requesting purge. async fn delete_table(&self, table: &TableIdent, purge: bool) -> Result<()> { let client = self.client().await?; diff --git a/crates/catalog/rest/src/scan_planning.rs b/crates/catalog/rest/src/scan_planning.rs index 3bd9216293..f0d9c24db9 100644 --- a/crates/catalog/rest/src/scan_planning.rs +++ b/crates/catalog/rest/src/scan_planning.rs @@ -222,7 +222,9 @@ pub struct WaitForPlanOptions { pub max_delay: Duration, /// Bound on the best-effort cancel after giving up. Zero uses 5s. pub cancel_grace_period: Duration, - /// Poll attempts after the first. Zero uses 10. + /// Poll attempts after the first. Zero uses 10 when [`Self::timeout`] is + /// `None`. When a timeout is set, zero means keep polling until the + /// deadline. pub max_retries: u32, /// Optional overall deadline. `None` relies on `max_retries`. pub timeout: Option, @@ -495,19 +497,24 @@ impl RestCatalog { self.require_endpoint(&V1_FETCH_PLAN_RESULT, "fetchPlanningResult") .await?; + let (_, _, grace, _, _) = resolve_wait_options(&opts); + let mut guard = + PlanAbandonGuard::new(self.clone_uninitialized(), table.clone(), plan_id, grace); + let work = self.wait_for_plan_loop(table, plan_id, opts.clone()); - if let Some(timeout) = opts.timeout { + let result = if let Some(timeout) = opts.timeout { match tokio::time::timeout(timeout, work).await { Ok(result) => result, Err(_) => { - self.abandon_plan(table, plan_id, resolve_wait_options(&opts).2) - .await; + self.abandon_plan(table, plan_id, grace).await; Err(Error::new(ErrorKind::Unexpected, MSG_PLAN_POLL_EXHAUSTED)) } } } else { work.await - } + }; + guard.disarm(); + result } async fn wait_for_plan_loop( @@ -516,7 +523,8 @@ impl RestCatalog { plan_id: &str, opts: WaitForPlanOptions, ) -> Result { - let (min_delay, max_delay, grace, max_retries) = resolve_wait_options(&opts); + let (min_delay, max_delay, grace, max_retries, clamp_retry_after) = + resolve_wait_options(&opts); let fetch_opts = FetchPlanningResultOptions { access_delegation: opts.access_delegation.clone(), }; @@ -545,13 +553,13 @@ impl RestCatalog { Err(PollError::Terminal(err)) => return Err(err), }; - if retries >= max_retries { + if max_retries > 0 && retries >= max_retries { self.abandon_plan(table, plan_id, grace).await; return Err(Error::new(ErrorKind::Unexpected, MSG_PLAN_POLL_EXHAUSTED)); } retries += 1; sleep = next_scan_plan_backoff(sleep, min_delay, max_delay); - sleep = apply_retry_after(sleep, retry_after, min_delay, max_delay); + sleep = apply_retry_after(sleep, retry_after, min_delay, max_delay, clamp_retry_after); tokio::time::sleep(sleep).await; } } @@ -653,6 +661,45 @@ impl RestCatalog { } } +/// Best-effort DELETE of a submitted plan if [`RestCatalog::wait_for_plan`] is +/// dropped (caller timeout, cancellation) before it returns. +struct PlanAbandonGuard { + catalog: Option, + table: TableIdent, + plan_id: String, + grace: Duration, +} + +impl PlanAbandonGuard { + fn new(catalog: RestCatalog, table: TableIdent, plan_id: &str, grace: Duration) -> Self { + Self { + catalog: Some(catalog), + table, + plan_id: plan_id.to_string(), + grace, + } + } + + fn disarm(&mut self) { + self.catalog = None; + } +} + +impl Drop for PlanAbandonGuard { + fn drop(&mut self) { + let Some(catalog) = self.catalog.take() else { + return; + }; + let table = self.table.clone(); + let plan_id = self.plan_id.clone(); + let grace = self.grace; + let runtime = catalog.runtime().clone(); + drop(runtime.io().spawn(async move { + let _ = tokio::time::timeout(grace, catalog.cancel_planning(&table, &plan_id)).await; + })); + } +} + enum PollError { Retry { retry_after: Option, @@ -774,7 +821,7 @@ fn escape_opaque_path_segment(s: &str) -> String { } } -fn resolve_wait_options(opts: &WaitForPlanOptions) -> (Duration, Duration, Duration, u32) { +fn resolve_wait_options(opts: &WaitForPlanOptions) -> (Duration, Duration, Duration, u32, bool) { let mut min_delay = opts.min_delay; let mut max_delay = opts.max_delay; let mut grace = opts.cancel_grace_period; @@ -791,10 +838,11 @@ fn resolve_wait_options(opts: &WaitForPlanOptions) -> (Duration, Duration, Durat if grace.is_zero() { grace = Duration::from_secs(5); } - if max_retries == 0 { + if max_retries == 0 && opts.timeout.is_none() { max_retries = 10; } - (min_delay, max_delay, grace, max_retries) + let clamp_retry_after = opts.timeout.is_none(); + (min_delay, max_delay, grace, max_retries, clamp_retry_after) } fn next_scan_plan_backoff(prev: Duration, min_delay: Duration, max_delay: Duration) -> Duration { @@ -821,17 +869,31 @@ fn apply_retry_after( retry_after: Option, min_delay: Duration, max_delay: Duration, + clamp_to_max: bool, ) -> Duration { let Some(retry_after) = retry_after.filter(|d| !d.is_zero()) else { return backoff; }; - retry_after.max(min_delay).min(max_delay) + let delay = retry_after.max(min_delay); + if clamp_to_max { + delay.min(max_delay) + } else { + delay + } } fn parse_retry_after(value: Option<&reqwest::header::HeaderValue>) -> Option { let value = value?.to_str().ok()?.trim(); - let seconds: i64 = value.parse().ok()?; - (seconds > 0).then(|| Duration::from_secs(seconds as u64)) + let seconds: u64 = value.parse().ok()?; + if seconds == 0 { + return None; + } + // Duration::from_secs panics when secs * 1e9 overflows u64. + const MAX_SECS: u64 = u64::MAX / 1_000_000_000; + if seconds > MAX_SECS { + return None; + } + Some(Duration::from_secs(seconds)) } fn require_plan_id(plan_id: &str) -> Result<()> { @@ -975,6 +1037,28 @@ mod tests { assert_eq!(escape_opaque_path_segment("plain"), "plain"); } + #[test] + fn parse_retry_after_ignores_zero_and_overflow() { + let zero = reqwest::header::HeaderValue::from_static("0"); + assert!(parse_retry_after(Some(&zero)).is_none()); + let huge = reqwest::header::HeaderValue::from_static("99999999999"); + assert!(parse_retry_after(Some(&huge)).is_none()); + let ok = reqwest::header::HeaderValue::from_static("2"); + assert_eq!(parse_retry_after(Some(&ok)), Some(Duration::from_secs(2))); + } + + #[test] + fn apply_retry_after_skips_max_clamp_when_a_deadline_is_set() { + let retry = Some(Duration::from_secs(30)); + let min = Duration::from_millis(100); + let max = Duration::from_secs(5); + assert_eq!(apply_retry_after(min, retry, min, max, true), max); + assert_eq!( + apply_retry_after(min, retry, min, max, false), + Duration::from_secs(30) + ); + } + #[tokio::test] async fn default_config_does_not_advertise_plan() { let mut server = Server::new_async().await; @@ -1276,6 +1360,153 @@ mod tests { cancel.assert_async().await; } + #[tokio::test] + async fn wait_for_plan_cancels_after_timeout() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let poll = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"submitted"}"#) + .expect_at_least(1) + .create_async() + .await; + let cancel = server + .mock("DELETE", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(204) + .expect(1) + .create_async() + .await; + let catalog = catalog(&server.url()); + let err = catalog + .wait_for_plan(&table(), "p1", WaitForPlanOptions { + min_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + timeout: Some(Duration::from_millis(80)), + max_retries: 0, + cancel_grace_period: Duration::from_secs(1), + ..Default::default() + }) + .await + .unwrap_err(); + assert!(is_plan_poll_exhausted(&err)); + config.assert_async().await; + poll.assert_async().await; + cancel.assert_async().await; + } + + #[tokio::test] + async fn wait_for_plan_cancels_when_future_is_dropped() { + let mut server = Server::new_async().await; + let body = json!({ + "overrides": {}, + "defaults": {}, + "endpoints": ALL_PLAN, + }); + let config = server + .mock("GET", "/v1/config") + .with_status(200) + .with_body(body.to_string()) + .expect_at_least(1) + .create_async() + .await; + let poll = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"submitted"}"#) + .expect_at_least(1) + .create_async() + .await; + let cancel = server + .mock("DELETE", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(204) + .expect(1) + .create_async() + .await; + let catalog = catalog(&server.url()); + let table = table(); + let wait = catalog.wait_for_plan(&table, "p1", WaitForPlanOptions { + min_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + max_retries: 50, + cancel_grace_period: Duration::from_secs(1), + ..Default::default() + }); + let _ = tokio::time::timeout(Duration::from_millis(80), wait).await; + tokio::time::sleep(Duration::from_millis(250)).await; + config.assert_async().await; + poll.assert_async().await; + cancel.assert_async().await; + } + + #[tokio::test] + async fn wait_for_plan_timeout_zero_retries_keeps_polling() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let pending = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"submitted"}"#) + .expect(11) + .create_async() + .await; + let done = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"completed"}"#) + .expect(1) + .create_async() + .await; + let catalog = catalog(&server.url()); + let result = catalog + .wait_for_plan(&table(), "p1", WaitForPlanOptions { + min_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + max_retries: 0, + timeout: Some(Duration::from_secs(5)), + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(result.status, PlanStatus::Completed); + config.assert_async().await; + pending.assert_async().await; + done.assert_async().await; + } + + #[tokio::test] + async fn wait_for_plan_huge_retry_after_does_not_panic() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let busy = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(503) + .with_header("Retry-After", "99999999999") + .expect(1) + .create_async() + .await; + let done = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"completed"}"#) + .expect(1) + .create_async() + .await; + let catalog = catalog(&server.url()); + let result = catalog + .wait_for_plan(&table(), "p1", WaitForPlanOptions { + min_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(result.status, PlanStatus::Completed); + config.assert_async().await; + busy.assert_async().await; + done.assert_async().await; + } + #[tokio::test] async fn wait_for_plan_rejects_empty_plan_id() { let catalog = catalog("http://127.0.0.1:1"); From 110578715832b4a0f7dc202f27a67a2273c0c528 Mon Sep 17 00:00:00 2001 From: David Dallakyan Date: Sun, 16 Aug 2026 17:08:56 +0400 Subject: [PATCH 3/6] fix(catalog-rest): honor HTTP-date Retry-After on plan polls RFC 9110 allows both delta-seconds and IMF-fixdate. Ignoring the date form made wait_for_plan retry immediately against a server that asked us to wait. Tests also cover cancelled/failed/expired poll outcomes and the namespace 404 split. --- crates/catalog/rest/src/scan_planning.rs | 218 +++++++++++++++++++++-- 1 file changed, 208 insertions(+), 10 deletions(-) diff --git a/crates/catalog/rest/src/scan_planning.rs b/crates/catalog/rest/src/scan_planning.rs index f0d9c24db9..bf5ba62471 100644 --- a/crates/catalog/rest/src/scan_planning.rs +++ b/crates/catalog/rest/src/scan_planning.rs @@ -884,16 +884,25 @@ fn apply_retry_after( fn parse_retry_after(value: Option<&reqwest::header::HeaderValue>) -> Option { let value = value?.to_str().ok()?.trim(); - let seconds: u64 = value.parse().ok()?; - if seconds == 0 { - return None; - } - // Duration::from_secs panics when secs * 1e9 overflows u64. - const MAX_SECS: u64 = u64::MAX / 1_000_000_000; - if seconds > MAX_SECS { - return None; - } - Some(Duration::from_secs(seconds)) + if let Ok(seconds) = value.parse::() { + if seconds == 0 { + return None; + } + // Duration::from_secs panics when secs * 1e9 overflows u64. + const MAX_SECS: u64 = u64::MAX / 1_000_000_000; + if seconds > MAX_SECS { + return None; + } + return Some(Duration::from_secs(seconds)); + } + // RFC 9110 HTTP-date, IMF-fixdate form (RFC 2822). Obsolete RFC 850 / + // asctime values are ignored rather than adding another parser. + let retry_at = chrono::DateTime::parse_from_rfc2822(value).ok()?; + retry_at + .signed_duration_since(chrono::Utc::now()) + .to_std() + .ok() + .filter(|d| !d.is_zero()) } fn require_plan_id(plan_id: &str) -> Result<()> { @@ -1045,6 +1054,22 @@ mod tests { assert!(parse_retry_after(Some(&huge)).is_none()); let ok = reqwest::header::HeaderValue::from_static("2"); assert_eq!(parse_retry_after(Some(&ok)), Some(Duration::from_secs(2))); + let padded = reqwest::header::HeaderValue::from_static(" 2 "); + assert_eq!( + parse_retry_after(Some(&padded)), + Some(Duration::from_secs(2)) + ); + let invalid = reqwest::header::HeaderValue::from_static("not-a-date"); + assert!(parse_retry_after(Some(&invalid)).is_none()); + } + + #[test] + fn parse_retry_after_accepts_http_date() { + let past = reqwest::header::HeaderValue::from_static("Wed, 21 Oct 2015 07:28:00 GMT"); + assert!(parse_retry_after(Some(&past)).is_none()); + let future = reqwest::header::HeaderValue::from_static("Sun, 16 Aug 2099 12:00:00 GMT"); + let delay = parse_retry_after(Some(&future)).unwrap(); + assert!(delay > Duration::from_secs(60 * 60 * 24 * 365)); } #[test] @@ -1164,6 +1189,68 @@ mod tests { missing.assert_async().await; } + #[tokio::test] + async fn plan_404_splits_namespace() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let missing = server + .mock("POST", "/v1/namespaces/ns/tables/tbl/plan") + .with_status(404) + .with_body( + r#"{"error":{"message":"gone","type":"NoSuchNamespaceException","code":404}}"#, + ) + .create_async() + .await; + let catalog = catalog(&server.url()); + let err = catalog + .plan_table_scan(&table(), PlanTableScanRequest::default()) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::NamespaceNotFound); + config.assert_async().await; + missing.assert_async().await; + } + + #[tokio::test] + async fn fetch_planning_result_cancelled_is_an_error() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let fetch = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"cancelled"}"#) + .create_async() + .await; + let catalog = catalog(&server.url()); + let err = catalog + .fetch_planning_result(&table(), "p1", FetchPlanningResultOptions::default()) + .await + .unwrap_err(); + assert!(is_plan_cancelled(&err)); + config.assert_async().await; + fetch.assert_async().await; + } + + #[tokio::test] + async fn fetch_planning_result_failed_is_an_error() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let fetch = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"failed","error":{"message":"boom","type":"IcebergException","code":500}}"#) + .create_async() + .await; + let catalog = catalog(&server.url()); + let err = catalog + .fetch_planning_result(&table(), "p1", FetchPlanningResultOptions::default()) + .await + .unwrap_err(); + assert!(is_plan_failed(&err)); + config.assert_async().await; + fetch.assert_async().await; + } + #[tokio::test] async fn fetch_result_expired_plan_id() { let mut server = Server::new_async().await; @@ -1293,6 +1380,117 @@ mod tests { second.assert_async().await; } + #[tokio::test] + async fn wait_for_plan_propagates_cancelled() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let poll = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"cancelled"}"#) + .expect(1) + .create_async() + .await; + let catalog = catalog(&server.url()); + let err = catalog + .wait_for_plan(&table(), "p1", WaitForPlanOptions { + min_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + ..Default::default() + }) + .await + .unwrap_err(); + assert!(is_plan_cancelled(&err)); + config.assert_async().await; + poll.assert_async().await; + } + + #[tokio::test] + async fn wait_for_plan_propagates_expired() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let poll = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(404) + .with_body( + r#"{"error":{"message":"expired","type":"NoSuchPlanIdException","code":404}}"#, + ) + .expect(1) + .create_async() + .await; + let catalog = catalog(&server.url()); + let err = catalog + .wait_for_plan(&table(), "p1", WaitForPlanOptions { + min_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + ..Default::default() + }) + .await + .unwrap_err(); + assert!(is_plan_expired(&err)); + config.assert_async().await; + poll.assert_async().await; + } + + #[tokio::test] + async fn wait_for_plan_propagates_failed() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let poll = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"failed","error":{"message":"boom","type":"IcebergException","code":500}}"#) + .expect(1) + .create_async() + .await; + let catalog = catalog(&server.url()); + let err = catalog + .wait_for_plan(&table(), "p1", WaitForPlanOptions { + min_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + ..Default::default() + }) + .await + .unwrap_err(); + assert!(is_plan_failed(&err)); + config.assert_async().await; + poll.assert_async().await; + } + + #[tokio::test] + async fn wait_for_plan_retries_java_idempotent_get_statuses() { + for status in [408, 429, 500, 502, 503, 504] { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let busy = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(status) + .expect(1) + .create_async() + .await; + let done = server + .mock("GET", "/v1/namespaces/ns/tables/tbl/plan/p1") + .with_status(200) + .with_body(r#"{"status":"completed"}"#) + .expect(1) + .create_async() + .await; + let catalog = catalog(&server.url()); + let result = catalog + .wait_for_plan(&table(), "p1", WaitForPlanOptions { + min_delay: Duration::from_millis(1), + max_delay: Duration::from_millis(1), + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(result.status, PlanStatus::Completed, "status {status}"); + config.assert_async().await; + busy.assert_async().await; + done.assert_async().await; + } + } + #[tokio::test] async fn wait_for_plan_retries_503_then_completes() { let mut server = Server::new_async().await; From f2cad848115ba72c74e93a143858f4499bd49ae1 Mon Sep 17 00:00:00 2001 From: David Dallakyan Date: Sun, 16 Aug 2026 17:41:57 +0400 Subject: [PATCH 4/6] feat(catalog-rest): wire TableScan to REST scan planning Decode REST content-file JSON into FileScanTask and auto-route TableScan::plan_files through a ScanPlanner with local fallback. --- crates/catalog/rest/public-api.txt | 3 + crates/catalog/rest/src/catalog.rs | 12 +- crates/catalog/rest/src/lib.rs | 1 + crates/catalog/rest/src/scan_decode.rs | 384 +++++++++++++++++++ crates/catalog/rest/src/scan_planning.rs | 467 ++++++++++++++++++++++- crates/iceberg/public-api.txt | 41 ++ crates/iceberg/src/scan/mod.rs | 291 +++++++++++++- crates/iceberg/src/scan/planner.rs | 97 +++++ crates/iceberg/src/table.rs | 19 +- 9 files changed, 1293 insertions(+), 22 deletions(-) create mode 100644 crates/catalog/rest/src/scan_decode.rs create mode 100644 crates/iceberg/src/scan/planner.rs diff --git a/crates/catalog/rest/public-api.txt b/crates/catalog/rest/public-api.txt index 8d5334b6e6..4ffaf78d24 100644 --- a/crates/catalog/rest/public-api.txt +++ b/crates/catalog/rest/public-api.txt @@ -398,6 +398,9 @@ pub fn iceberg_catalog_rest::RestCatalog::rename_table<'life0, 'life1, 'life2, ' pub fn iceberg_catalog_rest::RestCatalog::table_exists<'life0, 'life1, 'async_trait>(&'life0 self, table: &'life1 iceberg::catalog::TableIdent) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait pub fn iceberg_catalog_rest::RestCatalog::update_namespace<'life0, 'life1, 'async_trait>(&'life0 self, _namespace: &'life1 iceberg::catalog::NamespaceIdent, _properties: std::collections::hash::map::HashMap) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait pub fn iceberg_catalog_rest::RestCatalog::update_table<'life0, 'async_trait>(&'life0 self, commit: iceberg::catalog::TableCommit) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait +impl iceberg::scan::planner::ScanPlanner for iceberg_catalog_rest::RestCatalog +pub fn iceberg_catalog_rest::RestCatalog::plan_files<'life0, 'async_trait>(&'life0 self, request: iceberg::scan::planner::ScanPlanningRequest) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait +pub fn iceberg_catalog_rest::RestCatalog::supports_remote_scan_planning<'life0, 'async_trait>(&'life0 self) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait pub struct iceberg_catalog_rest::RestCatalogBuilder impl iceberg_catalog_rest::RestCatalogBuilder pub fn iceberg_catalog_rest::RestCatalogBuilder::with_auth_manager(self, auth_manager: alloc::sync::Arc) -> Self diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 9431a6c2ae..8c70b85262 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -1082,7 +1082,8 @@ impl Catalog for RestCatalog { .identifier(table_ident.clone()) .file_io(file_io) .metadata(response.metadata) - .runtime(self.runtime.clone()); + .runtime(self.runtime.clone()) + .scan_planner(Arc::new(self.clone_uninitialized())); if let Some(kms_client) = self.kms_client.clone() { table_builder = table_builder.kms_client(kms_client); } @@ -1143,7 +1144,8 @@ impl Catalog for RestCatalog { .identifier(table_ident.clone()) .file_io(file_io) .metadata(response.metadata) - .runtime(self.runtime.clone()); + .runtime(self.runtime.clone()) + .scan_planner(Arc::new(self.clone_uninitialized())); if let Some(kms_client) = self.kms_client.clone() { table_builder = table_builder.kms_client(kms_client); } @@ -1280,7 +1282,8 @@ impl Catalog for RestCatalog { .file_io(file_io) .metadata(response.metadata) .metadata_location(metadata_location.clone()) - .runtime(self.runtime.clone()); + .runtime(self.runtime.clone()) + .scan_planner(Arc::new(self.clone_uninitialized())); if let Some(kms_client) = self.kms_client.clone() { table_builder = table_builder.kms_client(kms_client); } @@ -1357,7 +1360,8 @@ impl Catalog for RestCatalog { .file_io(file_io) .metadata(response.metadata) .metadata_location(response.metadata_location) - .runtime(self.runtime.clone()); + .runtime(self.runtime.clone()) + .scan_planner(Arc::new(self.clone_uninitialized())); if let Some(kms_client) = self.kms_client.clone() { table_builder = table_builder.kms_client(kms_client); } diff --git a/crates/catalog/rest/src/lib.rs b/crates/catalog/rest/src/lib.rs index 465a4a10e7..93cb1b5626 100644 --- a/crates/catalog/rest/src/lib.rs +++ b/crates/catalog/rest/src/lib.rs @@ -58,6 +58,7 @@ pub use client::HttpClient; mod request; pub use request::{HttpRequest, HttpRequestBody}; mod endpoint; +mod scan_decode; mod scan_planning; mod types; diff --git a/crates/catalog/rest/src/scan_decode.rs b/crates/catalog/rest/src/scan_decode.rs new file mode 100644 index 0000000000..738ccb388a --- /dev/null +++ b/crates/catalog/rest/src/scan_decode.rs @@ -0,0 +1,384 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +//! Decode REST content-file JSON into [`FileScanTask`]s. +//! +//! REST `data-file` / `delete-files` entries are Iceberg content-file JSON +//! (`file-path`, `file-format`, kebab-case). That is not the Avro snake_case +//! payload [`iceberg::spec::deserialize_data_file_from_json`] expects. + +use iceberg::expr::BoundPredicate; +use iceberg::scan::{FileScanTask, FileScanTaskDeleteFile}; +use iceberg::spec::{ + DataContentType, DataFileFormat, Literal, NameMapping, SchemaRef, Struct, TableMetadataRef, +}; +use iceberg::{Error, ErrorKind, Result}; +use serde::Deserialize; +use serde_json::Value; + +use crate::scan_planning::RestFileScanTask; + +/// Per-scan context needed to materialize tasks. +pub(crate) struct ConvertContext { + pub(crate) metadata: TableMetadataRef, + pub(crate) snapshot_schema: SchemaRef, + pub(crate) project_field_ids: Vec, + pub(crate) case_sensitive: bool, + pub(crate) bound_filter: Option, + pub(crate) name_mapping: Option>, + pub(crate) unified_partition_type: Option>, +} + +#[derive(Debug, Clone, Deserialize)] +#[serde(rename_all = "kebab-case")] +struct RestContentFile { + file_path: String, + file_format: String, + file_size_in_bytes: u64, + #[serde(default)] + record_count: Option, + #[serde(default)] + content: Option, + #[serde(default)] + spec_id: i32, + #[serde(default)] + partition: Option, + #[serde(default)] + equality_ids: Option>, + #[serde(default)] + key_metadata: Option, + #[serde(default)] + first_row_id: Option, +} + +/// Decode expanded REST file-scan-tasks and delete-files into domain tasks. +pub(crate) fn decode_scan_tasks( + files: Vec, + delete_files: Vec, + ctx: &ConvertContext, +) -> Result> { + if files.is_empty() && delete_files.is_empty() { + return Ok(Vec::new()); + } + + let deletes = delete_files + .iter() + .map(to_delete_file) + .collect::>>()?; + + files + .into_iter() + .map(|task| to_file_scan_task(task, &deletes, ctx)) + .collect() +} + +fn parse_content_file(value: &Value) -> Result { + serde_json::from_value(value.clone()).map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + "failed to decode REST content-file JSON", + ) + .with_source(e) + }) +} + +fn to_delete_file(value: &Value) -> Result { + let rcf = parse_content_file(value)?; + let content = rcf + .content + .map(DataContentType::try_from) + .transpose()? + .unwrap_or(DataContentType::PositionDeletes); + Ok(FileScanTaskDeleteFile::builder() + .with_file_path(rcf.file_path) + .with_file_size_in_bytes(rcf.file_size_in_bytes) + .with_file_type(content) + .with_partition_spec_id(rcf.spec_id) + .with_equality_ids(rcf.equality_ids) + .with_key_metadata(decode_key_metadata(rcf.key_metadata.as_ref())) + .build()) +} + +fn to_file_scan_task( + task: RestFileScanTask, + all_deletes: &[FileScanTaskDeleteFile], + ctx: &ConvertContext, +) -> Result { + let rcf = parse_content_file(&task.data_file)?; + let spec = ctx + .metadata + .partition_spec_by_id(rcf.spec_id) + .cloned() + .ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!( + "scan plan referenced unknown partition spec id {}", + rcf.spec_id + ), + ) + })?; + + let data_file_format = rcf.file_format.parse::().map_err(|e| { + Error::new( + ErrorKind::DataInvalid, + format!( + "unsupported data file format {:?} in scan plan", + rcf.file_format + ), + ) + .with_source(e) + })?; + + let partition = decode_partition(&rcf, spec.as_ref(), &ctx.snapshot_schema)?; + let deletes = match task.delete_file_references { + Some(refs) => refs + .into_iter() + .map(|idx| { + let usize_idx = usize::try_from(idx).map_err(|_| { + Error::new( + ErrorKind::DataInvalid, + format!("delete-file-reference {idx} out of range"), + ) + })?; + all_deletes.get(usize_idx).cloned().ok_or_else(|| { + Error::new( + ErrorKind::DataInvalid, + format!("delete-file-reference {idx} out of range"), + ) + }) + }) + .collect::>>()?, + None => Vec::new(), + }; + + Ok(FileScanTask::builder() + .with_file_size_in_bytes(rcf.file_size_in_bytes) + .with_start(0) + .with_length(rcf.file_size_in_bytes) + .with_record_count(rcf.record_count) + .with_first_row_id(rcf.first_row_id) + .with_data_file_path(rcf.file_path) + .with_data_file_format(data_file_format) + .with_schema(ctx.snapshot_schema.clone()) + .with_project_field_ids(ctx.project_field_ids.clone()) + .with_predicate(ctx.bound_filter.clone()) + .with_deletes(deletes) + .with_partition(Some(partition)) + .with_partition_spec(Some(spec)) + .with_name_mapping(ctx.name_mapping.clone()) + .with_unified_partition_type(ctx.unified_partition_type.clone()) + .with_case_sensitive(ctx.case_sensitive) + .with_key_metadata(decode_key_metadata(rcf.key_metadata.as_ref())) + .build()) +} + +fn decode_partition( + rcf: &RestContentFile, + spec: &iceberg::spec::PartitionSpec, + schema: &SchemaRef, +) -> Result { + let partition_type = spec.partition_type(schema)?; + let fields = partition_type.fields(); + if fields.is_empty() { + return Ok(Struct::empty()); + } + + let Some(value) = rcf.partition.as_ref() else { + return Ok(Struct::from_iter(fields.iter().map(|_| None))); + }; + + let mut literals: Vec> = Vec::with_capacity(fields.len()); + match value { + Value::Array(values) => { + if values.len() != fields.len() { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "partition array has {} values but spec expects {}", + values.len(), + fields.len() + ), + )); + } + for (field, v) in fields.iter().zip(values) { + literals.push(Literal::try_from_json(v.clone(), &field.field_type)?); + } + } + Value::Object(map) => { + for field in fields { + let v = map + .get(&field.id.to_string()) + .cloned() + .unwrap_or(Value::Null); + literals.push(Literal::try_from_json(v, &field.field_type)?); + } + } + Value::Null => { + return Ok(Struct::from_iter(fields.iter().map(|_| None))); + } + other => { + return Err(Error::new( + ErrorKind::DataInvalid, + format!("unexpected partition encoding in scan plan: {other}"), + )); + } + } + + Ok(Struct::from_iter(literals)) +} + +fn decode_key_metadata(value: Option<&Value>) -> Option> { + match value { + Some(Value::Array(values)) => { + let bytes: Vec = values + .iter() + .filter_map(|v| v.as_u64().map(|n| n as u8)) + .collect(); + if bytes.is_empty() { + None + } else { + Some(bytes.into_boxed_slice()) + } + } + _ => None, + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use iceberg::spec::{DataContentType, DataFileFormat, TableMetadataRef}; + use serde_json::json; + + use super::*; + use crate::scan_planning::RestFileScanTask; + + fn unpartitioned_metadata() -> TableMetadataRef { + let json = r#"{ + "format-version": 2, + "table-uuid": "00000000-0000-0000-0000-000000000001", + "location": "s3://bucket/t", + "last-sequence-number": 1, + "last-updated-ms": 1, + "last-column-id": 1, + "current-schema-id": 0, + "schemas": [{ + "type": "struct", + "schema-id": 0, + "fields": [ + {"id": 1, "name": "id", "required": true, "type": "int"} + ] + }], + "default-spec-id": 0, + "partition-specs": [{"spec-id": 0, "fields": []}], + "last-partition-id": 999, + "default-sort-order-id": 0, + "sort-orders": [{"order-id": 0, "fields": []}], + "current-snapshot-id": -1, + "snapshots": [] + }"#; + Arc::new(serde_json::from_str(json).unwrap()) + } + + fn ctx() -> ConvertContext { + let metadata = unpartitioned_metadata(); + let schema: SchemaRef = metadata.current_schema().clone(); + ConvertContext { + metadata, + snapshot_schema: schema, + project_field_ids: vec![1], + case_sensitive: true, + bound_filter: None, + name_mapping: None, + unified_partition_type: None, + } + } + + #[test] + fn empty_plan_decodes_to_no_tasks() { + let tasks = decode_scan_tasks(vec![], vec![], &ctx()).unwrap(); + assert!(tasks.is_empty()); + } + + #[test] + fn data_file_json_becomes_file_scan_task() { + let file = RestFileScanTask { + data_file: json!({ + "content": 0, + "file-path": "s3://bucket/f.parquet", + "file-format": "PARQUET", + "spec-id": 0, + "partition": {}, + "record-count": 7, + "file-size-in-bytes": 128 + }), + delete_file_references: None, + residual_filter: None, + }; + let tasks = decode_scan_tasks(vec![file], vec![], &ctx()).unwrap(); + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0].data_file_path, "s3://bucket/f.parquet"); + assert_eq!(tasks[0].file_size_in_bytes, 128); + assert_eq!(tasks[0].length, 128); + assert_eq!(tasks[0].data_file_format, DataFileFormat::Parquet); + assert_eq!(tasks[0].record_count, Some(7)); + } + + #[test] + fn delete_file_references_are_resolved() { + let file = RestFileScanTask { + data_file: json!({ + "file-path": "s3://bucket/f.parquet", + "file-format": "parquet", + "file-size-in-bytes": 10 + }), + delete_file_references: Some(vec![0]), + residual_filter: None, + }; + let deletes = vec![json!({ + "content": 1, + "file-path": "s3://bucket/d.parquet", + "file-format": "parquet", + "file-size-in-bytes": 3 + })]; + let tasks = decode_scan_tasks(vec![file], deletes, &ctx()).unwrap(); + assert_eq!(tasks[0].deletes.len(), 1); + assert_eq!(tasks[0].deletes[0].file_path, "s3://bucket/d.parquet"); + assert_eq!( + tasks[0].deletes[0].file_type, + DataContentType::PositionDeletes + ); + } + + #[test] + fn unknown_spec_id_is_data_invalid() { + let file = RestFileScanTask { + data_file: json!({ + "file-path": "s3://bucket/f.parquet", + "file-format": "parquet", + "file-size-in-bytes": 10, + "spec-id": 99 + }), + delete_file_references: None, + residual_filter: None, + }; + let err = decode_scan_tasks(vec![file], vec![], &ctx()).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + } +} diff --git a/crates/catalog/rest/src/scan_planning.rs b/crates/catalog/rest/src/scan_planning.rs index bf5ba62471..57223619eb 100644 --- a/crates/catalog/rest/src/scan_planning.rs +++ b/crates/catalog/rest/src/scan_planning.rs @@ -17,13 +17,16 @@ //! REST server-side scan planning client. //! -//! Implements the plan / fetch-result / cancel / fetch-tasks endpoints and a -//! [`RestCatalog::wait_for_plan`] poller. Task decoding and `TableScan` -//! auto-routing are follow-ups: [`RestCatalog::supports_remote_scan_planning`] -//! stays `false` until those land. +//! Implements the plan / fetch-result / cancel / fetch-tasks endpoints, a +//! [`RestCatalog::wait_for_plan`] poller, content-file decoding, and +//! [`iceberg::scan::ScanPlanner`] so a table loaded from this catalog can +//! plan remotely when the server advertises the endpoints. +use std::collections::{HashSet, VecDeque}; use std::time::Duration; +use async_trait::async_trait; +use iceberg::scan::{ScanPlanner, ScanPlanningRequest, ScanPlanningResult}; use iceberg::{Error, ErrorKind, Result, TableIdent}; use rand::Rng; use reqwest::{Method, Response, StatusCode}; @@ -346,15 +349,11 @@ impl RestCatalog { /// Whether this catalog can complete a remote plan end-to-end. /// - /// Stays `false` until task decoding is wired into `TableScan`. Routing - /// auto-mode scans on endpoint capability alone would fail with - /// [`ErrorKind::FeatureUnsupported`] instead of falling back to local - /// planning. + /// True when the server advertises all four scan-planning endpoints. Task + /// decoding is implemented, so auto-mode [`iceberg::scan::TableScan`]s can + /// route here instead of falling back to local planning. pub async fn supports_remote_scan_planning(&self) -> Result { - // Touch config so a down server still surfaces as an error, but never - // claim end-to-end remote planning until TableScan can decode tasks. - let _ = self.supports_plan_table_scan().await?; - Ok(false) + self.supports_full_remote_scan_planning().await } /// Submits a server-side scan plan. @@ -659,6 +658,93 @@ impl RestCatalog { async fn abandon_plan(&self, table: &TableIdent, plan_id: &str, grace: Duration) { let _ = tokio::time::timeout(grace, self.cancel_planning(table, plan_id)).await; } + + async fn collect_scan_tasks( + &self, + table: &TableIdent, + tasks: ScanTasks, + ) -> Result<(Vec, Vec)> { + let mut files = tasks.file_scan_tasks; + let mut deletes = tasks.delete_files; + let mut queue: VecDeque = tasks.plan_tasks.into(); + let mut seen = HashSet::new(); + while let Some(handle) = queue.pop_front() { + if !seen.insert(handle.clone()) { + continue; + } + let resp = self + .fetch_scan_tasks(table, FetchScanTasksRequest { + idempotency_key: None, + plan_task: handle, + }) + .await?; + files.extend(resp.scan_tasks.file_scan_tasks); + deletes.extend(resp.scan_tasks.delete_files); + queue.extend(resp.scan_tasks.plan_tasks); + } + Ok((files, deletes)) + } + + async fn plan_scan(&self, request: ScanPlanningRequest) -> Result { + let ident = request.table_ident.clone(); + let resp = self + .plan_table_scan(&ident, PlanTableScanRequest { + snapshot_id: request.snapshot_id, + select: request.select.clone().unwrap_or_default(), + case_sensitive: Some(request.case_sensitive), + ..Default::default() + }) + .await?; + + let completed = match resp.status { + PlanStatus::Completed => CompletedPlanningResult { + status: PlanStatus::Completed, + scan_tasks: resp.scan_tasks, + storage_credentials: resp.storage_credentials, + }, + PlanStatus::Submitted => { + let plan_id = resp.plan_id.ok_or_else(|| { + Error::new(ErrorKind::DataInvalid, "submitted plan missing plan-id") + })?; + self.wait_for_plan(&ident, &plan_id, WaitForPlanOptions::default()) + .await? + } + PlanStatus::Failed => return Err(failed_plan_error(resp.error.as_ref())), + PlanStatus::Cancelled => { + return Err(Error::new( + ErrorKind::Unexpected, + "planTableScan response has invalid status cancelled", + )); + } + }; + + let (files, deletes) = self + .collect_scan_tasks(&ident, completed.scan_tasks) + .await?; + let ctx = crate::scan_decode::ConvertContext { + metadata: request.metadata, + snapshot_schema: request.snapshot_schema, + project_field_ids: request.project_field_ids, + case_sensitive: request.case_sensitive, + bound_filter: request.bound_filter, + name_mapping: request.name_mapping, + unified_partition_type: request.unified_partition_type, + }; + Ok(ScanPlanningResult { + tasks: crate::scan_decode::decode_scan_tasks(files, deletes, &ctx)?, + }) + } +} + +#[async_trait] +impl ScanPlanner for RestCatalog { + async fn supports_remote_scan_planning(&self) -> Result { + self.supports_full_remote_scan_planning().await + } + + async fn plan_files(&self, request: ScanPlanningRequest) -> Result { + self.plan_scan(request).await + } } /// Best-effort DELETE of a submitted plan if [`RestCatalog::wait_for_plan`] is @@ -915,7 +1001,12 @@ fn require_plan_id(plan_id: &str) -> Result<()> { #[cfg(test)] mod tests { - use iceberg::{NamespaceIdent, Runtime}; + use std::sync::Arc; + + use futures::TryStreamExt; + use iceberg::io::LocalFsStorageFactory; + use iceberg::scan::ScanPlanningMode; + use iceberg::{Catalog, NamespaceIdent, Runtime}; use mockito::Server; use serde_json::json; use uuid::Uuid; @@ -950,6 +1041,7 @@ mod tests { .mock("GET", "/v1/config") .with_status(200) .with_body(body.to_string()) + .expect_at_least(1) .create_async() .await } @@ -1097,13 +1189,13 @@ mod tests { } #[tokio::test] - async fn remote_scan_planning_stays_false_when_all_endpoints_are_advertised() { + async fn remote_scan_planning_is_true_when_all_endpoints_are_advertised() { let mut server = Server::new_async().await; let config = config_with_endpoints(&mut server, ALL_PLAN).await; let catalog = catalog(&server.url()); assert!(catalog.supports_plan_table_scan().await.unwrap()); assert!(catalog.supports_full_remote_scan_planning().await.unwrap()); - assert!(!catalog.supports_remote_scan_planning().await.unwrap()); + assert!(catalog.supports_remote_scan_planning().await.unwrap()); config.assert_async().await; } @@ -1749,4 +1841,349 @@ mod tests { config.assert_async().await; cancel.assert_async().await; } + + fn load_catalog(uri: &str) -> RestCatalog { + RestCatalog::new( + RestCatalogConfig::builder().uri(uri.to_string()).build(), + None, + Some(Arc::new(LocalFsStorageFactory)), + Runtime::current(), + None, + ) + } + + fn test1() -> TableIdent { + TableIdent::new(NamespaceIdent::new("ns1".into()), "test1".into()) + } + + fn data_file_json() -> serde_json::Value { + json!({ + "content": 0, + "file-path": "s3://warehouse/database/table/data/00000.parquet", + "file-format": "PARQUET", + "spec-id": 0, + "partition": {}, + "record-count": 1, + "file-size-in-bytes": 697 + }) + } + + async fn mock_load_table(server: &mut mockito::ServerGuard) -> mockito::Mock { + server + .mock("GET", "/v1/namespaces/ns1/tables/test1") + .with_status(200) + .with_body_from_file(format!( + "{}/testdata/load_table_response.json", + env!("CARGO_MANIFEST_DIR") + )) + .create_async() + .await + } + + fn empty_table_body() -> String { + json!({ + "metadata-location": "s3://warehouse/ns1/test1/metadata.json", + "metadata": { + "format-version": 2, + "table-uuid": "b55d9dda-6561-423a-8bfc-787980ce421f", + "location": "s3://warehouse/ns1/test1", + "last-sequence-number": 0, + "last-updated-ms": 1, + "last-column-id": 2, + "current-schema-id": 0, + "schemas": [{ + "type": "struct", + "schema-id": 0, + "fields": [ + {"id": 1, "name": "id", "required": false, "type": "int"}, + {"id": 2, "name": "data", "required": false, "type": "string"} + ] + }], + "default-spec-id": 0, + "partition-specs": [{"spec-id": 0, "fields": []}], + "last-partition-id": 999, + "default-sort-order-id": 0, + "sort-orders": [{"order-id": 0, "fields": []}], + "current-snapshot-id": -1, + "snapshots": [] + } + }) + .to_string() + } + + #[tokio::test] + async fn table_scan_completed_plan_yields_file_scan_task() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let load = mock_load_table(&mut server).await; + let plan = server + .mock("POST", "/v1/namespaces/ns1/tables/test1/plan") + .with_status(200) + .with_body( + json!({ + "status": "completed", + "plan-id": "p1", + "file-scan-tasks": [{"data-file": data_file_json()}] + }) + .to_string(), + ) + .expect(1) + .create_async() + .await; + let catalog = load_catalog(&server.url()); + let table = catalog.load_table(&test1()).await.unwrap(); + let tasks: Vec<_> = table + .scan() + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(tasks.len(), 1); + assert_eq!( + tasks[0].data_file_path, + "s3://warehouse/database/table/data/00000.parquet" + ); + assert_eq!(tasks[0].file_size_in_bytes, 697); + assert_eq!( + tasks[0].data_file_format, + iceberg::spec::DataFileFormat::Parquet + ); + config.assert_async().await; + load.assert_async().await; + plan.assert_async().await; + } + + #[tokio::test] + async fn table_scan_submitted_plan_waits_then_completes() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let load = mock_load_table(&mut server).await; + let plan = server + .mock("POST", "/v1/namespaces/ns1/tables/test1/plan") + .with_status(200) + .with_body(r#"{"status":"submitted","plan-id":"p1"}"#) + .expect(1) + .create_async() + .await; + let pending = server + .mock("GET", "/v1/namespaces/ns1/tables/test1/plan/p1") + .with_status(200) + .with_body(r#"{"status":"submitted"}"#) + .expect(1) + .create_async() + .await; + let done = server + .mock("GET", "/v1/namespaces/ns1/tables/test1/plan/p1") + .with_status(200) + .with_body( + json!({ + "status": "completed", + "file-scan-tasks": [{"data-file": data_file_json()}] + }) + .to_string(), + ) + .expect(1) + .create_async() + .await; + let catalog = load_catalog(&server.url()); + let table = catalog.load_table(&test1()).await.unwrap(); + let tasks: Vec<_> = table + .scan() + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0].file_size_in_bytes, 697); + config.assert_async().await; + load.assert_async().await; + plan.assert_async().await; + pending.assert_async().await; + done.assert_async().await; + } + + #[tokio::test] + async fn table_scan_expands_plan_task_handles() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let load = mock_load_table(&mut server).await; + let plan = server + .mock("POST", "/v1/namespaces/ns1/tables/test1/plan") + .with_status(200) + .with_body(r#"{"status":"completed","plan-id":"p1","plan-tasks":["h1"]}"#) + .expect(1) + .create_async() + .await; + let tasks_ep = server + .mock("POST", "/v1/namespaces/ns1/tables/test1/tasks") + .match_body(r#"{"plan-task":"h1"}"#) + .with_status(200) + .with_body( + json!({ + "file-scan-tasks": [{"data-file": data_file_json()}] + }) + .to_string(), + ) + .expect(1) + .create_async() + .await; + let catalog = load_catalog(&server.url()); + let table = catalog.load_table(&test1()).await.unwrap(); + let tasks: Vec<_> = table + .scan() + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(tasks.len(), 1); + assert_eq!( + tasks[0].data_file_path, + "s3://warehouse/database/table/data/00000.parquet" + ); + config.assert_async().await; + load.assert_async().await; + plan.assert_async().await; + tasks_ep.assert_async().await; + } + + #[tokio::test] + async fn table_scan_reissued_plan_task_does_not_loop() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let load = mock_load_table(&mut server).await; + let plan = server + .mock("POST", "/v1/namespaces/ns1/tables/test1/plan") + .with_status(200) + .with_body(r#"{"status":"completed","plan-id":"p1","plan-tasks":["h1"]}"#) + .expect(1) + .create_async() + .await; + let tasks_ep = server + .mock("POST", "/v1/namespaces/ns1/tables/test1/tasks") + .with_status(200) + .with_body(r#"{"plan-tasks":["h1"]}"#) + .expect(1) + .create_async() + .await; + let catalog = load_catalog(&server.url()); + let table = catalog.load_table(&test1()).await.unwrap(); + let tasks: Vec<_> = table + .scan() + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert!(tasks.is_empty()); + config.assert_async().await; + load.assert_async().await; + plan.assert_async().await; + tasks_ep.assert_async().await; + } + + #[tokio::test] + async fn table_scan_empty_plan_is_empty_not_error() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let load = mock_load_table(&mut server).await; + let plan = server + .mock("POST", "/v1/namespaces/ns1/tables/test1/plan") + .with_status(200) + .with_body(r#"{"status":"completed","plan-id":"p1"}"#) + .expect(1) + .create_async() + .await; + let catalog = load_catalog(&server.url()); + let table = catalog.load_table(&test1()).await.unwrap(); + let tasks: Vec<_> = table + .scan() + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert!(tasks.is_empty()); + config.assert_async().await; + load.assert_async().await; + plan.assert_async().await; + } + + #[tokio::test] + async fn auto_mode_without_plan_endpoints_does_not_post_plan() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, &[]).await; + let load = server + .mock("GET", "/v1/namespaces/ns1/tables/test1") + .with_status(200) + .with_body(empty_table_body()) + .create_async() + .await; + let plan = server + .mock("POST", "/v1/namespaces/ns1/tables/test1/plan") + .expect(0) + .create_async() + .await; + let catalog = load_catalog(&server.url()); + let table = catalog.load_table(&test1()).await.unwrap(); + let tasks: Vec<_> = table + .scan() + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert!(tasks.is_empty()); + config.assert_async().await; + load.assert_async().await; + plan.assert_async().await; + } + + #[tokio::test] + async fn remote_mode_without_endpoints_is_feature_unsupported() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, &[]).await; + let load = server + .mock("GET", "/v1/namespaces/ns1/tables/test1") + .with_status(200) + .with_body(empty_table_body()) + .create_async() + .await; + let catalog = load_catalog(&server.url()); + let table = catalog.load_table(&test1()).await.unwrap(); + let err = match table + .scan() + .with_scan_planning_mode(ScanPlanningMode::Remote) + .build() + .unwrap() + .plan_files() + .await + { + Ok(_) => panic!("expected FeatureUnsupported"), + Err(e) => e, + }; + assert_eq!(err.kind(), ErrorKind::FeatureUnsupported); + config.assert_async().await; + load.assert_async().await; + } } diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index bb34d0d620..0489a31537 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -1270,6 +1270,21 @@ pub const iceberg::puffin::APACHE_DATASKETCHES_THETA_V1: &str pub const iceberg::puffin::CREATED_BY_PROPERTY: &str pub const iceberg::puffin::DELETION_VECTOR_V1: &str pub mod iceberg::scan +pub enum iceberg::scan::ScanPlanningMode +pub iceberg::scan::ScanPlanningMode::Auto +pub iceberg::scan::ScanPlanningMode::Local +pub iceberg::scan::ScanPlanningMode::Remote +impl core::clone::Clone for iceberg::scan::ScanPlanningMode +pub fn iceberg::scan::ScanPlanningMode::clone(&self) -> iceberg::scan::ScanPlanningMode +impl core::cmp::Eq for iceberg::scan::ScanPlanningMode +impl core::cmp::PartialEq for iceberg::scan::ScanPlanningMode +pub fn iceberg::scan::ScanPlanningMode::eq(&self, other: &iceberg::scan::ScanPlanningMode) -> bool +impl core::default::Default for iceberg::scan::ScanPlanningMode +pub fn iceberg::scan::ScanPlanningMode::default() -> iceberg::scan::ScanPlanningMode +impl core::fmt::Debug for iceberg::scan::ScanPlanningMode +pub fn iceberg::scan::ScanPlanningMode::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::marker::Copy for iceberg::scan::ScanPlanningMode +impl core::marker::StructuralPartialEq for iceberg::scan::ScanPlanningMode pub struct iceberg::scan::FileScanTask pub iceberg::scan::FileScanTask::case_sensitive: bool pub iceberg::scan::FileScanTask::data_file_format: iceberg::spec::DataFileFormat @@ -1335,6 +1350,27 @@ impl core::clone::Clone for iceberg::scan::ScanMetrics pub fn iceberg::scan::ScanMetrics::clone(&self) -> iceberg::scan::ScanMetrics impl core::fmt::Debug for iceberg::scan::ScanMetrics pub fn iceberg::scan::ScanMetrics::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +pub struct iceberg::scan::ScanPlanningRequest +pub iceberg::scan::ScanPlanningRequest::bound_filter: core::option::Option +pub iceberg::scan::ScanPlanningRequest::case_sensitive: bool +pub iceberg::scan::ScanPlanningRequest::metadata: iceberg::spec::TableMetadataRef +pub iceberg::scan::ScanPlanningRequest::name_mapping: core::option::Option> +pub iceberg::scan::ScanPlanningRequest::project_field_ids: alloc::vec::Vec +pub iceberg::scan::ScanPlanningRequest::select: core::option::Option> +pub iceberg::scan::ScanPlanningRequest::snapshot_id: core::option::Option +pub iceberg::scan::ScanPlanningRequest::snapshot_schema: iceberg::spec::SchemaRef +pub iceberg::scan::ScanPlanningRequest::table_ident: iceberg::TableIdent +pub iceberg::scan::ScanPlanningRequest::unified_partition_type: core::option::Option> +impl core::clone::Clone for iceberg::scan::ScanPlanningRequest +pub fn iceberg::scan::ScanPlanningRequest::clone(&self) -> iceberg::scan::ScanPlanningRequest +impl core::fmt::Debug for iceberg::scan::ScanPlanningRequest +pub fn iceberg::scan::ScanPlanningRequest::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +pub struct iceberg::scan::ScanPlanningResult +pub iceberg::scan::ScanPlanningResult::tasks: alloc::vec::Vec +impl core::clone::Clone for iceberg::scan::ScanPlanningResult +pub fn iceberg::scan::ScanPlanningResult::clone(&self) -> iceberg::scan::ScanPlanningResult +impl core::fmt::Debug for iceberg::scan::ScanPlanningResult +pub fn iceberg::scan::ScanPlanningResult::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct iceberg::scan::ScanResult impl iceberg::scan::ScanResult pub fn iceberg::scan::ScanResult::metrics(&self) -> &iceberg::scan::ScanMetrics @@ -1362,6 +1398,10 @@ pub fn iceberg::scan::TableScanBuilder<'a>::with_filter(self, predicate: iceberg pub fn iceberg::scan::TableScanBuilder<'a>::with_manifest_entry_concurrency_limit(self, limit: usize) -> Self pub fn iceberg::scan::TableScanBuilder<'a>::with_row_group_filtering_enabled(self, row_group_filtering_enabled: bool) -> Self pub fn iceberg::scan::TableScanBuilder<'a>::with_row_selection_enabled(self, row_selection_enabled: bool) -> Self +pub fn iceberg::scan::TableScanBuilder<'a>::with_scan_planning_mode(self, mode: iceberg::scan::ScanPlanningMode) -> Self +pub trait iceberg::scan::ScanPlanner: core::fmt::Debug + core::marker::Send + core::marker::Sync +pub fn iceberg::scan::ScanPlanner::plan_files<'life0, 'async_trait>(&'life0 self, request: iceberg::scan::ScanPlanningRequest) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait +pub fn iceberg::scan::ScanPlanner::supports_remote_scan_planning<'life0, 'async_trait>(&'life0 self) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait pub type iceberg::scan::ArrowRecordBatchStream = futures_core::stream::BoxStream<'static, iceberg::Result> pub type iceberg::scan::FileScanTaskStream = futures_core::stream::BoxStream<'static, iceberg::Result> pub mod iceberg::sensitive @@ -3206,6 +3246,7 @@ pub fn iceberg::table::TableBuilder::metadata>(self, metadata_location: T) -> Self pub fn iceberg::table::TableBuilder::readonly(self, readonly: bool) -> Self pub fn iceberg::table::TableBuilder::runtime(self, runtime: iceberg::Runtime) -> Self +pub fn iceberg::table::TableBuilder::scan_planner(self, scan_planner: alloc::sync::Arc) -> Self pub mod iceberg::test_utils pub fn iceberg::test_utils::check_record_batches(record_batches: alloc::vec::Vec, expected_schema: expect_test::Expect, expected_data: expect_test::Expect, ignore_check_columns: &[&str], sort_column: core::option::Option<&str>) pub async fn iceberg::test_utils::make_encrypted_table() -> iceberg::table::Table diff --git a/crates/iceberg/src/scan/mod.rs b/crates/iceberg/src/scan/mod.rs index b7c7254ae2..e389ca50e2 100644 --- a/crates/iceberg/src/scan/mod.rs +++ b/crates/iceberg/src/scan/mod.rs @@ -21,6 +21,7 @@ mod cache; use cache::*; mod context; use context::*; +mod planner; mod task; use std::sync::Arc; @@ -29,6 +30,7 @@ use arrow_array::RecordBatch; use futures::channel::mpsc::{Sender, channel}; use futures::stream::BoxStream; use futures::{SinkExt, StreamExt, TryStreamExt}; +pub use planner::*; pub use task::*; use crate::arrow::ArrowReaderBuilder; @@ -45,7 +47,7 @@ use crate::runtime::Runtime; use crate::spec::{DEFAULT_SCHEMA_NAME_MAPPING, DataContentType, NameMapping, SnapshotRef}; use crate::table::Table; use crate::util::available_parallelism; -use crate::{Error, ErrorKind, Result}; +use crate::{Error, ErrorKind, Result, TableIdent}; /// A stream of arrow [`RecordBatch`]es. pub type ArrowRecordBatchStream = BoxStream<'static, Result>; @@ -64,6 +66,7 @@ pub struct TableScanBuilder<'a> { concurrency_limit_manifest_files: usize, row_group_filtering_enabled: bool, row_selection_enabled: bool, + planning_mode: ScanPlanningMode, } impl<'a> TableScanBuilder<'a> { @@ -82,6 +85,7 @@ impl<'a> TableScanBuilder<'a> { concurrency_limit_manifest_files: num_cpus, row_group_filtering_enabled: true, row_selection_enabled: false, + planning_mode: ScanPlanningMode::Auto, } } @@ -188,6 +192,15 @@ impl<'a> TableScanBuilder<'a> { self } + /// Sets whether this scan plans locally, remotely, or automatically. + /// + /// Defaults to [`ScanPlanningMode::Auto`]: remote planning when the table's + /// [`ScanPlanner`] advertises it, otherwise local manifest planning. + pub fn with_scan_planning_mode(mut self, mode: ScanPlanningMode) -> Self { + self.planning_mode = mode; + self + } + /// Build the table scan. pub fn build(self) -> Result { let snapshot = match self.snapshot_id { @@ -215,6 +228,9 @@ impl<'a> TableScanBuilder<'a> { row_group_filtering_enabled: self.row_group_filtering_enabled, row_selection_enabled: self.row_selection_enabled, runtime: self.table.runtime().clone(), + scan_planner: self.table.scan_planner(), + table_ident: self.table.identifier().clone(), + planning_mode: self.planning_mode, }); }; current_snapshot_id.clone() @@ -344,6 +360,9 @@ impl<'a> TableScanBuilder<'a> { row_group_filtering_enabled: self.row_group_filtering_enabled, row_selection_enabled: self.row_selection_enabled, runtime: self.table.runtime().clone(), + scan_planner: self.table.scan_planner(), + table_ident: self.table.identifier().clone(), + planning_mode: self.planning_mode, }) } } @@ -374,11 +393,80 @@ pub struct TableScan { row_selection_enabled: bool, runtime: Runtime, + scan_planner: Option>, + table_ident: TableIdent, + planning_mode: ScanPlanningMode, } impl TableScan { /// Returns a stream of [`FileScanTask`]s. pub async fn plan_files(&self) -> Result { + match self.planning_mode { + ScanPlanningMode::Local => self.plan_files_local().await, + ScanPlanningMode::Remote => self.plan_files_remote().await, + ScanPlanningMode::Auto => { + if self.remote_planning_available().await? { + match self.plan_files_remote().await { + Ok(tasks) => Ok(tasks), + Err(e) if e.kind() == ErrorKind::FeatureUnsupported => { + self.plan_files_local().await + } + Err(e) => Err(e), + } + } else { + self.plan_files_local().await + } + } + } + } + + async fn remote_planning_available(&self) -> Result { + match &self.scan_planner { + Some(planner) => planner.supports_remote_scan_planning().await, + None => Ok(false), + } + } + + async fn plan_files_remote(&self) -> Result { + let Some(planner) = &self.scan_planner else { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "remote scan planning is unavailable", + )); + }; + if !planner.supports_remote_scan_planning().await? { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "remote scan planning is unavailable", + )); + } + let Some(plan_context) = self.plan_context.as_ref() else { + return Ok(Box::pin(futures::stream::empty())); + }; + + let result = planner + .plan_files(ScanPlanningRequest { + table_ident: self.table_ident.clone(), + snapshot_id: Some(plan_context.snapshot.snapshot_id()), + select: self.column_names.clone(), + case_sensitive: plan_context.case_sensitive, + project_field_ids: plan_context.field_ids.as_ref().clone(), + metadata: plan_context.table_metadata.clone(), + snapshot_schema: plan_context.snapshot_schema.clone(), + bound_filter: plan_context + .snapshot_bound_predicate + .as_ref() + .map(|p| p.as_ref().clone()), + name_mapping: plan_context.name_mapping.clone(), + unified_partition_type: plan_context.unified_partition_type.clone(), + }) + .await?; + Ok(Box::pin(futures::stream::iter( + result.tasks.into_iter().map(Ok), + ))) + } + + async fn plan_files_local(&self) -> Result { let Some(plan_context) = self.plan_context.as_ref() else { return Ok(Box::pin(futures::stream::empty())); }; @@ -666,7 +754,9 @@ pub mod tests { RESERVED_COL_NAME_POS, RESERVED_COL_NAME_SPEC_ID, RESERVED_FIELD_ID_DELETE_FILE_PATH, RESERVED_FIELD_ID_DELETE_FILE_POS, RESERVED_FIELD_ID_POS, }; - use crate::scan::FileScanTask; + use crate::scan::{ + FileScanTask, ScanPlanner, ScanPlanningMode, ScanPlanningRequest, ScanPlanningResult, + }; use crate::spec::{ DEFAULT_SCHEMA_NAME_MAPPING, DataContentType, DataFileBuilder, DataFileFormat, Datum, FormatVersion, Literal, MAIN_BRANCH, ManifestEntry, ManifestListWriter, ManifestStatus, @@ -3686,4 +3776,201 @@ pub mod tests { .collect(); assert_eq!(y, (1100..1200).collect::>()); } + + #[derive(Debug)] + struct FakePlanner { + remote: bool, + kind: Option, + tasks: Vec, + } + + #[async_trait::async_trait] + impl ScanPlanner for FakePlanner { + async fn supports_remote_scan_planning(&self) -> crate::Result { + Ok(self.remote) + } + + async fn plan_files( + &self, + _request: ScanPlanningRequest, + ) -> crate::Result { + if let Some(kind) = self.kind { + return Err(crate::Error::new(kind, "fake planner error")); + } + Ok(ScanPlanningResult { + tasks: self.tasks.clone(), + }) + } + } + + fn remote_stub_task() -> FileScanTask { + let schema = Arc::new( + Schema::builder() + .with_fields(vec![ + NestedField::required(1, "x", Type::Primitive(PrimitiveType::Int)).into(), + ]) + .build() + .unwrap(), + ); + FileScanTask::builder() + .with_data_file_path("s3://bucket/remote.parquet".to_string()) + .with_file_size_in_bytes(42) + .with_start(0) + .with_length(42) + .with_project_field_ids(vec![1]) + .with_schema(schema) + .with_data_file_format(DataFileFormat::Parquet) + .with_case_sensitive(true) + .build() + } + + fn table_with_planner(fixture: &TableTestFixture, planner: Arc) -> Table { + Table::builder() + .metadata(fixture.table.metadata_ref()) + .identifier(fixture.table.identifier().clone()) + .file_io(fixture.table.file_io().clone()) + .metadata_location(fixture.table.metadata_location().unwrap()) + .runtime(fixture.table.runtime().clone()) + .scan_planner(planner) + .build() + .unwrap() + } + + #[tokio::test] + async fn auto_mode_without_planner_plans_locally() { + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + let tasks: Vec<_> = fixture + .table + .scan() + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(tasks.len(), 2); + } + + #[tokio::test] + async fn auto_mode_uses_remote_tasks_when_planner_supports_it() { + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + let table = table_with_planner( + &fixture, + Arc::new(FakePlanner { + remote: true, + kind: None, + tasks: vec![remote_stub_task()], + }), + ); + let tasks: Vec<_> = table + .scan() + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0].data_file_path, "s3://bucket/remote.parquet"); + assert_eq!(tasks[0].file_size_in_bytes, 42); + } + + #[tokio::test] + async fn auto_mode_falls_back_to_local_on_feature_unsupported() { + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + let table = table_with_planner( + &fixture, + Arc::new(FakePlanner { + remote: true, + kind: Some(ErrorKind::FeatureUnsupported), + tasks: vec![], + }), + ); + let tasks: Vec<_> = table + .scan() + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(tasks.len(), 2); + } + + #[tokio::test] + async fn auto_mode_does_not_fall_back_on_other_remote_errors() { + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + let table = table_with_planner( + &fixture, + Arc::new(FakePlanner { + remote: true, + kind: Some(ErrorKind::Unexpected), + tasks: vec![], + }), + ); + let err = match table.scan().build().unwrap().plan_files().await { + Ok(_) => panic!("expected remote planning to fail"), + Err(e) => e, + }; + assert_eq!(err.kind(), ErrorKind::Unexpected); + } + + #[tokio::test] + async fn remote_mode_without_planner_is_feature_unsupported() { + let fixture = TableTestFixture::new(); + let err = match fixture + .table + .scan() + .with_scan_planning_mode(ScanPlanningMode::Remote) + .build() + .unwrap() + .plan_files() + .await + { + Ok(_) => panic!("expected FeatureUnsupported"), + Err(e) => e, + }; + assert_eq!(err.kind(), ErrorKind::FeatureUnsupported); + } + + #[tokio::test] + async fn local_mode_ignores_a_capable_planner() { + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + let table = table_with_planner( + &fixture, + Arc::new(FakePlanner { + remote: true, + kind: None, + tasks: vec![remote_stub_task()], + }), + ); + let tasks: Vec<_> = table + .scan() + .with_scan_planning_mode(ScanPlanningMode::Local) + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(tasks.len(), 2); + assert!( + tasks + .iter() + .all(|t| t.data_file_path != "s3://bucket/remote.parquet") + ); + } } diff --git a/crates/iceberg/src/scan/planner.rs b/crates/iceberg/src/scan/planner.rs new file mode 100644 index 0000000000..a6ec6d8b92 --- /dev/null +++ b/crates/iceberg/src/scan/planner.rs @@ -0,0 +1,97 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +//! Optional remote scan-planning seam. +//! +//! A catalog that can plan scans server-side implements [`ScanPlanner`] and +//! injects it onto the [`Table`](crate::table::Table)s it loads. The core +//! [`Catalog`](crate::Catalog) trait is left untouched. [`TableScan::plan_files`] +//! (crate::scan::TableScan::plan_files) routes on [`ScanPlanningMode`]. + +use std::fmt::Debug; +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::expr::BoundPredicate; +use crate::scan::FileScanTask; +use crate::spec::{NameMapping, SchemaRef, StructType, TableMetadataRef}; +use crate::{Result, TableIdent}; + +/// How [`TableScan::plan_files`](crate::scan::TableScan::plan_files) chooses +/// between local manifest planning and a catalog-provided [`ScanPlanner`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum ScanPlanningMode { + /// Use remote planning when the table's planner advertises it, otherwise + /// plan locally. This is the default. + #[default] + Auto, + /// Always plan by reading manifests through the table's FileIO. + Local, + /// Require a planner that supports remote planning; error otherwise. + Remote, +} + +/// Input a [`TableScan`](crate::scan::TableScan) hands to a [`ScanPlanner`]. +#[derive(Debug, Clone)] +pub struct ScanPlanningRequest { + /// Identifier of the table being scanned. + pub table_ident: TableIdent, + /// Snapshot to scan. `None` lets the server use the current snapshot. + pub snapshot_id: Option, + /// Projected column names, or `None` to select all top-level columns. + pub select: Option>, + /// Case sensitivity for filter and projection binding. + pub case_sensitive: bool, + /// Resolved field ids to stamp onto every produced [`FileScanTask`]. + pub project_field_ids: Vec, + /// Table metadata at scan time, used to resolve partition specs. + pub metadata: TableMetadataRef, + /// Schema of the scanned snapshot. + pub snapshot_schema: SchemaRef, + /// The scan's bound filter, applied as the per-task row predicate. + pub bound_filter: Option, + /// Optional name mapping from table properties. + pub name_mapping: Option>, + /// Unified partition type when `_partition` is projected. + pub unified_partition_type: Option>, +} + +/// Tasks produced by a [`ScanPlanner`]. +#[derive(Debug, Clone)] +pub struct ScanPlanningResult { + /// Planned file scan tasks. + pub tasks: Vec, +} + +/// Catalog capability for planning scans remotely. +/// +/// REST catalogs implement this; other catalogs leave it unset and planning +/// stays local. [`ScanPlanner::supports_remote_scan_planning`] must only be +/// true when the implementation can decode the server's task payload into +/// [`FileScanTask`]s. +#[async_trait] +pub trait ScanPlanner: Debug + Send + Sync { + /// Whether this planner can complete a remote plan end-to-end. + /// + /// Auto mode routes on this flag. Returning true without a working decoder + /// would fail auto-mode scans instead of falling back to local planning. + async fn supports_remote_scan_planning(&self) -> Result; + + /// Plan a scan server-side and return [`FileScanTask`]s. + async fn plan_files(&self, request: ScanPlanningRequest) -> Result; +} diff --git a/crates/iceberg/src/table.rs b/crates/iceberg/src/table.rs index 31feade038..d72ea7a3f3 100644 --- a/crates/iceberg/src/table.rs +++ b/crates/iceberg/src/table.rs @@ -26,7 +26,7 @@ use crate::inspect::MetadataTable; use crate::io::FileIO; use crate::io::object_cache::ObjectCache; use crate::runtime::Runtime; -use crate::scan::TableScanBuilder; +use crate::scan::{ScanPlanner, TableScanBuilder}; use crate::spec::{ManifestListReader, SchemaRef, SnapshotRef, TableMetadata, TableMetadataRef}; use crate::{Error, ErrorKind, Result, TableIdent}; @@ -41,6 +41,7 @@ pub struct TableBuilder { disable_cache: bool, cache_size_bytes: Option, runtime: Option, + scan_planner: Option>, } impl TableBuilder { @@ -55,6 +56,7 @@ impl TableBuilder { disable_cache: false, cache_size_bytes: None, runtime: None, + scan_planner: None, } } @@ -118,6 +120,13 @@ impl TableBuilder { self } + /// optional - sets a [`ScanPlanner`] used by [`Table::scan`] when the + /// scan's planning mode selects remote planning. + pub fn scan_planner(mut self, scan_planner: Arc) -> Self { + self.scan_planner = Some(scan_planner); + self + } + /// build the Table pub fn build(self) -> Result { let Self { @@ -130,6 +139,7 @@ impl TableBuilder { disable_cache, cache_size_bytes, runtime, + scan_planner, } = self; let Some(file_io) = file_io else { @@ -190,6 +200,7 @@ impl TableBuilder { object_cache, runtime, encryption_manager, + scan_planner, }) } } @@ -205,6 +216,7 @@ pub struct Table { object_cache: Arc, runtime: Runtime, encryption_manager: Option>, + scan_planner: Option>, } impl Table { @@ -291,6 +303,11 @@ impl Table { &self.runtime } + /// Returns the optional [`ScanPlanner`] injected when this table was loaded. + pub(crate) fn scan_planner(&self) -> Option> { + self.scan_planner.clone() + } + /// Returns the flag indicating whether the `Table` is readonly or not pub fn readonly(&self) -> bool { self.readonly From 91b42c6a2d606dff05b53ae9ccd463c2f3aca8e4 Mon Sep 17 00:00:00 2001 From: David Dallakyan Date: Sun, 16 Aug 2026 18:55:08 +0400 Subject: [PATCH 5/6] fix(catalog-rest): decode REST content-file strings and honor local planning default REST content-file JSON uses kebab-case string discriminators and hex key-metadata. Decode each ScanTasks payload against its own delete-files. Default ScanPlanningMode to Local so existing scans do not POST /plan. Auto remotes only for scan-planning-mode=server. Share the live catalog client when injecting the planner. --- crates/catalog/rest/public-api.txt | 27 +- crates/catalog/rest/src/catalog.rs | 26 +- crates/catalog/rest/src/scan_decode.rs | 259 ++++++++------ crates/catalog/rest/src/scan_planning.rs | 432 +++++++++++++++++++++-- crates/iceberg/src/scan/mod.rs | 143 +++++++- crates/iceberg/src/scan/planner.rs | 12 +- 6 files changed, 757 insertions(+), 142 deletions(-) diff --git a/crates/catalog/rest/public-api.txt b/crates/catalog/rest/public-api.txt index 4ffaf78d24..5ab2241869 100644 --- a/crates/catalog/rest/public-api.txt +++ b/crates/catalog/rest/public-api.txt @@ -415,8 +415,31 @@ pub fn iceberg_catalog_rest::RestCatalogBuilder::load(self, name: impl core::con pub fn iceberg_catalog_rest::RestCatalogBuilder::with_kms_client_factory(self, kms_client_factory: alloc::sync::Arc) -> Self pub fn iceberg_catalog_rest::RestCatalogBuilder::with_runtime(self, runtime: iceberg::runtime::Runtime) -> Self pub fn iceberg_catalog_rest::RestCatalogBuilder::with_storage_factory(self, storage_factory: alloc::sync::Arc) -> Self +pub struct iceberg_catalog_rest::RestContentFile +pub iceberg_catalog_rest::RestContentFile::content: core::option::Option +pub iceberg_catalog_rest::RestContentFile::equality_ids: core::option::Option> +pub iceberg_catalog_rest::RestContentFile::file_format: alloc::string::String +pub iceberg_catalog_rest::RestContentFile::file_path: alloc::string::String +pub iceberg_catalog_rest::RestContentFile::file_size_in_bytes: u64 +pub iceberg_catalog_rest::RestContentFile::first_row_id: core::option::Option +pub iceberg_catalog_rest::RestContentFile::key_metadata: core::option::Option> +pub iceberg_catalog_rest::RestContentFile::partition: core::option::Option +pub iceberg_catalog_rest::RestContentFile::record_count: core::option::Option +pub iceberg_catalog_rest::RestContentFile::spec_id: i32 +impl core::clone::Clone for iceberg_catalog_rest::RestContentFile +pub fn iceberg_catalog_rest::RestContentFile::clone(&self) -> iceberg_catalog_rest::RestContentFile +impl core::cmp::Eq for iceberg_catalog_rest::RestContentFile +impl core::cmp::PartialEq for iceberg_catalog_rest::RestContentFile +pub fn iceberg_catalog_rest::RestContentFile::eq(&self, other: &iceberg_catalog_rest::RestContentFile) -> bool +impl core::fmt::Debug for iceberg_catalog_rest::RestContentFile +pub fn iceberg_catalog_rest::RestContentFile::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl core::marker::StructuralPartialEq for iceberg_catalog_rest::RestContentFile +impl serde_core::ser::Serialize for iceberg_catalog_rest::RestContentFile +pub fn iceberg_catalog_rest::RestContentFile::serialize<__S>(&self, __serializer: __S) -> core::result::Result<<__S as serde_core::ser::Serializer>::Ok, <__S as serde_core::ser::Serializer>::Error> where __S: serde_core::ser::Serializer +impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::RestContentFile +pub fn iceberg_catalog_rest::RestContentFile::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub struct iceberg_catalog_rest::RestFileScanTask -pub iceberg_catalog_rest::RestFileScanTask::data_file: serde_json::value::Value +pub iceberg_catalog_rest::RestFileScanTask::data_file: iceberg_catalog_rest::RestContentFile pub iceberg_catalog_rest::RestFileScanTask::delete_file_references: core::option::Option> pub iceberg_catalog_rest::RestFileScanTask::residual_filter: core::option::Option impl core::clone::Clone for iceberg_catalog_rest::RestFileScanTask @@ -432,7 +455,7 @@ pub fn iceberg_catalog_rest::RestFileScanTask::serialize<__S>(&self, __serialize impl<'de> serde_core::de::Deserialize<'de> for iceberg_catalog_rest::RestFileScanTask pub fn iceberg_catalog_rest::RestFileScanTask::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> pub struct iceberg_catalog_rest::ScanTasks -pub iceberg_catalog_rest::ScanTasks::delete_files: alloc::vec::Vec +pub iceberg_catalog_rest::ScanTasks::delete_files: alloc::vec::Vec pub iceberg_catalog_rest::ScanTasks::file_scan_tasks: alloc::vec::Vec pub iceberg_catalog_rest::ScanTasks::plan_tasks: alloc::vec::Vec impl core::clone::Clone for iceberg_catalog_rest::ScanTasks diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 8c70b85262..7a16e8c456 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -494,7 +494,7 @@ pub struct RestCatalog { /// /// It could be different from the config fetched from the server and used at runtime. user_config: RestCatalogConfig, - client: OnceCell, + client: Arc>, /// Storage factory for creating FileIO instances. storage_factory: Option>, runtime: Runtime, @@ -514,7 +514,7 @@ impl RestCatalog { Self { auth_manager, user_config: config, - client: OnceCell::new(), + client: Arc::new(OnceCell::new()), storage_factory, runtime, kms_client, @@ -533,6 +533,20 @@ impl RestCatalog { ) } + /// Shares the already-initialized HTTP client. Used when injecting this + /// catalog as a [`iceberg::scan::ScanPlanner`] so Auto/capability checks + /// do not repeat `GET /v1/config`. + pub(crate) fn clone_initialized(&self) -> Self { + Self { + auth_manager: self.auth_manager.clone(), + user_config: self.user_config.clone(), + client: Arc::clone(&self.client), + storage_factory: self.storage_factory.clone(), + runtime: self.runtime.clone(), + kms_client: self.kms_client.clone(), + } + } + pub(crate) fn runtime(&self) -> &Runtime { &self.runtime } @@ -1083,7 +1097,7 @@ impl Catalog for RestCatalog { .file_io(file_io) .metadata(response.metadata) .runtime(self.runtime.clone()) - .scan_planner(Arc::new(self.clone_uninitialized())); + .scan_planner(Arc::new(self.clone_initialized())); if let Some(kms_client) = self.kms_client.clone() { table_builder = table_builder.kms_client(kms_client); } @@ -1145,7 +1159,7 @@ impl Catalog for RestCatalog { .file_io(file_io) .metadata(response.metadata) .runtime(self.runtime.clone()) - .scan_planner(Arc::new(self.clone_uninitialized())); + .scan_planner(Arc::new(self.clone_initialized())); if let Some(kms_client) = self.kms_client.clone() { table_builder = table_builder.kms_client(kms_client); } @@ -1283,7 +1297,7 @@ impl Catalog for RestCatalog { .metadata(response.metadata) .metadata_location(metadata_location.clone()) .runtime(self.runtime.clone()) - .scan_planner(Arc::new(self.clone_uninitialized())); + .scan_planner(Arc::new(self.clone_initialized())); if let Some(kms_client) = self.kms_client.clone() { table_builder = table_builder.kms_client(kms_client); } @@ -1361,7 +1375,7 @@ impl Catalog for RestCatalog { .metadata(response.metadata) .metadata_location(response.metadata_location) .runtime(self.runtime.clone()) - .scan_planner(Arc::new(self.clone_uninitialized())); + .scan_planner(Arc::new(self.clone_initialized())); if let Some(kms_client) = self.kms_client.clone() { table_builder = table_builder.kms_client(kms_client); } diff --git a/crates/catalog/rest/src/scan_decode.rs b/crates/catalog/rest/src/scan_decode.rs index 738ccb388a..08d5813c49 100644 --- a/crates/catalog/rest/src/scan_decode.rs +++ b/crates/catalog/rest/src/scan_decode.rs @@ -27,10 +27,9 @@ use iceberg::spec::{ DataContentType, DataFileFormat, Literal, NameMapping, SchemaRef, Struct, TableMetadataRef, }; use iceberg::{Error, ErrorKind, Result}; -use serde::Deserialize; use serde_json::Value; -use crate::scan_planning::RestFileScanTask; +use crate::scan_planning::{RestContentFile, RestFileScanTask}; /// Per-scan context needed to materialize tasks. pub(crate) struct ConvertContext { @@ -43,32 +42,10 @@ pub(crate) struct ConvertContext { pub(crate) unified_partition_type: Option>, } -#[derive(Debug, Clone, Deserialize)] -#[serde(rename_all = "kebab-case")] -struct RestContentFile { - file_path: String, - file_format: String, - file_size_in_bytes: u64, - #[serde(default)] - record_count: Option, - #[serde(default)] - content: Option, - #[serde(default)] - spec_id: i32, - #[serde(default)] - partition: Option, - #[serde(default)] - equality_ids: Option>, - #[serde(default)] - key_metadata: Option, - #[serde(default)] - first_row_id: Option, -} - -/// Decode expanded REST file-scan-tasks and delete-files into domain tasks. +/// Decode one ScanTasks payload's file-scan-tasks against that payload's delete-files. pub(crate) fn decode_scan_tasks( files: Vec, - delete_files: Vec, + delete_files: Vec, ctx: &ConvertContext, ) -> Result> { if files.is_empty() && delete_files.is_empty() { @@ -76,7 +53,7 @@ pub(crate) fn decode_scan_tasks( } let deletes = delete_files - .iter() + .into_iter() .map(to_delete_file) .collect::>>()?; @@ -86,30 +63,27 @@ pub(crate) fn decode_scan_tasks( .collect() } -fn parse_content_file(value: &Value) -> Result { - serde_json::from_value(value.clone()).map_err(|e| { - Error::new( - ErrorKind::DataInvalid, - "failed to decode REST content-file JSON", - ) - .with_source(e) - }) -} - -fn to_delete_file(value: &Value) -> Result { - let rcf = parse_content_file(value)?; - let content = rcf - .content - .map(DataContentType::try_from) - .transpose()? - .unwrap_or(DataContentType::PositionDeletes); +fn to_delete_file(rcf: RestContentFile) -> Result { + let content = match rcf.content { + Some(DataContentType::PositionDeletes) => DataContentType::PositionDeletes, + Some(DataContentType::EqualityDeletes) => DataContentType::EqualityDeletes, + Some(DataContentType::Data) | None => { + return Err(Error::new( + ErrorKind::DataInvalid, + format!( + "delete file {} is missing a valid content type", + rcf.file_path + ), + )); + } + }; Ok(FileScanTaskDeleteFile::builder() .with_file_path(rcf.file_path) .with_file_size_in_bytes(rcf.file_size_in_bytes) .with_file_type(content) .with_partition_spec_id(rcf.spec_id) .with_equality_ids(rcf.equality_ids) - .with_key_metadata(decode_key_metadata(rcf.key_metadata.as_ref())) + .with_key_metadata(rcf.key_metadata) .build()) } @@ -118,7 +92,7 @@ fn to_file_scan_task( all_deletes: &[FileScanTaskDeleteFile], ctx: &ConvertContext, ) -> Result { - let rcf = parse_content_file(&task.data_file)?; + let rcf = task.data_file; let spec = ctx .metadata .partition_spec_by_id(rcf.spec_id) @@ -183,7 +157,7 @@ fn to_file_scan_task( .with_name_mapping(ctx.name_mapping.clone()) .with_unified_partition_type(ctx.unified_partition_type.clone()) .with_case_sensitive(ctx.case_sensitive) - .with_key_metadata(decode_key_metadata(rcf.key_metadata.as_ref())) + .with_key_metadata(rcf.key_metadata) .build()) } @@ -242,23 +216,6 @@ fn decode_partition( Ok(Struct::from_iter(literals)) } -fn decode_key_metadata(value: Option<&Value>) -> Option> { - match value { - Some(Value::Array(values)) => { - let bytes: Vec = values - .iter() - .filter_map(|v| v.as_u64().map(|n| n as u8)) - .collect(); - if bytes.is_empty() { - None - } else { - Some(bytes.into_boxed_slice()) - } - } - _ => None, - } -} - #[cfg(test)] mod tests { use std::sync::Arc; @@ -267,7 +224,23 @@ mod tests { use serde_json::json; use super::*; - use crate::scan_planning::RestFileScanTask; + use crate::scan_planning::{RestContentFile, RestFileScanTask}; + + fn content_file(value: Value) -> RestContentFile { + serde_json::from_value(value).unwrap() + } + + fn delete_files(values: Vec) -> Vec { + values.into_iter().map(content_file).collect() + } + + fn file_task(data_file: Value) -> RestFileScanTask { + RestFileScanTask { + data_file: content_file(data_file), + delete_file_references: None, + residual_filter: None, + } + } fn unpartitioned_metadata() -> TableMetadataRef { let json = r#"{ @@ -318,19 +291,15 @@ mod tests { #[test] fn data_file_json_becomes_file_scan_task() { - let file = RestFileScanTask { - data_file: json!({ - "content": 0, - "file-path": "s3://bucket/f.parquet", - "file-format": "PARQUET", - "spec-id": 0, - "partition": {}, - "record-count": 7, - "file-size-in-bytes": 128 - }), - delete_file_references: None, - residual_filter: None, - }; + let file = file_task(json!({ + "content": 0, + "file-path": "s3://bucket/f.parquet", + "file-format": "PARQUET", + "spec-id": 0, + "partition": {}, + "record-count": 7, + "file-size-in-bytes": 128 + })); let tasks = decode_scan_tasks(vec![file], vec![], &ctx()).unwrap(); assert_eq!(tasks.len(), 1); assert_eq!(tasks[0].data_file_path, "s3://bucket/f.parquet"); @@ -342,21 +311,18 @@ mod tests { #[test] fn delete_file_references_are_resolved() { - let file = RestFileScanTask { - data_file: json!({ - "file-path": "s3://bucket/f.parquet", - "file-format": "parquet", - "file-size-in-bytes": 10 - }), - delete_file_references: Some(vec![0]), - residual_filter: None, - }; - let deletes = vec![json!({ + let mut file = file_task(json!({ + "file-path": "s3://bucket/f.parquet", + "file-format": "parquet", + "file-size-in-bytes": 10 + })); + file.delete_file_references = Some(vec![0]); + let deletes = delete_files(vec![json!({ "content": 1, "file-path": "s3://bucket/d.parquet", "file-format": "parquet", "file-size-in-bytes": 3 - })]; + })]); let tasks = decode_scan_tasks(vec![file], deletes, &ctx()).unwrap(); assert_eq!(tasks[0].deletes.len(), 1); assert_eq!(tasks[0].deletes[0].file_path, "s3://bucket/d.parquet"); @@ -368,17 +334,112 @@ mod tests { #[test] fn unknown_spec_id_is_data_invalid() { - let file = RestFileScanTask { - data_file: json!({ - "file-path": "s3://bucket/f.parquet", - "file-format": "parquet", - "file-size-in-bytes": 10, - "spec-id": 99 - }), - delete_file_references: None, - residual_filter: None, - }; + let file = file_task(json!({ + "file-path": "s3://bucket/f.parquet", + "file-format": "parquet", + "file-size-in-bytes": 10, + "spec-id": 99 + })); let err = decode_scan_tasks(vec![file], vec![], &ctx()).unwrap_err(); assert_eq!(err.kind(), ErrorKind::DataInvalid); } + + #[test] + fn content_string_data_decodes() { + let file = file_task(json!({ + "content": "data", + "file-path": "s3://bucket/f.parquet", + "file-format": "PARQUET", + "spec-id": 0, + "file-size-in-bytes": 128 + })); + let tasks = decode_scan_tasks(vec![file], vec![], &ctx()).unwrap(); + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0].data_file_path, "s3://bucket/f.parquet"); + } + + #[test] + fn content_string_position_deletes_decodes() { + let mut file = file_task(json!({ + "file-path": "s3://bucket/f.parquet", + "file-format": "parquet", + "file-size-in-bytes": 10 + })); + file.delete_file_references = Some(vec![0]); + let deletes = delete_files(vec![json!({ + "content": "position-deletes", + "file-path": "s3://bucket/d.parquet", + "file-format": "parquet", + "file-size-in-bytes": 3 + })]); + let tasks = decode_scan_tasks(vec![file], deletes, &ctx()).unwrap(); + assert_eq!(tasks[0].deletes.len(), 1); + assert_eq!(tasks[0].deletes[0].file_path, "s3://bucket/d.parquet"); + assert_eq!( + tasks[0].deletes[0].file_type, + DataContentType::PositionDeletes + ); + } + + #[test] + fn content_integer_ordinals_still_decode() { + let mut file = file_task(json!({ + "content": 0, + "file-path": "s3://bucket/f.parquet", + "file-format": "parquet", + "file-size-in-bytes": 10 + })); + file.delete_file_references = Some(vec![0]); + let deletes = delete_files(vec![json!({ + "content": 1, + "file-path": "s3://bucket/d.parquet", + "file-format": "parquet", + "file-size-in-bytes": 3 + })]); + let tasks = decode_scan_tasks(vec![file], deletes, &ctx()).unwrap(); + assert_eq!( + tasks[0].deletes[0].file_type, + DataContentType::PositionDeletes + ); + } + + #[test] + fn delete_file_without_content_is_data_invalid() { + let mut file = file_task(json!({ + "file-path": "s3://bucket/f.parquet", + "file-format": "parquet", + "file-size-in-bytes": 10 + })); + file.delete_file_references = Some(vec![0]); + let deletes = delete_files(vec![json!({ + "file-path": "s3://bucket/d.parquet", + "file-format": "parquet", + "file-size-in-bytes": 3 + })]); + let err = decode_scan_tasks(vec![file], deletes, &ctx()).unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DataInvalid); + } + + #[test] + fn hex_key_metadata_decodes_to_bytes() { + let file = file_task(json!({ + "file-path": "s3://bucket/f.parquet", + "file-format": "parquet", + "file-size-in-bytes": 10, + "key-metadata": "00000000000000000000000000000000" + })); + let tasks = decode_scan_tasks(vec![file], vec![], &ctx()).unwrap(); + assert_eq!(tasks[0].key_metadata.as_deref(), Some([0u8; 16].as_slice())); + } + + #[test] + fn non_hex_key_metadata_is_rejected() { + let err = serde_json::from_value::(json!({ + "file-path": "s3://bucket/f.parquet", + "file-format": "parquet", + "file-size-in-bytes": 10, + "key-metadata": "not-hex" + })); + assert!(err.is_err()); + } } diff --git a/crates/catalog/rest/src/scan_planning.rs b/crates/catalog/rest/src/scan_planning.rs index 57223619eb..7c218d2d19 100644 --- a/crates/catalog/rest/src/scan_planning.rs +++ b/crates/catalog/rest/src/scan_planning.rs @@ -26,12 +26,13 @@ use std::collections::{HashSet, VecDeque}; use std::time::Duration; use async_trait::async_trait; -use iceberg::scan::{ScanPlanner, ScanPlanningRequest, ScanPlanningResult}; +use iceberg::scan::{FileScanTask, ScanPlanner, ScanPlanningRequest, ScanPlanningResult}; +use iceberg::spec::DataContentType; use iceberg::{Error, ErrorKind, Result, TableIdent}; use rand::Rng; use reqwest::{Method, Response, StatusCode}; use serde::de::{self, Deserializer}; -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Serialize, Serializer}; use uuid::{Uuid, Variant, Version}; use crate::catalog::RestCatalog; @@ -77,20 +78,63 @@ pub struct ScanTasks { /// Opaque plan-task handles that still need [`RestCatalog::fetch_scan_tasks`]. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub plan_tasks: Vec, - /// File scan tasks. `data-file` is left as JSON until a decoder lands. + /// File scan tasks. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub file_scan_tasks: Vec, - /// Delete files referenced by the scan tasks, as raw REST JSON. + /// Delete files referenced by indices in this payload's `file-scan-tasks`. #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub delete_files: Vec, + pub delete_files: Vec, } -/// REST `FileScanTask` wire payload. Nested content-files stay opaque. +/// REST ContentFile JSON. Unknown fields such as metrics maps are ignored. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub struct RestContentFile { + /// Absolute file path. + pub file_path: String, + /// File format string (`PARQUET`, `AVRO`, ...). + pub file_format: String, + /// File size in bytes. + pub file_size_in_bytes: u64, + /// Record count when the server provided it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub record_count: Option, + /// `data`, `position-deletes`, or `equality-deletes`. Avro ordinals 0/1/2 are also accepted. + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "deserialize_content", + serialize_with = "serialize_content" + )] + pub content: Option, + /// Partition spec id. Defaults to 0. + #[serde(default)] + pub spec_id: i32, + /// Partition values as a JSON array or object keyed by field id. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub partition: Option, + /// Equality field ids for equality deletes. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub equality_ids: Option>, + /// Encryption key metadata. OpenAPI hex text, or a JSON array of 0..=255 integers. + #[serde( + default, + skip_serializing_if = "Option::is_none", + deserialize_with = "deserialize_key_metadata", + serialize_with = "serialize_key_metadata" + )] + pub key_metadata: Option>, + /// First row id when the server provided it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub first_row_id: Option, +} + +/// REST `FileScanTask` wire payload. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "kebab-case")] pub struct RestFileScanTask { - /// REST ContentFile JSON for the data file. - pub data_file: serde_json::Value, + /// REST ContentFile for the data file. + pub data_file: RestContentFile, /// Indices into the sibling delete-files array. #[serde(default, skip_serializing_if = "Option::is_none")] pub delete_file_references: Option>, @@ -333,6 +377,135 @@ pub fn is_plan_poll_exhausted(err: &Error) -> bool { err.message() == MSG_PLAN_POLL_EXHAUSTED } +fn parse_content_type(value: &serde_json::Value) -> std::result::Result { + match value { + serde_json::Value::String(s) => match s.as_str() { + "data" => Ok(DataContentType::Data), + "position-deletes" => Ok(DataContentType::PositionDeletes), + "equality-deletes" => Ok(DataContentType::EqualityDeletes), + other => Err(format!("unknown REST content type {other:?}")), + }, + serde_json::Value::Number(n) => { + let ordinal = n + .as_i64() + .ok_or_else(|| format!("content ordinal is not an integer: {n}"))?; + let ordinal = i32::try_from(ordinal) + .map_err(|_| format!("content ordinal out of range: {ordinal}"))?; + DataContentType::try_from(ordinal).map_err(|e| e.to_string()) + } + other => Err(format!("unexpected REST content encoding: {other}")), + } +} + +fn deserialize_content<'de, D>( + deserializer: D, +) -> std::result::Result, D::Error> +where D: Deserializer<'de> { + match serde_json::Value::deserialize(deserializer)? { + serde_json::Value::Null => Ok(None), + value => parse_content_type(&value) + .map(Some) + .map_err(de::Error::custom), + } +} + +fn serialize_content( + value: &Option, + serializer: S, +) -> std::result::Result +where + S: Serializer, +{ + match value { + None => serializer.serialize_none(), + Some(DataContentType::Data) => serializer.serialize_str("data"), + Some(DataContentType::PositionDeletes) => serializer.serialize_str("position-deletes"), + Some(DataContentType::EqualityDeletes) => serializer.serialize_str("equality-deletes"), + } +} + +fn decode_hex_key_metadata(value: &str) -> std::result::Result, String> { + if !value.len().is_multiple_of(2) { + return Err(format!( + "key-metadata hex string must have an even number of characters: {value:?}" + )); + } + value + .as_bytes() + .chunks_exact(2) + .map(|chunk| { + let high = decode_hex_digit(chunk[0], value)?; + let low = decode_hex_digit(chunk[1], value)?; + Ok((high << 4) | low) + }) + .collect() +} + +fn decode_hex_digit(digit: u8, value: &str) -> std::result::Result { + match digit { + b'0'..=b'9' => Ok(digit - b'0'), + b'a'..=b'f' => Ok(digit - b'a' + 10), + b'A'..=b'F' => Ok(digit - b'A' + 10), + _ => Err(format!("key-metadata is not valid hex: {value:?}")), + } +} + +fn deserialize_key_metadata<'de, D>( + deserializer: D, +) -> std::result::Result>, D::Error> +where D: Deserializer<'de> { + match serde_json::Value::deserialize(deserializer)? { + serde_json::Value::Null => Ok(None), + serde_json::Value::String(s) => { + if s.is_empty() { + return Ok(None); + } + decode_hex_key_metadata(&s) + .map(|bytes| Some(bytes.into_boxed_slice())) + .map_err(de::Error::custom) + } + serde_json::Value::Array(values) => { + if values.is_empty() { + return Ok(None); + } + let mut bytes = Vec::with_capacity(values.len()); + for value in values { + let n = value.as_u64().ok_or_else(|| { + de::Error::custom("key-metadata array element is not an integer") + })?; + let byte = u8::try_from(n) + .map_err(|_| de::Error::custom("key-metadata array element is out of range"))?; + bytes.push(byte); + } + Ok(Some(bytes.into_boxed_slice())) + } + other => Err(de::Error::custom(format!( + "unexpected key-metadata encoding: {other}" + ))), + } +} + +fn serialize_key_metadata( + value: &Option>, + serializer: S, +) -> std::result::Result +where + S: Serializer, +{ + match value { + None => serializer.serialize_none(), + Some(bytes) => { + const HEX: &[u8; 16] = b"0123456789ABCDEF"; + let mut out = String::with_capacity(bytes.len() * 2); + for byte in bytes.iter() { + out.push(HEX[(byte >> 4) as usize] as char); + out.push(HEX[(byte & 0x0f) as usize] as char); + } + serializer.serialize_str(&out) + } + } +} + impl RestCatalog { /// Whether the server advertised the synchronous plan endpoint. pub async fn supports_plan_table_scan(&self) -> Result { @@ -663,9 +836,10 @@ impl RestCatalog { &self, table: &TableIdent, tasks: ScanTasks, - ) -> Result<(Vec, Vec)> { - let mut files = tasks.file_scan_tasks; - let mut deletes = tasks.delete_files; + ctx: &crate::scan_decode::ConvertContext, + ) -> Result> { + let mut out = + crate::scan_decode::decode_scan_tasks(tasks.file_scan_tasks, tasks.delete_files, ctx)?; let mut queue: VecDeque = tasks.plan_tasks.into(); let mut seen = HashSet::new(); while let Some(handle) = queue.pop_front() { @@ -678,11 +852,14 @@ impl RestCatalog { plan_task: handle, }) .await?; - files.extend(resp.scan_tasks.file_scan_tasks); - deletes.extend(resp.scan_tasks.delete_files); + out.extend(crate::scan_decode::decode_scan_tasks( + resp.scan_tasks.file_scan_tasks, + resp.scan_tasks.delete_files, + ctx, + )?); queue.extend(resp.scan_tasks.plan_tasks); } - Ok((files, deletes)) + Ok(out) } async fn plan_scan(&self, request: ScanPlanningRequest) -> Result { @@ -718,9 +895,6 @@ impl RestCatalog { } }; - let (files, deletes) = self - .collect_scan_tasks(&ident, completed.scan_tasks) - .await?; let ctx = crate::scan_decode::ConvertContext { metadata: request.metadata, snapshot_schema: request.snapshot_schema, @@ -731,7 +905,9 @@ impl RestCatalog { unified_partition_type: request.unified_partition_type, }; Ok(ScanPlanningResult { - tasks: crate::scan_decode::decode_scan_tasks(files, deletes, &ctx)?, + tasks: self + .collect_scan_tasks(&ident, completed.scan_tasks, &ctx) + .await?, }) } } @@ -1086,13 +1262,17 @@ mod tests { "status": "completed", "plan-id": "p1", "plan-tasks": ["t1"], - "file-scan-tasks": [{"data-file": {"file-path": "s3://b/f.parquet"}}], + "file-scan-tasks": [{"data-file": { + "file-path": "s3://b/f.parquet", + "file-format": "PARQUET", + "file-size-in-bytes": 1 + }}], "storage-credentials": [{"prefix": "s3://b/", "config": {"s3.access-key-id": "k"}}] }); let parsed: PlanTableScanResponse = serde_json::from_value(body.clone()).unwrap(); assert_eq!(parsed.scan_tasks.plan_tasks, ["t1"]); assert_eq!( - parsed.scan_tasks.file_scan_tasks[0].data_file["file-path"], + parsed.scan_tasks.file_scan_tasks[0].data_file.file_path, "s3://b/f.parquet" ); assert_eq!( @@ -1374,7 +1554,7 @@ mod tests { .match_body(r#"{"plan-task":"h1"}"#) .with_status(200) .with_body( - r#"{"plan-tasks":["h2"],"file-scan-tasks":[{"data-file":{"file-path":"f"}}]}"#, + r#"{"plan-tasks":["h2"],"file-scan-tasks":[{"data-file":{"file-path":"f","file-format":"PARQUET","file-size-in-bytes":1}}]}"#, ) .create_async() .await; @@ -1880,6 +2060,27 @@ mod tests { .await } + async fn mock_load_table_with_config( + server: &mut mockito::ServerGuard, + config: serde_json::Value, + ) -> mockito::Mock { + let mut body: serde_json::Value = serde_json::from_str( + &std::fs::read_to_string(format!( + "{}/testdata/load_table_response.json", + env!("CARGO_MANIFEST_DIR") + )) + .unwrap(), + ) + .unwrap(); + body["config"] = config; + server + .mock("GET", "/v1/namespaces/ns1/tables/test1") + .with_status(200) + .with_body(body.to_string()) + .create_async() + .await + } + fn empty_table_body() -> String { json!({ "metadata-location": "s3://warehouse/ns1/test1/metadata.json", @@ -1934,6 +2135,7 @@ mod tests { let table = catalog.load_table(&test1()).await.unwrap(); let tasks: Vec<_> = table .scan() + .with_scan_planning_mode(ScanPlanningMode::Remote) .build() .unwrap() .plan_files() @@ -1993,6 +2195,7 @@ mod tests { let table = catalog.load_table(&test1()).await.unwrap(); let tasks: Vec<_> = table .scan() + .with_scan_planning_mode(ScanPlanningMode::Remote) .build() .unwrap() .plan_files() @@ -2039,6 +2242,7 @@ mod tests { let table = catalog.load_table(&test1()).await.unwrap(); let tasks: Vec<_> = table .scan() + .with_scan_planning_mode(ScanPlanningMode::Remote) .build() .unwrap() .plan_files() @@ -2081,6 +2285,7 @@ mod tests { let table = catalog.load_table(&test1()).await.unwrap(); let tasks: Vec<_> = table .scan() + .with_scan_planning_mode(ScanPlanningMode::Remote) .build() .unwrap() .plan_files() @@ -2112,6 +2317,7 @@ mod tests { let table = catalog.load_table(&test1()).await.unwrap(); let tasks: Vec<_> = table .scan() + .with_scan_planning_mode(ScanPlanningMode::Remote) .build() .unwrap() .plan_files() @@ -2186,4 +2392,188 @@ mod tests { config.assert_async().await; load.assert_async().await; } + + #[tokio::test] + async fn default_plan_files_does_not_post_plan_when_endpoints_are_advertised() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let load = mock_load_table(&mut server).await; + let plan = server + .mock("POST", "/v1/namespaces/ns1/tables/test1/plan") + .expect(0) + .create_async() + .await; + let catalog = load_catalog(&server.url()); + let table = catalog.load_table(&test1()).await.unwrap(); + let _ = table.scan().build().unwrap().plan_files().await; + config.assert_async().await; + load.assert_async().await; + plan.assert_async().await; + } + + #[tokio::test] + async fn auto_mode_with_client_config_does_not_post_plan() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let load = + mock_load_table_with_config(&mut server, json!({"scan-planning-mode": "client"})).await; + let plan = server + .mock("POST", "/v1/namespaces/ns1/tables/test1/plan") + .expect(0) + .create_async() + .await; + let catalog = load_catalog(&server.url()); + let table = catalog.load_table(&test1()).await.unwrap(); + let _ = table + .scan() + .with_scan_planning_mode(ScanPlanningMode::Auto) + .build() + .unwrap() + .plan_files() + .await; + config.assert_async().await; + load.assert_async().await; + plan.assert_async().await; + } + + #[tokio::test] + async fn auto_mode_with_server_config_posts_plan() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let load = + mock_load_table_with_config(&mut server, json!({"scan-planning-mode": "server"})).await; + let plan = server + .mock("POST", "/v1/namespaces/ns1/tables/test1/plan") + .with_status(200) + .with_body( + json!({ + "status": "completed", + "plan-id": "p1", + "file-scan-tasks": [{"data-file": data_file_json()}] + }) + .to_string(), + ) + .expect(1) + .create_async() + .await; + let catalog = load_catalog(&server.url()); + let table = catalog.load_table(&test1()).await.unwrap(); + let tasks: Vec<_> = table + .scan() + .with_scan_planning_mode(ScanPlanningMode::Auto) + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(tasks.len(), 1); + config.assert_async().await; + load.assert_async().await; + plan.assert_async().await; + } + + #[tokio::test] + async fn delete_file_references_are_resolved_per_payload() { + let mut server = Server::new_async().await; + let config = config_with_endpoints(&mut server, ALL_PLAN).await; + let load = mock_load_table(&mut server).await; + let plan = server + .mock("POST", "/v1/namespaces/ns1/tables/test1/plan") + .with_status(200) + .with_body( + json!({ + "status": "completed", + "plan-id": "p1", + "plan-tasks": ["h1"], + "delete-files": [{ + "content": "position-deletes", + "file-path": "s3://warehouse/decoy.parquet", + "file-format": "PARQUET", + "spec-id": 0, + "file-size-in-bytes": 3 + }] + }) + .to_string(), + ) + .expect(1) + .create_async() + .await; + let tasks_ep = server + .mock("POST", "/v1/namespaces/ns1/tables/test1/tasks") + .match_body(r#"{"plan-task":"h1"}"#) + .with_status(200) + .with_body( + json!({ + "file-scan-tasks": [{ + "data-file": data_file_json(), + "delete-file-references": [0] + }], + "delete-files": [{ + "content": "position-deletes", + "file-path": "s3://warehouse/d1.parquet", + "file-format": "PARQUET", + "spec-id": 0, + "file-size-in-bytes": 3 + }] + }) + .to_string(), + ) + .expect(1) + .create_async() + .await; + let catalog = load_catalog(&server.url()); + let table = catalog.load_table(&test1()).await.unwrap(); + let tasks: Vec<_> = table + .scan() + .with_scan_planning_mode(ScanPlanningMode::Remote) + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(tasks.len(), 1); + assert_eq!(tasks[0].deletes.len(), 1); + assert_eq!(tasks[0].deletes[0].file_path, "s3://warehouse/d1.parquet"); + config.assert_async().await; + load.assert_async().await; + plan.assert_async().await; + tasks_ep.assert_async().await; + } + + #[tokio::test] + async fn load_table_then_auto_plan_hits_config_once() { + let mut server = Server::new_async().await; + let config = server + .mock("GET", "/v1/config") + .with_status(200) + .with_body( + json!({ + "overrides": {}, + "defaults": {}, + "endpoints": [] + }) + .to_string(), + ) + .expect(1) + .create_async() + .await; + let load = mock_load_table(&mut server).await; + let catalog = load_catalog(&server.url()); + let table = catalog.load_table(&test1()).await.unwrap(); + let _ = table + .scan() + .with_scan_planning_mode(ScanPlanningMode::Auto) + .build() + .unwrap() + .plan_files() + .await; + config.assert_async().await; + load.assert_async().await; + } } diff --git a/crates/iceberg/src/scan/mod.rs b/crates/iceberg/src/scan/mod.rs index e389ca50e2..68114a4f6d 100644 --- a/crates/iceberg/src/scan/mod.rs +++ b/crates/iceberg/src/scan/mod.rs @@ -85,7 +85,7 @@ impl<'a> TableScanBuilder<'a> { concurrency_limit_manifest_files: num_cpus, row_group_filtering_enabled: true, row_selection_enabled: false, - planning_mode: ScanPlanningMode::Auto, + planning_mode: ScanPlanningMode::Local, } } @@ -194,8 +194,10 @@ impl<'a> TableScanBuilder<'a> { /// Sets whether this scan plans locally, remotely, or automatically. /// - /// Defaults to [`ScanPlanningMode::Auto`]: remote planning when the table's - /// [`ScanPlanner`] advertises it, otherwise local manifest planning. + /// Defaults to [`ScanPlanningMode::Local`]. [`ScanPlanningMode::Auto`] + /// uses remote planning only when the planner advertises it and table + /// config is `scan-planning-mode=server` (case-insensitive). A missing key + /// is treated as client. pub fn with_scan_planning_mode(mut self, mode: ScanPlanningMode) -> Self { self.planning_mode = mode; self @@ -403,9 +405,17 @@ impl TableScan { pub async fn plan_files(&self) -> Result { match self.planning_mode { ScanPlanningMode::Local => self.plan_files_local().await, - ScanPlanningMode::Remote => self.plan_files_remote().await, + ScanPlanningMode::Remote => { + if self.table_requests_client_planning() { + return Err(Error::new( + ErrorKind::FeatureUnsupported, + "remote scan planning is disabled by scan-planning-mode=client", + )); + } + self.plan_files_remote().await + } ScanPlanningMode::Auto => { - if self.remote_planning_available().await? { + if self.remote_planning_available().await? && self.table_allows_server_planning() { match self.plan_files_remote().await { Ok(tasks) => Ok(tasks), Err(e) if e.kind() == ErrorKind::FeatureUnsupported => { @@ -420,6 +430,26 @@ impl TableScan { } } + fn scan_planning_mode_property(&self) -> Option<&str> { + let from_meta = self + .plan_context + .as_ref() + .and_then(|ctx| ctx.table_metadata.properties().get("scan-planning-mode")); + from_meta + .or_else(|| self.file_io.config().get("scan-planning-mode")) + .map(String::as_str) + } + + fn table_requests_client_planning(&self) -> bool { + self.scan_planning_mode_property() + .is_some_and(|value| value.eq_ignore_ascii_case("client")) + } + + fn table_allows_server_planning(&self) -> bool { + self.scan_planning_mode_property() + .is_some_and(|value| value.eq_ignore_ascii_case("server")) + } + async fn remote_planning_available(&self) -> Result { match &self.scan_planner { Some(planner) => planner.supports_remote_scan_planning().await, @@ -747,7 +777,7 @@ pub mod tests { use crate::arrow::ArrowReaderBuilder; use crate::expr::{BoundPredicate, Reference}; - use crate::io::{FileIO, OutputFile}; + use crate::io::{FileIO, FileIOBuilder, LocalFsStorageFactory, OutputFile}; use crate::metadata_columns::{ RESERVED_COL_NAME_DELETE_FILE_PATH, RESERVED_COL_NAME_DELETE_FILE_POS, RESERVED_COL_NAME_FILE, RESERVED_COL_NAME_LAST_UPDATED_SEQUENCE_NUMBER, @@ -3825,10 +3855,31 @@ pub mod tests { } fn table_with_planner(fixture: &TableTestFixture, planner: Arc) -> Table { + table_with_planner_and_mode(fixture, planner, None) + } + + fn table_with_server_planner( + fixture: &TableTestFixture, + planner: Arc, + ) -> Table { + table_with_planner_and_mode(fixture, planner, Some("server")) + } + + fn table_with_planner_and_mode( + fixture: &TableTestFixture, + planner: Arc, + scan_planning_mode: Option<&str>, + ) -> Table { + let file_io = match scan_planning_mode { + Some(mode) => FileIOBuilder::new(Arc::new(LocalFsStorageFactory)) + .with_prop("scan-planning-mode", mode) + .build(), + None => fixture.table.file_io().clone(), + }; Table::builder() .metadata(fixture.table.metadata_ref()) .identifier(fixture.table.identifier().clone()) - .file_io(fixture.table.file_io().clone()) + .file_io(file_io) .metadata_location(fixture.table.metadata_location().unwrap()) .runtime(fixture.table.runtime().clone()) .scan_planner(planner) @@ -3855,7 +3906,37 @@ pub mod tests { } #[tokio::test] - async fn auto_mode_uses_remote_tasks_when_planner_supports_it() { + async fn default_mode_is_local_even_with_a_capable_planner() { + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + let table = table_with_server_planner( + &fixture, + Arc::new(FakePlanner { + remote: true, + kind: None, + tasks: vec![remote_stub_task()], + }), + ); + let tasks: Vec<_> = table + .scan() + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(tasks.len(), 2); + assert!( + tasks + .iter() + .all(|t| t.data_file_path != "s3://bucket/remote.parquet") + ); + } + + #[tokio::test] + async fn auto_mode_without_server_property_plans_locally() { let mut fixture = TableTestFixture::new(); fixture.setup_manifest_files().await; let table = table_with_planner( @@ -3868,6 +3949,38 @@ pub mod tests { ); let tasks: Vec<_> = table .scan() + .with_scan_planning_mode(ScanPlanningMode::Auto) + .build() + .unwrap() + .plan_files() + .await + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(tasks.len(), 2); + assert!( + tasks + .iter() + .all(|t| t.data_file_path != "s3://bucket/remote.parquet") + ); + } + + #[tokio::test] + async fn auto_mode_uses_remote_tasks_when_planner_supports_it() { + let mut fixture = TableTestFixture::new(); + fixture.setup_manifest_files().await; + let table = table_with_server_planner( + &fixture, + Arc::new(FakePlanner { + remote: true, + kind: None, + tasks: vec![remote_stub_task()], + }), + ); + let tasks: Vec<_> = table + .scan() + .with_scan_planning_mode(ScanPlanningMode::Auto) .build() .unwrap() .plan_files() @@ -3885,7 +3998,7 @@ pub mod tests { async fn auto_mode_falls_back_to_local_on_feature_unsupported() { let mut fixture = TableTestFixture::new(); fixture.setup_manifest_files().await; - let table = table_with_planner( + let table = table_with_server_planner( &fixture, Arc::new(FakePlanner { remote: true, @@ -3895,6 +4008,7 @@ pub mod tests { ); let tasks: Vec<_> = table .scan() + .with_scan_planning_mode(ScanPlanningMode::Auto) .build() .unwrap() .plan_files() @@ -3910,7 +4024,7 @@ pub mod tests { async fn auto_mode_does_not_fall_back_on_other_remote_errors() { let mut fixture = TableTestFixture::new(); fixture.setup_manifest_files().await; - let table = table_with_planner( + let table = table_with_server_planner( &fixture, Arc::new(FakePlanner { remote: true, @@ -3918,7 +4032,14 @@ pub mod tests { tasks: vec![], }), ); - let err = match table.scan().build().unwrap().plan_files().await { + let err = match table + .scan() + .with_scan_planning_mode(ScanPlanningMode::Auto) + .build() + .unwrap() + .plan_files() + .await + { Ok(_) => panic!("expected remote planning to fail"), Err(e) => e, }; diff --git a/crates/iceberg/src/scan/planner.rs b/crates/iceberg/src/scan/planner.rs index a6ec6d8b92..2dff5ad996 100644 --- a/crates/iceberg/src/scan/planner.rs +++ b/crates/iceberg/src/scan/planner.rs @@ -36,13 +36,19 @@ use crate::{Result, TableIdent}; /// between local manifest planning and a catalog-provided [`ScanPlanner`]. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum ScanPlanningMode { - /// Use remote planning when the table's planner advertises it, otherwise - /// plan locally. This is the default. - #[default] + /// Use remote planning when the table's planner advertises it and the + /// table config is `scan-planning-mode=server`. Missing or `client` stays + /// local. Select this mode explicitly; the default is [`Self::Local`]. Auto, /// Always plan by reading manifests through the table's FileIO. + /// This is the default, matching Java `CLIENT` and Go `ScanPlanningLocal`. + #[default] Local, /// Require a planner that supports remote planning; error otherwise. + /// + /// An explicit `scan-planning-mode=client` table config still fails as + /// unsupported. A missing key does not; Remote is user opt-in and overrides + /// the client default. Remote, } From 6233b7c3405b28c60eb2cea153d9fbfebca2d975 Mon Sep 17 00:00:00 2001 From: David Dallakyan Date: Sun, 16 Aug 2026 19:15:33 +0400 Subject: [PATCH 6/6] fix(catalog-rest): send use-snapshot-schema for snapshot plans OpenAPI defaults the flag to false (current table schema). TableScan always binds the snapshot schema, so time travel would plan against the wrong schema without this flag. --- crates/catalog/rest/src/scan_planning.rs | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/crates/catalog/rest/src/scan_planning.rs b/crates/catalog/rest/src/scan_planning.rs index 7c218d2d19..eda1c059de 100644 --- a/crates/catalog/rest/src/scan_planning.rs +++ b/crates/catalog/rest/src/scan_planning.rs @@ -864,11 +864,16 @@ impl RestCatalog { async fn plan_scan(&self, request: ScanPlanningRequest) -> Result { let ident = request.table_ident.clone(); + // TableScan always binds the snapshot schema. OpenAPI defaults + // use-snapshot-schema to false (current table schema), which would + // plan time travel against the wrong schema. + let use_snapshot_schema = request.snapshot_id.map(|_| true); let resp = self .plan_table_scan(&ident, PlanTableScanRequest { snapshot_id: request.snapshot_id, select: request.select.clone().unwrap_or_default(), case_sensitive: Some(request.case_sensitive), + use_snapshot_schema, ..Default::default() }) .await?; @@ -2119,6 +2124,10 @@ mod tests { let load = mock_load_table(&mut server).await; let plan = server .mock("POST", "/v1/namespaces/ns1/tables/test1/plan") + .match_body(mockito::Matcher::PartialJson(json!({ + "snapshot-id": 3497810964824022504i64, + "use-snapshot-schema": true + }))) .with_status(200) .with_body( json!({