Skip to content

feat(catalog-rest): enforce negotiated endpoints for exists checks #2794

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
huan233usc wants to merge 1 commit into
base: main
Choose a base branch
from
+334 −38
Open

feat(catalog-rest): enforce negotiated endpoints for exists checks #2794

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
311 changes: 273 additions & 38 deletions crates/catalog/rest/src/catalog.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

//! This module contains the iceberg REST catalog implementation.

use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use std::future::Future;
use std::str::FromStr;
use std::sync::Arc;
Expand All @@ -41,6 +41,7 @@ use typed_builder::TypedBuilder;
use crate::client::{
HttpClient, deserialize_catalog_response, deserialize_unexpected_catalog_error,
};
use crate::endpoint::Endpoint;
use crate::types::{
CatalogConfig, CommitTableRequest, CommitTableResponse, CreateNamespaceRequest,
CreateTableRequest, ListNamespaceResponse, ListTablesResponse, LoadTableResult,
Expand Down Expand Up @@ -362,6 +363,8 @@ struct RestContext {
///
/// It's could be different from the user config.
config: RestCatalogConfig,
/// Capabilities the server advertises (see [`RestCatalog::supports_endpoint`]).
endpoints: HashSet<Endpoint>,
}

/// Rest catalog implementation.
Expand Down Expand Up @@ -431,14 +434,51 @@ impl RestCatalog {
.get_or_try_init(|| async {
let client = HttpClient::new(&self.user_config)?;
let catalog_config = RestCatalog::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 })
Ok(RestContext {
config,
client,
endpoints,
})
})
.await
}

/// Returns whether the server supports `endpoint`, per the `endpoints` it
/// advertised in `GET /v1/config` (or a default base set when it advertised
/// none).
pub(crate) async fn supports_endpoint(&self, endpoint: &Endpoint) -> Result<bool> {
Ok(self.context().await?.endpoints.contains(endpoint))
}

/// Issue a `HEAD` request to `url` and interpret it as an existence check:
/// `2xx` means it exists, `404` means it doesn't.
async fn check_exists_via_head(&self, url: String) -> Result<bool> {
let context = self.context().await?;
let request = context.client.request(Method::HEAD, url).build()?;
let http_response = context.client.query_catalog(request).await?;

match http_response.status() {
StatusCode::NO_CONTENT | StatusCode::OK => Ok(true),
StatusCode::NOT_FOUND => Ok(false),
_ => Err(deserialize_unexpected_catalog_error(
http_response,
context.client.disable_header_redaction(),
)
.await),
}
}

/// Load the runtime config from the server by `user_config`.
///
/// It's required for a REST catalog to update its config after creation.
Expand Down Expand Up @@ -648,24 +688,22 @@ impl Catalog for RestCatalog {
}

async fn namespace_exists(&self, ns: &NamespaceIdent) -> Result<bool> {
let context = self.context().await?;

let request = context
.client
.request(Method::HEAD, context.config.namespace_endpoint(ns))
.build()?;

let http_response = context.client.query_catalog(request).await?;

match http_response.status() {
StatusCode::NO_CONTENT | StatusCode::OK => Ok(true),
StatusCode::NOT_FOUND => Ok(false),
_ => Err(deserialize_unexpected_catalog_error(
http_response,
context.client.disable_header_redaction(),
)
.await),
let head_endpoint = Endpoint::new(Method::HEAD, "/v1/{prefix}/namespaces/{namespace}");

// 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(&head_endpoint).await? {
return match self.get_namespace(ns).await {
Ok(_) => Ok(true),
Err(e) if e.kind() == ErrorKind::NamespaceNotFound => Ok(false),
Err(e) => Err(e),
};
}

let url = self.context().await?.config.namespace_endpoint(ns);
self.check_exists_via_head(url).await
}

async fn update_namespace(
Expand Down Expand Up @@ -909,24 +947,24 @@ impl Catalog for RestCatalog {

/// Check if a table exists in the catalog.
async fn table_exists(&self, table: &TableIdent) -> Result<bool> {
let context = self.context().await?;

let request = context
.client
.request(Method::HEAD, context.config.table_endpoint(table))
.build()?;

let http_response = context.client.query_catalog(request).await?;

match http_response.status() {
StatusCode::NO_CONTENT | StatusCode::OK => Ok(true),
StatusCode::NOT_FOUND => Ok(false),
_ => Err(deserialize_unexpected_catalog_error(
http_response,
context.client.disable_header_redaction(),
)
.await),
let head_endpoint = Endpoint::new(
Method::HEAD,
"/v1/{prefix}/namespaces/{namespace}/tables/{table}",
);

// 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(&head_endpoint).await? {
return match self.load_table(table).await {
Ok(_) => Ok(true),
Err(e) if e.kind() == ErrorKind::TableNotFound => Ok(false),
Err(e) => Err(e),
};
}

let url = self.context().await?.config.table_endpoint(table);
self.check_exists_via_head(url).await
}

/// Rename a table in the catalog.
Expand Down Expand Up @@ -1183,6 +1221,127 @@ mod tests {
.await
}

/// Config mock that advertises the HEAD table/namespace-exists endpoints, so
/// `{table,namespace}_exists` take the HEAD path rather than the GET fallback.
async fn create_config_mock_with_exists_endpoints(server: &mut ServerGuard) -> Mock {
server
.mock("GET", "/v1/config")
.with_status(200)
.with_body(
r#"{
"overrides": { "warehouse": "s3://iceberg-catalog" },
"defaults": {},
"endpoints": [
"HEAD /v1/{prefix}/namespaces/{namespace}",
"HEAD /v1/{prefix}/namespaces/{namespace}/tables/{table}"
]
}"#,
)
.create_async()
.await
}

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

let config_mock = server
.mock("GET", "/v1/config")
.with_status(200)
.with_body(
r#"{
"overrides": {},
"defaults": {},
"endpoints": [
"GET /v1/{prefix}/namespaces",
"POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan"
]
}"#,
)
.create_async()
.await;

