feat(rest): support vended storage credentials with per-prefix FileIO #2651
+1,148
−47
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 |
|---|---|---|
|
|
@@ -49,7 +49,7 @@ use crate::response::HttpResponse; | |
| use crate::types::{ | ||
| CatalogConfig, CommitTableRequest, CommitTableResponse, CreateNamespaceRequest, | ||
| CreateTableRequest, ListNamespaceResponse, ListTablesResponse, LoadTableResult, | ||
| NamespaceResponse, RegisterTableRequest, RenameTableRequest, StorageCredential, | ||
| }; | ||
|
|
||
| /// REST catalog URI | ||
|
|
@@ -847,6 +847,7 @@ impl RestSessionCatalog { | |
| &self, | ||
| metadata_location: Option<&str>, | ||
| extra_config: Option<HashMap<String, String>>, | ||
| storage_credentials: Option<&[StorageCredential]>, | ||
| ) -> Result<FileIO> { | ||
| let mut props = self.client().await?.config.props.clone(); | ||
| if let Some(config) = extra_config { | ||
|
|
@@ -879,9 +880,20 @@ impl RestSessionCatalog { | |
| ) | ||
| })?; | ||
|
|
||
| let mut builder = FileIOBuilder::new(factory).with_props(props.clone()); | ||
|
|
||
| // Vended credentials are scoped per location prefix: give each its own | ||
| // storage. Paths under no vended prefix fall back to the default `props` | ||
| // above, which carry no credentials. | ||
| if let Some(creds) = storage_credentials { | ||
| for cred in creds { | ||
| let mut prefixed = props.clone(); | ||
| prefixed.extend(cred.config.clone()); | ||
| builder = builder.with_prefixed_props(cred.prefix.clone(), prefixed); | ||
| } | ||
| } | ||
|
|
||
| Ok(builder.build()) | ||
| } | ||
| } | ||
|
|
||
|
|
@@ -1181,14 +1193,15 @@ impl SessionCatalog for RestSessionCatalog { | |
| "Metadata location missing in `create_table` response!", | ||
| ))?; | ||
|
|
||
| let mut base_config = response.config.clone(); | ||
| base_config.extend(self.user_config.props.clone()); | ||
|
|
||
| let file_io = self | ||
| .load_file_io( | ||
| Some(metadata_location), | ||
| Some(base_config), | ||
| response.storage_credentials.as_deref(), | ||
| ) | ||
| .await?; | ||
|
|
||
| let mut table_builder = Table::builder() | ||
|
|
@@ -1219,6 +1232,9 @@ impl SessionCatalog for RestSessionCatalog { | |
| ) -> Result<Table> { | ||
| let client = self.client().await?; | ||
|
|
||
| // Vended credentials are opt-in via a `header.X-Iceberg-Access-Delegation` | ||
| // catalog property (applied to every request like the Iceberg Java client); | ||
| // any returned `storage_credentials` are wired into the FileIO below. | ||
| let request = HttpRequest::build( | ||
| client | ||
| .http_client | ||
|
|
@@ -1245,14 +1261,15 @@ impl SessionCatalog for RestSessionCatalog { | |
| } | ||
| }; | ||
|
|
||
| let mut base_config = response.config.clone(); | ||
| base_config.extend(self.user_config.props.clone()); | ||
|
|
||
| let file_io = self | ||
| .load_file_io( | ||
| response.metadata_location.as_deref(), | ||
| Some(base_config), | ||
| response.storage_credentials.as_deref(), | ||
| ) | ||
| .await?; | ||
|
|
||
| let mut table_builder = Table::builder() | ||
|
|
@@ -1391,7 +1408,16 @@ impl SessionCatalog for RestSessionCatalog { | |
| "Metadata location missing in `register_table` response!", | ||
| ))?; | ||
|
|
||
| let mut base_config = response.config.clone(); | ||
| base_config.extend(self.user_config.props.clone()); | ||
|
|
||
| let file_io = self | ||
| .load_file_io( | ||
| Some(metadata_location), | ||
| Some(base_config), | ||
| response.storage_credentials.as_deref(), | ||
| ) | ||
| .await?; | ||
|
|
||
| let mut table_builder = Table::builder() | ||
| .identifier(table_ident.clone()) | ||
|
|
@@ -1469,8 +1495,12 @@ impl SessionCatalog for RestSessionCatalog { | |
| } | ||
| }; | ||
|
|
||
| // The commit response carries no credentials, so this FileIO has only the | ||
| // catalog-level config. `Transaction::do_commit` swaps in the credentialed | ||
| // one from its pre-commit load, and it is the only caller there can be: | ||
| // `TableCommit` is buildable inside the `iceberg` crate alone. | ||
| let file_io = self | ||
| .load_file_io(Some(&response.metadata_location), None, None) | ||
|
Contributor
There was a problem hiding this comment. Wonder what the risks are here of not creating FileIo here without the credentials? Is that idea that after we have credential refresh we will just automatically refresh on the next read / write?
Member
Author
There was a problem hiding this comment. @xanderbailey Good question. No risk on the
Contributor
There was a problem hiding this comment. I guess in Java they have
Member
Author
There was a problem hiding this comment. @xanderbailey Good comparison — I think we're actually closer to it than it looks. On bullet-proofing: |
||
| .await?; | ||
|
|
||
| let mut table_builder = Table::builder() | ||
|
|
@@ -3794,6 +3824,114 @@ mod tests { | |
| rename_table_mock.assert_async().await; | ||
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn test_load_table_uses_vended_credentials() { | ||
| let mut server = Server::new_async().await; | ||
|
|
||
| let config_mock = create_config_mock(&mut server).await; | ||
|
|
||
| // Vended credentials are opt-in via a `header.*` catalog property (like the | ||
| // Java client). With it configured, the header is sent and the response's | ||
| // `storage-credentials` are accepted (the FileIO builds). | ||
| let load_table_mock = server | ||
| .mock("GET", "/v1/namespaces/ns1/tables/test1") | ||
| .match_header("x-iceberg-access-delegation", "vended-credentials") | ||
| .with_status(200) | ||
| .with_body_from_file(format!( | ||
| "{}/testdata/{}", | ||
| env!("CARGO_MANIFEST_DIR"), | ||
| "load_table_response_with_credentials.json" | ||
| )) | ||
| .create_async() | ||
| .await; | ||
|
|
||
| let props = HashMap::from([( | ||
| "header.X-Iceberg-Access-Delegation".to_string(), | ||
| "vended-credentials".to_string(), | ||
| )]); | ||
| let catalog = RestCatalog::new( | ||
| SessionContext::empty(), | ||
| RestCatalogConfig::builder() | ||
| .uri(server.url()) | ||
| .props(props) | ||
| .build(), | ||
| None, | ||
| Some(Arc::new(LocalFsStorageFactory)), | ||
| Runtime::current(), | ||
| None, | ||
| ); | ||
|
|
||
| let table = catalog | ||
| .load_table(&TableIdent::from_strs(["ns1", "test1"]).unwrap()) | ||
| .await | ||
| .unwrap(); | ||
|
|
||
| assert_eq!( | ||
| "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json", | ||
| table.metadata_location().unwrap() | ||
| ); | ||
|
|
||
| // The point of the feature: the vended credentials reach the FileIO, | ||
| // scoped to the prefix the server sent them for. | ||
| let file_io = table.file_io(); | ||
| let vended = file_io.config_for("s3://warehouse/database/table/data/f.parquet"); | ||
| assert_eq!( | ||
| vended.get("s3.access-key-id"), | ||
| Some(&"vended-key-id".to_string()) | ||
| ); | ||
| assert_eq!( | ||
| vended.get("s3.session-token"), | ||
| Some(&"vended-token".to_string()) | ||
| ); | ||
| // A path outside the prefix keeps the credential-free default. | ||
| assert_eq!( | ||
| file_io | ||
| .config_for("s3://warehouse/other/f.parquet") | ||
| .get("s3.access-key-id"), | ||
| None | ||
| ); | ||
|
|
||
| config_mock.assert_async().await; | ||
| load_table_mock.assert_async().await; | ||
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn test_load_table_omits_delegation_header_by_default() { | ||
| let mut server = Server::new_async().await; | ||
|
|
||
| let config_mock = create_config_mock(&mut server).await; | ||
|
|
||
| // No delegation header is hardcoded: without a `header.*` prop, none is sent. | ||
| let load_table_mock = server | ||
| .mock("GET", "/v1/namespaces/ns1/tables/test1") | ||
| .match_header("x-iceberg-access-delegation", mockito::Matcher::Missing) | ||
| .with_status(200) | ||
| .with_body_from_file(format!( | ||
| "{}/testdata/{}", | ||
| env!("CARGO_MANIFEST_DIR"), | ||
| "load_table_response.json" | ||
| )) | ||
| .create_async() | ||
| .await; | ||
|
|
||
| let catalog = RestCatalog::new( | ||
| SessionContext::empty(), | ||
| RestCatalogConfig::builder().uri(server.url()).build(), | ||
| None, | ||
| Some(Arc::new(LocalFsStorageFactory)), | ||
| Runtime::current(), | ||
| None, | ||
| ); | ||
|
|
||
| catalog | ||
| .load_table(&TableIdent::from_strs(["ns1", "test1"]).unwrap()) | ||
| .await | ||
| .unwrap(); | ||
|
|
||
| config_mock.assert_async().await; | ||
| load_table_mock.assert_async().await; | ||
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn test_create_table() { | ||
| let mut server = Server::new_async().await; | ||
|
|
@@ -4009,6 +4147,7 @@ mod tests { | |
|
|
||
| let config_mock = create_config_mock(&mut server).await; | ||
|
|
||
| // GET hit once: the transaction refreshes the table before committing. | ||
| let load_table_mock = server | ||
| .mock("GET", "/v1/namespaces/ns1/tables/test1") | ||
| .with_status(200) | ||
|
|
@@ -4017,6 +4156,7 @@ mod tests { | |
| env!("CARGO_MANIFEST_DIR"), | ||
| "load_table_response.json" | ||
| )) | ||
| .expect(1) | ||
| .create_async() | ||
| .await; | ||
|
|
||
|
|
@@ -4273,6 +4413,56 @@ mod tests { | |
| register_table_mock.assert_async().await; | ||
| } | ||
|
|
||
| /// The register response is a `LoadTableResult`, so the vended credentials | ||
| /// it carries have to reach the table's FileIO like `load_table`'s do. | ||
| #[tokio::test] | ||
| async fn test_register_table_uses_vended_credentials() { | ||
| let mut server = Server::new_async().await; | ||
| let config_mock = create_config_mock(&mut server).await; | ||
| let register_table_mock = server | ||
| .mock("POST", "/v1/namespaces/ns1/register") | ||
| .with_status(200) | ||
| .with_body_from_file(format!( | ||
| "{}/testdata/{}", | ||
| env!("CARGO_MANIFEST_DIR"), | ||
| "load_table_response_with_credentials.json" | ||
| )) | ||
| .create_async() | ||
| .await; | ||
|
|
||
| let catalog = session_catalog(RestCatalogConfig::builder().uri(server.url()).build()); | ||
| let table_ident = TableIdent::from_strs(["ns1", "test1"]).unwrap(); | ||
| let table = catalog | ||
| .register_table( | ||
| &SessionContext::empty(), | ||
| &table_ident, | ||
| "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json".to_string(), | ||
| ) | ||
| .await | ||
| .unwrap(); | ||
|
|
||
| let vended = table | ||
| .file_io() | ||
| .config_for("s3://warehouse/database/table/data/f.parquet"); | ||
| assert_eq!( | ||
| vended.get("s3.access-key-id").map(String::as_str), | ||
| Some("vended-key-id") | ||
| ); | ||
| // Scoped to the prefix the server sent them for. | ||
| let outside = table | ||
| .file_io() | ||
| .config_for("s3://other-bucket/data/f.parquet"); | ||
| assert_eq!(outside.get("s3.access-key-id"), None); | ||
| // The response's table config reaches the FileIO too, as in load_table. | ||
| assert_eq!( | ||
| table.file_io().config().get("region").map(String::as_str), | ||
| Some("us-west-2") | ||
| ); | ||
|
|
||
| config_mock.assert_async().await; | ||
| register_table_mock.assert_async().await; | ||
| } | ||
|
|
||
| #[tokio::test] | ||
| async fn test_register_table_404() { | ||
| let mut server = Server::new_async().await; | ||
|
|
||
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 |
|---|---|---|
| @@ -0,0 +1,78 @@ | ||
| { | ||
| "metadata-location": "s3://warehouse/database/table/metadata/00001-5f2f8166-244c-4eae-ac36-384ecdec81fc.gz.metadata.json", | ||
| "metadata": { | ||
| "format-version": 1, | ||
| "table-uuid": "b55d9dda-6561-423a-8bfc-787980ce421f", | ||
| "location": "s3://warehouse/database/table", | ||
| "last-updated-ms": 1646787054459, | ||
| "last-column-id": 2, | ||
| "schema": { | ||
| "type": "struct", | ||
| "schema-id": 0, | ||
| "fields": [ | ||
| {"id": 1, "name": "id", "required": false, "type": "int"}, | ||
| {"id": 2, "name": "data", "required": false, "type": "string"} | ||
| ] | ||
| }, | ||
| "current-schema-id": 0, | ||
| "schemas": [ | ||
| { | ||
| "type": "struct", | ||
| "schema-id": 0, | ||
| "fields": [ | ||
| {"id": 1, "name": "id", "required": false, "type": "int"}, | ||
| {"id": 2, "name": "data", "required": false, "type": "string"} | ||
| ] | ||
| } | ||
| ], | ||
| "partition-spec": [], | ||
| "default-spec-id": 0, | ||
| "partition-specs": [{"spec-id": 0, "fields": []}], | ||
| "last-partition-id": 999, | ||
| "default-sort-order-id": 0, | ||
| "sort-orders": [{"order-id": 0, "fields": []}], | ||
| "properties": {"owner": "bryan", "write.metadata.compression-codec": "gzip"}, | ||
| "current-snapshot-id": 3497810964824022504, | ||
| "refs": {"main": {"snapshot-id": 3497810964824022504, "type": "branch"}}, | ||
| "snapshots": [ | ||
| { | ||
| "snapshot-id": 3497810964824022504, | ||
| "timestamp-ms": 1646787054459, | ||
| "summary": { | ||
| "operation": "append", | ||
| "spark.app.id": "local-1646787004168", | ||
| "added-data-files": "1", | ||
| "added-records": "1", | ||
| "added-files-size": "697", | ||
| "changed-partition-count": "1", | ||
| "total-records": "1", | ||
| "total-files-size": "697", | ||
| "total-data-files": "1", | ||
| "total-delete-files": "0", | ||
| "total-position-deletes": "0", | ||
| "total-equality-deletes": "0" | ||
| }, | ||
| "manifest-list": "s3://warehouse/database/table/metadata/snap-3497810964824022504-1-c4f68204-666b-4e50-a9df-b10c34bf6b82.avro", | ||
| "schema-id": 0 | ||
| } | ||
| ], | ||
| "snapshot-log": [{"timestamp-ms": 1646787054459, "snapshot-id": 3497810964824022504}], | ||
| "metadata-log": [ | ||
| { | ||
| "timestamp-ms": 1646787031514, | ||
| "metadata-file": "s3://warehouse/database/table/metadata/00000-88484a1c-00e5-4a07-a787-c0e7aeffa805.gz.metadata.json" | ||
| } | ||
| ] | ||
| }, | ||
| "config": {"client.factory": "io.tabular.iceberg.catalog.TabularAwsClientFactory", "region": "us-west-2"}, | ||
| "storage-credentials": [ | ||
| { | ||
| "prefix": "s3://warehouse/database/table", | ||
| "config": { | ||
| "s3.access-key-id": "vended-key-id", | ||
| "s3.secret-access-key": "vended-secret", | ||
| "s3.session-token": "vended-token" | ||
| } | ||
| } | ||
| ] | ||
| } |
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.
qq: do we plan to handle credential refresh as follow up?
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.
@huan233usc Good question！— not in this PR. FileIO is built once from the LoadTable creds, so no refresh yet. Good follow-up: the opendal S3 backend already has a ProvideCredential hook (reqsign re-invokes on expires_in), so a provider re-fetching from loadCredentials before s3.session-token-expires-at-ms — like Java's VendedCredentialsProvider — fits there. Will track as a separate PR.
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.
hey @huan233usc @plusplusjiajia, cred refresh is pretty important for a workflow I'm working on, so I've created a PR for it here: #2932
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.
@zakariya-s Great to see this — credential refresh is the natural next step on top of the vended credentials here.