-
Notifications
You must be signed in to change notification settings - Fork 510
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.
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
apache:main
Choose a base branch
from
huan233usc:feat/rest-endpoint-enforcement
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+334
−38
Open
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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; | ||
|
|
@@ -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, | ||
|
|
@@ -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. | ||
|
|
@@ -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. | ||
|
|
@@ -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}"); | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The construction is per-call. Can we make this static? |
||
|
|
||
| // 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( | ||
|
|
@@ -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}", | ||
| ); | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Same here, can we make this static? |
||
|
|
||
| // 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. | ||
|
|
@@ -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 | ||
| } | ||
|
|
@@ -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") | ||
|
|
@@ -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; | ||
|
|
@@ -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") | ||
|
|
@@ -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; | ||
|
|
||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We can pass the context in directly. Otherwise we are fetching context twice: first to get url, then to get client