From ce8f335da5066459f83e92c52f97b838bfd70c4c Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Wed, 29 Jul 2026 00:13:51 +0200 Subject: [PATCH 1/6] Add SessionCatalog trait and SessionContext --- crates/iceberg/public-api.txt | 31 +++++ crates/iceberg/src/catalog/mod.rs | 2 + crates/iceberg/src/catalog/session.rs | 187 ++++++++++++++++++++++++++ 3 files changed, 220 insertions(+) create mode 100644 crates/iceberg/src/catalog/session.rs diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 610675c83d..80f2633b06 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -3658,6 +3658,16 @@ impl core::clone::Clone for iceberg::RuntimeHandle pub fn iceberg::RuntimeHandle::clone(&self) -> iceberg::RuntimeHandle impl core::fmt::Debug for iceberg::RuntimeHandle pub fn iceberg::RuntimeHandle::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +pub struct iceberg::SessionContext +impl iceberg::SessionContext +pub fn iceberg::SessionContext::empty() -> Self +pub fn iceberg::SessionContext::session_id(&self) -> &str +impl core::clone::Clone for iceberg::SessionContext +pub fn iceberg::SessionContext::clone(&self) -> iceberg::SessionContext +impl core::fmt::Debug for iceberg::SessionContext +pub fn iceberg::SessionContext::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl iceberg::SessionContext +pub fn iceberg::SessionContext::builder() -> SessionContextBuilder<(())> pub struct iceberg::TableCommit impl iceberg::TableCommit pub fn iceberg::TableCommit::apply(self, table: iceberg::table::Table) -> iceberg::Result @@ -3765,5 +3775,26 @@ pub fn iceberg::memory::MemoryCatalogBuilder::load(self, name: impl core::conver pub fn iceberg::memory::MemoryCatalogBuilder::with_kms_client_factory(self, kms_client_factory: alloc::sync::Arc) -> Self pub fn iceberg::memory::MemoryCatalogBuilder::with_runtime(self, runtime: iceberg::Runtime) -> Self pub fn iceberg::memory::MemoryCatalogBuilder::with_storage_factory(self, storage_factory: alloc::sync::Arc) -> Self +pub trait iceberg::SessionCatalog: core::fmt::Debug + core::marker::Send + core::marker::Sync +pub fn iceberg::SessionCatalog::create_namespace<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, namespace: &'life2 iceberg::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, 'life2: 'async_trait +pub fn iceberg::SessionCatalog::create_table<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, namespace: &'life2 iceberg::NamespaceIdent, creation: iceberg::TableCreation) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg::SessionCatalog::drop_namespace<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, namespace: &'life2 iceberg::NamespaceIdent) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg::SessionCatalog::drop_table<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, table: &'life2 iceberg::TableIdent) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg::SessionCatalog::get_namespace<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, namespace: &'life2 iceberg::NamespaceIdent) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg::SessionCatalog::list_namespaces<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, parent: core::option::Option<&'life2 iceberg::NamespaceIdent>) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg::SessionCatalog::list_tables<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, namespace: &'life2 iceberg::NamespaceIdent) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg::SessionCatalog::load_table<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, table: &'life2 iceberg::TableIdent) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg::SessionCatalog::namespace_exists<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, namespace: &'life2 iceberg::NamespaceIdent) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg::SessionCatalog::purge_table<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, table: &'life2 iceberg::TableIdent) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg::SessionCatalog::register_table<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, table: &'life2 iceberg::TableIdent, metadata_location: alloc::string::String) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg::SessionCatalog::rename_table<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, src: &'life2 iceberg::TableIdent, dest: &'life3 iceberg::TableIdent) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg::SessionCatalog::table_exists<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, table: &'life2 iceberg::TableIdent) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait +pub fn iceberg::SessionCatalog::update_namespace<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, namespace: &'life2 iceberg::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, 'life2: 'async_trait +pub fn iceberg::SessionCatalog::update_table<'life0, 'life1, 'async_trait>(&'life0 self, context: &'life1 iceberg::SessionContext, commit: iceberg::TableCommit) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait +pub trait iceberg::SessionCatalogBuilder: core::default::Default + core::fmt::Debug + core::marker::Send + core::marker::Sync +pub type iceberg::SessionCatalogBuilder::C: iceberg::SessionCatalog +pub fn iceberg::SessionCatalogBuilder::load(self, name: impl core::convert::Into, props: std::collections::hash::map::HashMap) -> impl core::future::future::Future> + core::marker::Send +pub fn iceberg::SessionCatalogBuilder::with_runtime(self, runtime: iceberg::Runtime) -> Self +pub fn iceberg::SessionCatalogBuilder::with_storage_factory(self, storage_factory: alloc::sync::Arc) -> Self pub async fn iceberg::drop_table_data(table_info: &iceberg::table::Table) -> iceberg::Result<()> pub type iceberg::Result = core::result::Result diff --git a/crates/iceberg/src/catalog/mod.rs b/crates/iceberg/src/catalog/mod.rs index 64cf2617e5..d2ecc7d51e 100644 --- a/crates/iceberg/src/catalog/mod.rs +++ b/crates/iceberg/src/catalog/mod.rs @@ -19,6 +19,7 @@ pub mod memory; mod metadata_location; +mod session; pub(crate) mod utils; use std::collections::HashMap; @@ -36,6 +37,7 @@ pub use metadata_location::*; #[cfg(test)] use mockall::automock; use serde_derive::{Deserialize, Serialize}; +pub use session::*; use typed_builder::TypedBuilder; use uuid::Uuid; diff --git a/crates/iceberg/src/catalog/session.rs b/crates/iceberg/src/catalog/session.rs new file mode 100644 index 0000000000..b13baa081b --- /dev/null +++ b/crates/iceberg/src/catalog/session.rs @@ -0,0 +1,187 @@ +// 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. + +//! Session catalog API for Apache Iceberg. + +use std::collections::HashMap; +use std::fmt::Debug; +use std::future::Future; +use std::sync::Arc; + +use async_trait::async_trait; +use typed_builder::TypedBuilder; +use uuid::Uuid; + +use crate::io::StorageFactory; +use crate::runtime::Runtime; +use crate::table::Table; +use crate::{Namespace, NamespaceIdent, Result, TableCommit, TableCreation, TableIdent}; + +/// Context for a session. +#[derive(Debug, Clone, TypedBuilder)] +pub struct SessionContext { + session_id: String, +} + +impl SessionContext { + /// Creates a new unique but empty session. + pub fn empty() -> Self { + Self { + session_id: Uuid::new_v4().to_string(), + } + } + + /// Returns a string that identifies this session. + /// + /// This can be used for caching state within a session. + pub fn session_id(&self) -> &str { + &self.session_id + } +} + +/// The catalog API for Iceberg Rust that includes session handling. +#[async_trait] +pub trait SessionCatalog: Debug + Send + Sync { + /// List namespaces inside the catalog. + async fn list_namespaces( + &self, + context: &SessionContext, + parent: Option<&NamespaceIdent>, + ) -> Result>; + + /// Create a new namespace inside the catalog. + async fn create_namespace( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result; + + /// Get a namespace information from the catalog. + async fn get_namespace( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + ) -> Result; + + /// Check if namespace exists in catalog. + async fn namespace_exists( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + ) -> Result; + + /// Update a namespace inside the catalog. + /// + /// # Behavior + /// + /// The properties must be the full set of namespace. + async fn update_namespace( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result<()>; + + /// Drop a namespace from the catalog, or returns error if it doesn't exist. + async fn drop_namespace( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + ) -> Result<()>; + + /// List tables from namespace. + async fn list_tables( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + ) -> Result>; + + /// Create a new table inside the namespace. + async fn create_table( + &self, + context: &SessionContext, + namespace: &NamespaceIdent, + creation: TableCreation, + ) -> Result; + + /// Load table from the catalog. + async fn load_table(&self, context: &SessionContext, table: &TableIdent) -> Result
; + + /// Drop a table from the catalog, or returns error if it doesn't exist. + async fn drop_table(&self, context: &SessionContext, table: &TableIdent) -> Result<()>; + + /// Drop a table from the catalog and delete the underlying table data. + /// + /// Implementations should load the table metadata, drop the table + /// from the catalog, then delete all associated data and metadata files. + /// The [`drop_table_data`](super::utils::drop_table_data) utility function can + /// be used for the file cleanup step. + async fn purge_table(&self, context: &SessionContext, table: &TableIdent) -> Result<()>; + + /// Check if a table exists in the catalog. + async fn table_exists(&self, context: &SessionContext, table: &TableIdent) -> Result; + + /// Rename a table in the catalog. + async fn rename_table( + &self, + context: &SessionContext, + src: &TableIdent, + dest: &TableIdent, + ) -> Result<()>; + + /// Register an existing table to the catalog. + async fn register_table( + &self, + context: &SessionContext, + table: &TableIdent, + metadata_location: String, + ) -> Result
; + + /// Update a table to the catalog. + async fn update_table(&self, context: &SessionContext, commit: TableCommit) -> Result
; +} + +/// Common interface for all session catalog builders. +pub trait SessionCatalogBuilder: Default + Debug + Send + Sync { + /// The session catalog type that this builder creates. + type C: SessionCatalog; + + /// Set a custom StorageFactory to use for storage operations. + /// + /// When a StorageFactory is provided, the catalog will use it to build FileIO + /// instances for all storage operations instead of using the default factory. + /// + /// # Arguments + /// + /// * `storage_factory` - The StorageFactory to use for creating storage instances + fn with_storage_factory(self, storage_factory: Arc) -> Self; + + /// Set a custom tokio Runtime to use for spawning async tasks. + /// + /// When a Runtime is provided, the catalog will propagate it to all tables + /// it creates. Tasks such as scan planning and delete file processing + /// will be spawned on this runtime. + fn with_runtime(self, runtime: Runtime) -> Self; + + /// Create a new session catalog instance. + fn load( + self, + name: impl Into, + props: HashMap, + ) -> impl Future> + Send; +} From a8c8a61ba661cb2069f5e9dd3817cb7872835fd7 Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sun, 26 Jul 2026 19:45:30 +0200 Subject: [PATCH 2/6] Add remaining SessionContext fields --- crates/iceberg/public-api.txt | 5 +++- crates/iceberg/src/catalog/session.rs | 37 +++++++++++++++++++++++++-- 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 80f2633b06..6ba36d73ac 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -3660,14 +3660,17 @@ impl core::fmt::Debug for iceberg::RuntimeHandle pub fn iceberg::RuntimeHandle::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct iceberg::SessionContext impl iceberg::SessionContext +pub fn iceberg::SessionContext::credentials(&self) -> &std::collections::hash::map::HashMap pub fn iceberg::SessionContext::empty() -> Self +pub fn iceberg::SessionContext::identity(&self) -> core::option::Option<&str> +pub fn iceberg::SessionContext::properties(&self) -> &std::collections::hash::map::HashMap pub fn iceberg::SessionContext::session_id(&self) -> &str impl core::clone::Clone for iceberg::SessionContext pub fn iceberg::SessionContext::clone(&self) -> iceberg::SessionContext impl core::fmt::Debug for iceberg::SessionContext pub fn iceberg::SessionContext::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl iceberg::SessionContext -pub fn iceberg::SessionContext::builder() -> SessionContextBuilder<(())> +pub fn iceberg::SessionContext::builder() -> SessionContextBuilder<((), (), (), ())> pub struct iceberg::TableCommit impl iceberg::TableCommit pub fn iceberg::TableCommit::apply(self, table: iceberg::table::Table) -> iceberg::Result diff --git a/crates/iceberg/src/catalog/session.rs b/crates/iceberg/src/catalog/session.rs index b13baa081b..0e4b9077c4 100644 --- a/crates/iceberg/src/catalog/session.rs +++ b/crates/iceberg/src/catalog/session.rs @@ -34,7 +34,22 @@ use crate::{Namespace, NamespaceIdent, Result, TableCommit, TableCreation, Table /// Context for a session. #[derive(Debug, Clone, TypedBuilder)] pub struct SessionContext { + /// The unique identifier for this session. + /// + /// Note that the session_id may be used for caching session-scoped state + /// and re-use of a session_id with different session context may result in + /// unexpected behavior. session_id: String, + + /// An optional user or principal associated with the session. + #[builder(default, setter(strip_option))] + identity: Option, + + #[builder(default)] + properties: HashMap, + + #[builder(default)] + credentials: HashMap, } impl SessionContext { @@ -42,15 +57,33 @@ impl SessionContext { pub fn empty() -> Self { Self { session_id: Uuid::new_v4().to_string(), + identity: None, + properties: HashMap::new(), + credentials: HashMap::new(), } } - /// Returns a string that identifies this session. + /// Returns the identifier for this session. /// - /// This can be used for caching state within a session. + /// The identifier may be used for caching state within a session. pub fn session_id(&self) -> &str { &self.session_id } + + /// Returns a string that identifies the current user or principal. + pub fn identity(&self) -> Option<&str> { + self.identity.as_deref() + } + + /// Returns a map of properties currently set for the session. + pub fn properties(&self) -> &HashMap { + &self.properties + } + + /// Returns the session's credential map. + pub fn credentials(&self) -> &HashMap { + &self.credentials + } } /// The catalog API for Iceberg Rust that includes session handling. From ad60abbb5b1aa0ec1d506b55a1bf980fcf7cb11d Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Sun, 26 Jul 2026 20:17:45 +0200 Subject: [PATCH 3/6] Use safe Credential type in SessionContext --- crates/iceberg/public-api.txt | 12 ++++++- crates/iceberg/src/catalog/session.rs | 48 +++++++++++++++++++++++++-- 2 files changed, 57 insertions(+), 3 deletions(-) diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 6ba36d73ac..0404f35597 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -3510,6 +3510,16 @@ impl serde_core::ser::Serialize for iceberg::ViewUpdate pub fn iceberg::ViewUpdate::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::ViewUpdate pub fn iceberg::ViewUpdate::deserialize<__D>(__deserializer: __D) -> core::result::Result::Error> where __D: serde_core::de::Deserializer<'de> +pub struct iceberg::Credential(_) +impl iceberg::Credential +pub fn iceberg::Credential::expose(&self) -> &str +pub fn iceberg::Credential::new(value: alloc::string::String) -> Self +impl core::clone::Clone for iceberg::Credential +pub fn iceberg::Credential::clone(&self) -> iceberg::Credential +impl core::convert::From for iceberg::Credential +pub fn iceberg::Credential::from(value: alloc::string::String) -> Self +impl core::fmt::Debug for iceberg::Credential +pub fn iceberg::Credential::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct iceberg::Error impl iceberg::Error pub fn iceberg::Error::backtrace(&self) -> &std::backtrace::Backtrace @@ -3660,7 +3670,7 @@ impl core::fmt::Debug for iceberg::RuntimeHandle pub fn iceberg::RuntimeHandle::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result pub struct iceberg::SessionContext impl iceberg::SessionContext -pub fn iceberg::SessionContext::credentials(&self) -> &std::collections::hash::map::HashMap +pub fn iceberg::SessionContext::credentials(&self) -> &std::collections::hash::map::HashMap pub fn iceberg::SessionContext::empty() -> Self pub fn iceberg::SessionContext::identity(&self) -> core::option::Option<&str> pub fn iceberg::SessionContext::properties(&self) -> &std::collections::hash::map::HashMap diff --git a/crates/iceberg/src/catalog/session.rs b/crates/iceberg/src/catalog/session.rs index 0e4b9077c4..2836546753 100644 --- a/crates/iceberg/src/catalog/session.rs +++ b/crates/iceberg/src/catalog/session.rs @@ -25,6 +25,7 @@ use std::sync::Arc; use async_trait::async_trait; use typed_builder::TypedBuilder; use uuid::Uuid; +use zeroize::Zeroizing; use crate::io::StorageFactory; use crate::runtime::Runtime; @@ -49,7 +50,7 @@ pub struct SessionContext { properties: HashMap, #[builder(default)] - credentials: HashMap, + credentials: HashMap, } impl SessionContext { @@ -81,11 +82,41 @@ impl SessionContext { } /// Returns the session's credential map. - pub fn credentials(&self) -> &HashMap { + pub fn credentials(&self) -> &HashMap { &self.credentials } } +/// A string-like type containing sensitive information such as passwords or tokens. +/// +/// It is redacted from logs and automatically zeroized. +#[derive(Clone)] +pub struct Credential(Zeroizing); + +impl Credential { + /// Creates a new [`Credential`] from a string. + pub fn new(value: String) -> Self { + Self(Zeroizing::new(value)) + } + + /// Returns the raw value of the credential. + pub fn expose(&self) -> &str { + &self.0 + } +} + +impl Debug for Credential { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("Credential([REDACTED])") + } +} + +impl From for Credential { + fn from(value: String) -> Self { + Self::new(value) + } +} + /// The catalog API for Iceberg Rust that includes session handling. #[async_trait] pub trait SessionCatalog: Debug + Send + Sync { @@ -218,3 +249,16 @@ pub trait SessionCatalogBuilder: Default + Debug + Send + Sync { props: HashMap, ) -> impl Future> + Send; } + +#[cfg(test)] +mod tests { + use crate::Credential; + + #[test] + fn test_credential_redacts_value() { + let sensitive_value = "my-pw-12346"; + + let logged = format!("{:?}", Credential::new(sensitive_value.to_string())); + assert!(!logged.contains(sensitive_value)); + } +} From 0dea3836f90c061a814e8f626e359ed895c8839e Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Wed, 29 Jul 2026 00:14:18 +0200 Subject: [PATCH 4/6] Add RestSessionCatalog --- crates/catalog/rest/public-api.txt | 24 +- crates/catalog/rest/src/catalog.rs | 652 ++++++++++++++++++++++++----- 2 files changed, 561 insertions(+), 115 deletions(-) diff --git a/crates/catalog/rest/public-api.txt b/crates/catalog/rest/public-api.txt index 776b11c40a..22a458fa85 100644 --- a/crates/catalog/rest/public-api.txt +++ b/crates/catalog/rest/public-api.txt @@ -226,13 +226,14 @@ pub fn iceberg_catalog_rest::RestCatalog::purge_table<'life0, 'life1, 'async_tra pub fn iceberg_catalog_rest::RestCatalog::register_table<'life0, 'life1, 'async_trait>(&'life0 self, table_ident: &'life1 iceberg::catalog::TableIdent, metadata_location: alloc::string::String) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait pub fn iceberg_catalog_rest::RestCatalog::rename_table<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, src: &'life1 iceberg::catalog::TableIdent, dest: &'life2 iceberg::catalog::TableIdent) -> 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::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_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 pub struct iceberg_catalog_rest::RestCatalogBuilder impl iceberg_catalog_rest::RestCatalogBuilder pub fn iceberg_catalog_rest::RestCatalogBuilder::with_client(self, client: reqwest::async_impl::client::Client) -> Self +pub fn iceberg_catalog_rest::RestCatalogBuilder::with_session(self, session: iceberg::catalog::session::SessionContext) -> Self impl core::default::Default for iceberg_catalog_rest::RestCatalogBuilder -pub fn iceberg_catalog_rest::RestCatalogBuilder::default() -> Self +pub fn iceberg_catalog_rest::RestCatalogBuilder::default() -> iceberg_catalog_rest::RestCatalogBuilder impl core::fmt::Debug for iceberg_catalog_rest::RestCatalogBuilder pub fn iceberg_catalog_rest::RestCatalogBuilder::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl iceberg::catalog::CatalogBuilder for iceberg_catalog_rest::RestCatalogBuilder @@ -241,6 +242,25 @@ 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::RestSessionCatalog +impl core::fmt::Debug for iceberg_catalog_rest::RestSessionCatalog +pub fn iceberg_catalog_rest::RestSessionCatalog::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl iceberg::catalog::session::SessionCatalog for iceberg_catalog_rest::RestSessionCatalog +pub fn iceberg_catalog_rest::RestSessionCatalog::create_namespace<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _session: &'life1 iceberg::catalog::session::SessionContext, namespace: &'life2 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, 'life2: 'async_trait +pub fn iceberg_catalog_rest::RestSessionCatalog::create_table<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _session: &'life1 iceberg::catalog::session::SessionContext, namespace: &'life2 iceberg::catalog::NamespaceIdent, creation: iceberg::catalog::TableCreation) -> 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::RestSessionCatalog::drop_namespace<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _session: &'life1 iceberg::catalog::session::SessionContext, namespace: &'life2 iceberg::catalog::NamespaceIdent) -> 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::RestSessionCatalog::drop_table<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, session: &'life1 iceberg::catalog::session::SessionContext, table: &'life2 iceberg::catalog::TableIdent) -> 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::RestSessionCatalog::get_namespace<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _session: &'life1 iceberg::catalog::session::SessionContext, namespace: &'life2 iceberg::catalog::NamespaceIdent) -> 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::RestSessionCatalog::list_namespaces<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _session: &'life1 iceberg::catalog::session::SessionContext, parent: core::option::Option<&'life2 iceberg::catalog::NamespaceIdent>) -> 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::RestSessionCatalog::list_tables<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _session: &'life1 iceberg::catalog::session::SessionContext, namespace: &'life2 iceberg::catalog::NamespaceIdent) -> 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::RestSessionCatalog::load_table<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _session: &'life1 iceberg::catalog::session::SessionContext, table_ident: &'life2 iceberg::catalog::TableIdent) -> 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::RestSessionCatalog::namespace_exists<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, session: &'life1 iceberg::catalog::session::SessionContext, ns: &'life2 iceberg::catalog::NamespaceIdent) -> 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::RestSessionCatalog::purge_table<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, session: &'life1 iceberg::catalog::session::SessionContext, table: &'life2 iceberg::catalog::TableIdent) -> 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::RestSessionCatalog::register_table<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _session: &'life1 iceberg::catalog::session::SessionContext, table_ident: &'life2 iceberg::catalog::TableIdent, metadata_location: alloc::string::String) -> 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::RestSessionCatalog::rename_table<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _session: &'life1 iceberg::catalog::session::SessionContext, src: &'life2 iceberg::catalog::TableIdent, dest: &'life3 iceberg::catalog::TableIdent) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg_catalog_rest::RestSessionCatalog::table_exists<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, session: &'life1 iceberg::catalog::session::SessionContext, table: &'life2 iceberg::catalog::TableIdent) -> 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::RestSessionCatalog::update_namespace<'life0, 'life1, 'life2, 'async_trait>(&'life0 self, _session: &'life1 iceberg::catalog::session::SessionContext, _namespace: &'life2 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, 'life2: 'async_trait +pub fn iceberg_catalog_rest::RestSessionCatalog::update_table<'life0, 'life1, 'async_trait>(&'life0 self, _session: &'life1 iceberg::catalog::session::SessionContext, commit: iceberg::catalog::TableCommit) -> core::pin::Pin> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait 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 diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 8642b32d22..375ac6a156 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -28,7 +28,7 @@ use iceberg::io::{FileIO, FileIOBuilder, StorageFactory}; use iceberg::table::Table; use iceberg::{ Catalog, CatalogBuilder, Error, ErrorKind, Namespace, NamespaceIdent, Result, Runtime, - TableCommit, TableCreation, TableIdent, + SessionCatalog, SessionCatalogBuilder, SessionContext, TableCommit, TableCreation, TableIdent, }; use itertools::Itertools; use reqwest::header::{ @@ -60,97 +60,41 @@ const CARGO_PKG_VERSION: &str = env!("CARGO_PKG_VERSION"); const PATH_V1: &str = "v1"; /// Builder for [`RestCatalog`]. -#[derive(Debug)] +#[derive(Debug, Default)] pub struct RestCatalogBuilder { - config: RestCatalogConfig, - storage_factory: Option>, - kms_client_factory: Option>, - runtime: Option, -} - -impl Default for RestCatalogBuilder { - fn default() -> Self { - Self { - config: RestCatalogConfig { - name: None, - uri: "".to_string(), - warehouse: None, - props: HashMap::new(), - client: None, - }, - storage_factory: None, - kms_client_factory: None, - runtime: None, - } - } + session: Option, + inner: RestSessionCatalogBuilder, } impl CatalogBuilder for RestCatalogBuilder { type C = RestCatalog; fn with_storage_factory(mut self, storage_factory: Arc) -> Self { - self.storage_factory = Some(storage_factory); + self.inner = self.inner.with_storage_factory(storage_factory); self } fn with_kms_client_factory(mut self, kms_client_factory: Arc) -> Self { - self.kms_client_factory = Some(kms_client_factory); + self.inner = self.inner.with_kms_client_factory(kms_client_factory); self } fn with_runtime(mut self, runtime: Runtime) -> Self { - self.runtime = Some(runtime); + self.inner = self.inner.with_runtime(runtime); self } fn load( - mut self, + self, name: impl Into, props: HashMap, ) -> impl Future> + Send { - self.config.name = Some(name.into()); - - if props.contains_key(REST_CATALOG_PROP_URI) { - self.config.uri = props - .get(REST_CATALOG_PROP_URI) - .cloned() - .unwrap_or_default(); - } - - if props.contains_key(REST_CATALOG_PROP_WAREHOUSE) { - self.config.warehouse = props.get(REST_CATALOG_PROP_WAREHOUSE).cloned() - } - - // Collect other remaining properties - self.config.props = props - .into_iter() - .filter(|(k, _)| k != REST_CATALOG_PROP_URI && k != REST_CATALOG_PROP_WAREHOUSE) - .collect(); - + let name = name.into(); async move { - if self.config.name.is_none() { - Err(Error::new( - ErrorKind::DataInvalid, - "Catalog name is required", - )) - } else if self.config.uri.is_empty() { - Err(Error::new( - ErrorKind::DataInvalid, - "Catalog uri is required", - )) - } else { - let runtime = self.runtime.unwrap_or_else(Runtime::current); - let kms_client = match self.kms_client_factory { - Some(factory) => Some(factory.create_kms_client(&self.config.props).await?), - None => None, - }; - Ok(RestCatalog::new( - self.config, - self.storage_factory, - runtime, - kms_client, - )) - } + let session = self.session.unwrap_or_else(SessionContext::empty); + let session_catalog = Arc::new(self.inner.load(name, props).await?); + + Ok(RestCatalog::from_session_catalog(session, session_catalog)) } } } @@ -158,7 +102,14 @@ impl CatalogBuilder for RestCatalogBuilder { impl RestCatalogBuilder { /// Configures the catalog with a custom HTTP client. pub fn with_client(mut self, client: Client) -> Self { - self.config.client = Some(client); + self.inner = self.inner.with_client(client); + self + } + + /// Configures the session that will be used with this catalog. + /// Overwrites the default empty session from SessionContext::empty(). + pub fn with_session(mut self, session: SessionContext) -> Self { + self.session = Some(session); self } } @@ -370,6 +321,191 @@ struct RestContext { /// Rest catalog implementation. #[derive(Debug)] pub struct RestCatalog { + session: SessionContext, + session_catalog: Arc, +} + +impl RestCatalog { + /// Creates a `RestCatalog` from a [`RestCatalogConfig`]. + #[cfg(test)] + fn new( + session: SessionContext, + config: RestCatalogConfig, + storage_factory: Option>, + runtime: Runtime, + kms_client: Option>, + ) -> Self { + let session_catalog = Arc::new(RestSessionCatalog::new( + config, + storage_factory, + runtime, + kms_client, + )); + + Self::from_session_catalog(session, session_catalog) + } + + fn from_session_catalog( + session: SessionContext, + session_catalog: Arc, + ) -> Self { + Self { + session, + session_catalog, + } + } + + /// Invalidate the current token without generating a new one. On the next request, the client + /// will attempt to generate a new token. + pub async fn invalidate_token(&self) -> Result<()> { + self.session_catalog.invalidate_token(&self.session).await + } + + /// Invalidate the current token and set a new one. Generates a new token before invalidating + /// the current token, meaning the old token will be used until this function acquires the lock + /// and overwrites the token. + /// + /// If credential is invalid, or the request fails, this method will return an error and leave + /// the current token unchanged. + pub async fn regenerate_token(&self) -> Result<()> { + self.session_catalog.regenerate_token(&self.session).await + } + + #[cfg(test)] + async fn context(&self) -> Result<&RestContext> { + self.session_catalog.context().await + } +} + +/// All requests and expected responses are derived from the REST catalog API spec: +/// https://github.com/apache/iceberg/blob/main/open-api/rest-catalog-open-api.yaml +#[async_trait] +impl Catalog for RestCatalog { + async fn list_namespaces( + &self, + parent: Option<&NamespaceIdent>, + ) -> Result> { + self.session_catalog + .list_namespaces(&self.session, parent) + .await + } + + async fn create_namespace( + &self, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result { + self.session_catalog + .create_namespace(&self.session, namespace, properties) + .await + } + + async fn get_namespace(&self, namespace: &NamespaceIdent) -> Result { + self.session_catalog + .get_namespace(&self.session, namespace) + .await + } + + async fn namespace_exists(&self, ns: &NamespaceIdent) -> Result { + self.session_catalog + .namespace_exists(&self.session, ns) + .await + } + + async fn update_namespace( + &self, + namespace: &NamespaceIdent, + properties: HashMap, + ) -> Result<()> { + self.session_catalog + .update_namespace(&self.session, namespace, properties) + .await + } + + async fn drop_namespace(&self, namespace: &NamespaceIdent) -> Result<()> { + self.session_catalog + .drop_namespace(&self.session, namespace) + .await + } + + async fn list_tables(&self, namespace: &NamespaceIdent) -> Result> { + self.session_catalog + .list_tables(&self.session, namespace) + .await + } + + /// Create a new table inside the namespace. + /// + /// In the resulting table, if there are any config properties that + /// are present in both the response from the REST server and the + /// config provided when creating this `RestCatalog` instance then + /// the value provided locally to the `RestCatalog` will take precedence. + async fn create_table( + &self, + namespace: &NamespaceIdent, + creation: TableCreation, + ) -> Result
{ + self.session_catalog + .create_table(&self.session, namespace, creation) + .await + } + + /// Load table from the catalog. + /// + /// If there are any config properties that are present in both the response from the REST + /// server and the config provided when creating this `RestCatalog` instance, then the value + /// provided locally to the `RestCatalog` will take precedence. + async fn load_table(&self, table_ident: &TableIdent) -> Result
{ + self.session_catalog + .load_table(&self.session, table_ident) + .await + } + + /// Drop a table from the catalog. + async fn drop_table(&self, table: &TableIdent) -> Result<()> { + self.session_catalog.drop_table(&self.session, table).await + } + + /// Drop a table from the catalog and purge its data by sending + /// `purgeRequested=true` to the REST server. + async fn purge_table(&self, table: &TableIdent) -> Result<()> { + self.session_catalog.purge_table(&self.session, table).await + } + + /// Check if a table exists in the catalog. + async fn table_exists(&self, table: &TableIdent) -> Result { + self.session_catalog + .table_exists(&self.session, table) + .await + } + + /// Rename a table in the catalog. + async fn rename_table(&self, src: &TableIdent, dest: &TableIdent) -> Result<()> { + self.session_catalog + .rename_table(&self.session, src, dest) + .await + } + + async fn register_table( + &self, + table_ident: &TableIdent, + metadata_location: String, + ) -> Result
{ + self.session_catalog + .register_table(&self.session, table_ident, metadata_location) + .await + } + + async fn update_table(&self, commit: TableCommit) -> Result
{ + self.session_catalog + .update_table(&self.session, commit) + .await + } +} + +/// Rest catalog implementation. +#[derive(Debug)] +pub struct RestSessionCatalog { /// User config is stored as-is and never be changed. /// /// It could be different from the config fetched from the server and used at runtime. @@ -382,8 +518,8 @@ pub struct RestCatalog { kms_client: Option>, } -impl RestCatalog { - /// Creates a `RestCatalog` from a [`RestCatalogConfig`]. +impl RestSessionCatalog { + /// Creates a `RestSessionCatalog` from a [`RestCatalogConfig`]. fn new( config: RestCatalogConfig, storage_factory: Option>, @@ -399,8 +535,29 @@ impl RestCatalog { } } + /// Invalidate the current token without generating a new one. On the next request, the client + /// will attempt to generate a new token. + async fn invalidate_token(&self, _session: &SessionContext) -> Result<()> { + self.context().await?.client.invalidate_token().await + } + + /// Invalidate the current token and set a new one. Generates a new token before invalidating + /// the current token, meaning the old token will be used until this function acquires the lock + /// and overwrites the token. + /// + /// If credential is invalid, or the request fails, this method will return an error and leave + /// the current token unchanged. + async fn regenerate_token(&self, _session: &SessionContext) -> Result<()> { + self.context().await?.client.regenerate_token().await + } + /// Sends a DELETE request for the given table, optionally requesting purge. - async fn delete_table(&self, table: &TableIdent, purge: bool) -> Result<()> { + async fn delete_table( + &self, + _session: &SessionContext, + table: &TableIdent, + purge: bool, + ) -> Result<()> { let context = self.context().await?; let mut request_builder = context @@ -433,7 +590,8 @@ impl RestCatalog { self.ctx .get_or_try_init(|| async { let client = HttpClient::new(&self.user_config)?; - let catalog_config = RestCatalog::load_config(&client, &self.user_config).await?; + let catalog_config = + RestSessionCatalog::load_config(&client, &self.user_config).await?; // Use the advertised endpoints as-is, falling back to // `DEFAULT_ENDPOINTS` when absent or empty. let endpoints = match &catalog_config.endpoints { @@ -545,30 +703,13 @@ impl RestCatalog { Ok(file_io) } - - /// Invalidate the current token without generating a new one. On the next request, the client - /// will attempt to generate a new token. - pub async fn invalidate_token(&self) -> Result<()> { - self.context().await?.client.invalidate_token().await - } - - /// Invalidate the current token and set a new one. Generates a new token before invalidating - /// the current token, meaning the old token will be used until this function acquires the lock - /// and overwrites the token. - /// - /// If credential is invalid, or the request fails, this method will return an error and leave - /// the current token unchanged. - pub async fn regenerate_token(&self) -> Result<()> { - self.context().await?.client.regenerate_token().await - } } -/// All requests and expected responses are derived from the REST catalog API spec: -/// https://github.com/apache/iceberg/blob/main/open-api/rest-catalog-open-api.yaml #[async_trait] -impl Catalog for RestCatalog { +impl SessionCatalog for RestSessionCatalog { async fn list_namespaces( &self, + _session: &SessionContext, parent: Option<&NamespaceIdent>, ) -> Result> { let context = self.context().await?; @@ -624,6 +765,7 @@ impl Catalog for RestCatalog { async fn create_namespace( &self, + _session: &SessionContext, namespace: &NamespaceIdent, properties: HashMap, ) -> Result { @@ -658,7 +800,11 @@ impl Catalog for RestCatalog { } } - async fn get_namespace(&self, namespace: &NamespaceIdent) -> Result { + async fn get_namespace( + &self, + _session: &SessionContext, + namespace: &NamespaceIdent, + ) -> Result { let context = self.context().await?; let request = context @@ -686,13 +832,17 @@ impl Catalog for RestCatalog { } } - async fn namespace_exists(&self, ns: &NamespaceIdent) -> Result { + async fn namespace_exists( + &self, + session: &SessionContext, + ns: &NamespaceIdent, + ) -> Result { // Prefer a cheap HEAD when the server advertises it; otherwise fall back // to loading the namespace (GET) and treating a missing namespace as // `false`, so this still works against servers that don't advertise the // HEAD route. if !self.supports_endpoint(&V1_NAMESPACE_EXISTS).await? { - return match self.get_namespace(ns).await { + return match self.get_namespace(session, ns).await { Ok(_) => Ok(true), Err(e) if e.kind() == ErrorKind::NamespaceNotFound => Ok(false), Err(e) => Err(e), @@ -706,6 +856,7 @@ impl Catalog for RestCatalog { async fn update_namespace( &self, + _session: &SessionContext, _namespace: &NamespaceIdent, _properties: HashMap, ) -> Result<()> { @@ -715,7 +866,11 @@ impl Catalog for RestCatalog { )) } - async fn drop_namespace(&self, namespace: &NamespaceIdent) -> Result<()> { + async fn drop_namespace( + &self, + _session: &SessionContext, + namespace: &NamespaceIdent, + ) -> Result<()> { let context = self.context().await?; let request = context @@ -739,7 +894,11 @@ impl Catalog for RestCatalog { } } - async fn list_tables(&self, namespace: &NamespaceIdent) -> Result> { + async fn list_tables( + &self, + _session: &SessionContext, + namespace: &NamespaceIdent, + ) -> Result> { let context = self.context().await?; let endpoint = context.config.tables_endpoint(namespace); let mut identifiers = Vec::new(); @@ -793,6 +952,7 @@ impl Catalog for RestCatalog { /// the value provided locally to the `RestCatalog` will take precedence. async fn create_table( &self, + _session: &SessionContext, namespace: &NamespaceIdent, creation: TableCreation, ) -> Result
{ @@ -877,7 +1037,11 @@ impl Catalog for RestCatalog { /// If there are any config properties that are present in both the response from the REST /// server and the config provided when creating this `RestCatalog` instance, then the value /// provided locally to the `RestCatalog` will take precedence. - async fn load_table(&self, table_ident: &TableIdent) -> Result
{ + async fn load_table( + &self, + _session: &SessionContext, + table_ident: &TableIdent, + ) -> Result
{ let context = self.context().await?; let request = context @@ -933,23 +1097,23 @@ impl Catalog for RestCatalog { } /// Drop a table from the catalog. - async fn drop_table(&self, table: &TableIdent) -> Result<()> { - self.delete_table(table, false).await + async fn drop_table(&self, session: &SessionContext, table: &TableIdent) -> Result<()> { + self.delete_table(session, table, false).await } /// Drop a table from the catalog and purge its data by sending /// `purgeRequested=true` to the REST server. - async fn purge_table(&self, table: &TableIdent) -> Result<()> { - self.delete_table(table, true).await + async fn purge_table(&self, session: &SessionContext, table: &TableIdent) -> Result<()> { + self.delete_table(session, table, true).await } /// Check if a table exists in the catalog. - async fn table_exists(&self, table: &TableIdent) -> Result { + async fn table_exists(&self, session: &SessionContext, table: &TableIdent) -> Result { // Prefer a cheap HEAD when the server advertises it; otherwise fall back // to loading the table (GET) and treating a missing table as `false`, so // this still works against servers that don't advertise the HEAD route. if !self.supports_endpoint(&V1_TABLE_EXISTS).await? { - return match self.load_table(table).await { + return match self.load_table(session, table).await { Ok(_) => Ok(true), Err(e) if e.kind() == ErrorKind::TableNotFound => Ok(false), Err(e) => Err(e), @@ -962,7 +1126,12 @@ impl Catalog for RestCatalog { } /// Rename a table in the catalog. - async fn rename_table(&self, src: &TableIdent, dest: &TableIdent) -> Result<()> { + async fn rename_table( + &self, + _session: &SessionContext, + src: &TableIdent, + dest: &TableIdent, + ) -> Result<()> { let context = self.context().await?; let request = context @@ -996,6 +1165,7 @@ impl Catalog for RestCatalog { async fn register_table( &self, + _session: &SessionContext, table_ident: &TableIdent, metadata_location: String, ) -> Result
{ @@ -1062,7 +1232,11 @@ impl Catalog for RestCatalog { table_builder.build() } - async fn update_table(&self, mut commit: TableCommit) -> Result
{ + async fn update_table( + &self, + _session: &SessionContext, + mut commit: TableCommit, + ) -> Result
{ let context = self.context().await?; let request = context @@ -1139,6 +1313,115 @@ impl Catalog for RestCatalog { } } +#[derive(Debug)] +struct RestSessionCatalogBuilder { + config: RestCatalogConfig, + storage_factory: Option>, + kms_client_factory: Option>, + runtime: Option, +} + +impl Default for RestSessionCatalogBuilder { + fn default() -> Self { + Self { + config: RestCatalogConfig { + name: None, + uri: "".to_string(), + warehouse: None, + props: HashMap::new(), + client: None, + }, + storage_factory: None, + kms_client_factory: None, + runtime: None, + } + } +} + +impl RestSessionCatalogBuilder { + /// Configures the catalog with a custom HTTP client. + pub fn with_client(mut self, client: Client) -> Self { + self.config.client = Some(client); + self + } + + /// Configures the catalog with a KMS client factory used to build clients + /// for encrypted tables. + pub fn with_kms_client_factory( + mut self, + kms_client_factory: Arc, + ) -> Self { + self.kms_client_factory = Some(kms_client_factory); + self + } +} + +impl SessionCatalogBuilder for RestSessionCatalogBuilder { + type C = RestSessionCatalog; + + fn with_storage_factory(mut self, storage_factory: Arc) -> Self { + self.storage_factory = Some(storage_factory); + self + } + + fn with_runtime(mut self, runtime: Runtime) -> Self { + self.runtime = Some(runtime); + self + } + + fn load( + mut self, + name: impl Into, + props: HashMap, + ) -> impl Future> + Send { + self.config.name = Some(name.into()); + + if props.contains_key(REST_CATALOG_PROP_URI) { + self.config.uri = props + .get(REST_CATALOG_PROP_URI) + .cloned() + .unwrap_or_default(); + } + + if props.contains_key(REST_CATALOG_PROP_WAREHOUSE) { + self.config.warehouse = props.get(REST_CATALOG_PROP_WAREHOUSE).cloned() + } + + // Collect other remaining properties + self.config.props = props + .into_iter() + .filter(|(k, _)| k != REST_CATALOG_PROP_URI && k != REST_CATALOG_PROP_WAREHOUSE) + .collect(); + + async move { + if self.config.name.is_none() { + Err(Error::new( + ErrorKind::DataInvalid, + "Catalog name is required", + )) + } else if self.config.uri.is_empty() { + Err(Error::new( + ErrorKind::DataInvalid, + "Catalog uri is required", + )) + } else { + let runtime = self.runtime.unwrap_or_else(Runtime::current); + let kms_client = match self.kms_client_factory { + Some(factory) => Some(factory.create_kms_client(&self.config.props).await?), + None => None, + }; + + Ok(RestSessionCatalog::new( + self.config, + self.storage_factory, + runtime, + kms_client, + )) + } + } + } +} + #[cfg(test)] mod tests { use std::fs::File; @@ -1179,6 +1462,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -1256,6 +1540,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -1265,12 +1550,24 @@ mod tests { let plan = "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan" .parse::() .unwrap(); - assert!(catalog.supports_endpoint(&plan).await.unwrap()); + assert!( + catalog + .session_catalog + .supports_endpoint(&plan) + .await + .unwrap() + ); // Advertised list is present but does not include this route. let delete_ns = "DELETE /v1/{prefix}/namespaces/{namespace}" .parse::() .unwrap(); - assert!(!catalog.supports_endpoint(&delete_ns).await.unwrap()); + assert!( + !catalog + .session_catalog + .supports_endpoint(&delete_ns) + .await + .unwrap() + ); config_mock.assert_async().await; } @@ -1287,6 +1584,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -1298,12 +1596,24 @@ mod tests { let load_table = "GET /v1/{prefix}/namespaces/{namespace}/tables/{table}" .parse::() .unwrap(); - assert!(catalog.supports_endpoint(&load_table).await.unwrap()); + assert!( + catalog + .session_catalog + .supports_endpoint(&load_table) + .await + .unwrap() + ); // But not an optional endpoint that must be advertised. let plan = "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan" .parse::() .unwrap(); - assert!(!catalog.supports_endpoint(&plan).await.unwrap()); + assert!( + !catalog + .session_catalog + .supports_endpoint(&plan) + .await + .unwrap() + ); config_mock.assert_async().await; } @@ -1322,6 +1632,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -1331,7 +1642,13 @@ mod tests { let load_table = "GET /v1/{prefix}/namespaces/{namespace}/tables/{table}" .parse::() .unwrap(); - assert!(catalog.supports_endpoint(&load_table).await.unwrap()); + assert!( + catalog + .session_catalog + .supports_endpoint(&load_table) + .await + .unwrap() + ); config_mock.assert_async().await; } @@ -1373,6 +1690,7 @@ mod tests { props.insert("credential".to_string(), "client1:secret1".to_string()); let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder() .uri(server.url()) .props(props) @@ -1422,6 +1740,7 @@ mod tests { let config_mock = create_config_mock(&mut server).await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder() .uri(server.url()) .props(props) @@ -1448,6 +1767,7 @@ mod tests { props.insert("credential".to_string(), "client1:secret1".to_string()); let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder() .uri(server.url()) .props(props) @@ -1481,6 +1801,7 @@ mod tests { props.insert("credential".to_string(), "client1:secret1".to_string()); let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder() .uri(server.url()) .props(props) @@ -1514,6 +1835,7 @@ mod tests { props.insert("credential".to_string(), "client1:secret1".to_string()); let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder() .uri(server.url()) .props(props) @@ -1547,6 +1869,7 @@ mod tests { props.insert("credential".to_string(), "client1:secret1".to_string()); let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder() .uri(server.url()) .props(props) @@ -1662,6 +1985,7 @@ mod tests { ); let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder() .uri(server.url()) .props(props) @@ -1714,6 +2038,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -1746,6 +2071,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -1799,6 +2125,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -1900,6 +2227,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -1955,6 +2283,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2000,6 +2329,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2035,6 +2365,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2072,6 +2403,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2102,6 +2434,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2144,6 +2477,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2214,6 +2548,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2347,6 +2682,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2393,6 +2729,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2424,6 +2761,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2463,6 +2801,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2496,6 +2835,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2532,6 +2872,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2651,6 +2992,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2689,6 +3031,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2840,6 +3183,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2911,6 +3255,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -3057,6 +3402,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -3124,6 +3470,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -3177,6 +3524,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -3218,7 +3566,8 @@ mod tests { assert!(catalog.is_ok()); - let catalog_config = catalog.unwrap().user_config; + let catalog = catalog.unwrap(); + let catalog_config = &catalog.session_catalog.user_config; assert_eq!(catalog_config.name.as_deref(), Some("test")); assert_eq!(catalog_config.uri, "http://localhost:8080"); assert_eq!(catalog_config.warehouse, None); @@ -3248,4 +3597,81 @@ mod tests { assert_eq!(err.message(), "Catalog uri is required"); } } + + #[tokio::test] + async fn test_create_session_catalog() { + let builder = RestSessionCatalogBuilder::default(); + + let result = builder + .load( + "test", + HashMap::from([ + ( + REST_CATALOG_PROP_URI.to_string(), + "http://localhost:8080".to_string(), + ), + ("a".to_string(), "b".to_string()), + ]), + ) + .await; + + assert!(result.is_ok()); + + let catalog = result.unwrap(); + + let catalog_config = catalog.user_config; + assert_eq!(catalog_config.name.as_deref(), Some("test")); + assert_eq!(catalog_config.uri, "http://localhost:8080"); + assert_eq!(catalog_config.warehouse, None); + // The default builder sets no client (only `with_client` does). + assert!(catalog_config.client.is_none()); + + // `uri` is consumed into its own field; other props are retained. + assert_eq!(catalog_config.props.get("a"), Some(&"b".to_string())); + assert!(!catalog_config.props.contains_key(REST_CATALOG_PROP_URI)); + } + + #[tokio::test] + async fn test_create_rest_catalog_with_session() { + let session = SessionContext::builder() + .session_id("test-id".to_string()) + .build(); + + let result = RestCatalogBuilder::default() + .with_session(session) + .load( + "test", + HashMap::from([( + REST_CATALOG_PROP_URI.to_string(), + "http://localhost:8080".to_string(), + )]), + ) + .await; + + assert!(result.is_ok()); + + // The session passed to `with_session` is the one the catalog is bound to. + let catalog = result.unwrap(); + assert_eq!(catalog.session.session_id(), "test-id"); + } + + #[tokio::test] + async fn test_create_rest_catalog_default_session() { + let result = RestCatalogBuilder::default() + .load( + "test", + HashMap::from([( + REST_CATALOG_PROP_URI.to_string(), + "http://localhost:8080".to_string(), + )]), + ) + .await; + + assert!(result.is_ok()); + + // Without `with_session`, the catalog falls back to `SessionContext::empty()`, + // which assigns a fresh v4 UUID. + let catalog = result.unwrap(); + assert!(uuid::Uuid::parse_str(catalog.session.session_id()).is_ok()); + } } From a349a8075387ca1757d373956263c636fe01bfa0 Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Thu, 16 Jul 2026 20:06:55 +0200 Subject: [PATCH 5/6] Extract RestContext::init() --- crates/catalog/rest/src/catalog.rs | 45 ++++++++++++++++-------------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 375ac6a156..953423bc48 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -318,6 +318,29 @@ struct RestContext { endpoints: HashSet, } +impl RestContext { + async fn init(config: &RestCatalogConfig) -> Result { + let client = HttpClient::new(config)?; + + let catalog_config = RestSessionCatalog::load_config(&client, config).await?; + // Use the advertised endpoints as-is, falling back to + // `DEFAULT_ENDPOINTS` when absent or empty. + let endpoints = match &catalog_config.endpoints { + Some(advertised) if !advertised.is_empty() => advertised.iter().cloned().collect(), + _ => crate::endpoint::DEFAULT_ENDPOINTS.clone(), + }; + let config = config.clone().merge_with_config(catalog_config); + + let client = client.update_with(&config)?; + + Ok(RestContext { + config, + client, + endpoints, + }) + } +} + /// Rest catalog implementation. #[derive(Debug)] pub struct RestCatalog { @@ -588,27 +611,7 @@ impl RestSessionCatalog { /// Gets the [`RestContext`] from the catalog. async fn context(&self) -> Result<&RestContext> { self.ctx - .get_or_try_init(|| async { - let client = HttpClient::new(&self.user_config)?; - let catalog_config = - RestSessionCatalog::load_config(&client, &self.user_config).await?; - // Use the advertised endpoints as-is, falling back to - // `DEFAULT_ENDPOINTS` when absent or empty. - let endpoints = match &catalog_config.endpoints { - Some(advertised) if !advertised.is_empty() => { - advertised.iter().cloned().collect() - } - _ => crate::endpoint::DEFAULT_ENDPOINTS.clone(), - }; - let config = self.user_config.clone().merge_with_config(catalog_config); - let client = client.update_with(&config)?; - - Ok(RestContext { - config, - client, - endpoints, - }) - }) + .get_or_try_init(async || RestContext::init(&self.user_config).await) .await } From 75caec7f6bcf89abb7cbe1b976a7d8d5fa7dac02 Mon Sep 17 00:00:00 2001 From: Jannik Steinmann Date: Fri, 17 Jul 2026 01:10:29 +0200 Subject: [PATCH 6/6] Refactor test suite to target RestSessionCatalog --- crates/catalog/rest/src/catalog.rs | 406 ++++++++++++++--------------- 1 file changed, 190 insertions(+), 216 deletions(-) diff --git a/crates/catalog/rest/src/catalog.rs b/crates/catalog/rest/src/catalog.rs index 953423bc48..a8e9970d73 100644 --- a/crates/catalog/rest/src/catalog.rs +++ b/crates/catalog/rest/src/catalog.rs @@ -1446,6 +1446,16 @@ mod tests { use super::*; + /// Builds a [`RestSessionCatalog`] with the default test storage factory and runtime. + fn session_catalog(config: RestCatalogConfig) -> RestSessionCatalog { + RestSessionCatalog::new( + config, + Some(Arc::new(LocalFsStorageFactory)), + Runtime::current(), + None, + ) + } + #[tokio::test] async fn test_update_config() { let mut server = Server::new_async().await; @@ -1464,13 +1474,7 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); assert_eq!( catalog @@ -1692,15 +1696,11 @@ mod tests { let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); - let catalog = RestCatalog::new( - SessionContext::empty(), + let catalog = session_catalog( RestCatalogConfig::builder() .uri(server.url()) .props(props) .build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -1742,15 +1742,11 @@ mod tests { let config_mock = create_config_mock(&mut server).await; - let catalog = RestCatalog::new( - SessionContext::empty(), + let catalog = session_catalog( RestCatalogConfig::builder() .uri(server.url()) .props(props) .build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -1769,15 +1765,11 @@ mod tests { let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); - let catalog = RestCatalog::new( - SessionContext::empty(), + let catalog = session_catalog( RestCatalogConfig::builder() .uri(server.url()) .props(props) .build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -1788,7 +1780,10 @@ mod tests { let oauth_mock = create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "ey000000000001", 200) .await; - catalog.invalidate_token().await.unwrap(); + catalog + .invalidate_token(&SessionContext::empty()) + .await + .unwrap(); let token = catalog.context().await.unwrap().client.token().await; oauth_mock.assert_async().await; assert_eq!(token, Some("ey000000000001".to_string())); @@ -1803,15 +1798,11 @@ mod tests { let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); - let catalog = RestCatalog::new( - SessionContext::empty(), + let catalog = session_catalog( RestCatalogConfig::builder() .uri(server.url()) .props(props) .build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -1822,7 +1813,10 @@ mod tests { let oauth_mock = create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "ey000000000001", 500) .await; - catalog.invalidate_token().await.unwrap(); + catalog + .invalidate_token(&SessionContext::empty()) + .await + .unwrap(); let token = catalog.context().await.unwrap().client.token().await; oauth_mock.assert_async().await; assert_eq!(token, None); @@ -1837,15 +1831,11 @@ mod tests { let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); - let catalog = RestCatalog::new( - SessionContext::empty(), + let catalog = session_catalog( RestCatalogConfig::builder() .uri(server.url()) .props(props) .build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -1856,7 +1846,10 @@ mod tests { let oauth_mock = create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "ey000000000001", 200) .await; - catalog.regenerate_token().await.unwrap(); + catalog + .regenerate_token(&SessionContext::empty()) + .await + .unwrap(); oauth_mock.assert_async().await; let token = catalog.context().await.unwrap().client.token().await; assert_eq!(token, Some("ey000000000001".to_string())); @@ -1871,15 +1864,11 @@ mod tests { let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); - let catalog = RestCatalog::new( - SessionContext::empty(), + let catalog = session_catalog( RestCatalogConfig::builder() .uri(server.url()) .props(props) .build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -1890,7 +1879,7 @@ mod tests { let oauth_mock = create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "ey000000000001", 500) .await; - let invalidate_result = catalog.regenerate_token().await; + let invalidate_result = catalog.regenerate_token(&SessionContext::empty()).await; assert!(invalidate_result.is_err()); oauth_mock.assert_async().await; let token = catalog.context().await.unwrap().client.token().await; @@ -1987,15 +1976,11 @@ mod tests { format!("{}{}", auth_server.url(), auth_server_path).to_string(), ); - let catalog = RestCatalog::new( - SessionContext::empty(), + let catalog = session_catalog( RestCatalogConfig::builder() .uri(server.url()) .props(props) .build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, ); let token = catalog.context().await.unwrap().client.token().await; @@ -2040,15 +2025,12 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); - let _namespaces = catalog.list_namespaces(None).await.unwrap(); + let _namespaces = catalog + .list_namespaces(&SessionContext::empty(), None) + .await + .unwrap(); config_mock.assert_async().await; list_ns_mock.assert_async().await; @@ -2073,15 +2055,12 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); - let namespaces = catalog.list_namespaces(None).await.unwrap(); + let namespaces = catalog + .list_namespaces(&SessionContext::empty(), None) + .await + .unwrap(); let expected_ns = vec![ NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(), @@ -2127,15 +2106,12 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); - let namespaces = catalog.list_namespaces(None).await.unwrap(); + let namespaces = catalog + .list_namespaces(&SessionContext::empty(), None) + .await + .unwrap(); let expected_ns = vec![ NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(), @@ -2229,15 +2205,12 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); - let namespaces = catalog.list_namespaces(None).await.unwrap(); + let namespaces = catalog + .list_namespaces(&SessionContext::empty(), None) + .await + .unwrap(); let expected_ns = vec![ NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(), @@ -2285,16 +2258,11 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); let namespaces = catalog .create_namespace( + &SessionContext::empty(), &NamespaceIdent::from_vec(vec!["ns1".to_string(), "ns11".to_string()]).unwrap(), HashMap::from([("key1".to_string(), "value1".to_string())]), ) @@ -2331,16 +2299,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); let namespaces = catalog - .get_namespace(&NamespaceIdent::new("ns1".to_string())) + .get_namespace( + &SessionContext::empty(), + &NamespaceIdent::new("ns1".to_string()), + ) .await .unwrap(); @@ -2367,17 +2332,14 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); assert!( catalog - .namespace_exists(&NamespaceIdent::new("ns1".to_string())) + .namespace_exists( + &SessionContext::empty(), + &NamespaceIdent::new("ns1".to_string()) + ) .await .unwrap() ); @@ -2436,16 +2398,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); catalog - .drop_namespace(&NamespaceIdent::new("ns1".to_string())) + .drop_namespace( + &SessionContext::empty(), + &NamespaceIdent::new("ns1".to_string()), + ) .await .unwrap(); @@ -2479,16 +2438,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); let tables = catalog - .list_tables(&NamespaceIdent::new("ns1".to_string())) + .list_tables( + &SessionContext::empty(), + &NamespaceIdent::new("ns1".to_string()), + ) .await .unwrap(); @@ -2550,16 +2506,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); let tables = catalog - .list_tables(&NamespaceIdent::new("ns1".to_string())) + .list_tables( + &SessionContext::empty(), + &NamespaceIdent::new("ns1".to_string()), + ) .await .unwrap(); @@ -2684,16 +2637,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); let tables = catalog - .list_tables(&NamespaceIdent::new("ns1".to_string())) + .list_tables( + &SessionContext::empty(), + &NamespaceIdent::new("ns1".to_string()), + ) .await .unwrap(); @@ -2731,19 +2681,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); catalog - .drop_table(&TableIdent::new( - NamespaceIdent::new("ns1".to_string()), - "table1".to_string(), - )) + .drop_table( + &SessionContext::empty(), + &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string()), + ) .await .unwrap(); @@ -2763,20 +2707,14 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); assert!( catalog - .table_exists(&TableIdent::new( - NamespaceIdent::new("ns1".to_string()), - "table1".to_string(), - )) + .table_exists( + &SessionContext::empty(), + &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string(),), + ) .await .unwrap() ); @@ -2837,16 +2775,11 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); catalog .rename_table( + &SessionContext::empty(), &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table1".to_string()), &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "table2".to_string()), ) @@ -2874,19 +2807,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); let table = catalog - .load_table(&TableIdent::new( - NamespaceIdent::new("ns1".to_string()), - "test1".to_string(), - )) + .load_table( + &SessionContext::empty(), + &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "test1".to_string()), + ) .await .unwrap(); @@ -2994,19 +2921,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); let table = catalog - .load_table(&TableIdent::new( - NamespaceIdent::new("ns1".to_string()), - "test1".to_string(), - )) + .load_table( + &SessionContext::empty(), + &TableIdent::new(NamespaceIdent::new("ns1".to_string()), "test1".to_string()), + ) .await; assert!(table.is_err()); @@ -3033,13 +2954,7 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); let table_creation = TableCreation::builder() .name("test1".to_string()) @@ -3086,7 +3001,11 @@ mod tests { .build(); let table = catalog - .create_table(&NamespaceIdent::from_strs(["ns1"]).unwrap(), table_creation) + .create_table( + &SessionContext::empty(), + &NamespaceIdent::from_strs(["ns1"]).unwrap(), + table_creation, + ) .await .unwrap(); @@ -3185,13 +3104,7 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); let table_creation = TableCreation::builder() .name("test1".to_string()) @@ -3213,7 +3126,11 @@ mod tests { .build(); let table_result = catalog - .create_table(&NamespaceIdent::from_strs(["ns1"]).unwrap(), table_creation) + .create_table( + &SessionContext::empty(), + &NamespaceIdent::from_strs(["ns1"]).unwrap(), + table_creation, + ) .await; assert!(table_result.is_err()); @@ -3472,13 +3389,7 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); let table_ident = TableIdent::new(NamespaceIdent::new("ns1".to_string()), "test1".to_string()); let metadata_location = String::from( @@ -3486,7 +3397,7 @@ mod tests { ); let table = catalog - .register_table(&table_ident, metadata_location) + .register_table(&SessionContext::empty(), &table_ident, metadata_location) .await .unwrap(); @@ -3526,13 +3437,7 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - SessionContext::empty(), - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); let table_ident = TableIdent::new(NamespaceIdent::new("ns1".to_string()), "test1".to_string()); @@ -3540,7 +3445,7 @@ mod tests { "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json", ); let table = catalog - .register_table(&table_ident, metadata_location) + .register_table(&SessionContext::empty(), &table_ident, metadata_location) .await; assert!(table.is_err()); @@ -3677,4 +3582,73 @@ mod tests { let catalog = result.unwrap(); assert!(uuid::Uuid::parse_str(catalog.session.session_id()).is_ok()); } + + /// Smoke test: a [`Catalog`] trait method delegates through the façade to `RestSessionCatalog`. + #[tokio::test] + async fn test_rest_catalog_delegates() { + let mut server = Server::new_async().await; + let config_mock = create_config_mock(&mut server).await; + let list_ns_mock = server + .mock("GET", "/v1/namespaces") + .with_body( + r#"{ + "namespaces": [["ns1"]] + }"#, + ) + .create_async() + .await; + + let catalog = RestCatalog::new( + SessionContext::empty(), + RestCatalogConfig::builder().uri(server.url()).build(), + Some(Arc::new(LocalFsStorageFactory)), + Runtime::current(), + None, + ); + + let namespaces = catalog.list_namespaces(None).await.unwrap(); + + assert_eq!(namespaces, vec![ + NamespaceIdent::from_vec(vec!["ns1".to_string()]).unwrap() + ]); + + config_mock.assert_async().await; + list_ns_mock.assert_async().await; + } + + /// Smoke test: the inherent `pub` token methods delegate through the façade to `RestSessionCatalog`. + #[tokio::test] + async fn test_rest_catalog_delegates_token() { + let mut server = Server::new_async().await; + let oauth_mock = create_oauth_mock(&mut server).await; + let config_mock = create_config_mock(&mut server).await; + + let mut props = HashMap::new(); + props.insert("credential".to_string(), "client1:secret1".to_string()); + + let catalog = RestCatalog::new( + SessionContext::empty(), + RestCatalogConfig::builder() + .uri(server.url()) + .props(props) + .build(), + Some(Arc::new(LocalFsStorageFactory)), + Runtime::current(), + None, + ); + + let token = catalog.context().await.unwrap().client.token().await; + assert_eq!(token, Some("ey000000000000".to_string())); + + let oauth_mock_2 = + create_oauth_mock_with_path(&mut server, "/v1/oauth/tokens", "ey000000000001", 200) + .await; + catalog.invalidate_token().await.unwrap(); + let token = catalog.context().await.unwrap().client.token().await; + assert_eq!(token, Some("ey000000000001".to_string())); + + oauth_mock.assert_async().await; + config_mock.assert_async().await; + oauth_mock_2.assert_async().await; + } }