let catalog = RestCatalog::new(
RestCatalogConfig::builder().uri(server.url()).build(),
Some(Arc::new(LocalFsStorageFactory)),
Runtime::current(),
None,
);

let plan = "POST /v1/{prefix}/namespaces/{namespace}/tables/{table}/plan"
.parse::<Endpoint>()
.unwrap();
assert!(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::<Endpoint>()
.unwrap();
assert!(!catalog.supports_endpoint(&delete_ns).await.unwrap());

config_mock.assert_async().await;
}

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

let config_mock = server
.mock("GET", "/v1/config")
.with_status(200)
.with_body(r#"{ "overrides": {}, "defaults": {} }"#)
.create_async()
.await;

let catalog = RestCatalog::new(
RestCatalogConfig::builder().uri(server.url()).build(),
Some(Arc::new(LocalFsStorageFactory)),
Runtime::current(),
None,
);

// A server that omits the `endpoints` field is assumed to support the
// standard base operations.
let load_table = "GET /v1/{prefix}/namespaces/{namespace}/tables/{table}"
.parse::<Endpoint>()
.unwrap();
assert!(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::<Endpoint>()
.unwrap();
assert!(!catalog.supports_endpoint(&plan).await.unwrap());

config_mock.assert_async().await;
}

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

// An explicit empty list is treated the same as an absent field: fall
// back to the standard base set.
let config_mock = server
.mock("GET", "/v1/config")
.with_status(200)
.with_body(r#"{ "overrides": {}, "defaults": {}, "endpoints": [] }"#)
.create_async()
.await;

let catalog = RestCatalog::new(
RestCatalogConfig::builder().uri(server.url()).build(),
Some(Arc::new(LocalFsStorageFactory)),
Runtime::current(),
None,
);

let load_table = "GET /v1/{prefix}/namespaces/{namespace}/tables/{table}"
.parse::<Endpoint>()
.unwrap();
assert!(catalog.supports_endpoint(&load_table).await.unwrap());

config_mock.assert_async().await;
}

async fn create_oauth_mock(server: &mut ServerGuard) -> Mock {
create_oauth_mock_with_path(server, "/v1/oauth/tokens", "ey000000000000", 200).await
}
Expand Down Expand Up @@ -1873,7 +2032,7 @@ mod tests {
async fn check_namespace_exists() {
let mut server = Server::new_async().await;

let config_mock = create_config_mock(&mut server).await;
let config_mock = create_config_mock_with_exists_endpoints(&mut server).await;

let get_ns_mock = server
.mock("HEAD", "/v1/namespaces/ns1")
Expand All @@ -1899,6 +2058,43 @@ mod tests {
get_ns_mock.assert_async().await;
}

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

// No `endpoints` advertised, and the default set has no HEAD namespace
// route, so `namespace_exists` falls back to a GET load-namespace.
let config_mock = create_config_mock(&mut server).await;
let get_ns_mock = server
.mock("GET", "/v1/namespaces/ns1")
.with_status(200)
.with_body(
r#"{
"namespace": ["ns1"],
"properties": {}
}"#,
)
.create_async()
.await;

let catalog = RestCatalog::new(
RestCatalogConfig::builder().uri(server.url()).build(),
Some(Arc::new(LocalFsStorageFactory)),
Runtime::current(),
None,
);

assert!(
catalog
.namespace_exists(&NamespaceIdent::new("ns1".to_string()))
.await
.unwrap()
);

config_mock.assert_async().await;
get_ns_mock.assert_async().await;
}

#[tokio::test]
async fn test_drop_namespace() {
let mut server = Server::new_async().await;
Expand Down Expand Up @@ -2225,7 +2421,7 @@ mod tests {
async fn test_check_table_exists() {
let mut server = Server::new_async().await;

let config_mock = create_config_mock(&mut server).await;
let config_mock = create_config_mock_with_exists_endpoints(&mut server).await;

let check_table_exists_mock = server
.mock("HEAD", "/v1/namespaces/ns1/tables/table1")
Expand Down Expand Up @@ -2254,6 +2450,45 @@ mod tests {
check_table_exists_mock.assert_async().await;
}

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

// No `endpoints` advertised, and the default set has no HEAD table
// route, so `table_exists` falls back to a GET load-table.
let config_mock = create_config_mock(&mut server).await;
let load_table_mock = server
.mock("GET", "/v1/namespaces/ns1/tables/table1")
.with_status(200)
.with_body_from_file(format!(
"{}/testdata/{}",
env!("CARGO_MANIFEST_DIR"),
"load_table_response.json"
))
.create_async()
.await;

let catalog = RestCatalog::new(
RestCatalogConfig::builder().uri(server.url()).build(),
Some(Arc::new(LocalFsStorageFactory)),
Runtime::current(),
None,
);

assert!(
catalog
.table_exists(&TableIdent::new(
NamespaceIdent::new("ns1".to_string()),
"table1".to_string(),
))
.await
.unwrap()
);

config_mock.assert_async().await;
load_table_mock.assert_async().await;
}

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