Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/catalog/rest/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ chrono = { workspace = true }
http = { workspace = true }
iceberg = { workspace = true }
itertools = { workspace = true }
moka = { version = "0.12.10", features = ["sync"] }
reqwest = { workspace = true }
serde = { workspace = true }
serde_derive = { workspace = true }
Expand Down
2 changes: 2 additions & 0 deletions crates/catalog/rest/public-api.txt
Original file line number Diff line number Diff line change
Expand Up @@ -381,6 +381,8 @@ pub const iceberg_catalog_rest::AUTH_TYPE_NONE: &str
pub const iceberg_catalog_rest::AUTH_TYPE_OAUTH2: &str
pub const iceberg_catalog_rest::REST_CATALOG_PROP_AUTH_TYPE: &str
pub const iceberg_catalog_rest::REST_CATALOG_PROP_DISABLE_HEADER_REDACTION: &str
pub const iceberg_catalog_rest::REST_CATALOG_PROP_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS: &str
pub const iceberg_catalog_rest::REST_CATALOG_PROP_TABLE_CACHE_MAX_ENTRIES: &str
pub const iceberg_catalog_rest::REST_CATALOG_PROP_URI: &str
pub const iceberg_catalog_rest::REST_CATALOG_PROP_WAREHOUSE: &str
pub trait iceberg_catalog_rest::AuthManager: core::fmt::Debug + core::marker::Send + core::marker::Sync
Expand Down
211 changes: 203 additions & 8 deletions crates/catalog/rest/src/catalog.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ use std::fmt::{Debug, Formatter};
use std::future::Future;
use std::str::FromStr;
use std::sync::{Arc, OnceLock};
use std::time::Duration;

use async_trait::async_trait;
use iceberg::encryption::kms::{KeyManagementClient, KmsClientFactory};
Expand Down Expand Up @@ -65,6 +66,17 @@ pub const REST_CATALOG_PROP_DISABLE_HEADER_REDACTION: &str = "disable-header-red
/// if a `token`, `credential` or `oauth2-server-uri` is configured, `none`
/// otherwise.
pub const REST_CATALOG_PROP_AUTH_TYPE: &str = "rest.auth.type";
/// Maximum number of loaded tables kept for freshness-aware loading (defaults
/// to 100; 0 disables the cache).
pub const REST_CATALOG_PROP_TABLE_CACHE_MAX_ENTRIES: &str = "rest-table-cache.max-entries";
/// How long a loaded table stays cached for freshness-aware loading, in
/// milliseconds (defaults to 5 minutes; must be positive and at most 1000 years).
pub const REST_CATALOG_PROP_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS: &str =
"rest-table-cache.expire-after-write-ms";

/// Tables loaded with an `ETag`, keyed by session id and identifier. Stale
/// entries are harmless: the server answers `304` only while the `ETag` matches.
type TableCache = moka::sync::Cache<(String, TableIdent), (String, Table)>;

const ICEBERG_REST_SPEC_VERSION: &str = "0.14.1";
const CARGO_PKG_VERSION: &str = env!("CARGO_PKG_VERSION");
Expand Down Expand Up @@ -297,6 +309,36 @@ impl RestCatalogConfig {
.unwrap_or(false)
}

/// Builds the table cache for freshness-aware loading.
fn table_cache(&self) -> Result<TableCache> {
let parse = |key: &str, default: u64| -> Result<u64> {
self.props.get(key).map_or(Ok(default), |v| {
v.parse().map_err(|e| {
Error::new(ErrorKind::DataInvalid, format!("Invalid {key}: {v}")).with_source(e)
})
})
};
let expire_after_write = Duration::from_millis(parse(
REST_CATALOG_PROP_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS,
300_000,
)?);
// Zero is ambiguous (disable or never expire), and moka panics beyond 1000 years.
if expire_after_write.is_zero()
|| expire_after_write > Duration::from_secs(1000 * 365 * 24 * 3600)
{
return Err(Error::new(
ErrorKind::DataInvalid,
format!(
"Invalid {REST_CATALOG_PROP_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS}: must be between 1 ms and 1000 years"
),
));
}
Ok(TableCache::builder()
.max_capacity(parse(REST_CATALOG_PROP_TABLE_CACHE_MAX_ENTRIES, 100)?)
.time_to_live(expire_after_write)
.build())
}

