diff --git a/crates/catalog/rest/public-api.txt b/crates/catalog/rest/public-api.txt index 776b11c40a..ab06ea50b2 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 da05fc7e3b..64d633eec3 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::{ @@ -59,97 +59,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)) } } } @@ -157,7 +101,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 } } @@ -364,9 +315,207 @@ struct RestContext { config: RestCatalogConfig, } +impl RestContext { + async fn init(config: &RestCatalogConfig) -> Result { + let client = HttpClient::new(config)?; + + let catalog_config = RestSessionCatalog::load_config(&client, config).await?; + let config = config.clone().merge_with_config(catalog_config); + + let client = client.update_with(&config)?; + + Ok(RestContext { config, client }) + } +} + /// 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. @@ -379,8 +528,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>, @@ -396,8 +545,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 @@ -428,44 +598,10 @@ impl RestCatalog { /// 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 = RestCatalog::load_config(&client, &self.user_config).await?; - let config = self.user_config.clone().merge_with_config(catalog_config); - let client = client.update_with(&config)?; - - Ok(RestContext { config, client }) - }) + .get_or_try_init(async || RestContext::init(&self.user_config).await) .await } - /// Load the runtime config from the server by `user_config`. - /// - /// It's required for a REST catalog to update its config after creation. - async fn load_config( - client: &HttpClient, - user_config: &RestCatalogConfig, - ) -> Result { - let mut request_builder = client.request(Method::GET, user_config.config_endpoint()); - - if let Some(warehouse_location) = &user_config.warehouse { - request_builder = request_builder.query(&[("warehouse", warehouse_location)]); - } - - let request = request_builder.build()?; - - let http_response = client.query_catalog(request).await?; - - match http_response.status() { - StatusCode::OK => deserialize_catalog_response(http_response).await, - _ => Err(deserialize_unexpected_catalog_error( - http_response, - client.disable_header_redaction(), - ) - .await), - } - } - async fn load_file_io( &self, metadata_location: Option<&str>, @@ -507,29 +643,39 @@ 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. + /// Load the runtime config from the server by `user_config`. /// - /// 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 + /// It's required for a REST catalog to update its config after creation. + async fn load_config( + client: &HttpClient, + user_config: &RestCatalogConfig, + ) -> Result { + let mut request_builder = client.request(Method::GET, user_config.config_endpoint()); + + if let Some(warehouse_location) = &user_config.warehouse { + request_builder = request_builder.query(&[("warehouse", warehouse_location)]); + } + + let request = request_builder.build()?; + + let http_response = client.query_catalog(request).await?; + + match http_response.status() { + StatusCode::OK => deserialize_catalog_response(http_response).await, + _ => Err(deserialize_unexpected_catalog_error( + http_response, + client.disable_header_redaction(), + ) + .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?; @@ -585,6 +731,7 @@ impl Catalog for RestCatalog { async fn create_namespace( &self, + _session: &SessionContext, namespace: &NamespaceIdent, properties: HashMap, ) -> Result { @@ -619,7 +766,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 @@ -647,7 +798,11 @@ impl Catalog for RestCatalog { } } - async fn namespace_exists(&self, ns: &NamespaceIdent) -> Result { + async fn namespace_exists( + &self, + _session: &SessionContext, + ns: &NamespaceIdent, + ) -> Result { let context = self.context().await?; let request = context @@ -670,6 +825,7 @@ impl Catalog for RestCatalog { async fn update_namespace( &self, + _session: &SessionContext, _namespace: &NamespaceIdent, _properties: HashMap, ) -> Result<()> { @@ -679,7 +835,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 @@ -703,7 +863,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(); @@ -757,6 +921,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
{ @@ -841,7 +1006,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 @@ -897,18 +1066,18 @@ 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 { let context = self.context().await?; let request = context @@ -930,7 +1099,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 @@ -964,6 +1138,7 @@ impl Catalog for RestCatalog { async fn register_table( &self, + _session: &SessionContext, table_ident: &TableIdent, metadata_location: String, ) -> Result
{ @@ -1030,7 +1205,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 @@ -1107,6 +1286,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; @@ -1128,6 +1416,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; @@ -1146,12 +1444,7 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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 @@ -1219,14 +1512,11 @@ mod tests { let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); - let catalog = RestCatalog::new( + 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; @@ -1268,14 +1558,11 @@ mod tests { let config_mock = create_config_mock(&mut server).await; - let catalog = RestCatalog::new( + 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; @@ -1294,14 +1581,11 @@ mod tests { let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); - let catalog = RestCatalog::new( + 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; @@ -1312,7 +1596,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())); @@ -1327,14 +1614,11 @@ mod tests { let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); - let catalog = RestCatalog::new( + 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; @@ -1345,7 +1629,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); @@ -1360,14 +1647,11 @@ mod tests { let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); - let catalog = RestCatalog::new( + 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; @@ -1378,7 +1662,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())); @@ -1393,14 +1680,11 @@ mod tests { let mut props = HashMap::new(); props.insert("credential".to_string(), "client1:secret1".to_string()); - let catalog = RestCatalog::new( + 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; @@ -1411,7 +1695,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; @@ -1508,14 +1792,11 @@ mod tests { format!("{}{}", auth_server.url(), auth_server_path).to_string(), ); - let catalog = RestCatalog::new( + 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; @@ -1555,19 +1836,17 @@ mod tests { .with_body( r#"{ "namespaces": [] - }"#, - ) - .create_async() - .await; - - let catalog = RestCatalog::new( - RestCatalogConfig::builder().uri(server.url()).build(), - Some(Arc::new(LocalFsStorageFactory)), - Runtime::current(), - None, - ); + }"#, + ) + .create_async() + .await; - let _namespaces = catalog.list_namespaces(None).await.unwrap(); + let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); + + let _namespaces = catalog + .list_namespaces(&SessionContext::empty(), None) + .await + .unwrap(); config_mock.assert_async().await; list_ns_mock.assert_async().await; @@ -1592,14 +1871,12 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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(), @@ -1645,14 +1922,12 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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(), @@ -1746,14 +2021,12 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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(), @@ -1801,15 +2074,11 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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())]), ) @@ -1846,15 +2115,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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(); @@ -1881,16 +2148,14 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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() ); @@ -1911,15 +2176,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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(); @@ -1953,15 +2216,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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(); @@ -2023,15 +2284,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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(); @@ -2156,15 +2415,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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(); @@ -2202,18 +2459,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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(); @@ -2233,19 +2485,14 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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() ); @@ -2266,15 +2513,11 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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()), ) @@ -2302,18 +2545,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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(); @@ -2421,18 +2659,13 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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()); @@ -2459,12 +2692,7 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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()) @@ -2511,7 +2739,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(); @@ -2610,12 +2842,7 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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()) @@ -2637,7 +2864,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()); @@ -2682,6 +2913,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2828,6 +3060,7 @@ mod tests { .await; let catalog = RestCatalog::new( + SessionContext::empty(), RestCatalogConfig::builder().uri(server.url()).build(), Some(Arc::new(LocalFsStorageFactory)), Runtime::current(), @@ -2894,12 +3127,7 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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( @@ -2907,7 +3135,7 @@ mod tests { ); let table = catalog - .register_table(&table_ident, metadata_location) + .register_table(&SessionContext::empty(), &table_ident, metadata_location) .await .unwrap(); @@ -2947,12 +3175,7 @@ mod tests { .create_async() .await; - let catalog = RestCatalog::new( - 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()); @@ -2960,7 +3183,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()); @@ -2989,7 +3212,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); @@ -3019,4 +3243,150 @@ 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()); + } + + /// 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; + } } diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index fcf3cb99e6..4ceec5e441 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -3636,6 +3636,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 @@ -3743,5 +3753,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 d2dc1ac1a0..d4cc80c8f4 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; +}