/// Merge the `RestCatalogConfig` with the a [`CatalogConfig`] (fetched from the REST server).
pub(crate) fn merge_with_config(mut self, mut config: CatalogConfig) -> Self {
if let Some(uri) = config.overrides.remove(REST_CATALOG_PROP_URI) {
Expand Down Expand Up @@ -433,6 +475,8 @@ struct RestCatalogClient {
config: RestCatalogConfig,
/// Capabilities the server advertises (see [`RestSessionCatalog::supports_endpoint`]).
endpoints: HashSet<Endpoint>,
/// Tables reused when the server answers a load with `304 Not Modified`.
table_cache: TableCache,
}

impl Debug for RestCatalogClient {
Expand Down Expand Up @@ -481,6 +525,7 @@ impl RestCatalogClient {
Ok(Self {
auth_manager,
catalog_session,
table_cache: config.table_cache()?,
config,
http_client,
endpoints,
Expand Down Expand Up @@ -1250,18 +1295,33 @@ impl SessionCatalog for RestSessionCatalog {
table_ident: &TableIdent,
) -> Result<Table> {
let client = self.client().await?;
let cache_key = (context.session_id().to_string(), table_ident.clone());
let cached = client.table_cache.get(&cache_key);

let request = HttpRequest::build(
client
.http_client
.request(Method::GET, client.config.table_endpoint(table_ident)),
)?;
let mut request_builder = client
.http_client
.request(Method::GET, client.config.table_endpoint(table_ident));
if let Some((etag, _)) = &cached {
request_builder = request_builder.header(header::IF_NONE_MATCH, etag);
}
let request = HttpRequest::build(request_builder)?;

let http_response = client.query_catalog(context, request).await?;
let etag = http_response
.headers()
.get(header::ETAG)
.and_then(|v| v.to_str().ok())
.map(str::to_string);

let response = match http_response.status() {
StatusCode::OK | StatusCode::NOT_MODIFIED => {
deserialize_catalog_response::<LoadTableResult>(http_response)?
StatusCode::OK => deserialize_catalog_response::<LoadTableResult>(http_response)?,
StatusCode::NOT_MODIFIED => {
return cached.map(|(_, table)| table).ok_or_else(|| {
Error::new(
ErrorKind::Unexpected,
"Received 304 Not Modified for a table load without If-None-Match",
)
});
}
StatusCode::NOT_FOUND => {
return Err(Error::new(
Expand Down Expand Up @@ -1296,11 +1356,16 @@ impl SessionCatalog for RestSessionCatalog {
table_builder = table_builder.kms_client(kms_client);
}

if let Some(metadata_location) = response.metadata_location {
let table = if let Some(metadata_location) = response.metadata_location {
table_builder.metadata_location(metadata_location).build()
} else {
table_builder.build()
}?;

if let Some(etag) = etag {
client.table_cache.insert(cache_key, (etag, table.clone()));
}
Ok(table)
}

/// Drop a table from the catalog.
Expand Down Expand Up @@ -4023,6 +4088,136 @@ mod tests {
rename_table_mock.assert_async().await;
}

/// Mocks one `loadTable` of `ns1.test1` that expects `if_none_match` and
/// answers `status`, carrying `etag` if given.
async fn load_table_mock(
server: &mut ServerGuard,
if_none_match: mockito::Matcher,
status: usize,
etag: Option<&str>,
) -> Mock {
let mut mock = server
.mock("GET", "/v1/namespaces/ns1/tables/test1")
.match_header("if-none-match", if_none_match)
.with_status(status);
if let Some(etag) = etag {
mock = mock.with_header("etag", etag);
}
if status == 200 {
mock = mock.with_body_from_file(format!(
"{}/testdata/load_table_response.json",
env!("CARGO_MANIFEST_DIR")
));
}
mock.create_async().await
}

fn test1() -> TableIdent {
TableIdent::from_strs(["ns1", "test1"]).unwrap()
}

#[tokio::test]
async fn test_load_table_revalidates_cached_table_with_etag() {
let mut server = Server::new_async().await;
let _config = create_config_mock(&mut server).await;
let v1 = load_table_mock(&mut server, mockito::Matcher::Missing, 200, Some("\"v1\"")).await;
let reuse_v1 = load_table_mock(&mut server, "\"v1\"".into(), 304, None).await;
let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());
let context = SessionContext::empty();

let loaded = catalog.load_table(&context, &test1()).await.unwrap();
let reused = catalog.load_table(&context, &test1()).await.unwrap();
assert_eq!(loaded.metadata_location(), reused.metadata_location());
v1.assert_async().await;
reuse_v1.assert_async().await;

reuse_v1.remove_async().await;
let v2 = load_table_mock(&mut server, "\"v1\"".into(), 200, Some("\"v2\"")).await;
let reuse_v2 = load_table_mock(&mut server, "\"v2\"".into(), 304, None).await;
catalog.load_table(&context, &test1()).await.unwrap();
catalog.load_table(&context, &test1()).await.unwrap();
v2.assert_async().await;
reuse_v2.assert_async().await;
}

#[tokio::test]
async fn test_load_table_not_modified_without_cached_table() {
let mut server = Server::new_async().await;
let _config = create_config_mock(&mut server).await;
let _load = load_table_mock(&mut server, mockito::Matcher::Missing, 304, None).await;
let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());

let err = catalog
.load_table(&SessionContext::empty(), &test1())
.await
.unwrap_err();

assert_eq!(err.kind(), ErrorKind::Unexpected);
}

#[tokio::test]
async fn test_table_cache_is_per_session() {
let mut server = Server::new_async().await;
let _config = create_config_mock(&mut server).await;
let load = load_table_mock(&mut server, mockito::Matcher::Missing, 200, Some("\"v1\""))
.await
.expect(2);
let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build());

catalog
.load_table(&SessionContext::empty(), &test1())
.await
.unwrap();
catalog
.load_table(&SessionContext::empty(), &test1())
.await
.unwrap();

load.assert_async().await;
}

#[tokio::test]
async fn test_table_cache_disabled_with_zero_max_entries() {
let mut server = Server::new_async().await;
let _config = create_config_mock(&mut server).await;
let load = load_table_mock(&mut server, mockito::Matcher::Missing, 200, Some("\"v1\""))
.await
.expect(2);
let catalog = session_catalog(
RestCatalogConfig::builder()
.uri(server.url())
.props(HashMap::from([(
REST_CATALOG_PROP_TABLE_CACHE_MAX_ENTRIES.to_string(),
"0".to_string(),
)]))
.build(),
);
let context = SessionContext::empty();

catalog.load_table(&context, &test1()).await.unwrap();
catalog.load_table(&context, &test1()).await.unwrap();

load.assert_async().await;
}

#[test]
fn test_table_cache_rejects_out_of_range_expiry() {
for value in ["0", "31536000000001"] {
let config = RestCatalogConfig::builder()
.uri("http://localhost".to_string())
.props(HashMap::from([(
REST_CATALOG_PROP_TABLE_CACHE_EXPIRE_AFTER_WRITE_MS.to_string(),
value.to_string(),
)]))
.build();

assert_eq!(
config.table_cache().unwrap_err().kind(),
ErrorKind::DataInvalid
);
}
}

#[tokio::test]
async fn test_load_table_404() {
let mut server = Server::new_async().await;
Expand Down
Loading