From b6b6c190bed5528ca3697b55a31267af57499f1c Mon Sep 17 00:00:00 2001 From: Patrick Butler Date: Tue, 25 Aug 2026 11:59:46 -0400 Subject: [PATCH 1/2] storage: Refresh catalog-vended storage credentials Vended credentials expire, and nothing refreshed them: `loadTable` hands their access keys to the FileIO once and OpenDAL keeps signing with them until the dataflow restarts. Adds `VendedCredentialLoader`, a `ProvideCredential` implementation that re-fetches from the catalog's `loadCredentials` endpoint. OpenDAL rebuilds its `Operator` for every file operation, so reqsign's own credential cache never survives one call; the loader therefore caches with its own expiry deadline, taken from `s3.session-token-expires-at-ms` where the catalog reports one and a short interval where it does not. On a 401 or 403 it invalidates the catalog token so the next attempt mints a fresh one. Materialize now builds the OAuth2 provider itself rather than passing a `credential` catalog property, so one token object serves both catalog requests and credential refreshes. The catalog client rejects a custom authenticator combined with that property, so the two cannot coexist. `connect` grows a table argument because the credentials endpoint is table-scoped. Installing a loader also means it alone supplies S3 credentials, since OpenDAL replaces its whole provider chain, so one is installed only when the connection asked for delegation. --- src/sql/src/pure.rs | 4 +- src/storage-types/src/connections.rs | 417 +++++++++++++++++++++++++-- src/storage/src/sink/iceberg.rs | 18 +- 3 files changed, 411 insertions(+), 28 deletions(-) diff --git a/src/sql/src/pure.rs b/src/sql/src/pure.rs index 82d6f4b79f05a..25c982728708a 100644 --- a/src/sql/src/pure.rs +++ b/src/sql/src/pure.rs @@ -625,8 +625,10 @@ async fn purify_create_sink( // Now that we've validated the sink's storage creds (if they exist) // we _could_ use them to build a complete Iceberg client (both catalog and storage). // TODO(kynan): Actually use those sink-specific creds here instead of ignoring them. + // Purification only proves the catalog is reachable, so it needs no table-scoped + // storage credentials. let _catalog = connection - .connect(storage_configuration, InTask::No) + .connect(storage_configuration, InTask::No, None) .await .map_err(|e| IcebergSinkPurificationError::CatalogError(Arc::new(e)))?; } diff --git a/src/storage-types/src/connections.rs b/src/storage-types/src/connections.rs index 67f4bc44d4623..e969418c09f39 100644 --- a/src/storage-types/src/connections.rs +++ b/src/storage-types/src/connections.rs @@ -14,7 +14,7 @@ use std::collections::{BTreeMap, BTreeSet}; use std::fmt; use std::net::SocketAddr; use std::sync::Arc; -use std::time::SystemTime; +use std::time::{Duration, Instant, SystemTime}; use anyhow::{Context, anyhow}; use async_trait::async_trait; @@ -24,15 +24,17 @@ use aws_sigv4::sign::v4; // Aliased to avoid colliding with `mz_ccsr::tls::Identity`. use aws_smithy_runtime_api::client::identity::Identity as AwsIdentity; use base64::Engine; -use http::{HeaderName, HeaderValue}; +use http::{HeaderMap, HeaderName, HeaderValue}; use iceberg::Catalog; use iceberg::CatalogBuilder; +use iceberg::TableIdent; use iceberg::io::{ GCS_CREDENTIALS_JSON, GCS_DISABLE_CONFIG_LOAD, GCS_DISABLE_VM_METADATA, GCS_USER_PROJECT, - S3_ACCESS_KEY_ID, S3_DISABLE_EC2_METADATA, S3_REGION, S3_SECRET_ACCESS_KEY, + S3_ACCESS_KEY_ID, S3_DISABLE_EC2_METADATA, S3_REGION, S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN, }; use iceberg_catalog_rest::{ - REST_CATALOG_PROP_URI, REST_CATALOG_PROP_WAREHOUSE, RequestAuthenticator, RestCatalogBuilder, + OAuth2TokenProvider, REST_CATALOG_PROP_URI, REST_CATALOG_PROP_WAREHOUSE, RequestAuthenticator, + RestCatalogBuilder, StorageCredential, TokenProvider, }; use iceberg_storage_opendal::{ AwsCredential, CustomAwsCredentialLoader, OpenDalStorageFactory, ProvideCredential, @@ -64,10 +66,11 @@ use rdkafka::config::FromClientConfigAndContext; use rdkafka::consumer::{BaseConsumer, Consumer}; use regex::Regex; use reqsign_core::time::Timestamp; -use reqwest::Request; +use reqwest::{Request, StatusCode}; use serde::{Deserialize, Deserializer, Serialize}; use tokio::net; use tokio::runtime::Handle; +use tokio::sync::Mutex; use tokio_postgres::config::SslMode; use tracing::{debug, info, warn}; use url::Url; @@ -92,15 +95,35 @@ pub mod gcp; pub mod inline; pub mod string_or_secret; -const REST_CATALOG_PROP_SCOPE: &str = "scope"; -const REST_CATALOG_PROP_CREDENTIAL: &str = "credential"; -/// Overrides the OAuth2 token endpoint. Spelled `uri` because that is the property name the -/// Iceberg REST clients agree on, even though the SQL option says `URL`. -const REST_CATALOG_PROP_OAUTH2_SERVER_URI: &str = "oauth2-server-uri"; +/// The OAuth2 form field naming the scopes a token is requested for. +/// +/// Materialize drives the OAuth2 exchange itself rather than through the `credential`, +/// `oauth2-server-uri`, and `scope` catalog properties, so that one token object serves both +/// catalog requests and storage-credential refreshes. +const OAUTH2_PARAM_SCOPE: &str = "scope"; /// Requests catalog-vended storage credentials. `iceberg-rust` turns `header.*` catalog /// properties into headers on every REST request, the same way the Iceberg Java client /// carries this one. const REST_CATALOG_PROP_ACCESS_DELEGATION: &str = "header.X-Iceberg-Access-Delegation"; +/// The same header, spelled for requests Materialize issues itself rather than through the +/// catalog client. +const ICEBERG_ACCESS_DELEGATION_HEADER: &str = "X-Iceberg-Access-Delegation"; + +/// Reports when vended S3 credentials expire, as spelled by the Iceberg Java client. Catalogs are +/// not required to send it. +const S3_SESSION_TOKEN_EXPIRES_AT_MS: &str = "s3.session-token-expires-at-ms"; + +/// How far ahead of a reported expiry to re-fetch a vended credential. +const VENDED_CREDENTIAL_REFRESH_BUFFER: Duration = Duration::from_secs(120); + +/// How long to trust a vended credential that reports no expiry. +/// +/// A catalog that omits [`S3_SESSION_TOKEN_EXPIRES_AT_MS`] leaves nothing to schedule against, and +/// a credential held past its real lifetime fails every S3 request until the dataflow restarts. So +/// re-fetch on a short interval instead: each one is a single REST call against a credential the +/// sink is already using. +// TODO: make this a dyncfg once we know what expiries real catalogs report. +const VENDED_CREDENTIAL_DEFAULT_TTL: Duration = Duration::from_secs(300); /// A credential loader that wraps an aws-sdk-rust credentials provider for use with /// iceberg/OpenDAL. This allows us to provide refreshable credentials from the AWS SDK @@ -155,6 +178,214 @@ impl ProvideCredential for AwsSdkCredentialLoader { } } +#[derive(Debug)] +struct VendedCredentialLoader { + client: reqwest::Client, + credential_endpoint: Url, + token: Arc, + cached: Mutex>, +} + +impl VendedCredentialLoader { + fn new( + client: reqwest::Client, + credential_endpoint: Url, + token: Arc, + ) -> Self { + Self { + client, + credential_endpoint, + token, + cached: Mutex::new(None), + } + } +} + +impl VendedCredentialLoader { + /// Fetches a fresh credential from the catalog, paired with the instant at which it should be + /// re-fetched. + async fn fetch(&self) -> reqsign_core::Result<(AwsCredential, Instant)> { + let token = self.token.token().await.map_err(|e| { + reqsign_core::Error::credential_invalid( + "failed to obtain a catalog token for vended Iceberg storage credentials", + ) + .with_source(e) + })?; + + let response = self + .client + .get(self.credential_endpoint.clone()) + .bearer_auth(token) + .header( + ICEBERG_ACCESS_DELEGATION_HEADER, + IcebergAccessDelegation::VendedCredentials.as_header_value(), + ) + .send() + .await + .map_err(|e| { + reqsign_core::Error::unexpected(format!( + "failed to request vended Iceberg storage credentials from {}", + self.credential_endpoint + )) + .with_source(e) + })?; + + let status = response.status(); + if !status.is_success() { + // A rejected token stays rejected until something re-mints it, and nothing else on + // this path does: the REST client's own 401 handling covers catalog requests, not + // ours. Drop it so the next attempt fetches a new one. + if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { + if let Err(e) = self.token.invalidate().await { + warn!( + error = %e.display_with_causes(), + "failed to invalidate catalog token after {status} from the \ + Iceberg credentials endpoint" + ); + } + } + // Safe to surface the body: only success responses carry credentials. + let body = response.text().await.unwrap_or_default(); + return Err(reqsign_core::Error::unexpected(format!( + "Iceberg catalog returned {status} for vended storage credentials at {}: {body}", + self.credential_endpoint + ))); + } + + let response: LoadCredentialsResponse = response.json().await.map_err(|e| { + reqsign_core::Error::unexpected( + "failed to parse the Iceberg catalog's vended storage credentials", + ) + .with_source(e) + })?; + + // `provide_credential` is handed no path, so the longest-prefix match the Iceberg spec + // describes is not available to us. This endpoint is scoped to a single table and in + // practice returns one credential; if a catalog returns several, the most specific + // prefix is the closest thing to a safe default. + if response.storage_credentials.len() > 1 { + debug!( + endpoint = %self.credential_endpoint, + count = response.storage_credentials.len(), + "Iceberg catalog vended multiple storage credentials; using the longest prefix" + ); + } + let credential = response + .storage_credentials + .into_iter() + .max_by_key(|credential| credential.prefix.len()) + .ok_or_else(|| { + reqsign_core::Error::credential_invalid(format!( + "Iceberg catalog vended no storage credentials at {}", + self.credential_endpoint + )) + })?; + + let missing = |prop: &str| { + reqsign_core::Error::credential_invalid(format!( + "vended Iceberg storage credential for prefix {} is missing {prop}", + credential.prefix + )) + }; + let access_key_id = credential + .config + .get(S3_ACCESS_KEY_ID) + .ok_or_else(|| missing(S3_ACCESS_KEY_ID))? + .clone(); + let secret_access_key = credential + .config + .get(S3_SECRET_ACCESS_KEY) + .ok_or_else(|| missing(S3_SECRET_ACCESS_KEY))? + .clone(); + + let expires_in = credential + .config + .get(S3_SESSION_TOKEN_EXPIRES_AT_MS) + .map(|raw| { + let millis = raw.parse::().map_err(|e| { + reqsign_core::Error::credential_invalid(format!( + "vended Iceberg storage credential for prefix {} has an unparseable \ + {S3_SESSION_TOKEN_EXPIRES_AT_MS}", + credential.prefix + )) + .with_source(e) + })?; + Timestamp::from_millisecond(millis) + }) + .transpose()?; + + Ok(( + AwsCredential { + access_key_id, + secret_access_key, + session_token: credential.config.get(S3_SESSION_TOKEN).cloned(), + expires_in, + }, + refresh_deadline(expires_in), + )) + } +} + +impl ProvideCredential for VendedCredentialLoader { + type Credential = AwsCredential; + + async fn provide_credential( + &self, + _ctx: &reqsign_core::Context, + ) -> reqsign_core::Result> { + // The lock is deliberately held across the fetch. `create_operator` builds a fresh + // OpenDAL `Operator` for every file operation, so reqsign's own credential cache never + // outlives a single call and this cache is all that stands between the sink and one + // catalog round trip per S3 request. Serializing here means a stale entry costs one + // refetch rather than one per in-flight operation. + let mut cached = self.cached.lock().await; + + if let Some((credential, refresh_at)) = &*cached + && Instant::now() < *refresh_at + { + return Ok(Some(credential.clone())); + } + + let (credential, refresh_at) = self.fetch().await?; + *cached = Some((credential.clone(), refresh_at)); + Ok(Some(credential)) + } +} + +/// Returns when a vended credential expiring at `expires_in` should be re-fetched. +/// +/// Refreshes [`VENDED_CREDENTIAL_REFRESH_BUFFER`] early to absorb clock skew and the latency of +/// the fetch itself. A credential that expires within the buffer, or has expired already, is +/// re-fetched on the next call. +fn refresh_deadline(expires_in: Option) -> Instant { + let now = Instant::now(); + let Some(expires_in) = expires_in else { + return now + VENDED_CREDENTIAL_DEFAULT_TTL; + }; + let remaining = expires_in + .as_system_time() + .duration_since(SystemTime::now()) + .unwrap_or_default(); + now + remaining.saturating_sub(VENDED_CREDENTIAL_REFRESH_BUFFER) +} + +/// The `loadCredentials` response envelope. `iceberg-rust` models the credential entries but not +/// this wrapper, because it only ever reads them out of a `loadTable` response. +#[derive(Debug, Deserialize)] +struct LoadCredentialsResponse { + #[serde(default, rename = "storage-credentials")] + storage_credentials: Vec, +} + +/// The part of the Iceberg REST `config` response Materialize reads for itself. +#[derive(Debug, Default, Deserialize)] +struct CatalogConfigResponse { + #[serde(default)] + defaults: BTreeMap, + #[serde(default)] + overrides: BTreeMap, +} + /// Converts an AWS SDK credential expiry into reqsign's [`Timestamp`]. /// /// Both failure modes require a nonsensical expiry (before the Unix epoch, or beyond year @@ -793,18 +1024,27 @@ impl IcebergCatalogConnection { } impl IcebergCatalogConnection { + /// Connects to the catalog. + /// + /// `table` names the table this handle will be used against. It is needed only to keep + /// catalog-vended storage credentials refreshed, which the REST specification scopes to a + /// single table. Passing `None` leaves the connection on whatever credentials the catalog + /// supplies at `loadTable` time, which expire. pub async fn connect( &self, storage_configuration: &StorageConfiguration, in_task: InTask, + table: Option<&TableIdent>, ) -> Result, anyhow::Error> { match self.catalog { IcebergCatalogImpl::S3TablesRest(ref s3tables) => { + // S3 Tables signs every request with SigV4 off a refreshable AWS provider, so it + // has no vended credential to keep alive. self.connect_s3tables(s3tables, storage_configuration, in_task) .await } IcebergCatalogImpl::Rest(ref rest) => { - self.connect_rest(rest, storage_configuration, in_task) + self.connect_rest(rest, storage_configuration, in_task, table) .await } } @@ -972,6 +1212,7 @@ impl IcebergCatalogConnection { rest: &RestIcebergCatalog, storage_configuration: &StorageConfiguration, in_task: InTask, + table: Option<&TableIdent>, ) -> Result, anyhow::Error> { let mut props = BTreeMap::from([( REST_CATALOG_PROP_URI.to_string(), @@ -982,6 +1223,10 @@ impl IcebergCatalogConnection { props.insert(REST_CATALOG_PROP_WAREHOUSE.to_string(), warehouse.clone()); } + // One client for catalog requests, OAuth token requests, and credential refreshes, so all + // three share a connection pool. `iceberg-rust` would otherwise default to its own. + let client = reqwest::Client::new(); + // Catalog auth is configured through a combination of `props` and `.with_authenticator(...)`, // which happen at different stages of the [`RestCatalogBuilder`] -> [`RestCatalog`] // construction pipeline. @@ -998,7 +1243,6 @@ impl IcebergCatalogConnection { ) .await .map_err(|e| anyhow!("failed to read Iceberg catalog credential: {e}"))?; - props.insert(REST_CATALOG_PROP_CREDENTIAL.to_string(), credential); if let Some(server_url) = server_url { // The OAuth2 exchange POSTs the catalog credential to this URL, so a URL @@ -1030,10 +1274,64 @@ impl IcebergCatalogConnection { server_url.clone(), ); } + // Materialize builds the OAuth2 provider rather than handing `iceberg-rust` a + // `credential` prop, so one token object backs both catalog requests and the + // vended-credential refresh below. The two configurations are mutually exclusive: + // the catalog client rejects a custom authenticator combined with a `credential` + // prop, so the props that would drive its own provider are set here instead. + let token_endpoint = match server_url { + Some(server_url) => server_url.clone(), + // Matches `iceberg-rust`'s default when no `oauth2-server-uri` is configured. + None => format!( + "{}/v1/oauth/tokens", + self.uri.as_str().trim_end_matches('/') + ), + }; + let (client_id, client_secret) = match credential.split_once(':') { + Some((client_id, client_secret)) => { + (Some(client_id.to_string()), client_secret.to_string()) + } + None => (None, credential), + }; + let oauth_params = BTreeMap::from([( + OAUTH2_PARAM_SCOPE.to_string(), + // The default `iceberg-rust` applies when the connection names no scope. + scope.clone().unwrap_or_else(|| "catalog".to_string()), + )]); + let token: Arc = Arc::new(OAuth2TokenProvider::new( + client.clone(), + client_id, + client_secret, + token_endpoint, + // The token request needs none of the catalog's headers, and + // `OAuth2TokenProvider` sets the form content type itself. + HeaderMap::new(), + oauth_params.into_iter().collect(), + )); - if let Some(scope) = scope { - props.insert(REST_CATALOG_PROP_SCOPE.to_string(), scope.clone()); - } + // Installing a loader hands it sole responsibility for S3 credentials: OpenDAL + // replaces its whole provider chain, including the static keys parsed out of the + // catalog's vended `storage-credentials` props. So only install one when the + // connection asked for delegation and we know which table to refresh, and let the + // static props serve every other case. + let customized_credential_load = match (&rest.access_delegation, table) { + (Some(IcebergAccessDelegation::VendedCredentials), Some(table)) => { + let endpoint = self + .table_credentials_endpoint( + &client, + &token, + rest.warehouse.as_deref(), + table, + ) + .await?; + Some(CustomAwsCredentialLoader::new(VendedCredentialLoader::new( + client.clone(), + endpoint, + Arc::clone(&token), + ))) + } + _ => None, + }; ( OpenDalStorageFactory::S3 { @@ -1043,9 +1341,9 @@ impl IcebergCatalogConnection { // vends instead, it returns per-table `storage-credentials` that // `iceberg-rust` wires into the same FileIO. // N.B. This is not confirmed to work with other catalog & storage implementations. - customized_credential_load: None, + customized_credential_load, }, - None, + Some(iceberg_catalog_rest::BearerTokenAuthenticator::new(token)), ) } IcebergCatalogAuth::Gcp(gcp_connection_reference) => { @@ -1089,8 +1387,9 @@ impl IcebergCatalogConnection { ); } - let mut catalog = - RestCatalogBuilder::default().with_storage_factory(Arc::new(storage_factory)); + let mut catalog = RestCatalogBuilder::default() + .with_storage_factory(Arc::new(storage_factory)) + .with_client(client); if let Some(auth) = custom_authenticator { catalog = catalog.with_authenticator(Arc::new(auth)); } @@ -1101,13 +1400,91 @@ impl IcebergCatalogConnection { Ok(Arc::new(catalog)) } + /// Builds the REST endpoint that vends storage credentials for `table`. + /// + /// The path carries the catalog's request prefix, which the REST specification has servers + /// announce from their `config` endpoint (`catalogs/` for Unity Catalog, absent for + /// others). `iceberg-rust` resolves the same value when it builds the catalog but keeps it + /// private, so this asks the server for it directly. + async fn table_credentials_endpoint( + &self, + client: &reqwest::Client, + token: &Arc, + warehouse: Option<&str>, + table: &TableIdent, + ) -> Result { + let mut config_endpoint = self.uri.clone(); + config_endpoint + .path_segments_mut() + .map_err(|_| anyhow!("Iceberg catalog URI cannot be a base: {}", self.uri))? + .pop_if_empty() + .extend(["v1", "config"]); + if let Some(warehouse) = warehouse { + config_endpoint + .query_pairs_mut() + .append_pair("warehouse", warehouse); + } + + let bearer = token + .token() + .await + .map_err(|e| anyhow!("failed to obtain an Iceberg catalog token: {e}"))?; + let response = client + .get(config_endpoint.clone()) + .bearer_auth(bearer) + .send() + .await + .with_context(|| { + format!("failed to request Iceberg catalog config at {config_endpoint}") + })?; + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + return Err(anyhow!( + "Iceberg catalog returned {status} for its config at {config_endpoint}: {body}" + )); + } + let config: CatalogConfigResponse = response.json().await.with_context(|| { + format!("failed to parse Iceberg catalog config from {config_endpoint}") + })?; + + // Overrides take precedence over defaults, matching how the catalog client merges them. + let prefix = config + .overrides + .get("prefix") + .or_else(|| config.defaults.get("prefix")); + + let mut endpoint = self.uri.clone(); + { + let mut segments = endpoint + .path_segments_mut() + .map_err(|_| anyhow!("Iceberg catalog URI cannot be a base: {}", self.uri))?; + segments.pop_if_empty().push("v1"); + // A prefix is a path fragment rather than a single segment, so it is split out + // instead of pushed whole, which would percent-encode its separators. + if let Some(prefix) = prefix { + segments.extend(prefix.split('/').filter(|s| !s.is_empty())); + } + // `to_url_string` joins multi-level namespaces with the unit separator the + // specification mandates; pushing it as one segment percent-encodes it to `%1F`. + segments + .push("namespaces") + .push(&table.namespace().to_url_string()) + .push("tables") + .push(&table.name) + .push("credentials"); + } + Ok(endpoint) + } + async fn validate( &self, _id: CatalogItemId, storage_configuration: &StorageConfiguration, ) -> Result<(), ConnectionValidationError> { + // Validation only lists namespaces, so it needs no table-scoped credentials. let catalog = self - .connect(storage_configuration, InTask::No) + .connect(storage_configuration, InTask::No, None) .await .map_err(|e| { ConnectionValidationError::Other(anyhow!("failed to connect to catalog: {e}")) diff --git a/src/storage/src/sink/iceberg.rs b/src/storage/src/sink/iceberg.rs index 50ed536bda5ab..c416e450ee33d 100644 --- a/src/storage/src/sink/iceberg.rs +++ b/src/storage/src/sink/iceberg.rs @@ -1135,9 +1135,13 @@ fn mint_batch_descriptions<'scope>( return Ok(()); } + let table_ident = TableIdent::new( + NamespaceIdent::new(connection.namespace.clone()), + connection.table.clone(), + ); let catalog = connection .catalog_connection - .connect(&storage_configuration, InTask::Yes) + .connect(&storage_configuration, InTask::Yes, Some(&table_ident)) .await .with_context(|| { format!( @@ -1542,9 +1546,11 @@ fn write_data_files<'scope, H: EnvelopeHandler + 'static>( .build_fallible(move |caps| { Box::pin(async move { let [capset]: &mut [_; 1] = caps.try_into().unwrap(); + let namespace_ident = NamespaceIdent::new(connection.namespace.clone()); + let table_ident = TableIdent::new(namespace_ident, connection.table.clone()); let catalog = connection .catalog_connection - .connect(&storage_configuration, InTask::Yes) + .connect(&storage_configuration, InTask::Yes, Some(&table_ident)) .await .with_context(|| { format!( @@ -1555,8 +1561,6 @@ fn write_data_files<'scope, H: EnvelopeHandler + 'static>( ) })?; - let namespace_ident = NamespaceIdent::new(connection.namespace.clone()); - let table_ident = TableIdent::new(namespace_ident, connection.table.clone()); while let Some(_) = table_ready_input.next().await { // Wait for table to be ready } @@ -2508,9 +2512,11 @@ fn commit_to_iceberg<'scope>( return Ok(()); } + let namespace_ident = NamespaceIdent::new(connection.namespace.clone()); + let table_ident = TableIdent::new(namespace_ident, connection.table.clone()); let catalog = connection .catalog_connection - .connect(&storage_configuration, InTask::Yes) + .connect(&storage_configuration, InTask::Yes, Some(&table_ident)) .await .with_context(|| { format!( @@ -2521,8 +2527,6 @@ fn commit_to_iceberg<'scope>( let mut write_handle = write_handle.await?; - let namespace_ident = NamespaceIdent::new(connection.namespace.clone()); - let table_ident = TableIdent::new(namespace_ident, connection.table.clone()); while let Some(_) = table_ready_input.next().await { // Wait for table to be ready } From 13401858b648513bb06ca44bd26d01ee6be20c60 Mon Sep 17 00:00:00 2001 From: Patrick Butler Date: Thu, 27 Aug 2026 18:08:50 -0400 Subject: [PATCH 2/2] address various comments --- src/storage-types/src/connections.rs | 355 +----------- .../src/connections/iceberg_credentials.rs | 503 ++++++++++++++++++ 2 files changed, 531 insertions(+), 327 deletions(-) create mode 100644 src/storage-types/src/connections/iceberg_credentials.rs diff --git a/src/storage-types/src/connections.rs b/src/storage-types/src/connections.rs index e969418c09f39..aa48f6a599c23 100644 --- a/src/storage-types/src/connections.rs +++ b/src/storage-types/src/connections.rs @@ -14,7 +14,7 @@ use std::collections::{BTreeMap, BTreeSet}; use std::fmt; use std::net::SocketAddr; use std::sync::Arc; -use std::time::{Duration, Instant, SystemTime}; +use std::time::SystemTime; use anyhow::{Context, anyhow}; use async_trait::async_trait; @@ -30,11 +30,11 @@ use iceberg::CatalogBuilder; use iceberg::TableIdent; use iceberg::io::{ GCS_CREDENTIALS_JSON, GCS_DISABLE_CONFIG_LOAD, GCS_DISABLE_VM_METADATA, GCS_USER_PROJECT, - S3_ACCESS_KEY_ID, S3_DISABLE_EC2_METADATA, S3_REGION, S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN, + S3_ACCESS_KEY_ID, S3_DISABLE_EC2_METADATA, S3_REGION, S3_SECRET_ACCESS_KEY, }; use iceberg_catalog_rest::{ OAuth2TokenProvider, REST_CATALOG_PROP_URI, REST_CATALOG_PROP_WAREHOUSE, RequestAuthenticator, - RestCatalogBuilder, StorageCredential, TokenProvider, + RestCatalogBuilder, TokenProvider, }; use iceberg_storage_opendal::{ AwsCredential, CustomAwsCredentialLoader, OpenDalStorageFactory, ProvideCredential, @@ -66,11 +66,10 @@ use rdkafka::config::FromClientConfigAndContext; use rdkafka::consumer::{BaseConsumer, Consumer}; use regex::Regex; use reqsign_core::time::Timestamp; -use reqwest::{Request, StatusCode}; +use reqwest::Request; use serde::{Deserialize, Deserializer, Serialize}; use tokio::net; use tokio::runtime::Handle; -use tokio::sync::Mutex; use tokio_postgres::config::SslMode; use tracing::{debug, info, warn}; use url::Url; @@ -92,6 +91,7 @@ use crate::errors::{ContextCreationError, CsrConnectError}; pub mod aws; pub mod gcp; +mod iceberg_credentials; pub mod inline; pub mod string_or_secret; @@ -101,29 +101,12 @@ pub mod string_or_secret; /// `oauth2-server-uri`, and `scope` catalog properties, so that one token object serves both /// catalog requests and storage-credential refreshes. const OAUTH2_PARAM_SCOPE: &str = "scope"; + +const REST_CATALOG_PROP_OAUTH2_SERVER_URI: &str = "oauth2-server-uri"; /// Requests catalog-vended storage credentials. `iceberg-rust` turns `header.*` catalog /// properties into headers on every REST request, the same way the Iceberg Java client /// carries this one. const REST_CATALOG_PROP_ACCESS_DELEGATION: &str = "header.X-Iceberg-Access-Delegation"; -/// The same header, spelled for requests Materialize issues itself rather than through the -/// catalog client. -const ICEBERG_ACCESS_DELEGATION_HEADER: &str = "X-Iceberg-Access-Delegation"; - -/// Reports when vended S3 credentials expire, as spelled by the Iceberg Java client. Catalogs are -/// not required to send it. -const S3_SESSION_TOKEN_EXPIRES_AT_MS: &str = "s3.session-token-expires-at-ms"; - -/// How far ahead of a reported expiry to re-fetch a vended credential. -const VENDED_CREDENTIAL_REFRESH_BUFFER: Duration = Duration::from_secs(120); - -/// How long to trust a vended credential that reports no expiry. -/// -/// A catalog that omits [`S3_SESSION_TOKEN_EXPIRES_AT_MS`] leaves nothing to schedule against, and -/// a credential held past its real lifetime fails every S3 request until the dataflow restarts. So -/// re-fetch on a short interval instead: each one is a single REST call against a credential the -/// sink is already using. -// TODO: make this a dyncfg once we know what expiries real catalogs report. -const VENDED_CREDENTIAL_DEFAULT_TTL: Duration = Duration::from_secs(300); /// A credential loader that wraps an aws-sdk-rust credentials provider for use with /// iceberg/OpenDAL. This allows us to provide refreshable credentials from the AWS SDK @@ -178,214 +161,6 @@ impl ProvideCredential for AwsSdkCredentialLoader { } } -#[derive(Debug)] -struct VendedCredentialLoader { - client: reqwest::Client, - credential_endpoint: Url, - token: Arc, - cached: Mutex>, -} - -impl VendedCredentialLoader { - fn new( - client: reqwest::Client, - credential_endpoint: Url, - token: Arc, - ) -> Self { - Self { - client, - credential_endpoint, - token, - cached: Mutex::new(None), - } - } -} - -impl VendedCredentialLoader { - /// Fetches a fresh credential from the catalog, paired with the instant at which it should be - /// re-fetched. - async fn fetch(&self) -> reqsign_core::Result<(AwsCredential, Instant)> { - let token = self.token.token().await.map_err(|e| { - reqsign_core::Error::credential_invalid( - "failed to obtain a catalog token for vended Iceberg storage credentials", - ) - .with_source(e) - })?; - - let response = self - .client - .get(self.credential_endpoint.clone()) - .bearer_auth(token) - .header( - ICEBERG_ACCESS_DELEGATION_HEADER, - IcebergAccessDelegation::VendedCredentials.as_header_value(), - ) - .send() - .await - .map_err(|e| { - reqsign_core::Error::unexpected(format!( - "failed to request vended Iceberg storage credentials from {}", - self.credential_endpoint - )) - .with_source(e) - })?; - - let status = response.status(); - if !status.is_success() { - // A rejected token stays rejected until something re-mints it, and nothing else on - // this path does: the REST client's own 401 handling covers catalog requests, not - // ours. Drop it so the next attempt fetches a new one. - if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { - if let Err(e) = self.token.invalidate().await { - warn!( - error = %e.display_with_causes(), - "failed to invalidate catalog token after {status} from the \ - Iceberg credentials endpoint" - ); - } - } - // Safe to surface the body: only success responses carry credentials. - let body = response.text().await.unwrap_or_default(); - return Err(reqsign_core::Error::unexpected(format!( - "Iceberg catalog returned {status} for vended storage credentials at {}: {body}", - self.credential_endpoint - ))); - } - - let response: LoadCredentialsResponse = response.json().await.map_err(|e| { - reqsign_core::Error::unexpected( - "failed to parse the Iceberg catalog's vended storage credentials", - ) - .with_source(e) - })?; - - // `provide_credential` is handed no path, so the longest-prefix match the Iceberg spec - // describes is not available to us. This endpoint is scoped to a single table and in - // practice returns one credential; if a catalog returns several, the most specific - // prefix is the closest thing to a safe default. - if response.storage_credentials.len() > 1 { - debug!( - endpoint = %self.credential_endpoint, - count = response.storage_credentials.len(), - "Iceberg catalog vended multiple storage credentials; using the longest prefix" - ); - } - let credential = response - .storage_credentials - .into_iter() - .max_by_key(|credential| credential.prefix.len()) - .ok_or_else(|| { - reqsign_core::Error::credential_invalid(format!( - "Iceberg catalog vended no storage credentials at {}", - self.credential_endpoint - )) - })?; - - let missing = |prop: &str| { - reqsign_core::Error::credential_invalid(format!( - "vended Iceberg storage credential for prefix {} is missing {prop}", - credential.prefix - )) - }; - let access_key_id = credential - .config - .get(S3_ACCESS_KEY_ID) - .ok_or_else(|| missing(S3_ACCESS_KEY_ID))? - .clone(); - let secret_access_key = credential - .config - .get(S3_SECRET_ACCESS_KEY) - .ok_or_else(|| missing(S3_SECRET_ACCESS_KEY))? - .clone(); - - let expires_in = credential - .config - .get(S3_SESSION_TOKEN_EXPIRES_AT_MS) - .map(|raw| { - let millis = raw.parse::().map_err(|e| { - reqsign_core::Error::credential_invalid(format!( - "vended Iceberg storage credential for prefix {} has an unparseable \ - {S3_SESSION_TOKEN_EXPIRES_AT_MS}", - credential.prefix - )) - .with_source(e) - })?; - Timestamp::from_millisecond(millis) - }) - .transpose()?; - - Ok(( - AwsCredential { - access_key_id, - secret_access_key, - session_token: credential.config.get(S3_SESSION_TOKEN).cloned(), - expires_in, - }, - refresh_deadline(expires_in), - )) - } -} - -impl ProvideCredential for VendedCredentialLoader { - type Credential = AwsCredential; - - async fn provide_credential( - &self, - _ctx: &reqsign_core::Context, - ) -> reqsign_core::Result> { - // The lock is deliberately held across the fetch. `create_operator` builds a fresh - // OpenDAL `Operator` for every file operation, so reqsign's own credential cache never - // outlives a single call and this cache is all that stands between the sink and one - // catalog round trip per S3 request. Serializing here means a stale entry costs one - // refetch rather than one per in-flight operation. - let mut cached = self.cached.lock().await; - - if let Some((credential, refresh_at)) = &*cached - && Instant::now() < *refresh_at - { - return Ok(Some(credential.clone())); - } - - let (credential, refresh_at) = self.fetch().await?; - *cached = Some((credential.clone(), refresh_at)); - Ok(Some(credential)) - } -} - -/// Returns when a vended credential expiring at `expires_in` should be re-fetched. -/// -/// Refreshes [`VENDED_CREDENTIAL_REFRESH_BUFFER`] early to absorb clock skew and the latency of -/// the fetch itself. A credential that expires within the buffer, or has expired already, is -/// re-fetched on the next call. -fn refresh_deadline(expires_in: Option) -> Instant { - let now = Instant::now(); - let Some(expires_in) = expires_in else { - return now + VENDED_CREDENTIAL_DEFAULT_TTL; - }; - let remaining = expires_in - .as_system_time() - .duration_since(SystemTime::now()) - .unwrap_or_default(); - now + remaining.saturating_sub(VENDED_CREDENTIAL_REFRESH_BUFFER) -} - -/// The `loadCredentials` response envelope. `iceberg-rust` models the credential entries but not -/// this wrapper, because it only ever reads them out of a `loadTable` response. -#[derive(Debug, Deserialize)] -struct LoadCredentialsResponse { - #[serde(default, rename = "storage-credentials")] - storage_credentials: Vec, -} - -/// The part of the Iceberg REST `config` response Materialize reads for itself. -#[derive(Debug, Default, Deserialize)] -struct CatalogConfigResponse { - #[serde(default)] - defaults: BTreeMap, - #[serde(default)] - overrides: BTreeMap, -} - /// Converts an AWS SDK credential expiry into reqsign's [`Timestamp`]. /// /// Both failure modes require a nonsensical expiry (before the Unix epoch, or beyond year @@ -1274,11 +1049,9 @@ impl IcebergCatalogConnection { server_url.clone(), ); } - // Materialize builds the OAuth2 provider rather than handing `iceberg-rust` a - // `credential` prop, so one token object backs both catalog requests and the - // vended-credential refresh below. The two configurations are mutually exclusive: - // the catalog client rejects a custom authenticator combined with a `credential` - // prop, so the props that would drive its own provider are set here instead. + + // Materialize builds an OAuth2 provider shared across both catalog requests + // and the vended credentials refresh below. let token_endpoint = match server_url { Some(server_url) => server_url.clone(), // Matches `iceberg-rust`'s default when no `oauth2-server-uri` is configured. @@ -1316,19 +1089,21 @@ impl IcebergCatalogConnection { // static props serve every other case. let customized_credential_load = match (&rest.access_delegation, table) { (Some(IcebergAccessDelegation::VendedCredentials), Some(table)) => { - let endpoint = self - .table_credentials_endpoint( - &client, - &token, - rest.warehouse.as_deref(), - table, - ) - .await?; - Some(CustomAwsCredentialLoader::new(VendedCredentialLoader::new( - client.clone(), - endpoint, - Arc::clone(&token), - ))) + let endpoint = iceberg_credentials::table_credentials_endpoint( + &self.uri, + &client, + &token, + rest.warehouse.as_deref(), + table, + ) + .await?; + Some(CustomAwsCredentialLoader::new( + iceberg_credentials::VendedCredentialLoader::new( + client.clone(), + endpoint, + Arc::clone(&token), + ), + )) } _ => None, }; @@ -1343,6 +1118,9 @@ impl IcebergCatalogConnection { // N.B. This is not confirmed to work with other catalog & storage implementations. customized_credential_load, }, + // NOTE: We construct our own OAuth authenticator for the Catalog client instead of using the one built in. + // This means we ignore auth overrides from `/v1/config` (e.g. `oauth2-server-uri`). + // This is okay because users can set these configs from Mz SQL. Some(iceberg_catalog_rest::BearerTokenAuthenticator::new(token)), ) } @@ -1400,83 +1178,6 @@ impl IcebergCatalogConnection { Ok(Arc::new(catalog)) } - /// Builds the REST endpoint that vends storage credentials for `table`. - /// - /// The path carries the catalog's request prefix, which the REST specification has servers - /// announce from their `config` endpoint (`catalogs/` for Unity Catalog, absent for - /// others). `iceberg-rust` resolves the same value when it builds the catalog but keeps it - /// private, so this asks the server for it directly. - async fn table_credentials_endpoint( - &self, - client: &reqwest::Client, - token: &Arc, - warehouse: Option<&str>, - table: &TableIdent, - ) -> Result { - let mut config_endpoint = self.uri.clone(); - config_endpoint - .path_segments_mut() - .map_err(|_| anyhow!("Iceberg catalog URI cannot be a base: {}", self.uri))? - .pop_if_empty() - .extend(["v1", "config"]); - if let Some(warehouse) = warehouse { - config_endpoint - .query_pairs_mut() - .append_pair("warehouse", warehouse); - } - - let bearer = token - .token() - .await - .map_err(|e| anyhow!("failed to obtain an Iceberg catalog token: {e}"))?; - let response = client - .get(config_endpoint.clone()) - .bearer_auth(bearer) - .send() - .await - .with_context(|| { - format!("failed to request Iceberg catalog config at {config_endpoint}") - })?; - let status = response.status(); - if !status.is_success() { - let body = response.text().await.unwrap_or_default(); - return Err(anyhow!( - "Iceberg catalog returned {status} for its config at {config_endpoint}: {body}" - )); - } - let config: CatalogConfigResponse = response.json().await.with_context(|| { - format!("failed to parse Iceberg catalog config from {config_endpoint}") - })?; - - // Overrides take precedence over defaults, matching how the catalog client merges them. - let prefix = config - .overrides - .get("prefix") - .or_else(|| config.defaults.get("prefix")); - - let mut endpoint = self.uri.clone(); - { - let mut segments = endpoint - .path_segments_mut() - .map_err(|_| anyhow!("Iceberg catalog URI cannot be a base: {}", self.uri))?; - segments.pop_if_empty().push("v1"); - // A prefix is a path fragment rather than a single segment, so it is split out - // instead of pushed whole, which would percent-encode its separators. - if let Some(prefix) = prefix { - segments.extend(prefix.split('/').filter(|s| !s.is_empty())); - } - // `to_url_string` joins multi-level namespaces with the unit separator the - // specification mandates; pushing it as one segment percent-encodes it to `%1F`. - segments - .push("namespaces") - .push(&table.namespace().to_url_string()) - .push("tables") - .push(&table.name) - .push("credentials"); - } - Ok(endpoint) - } - async fn validate( &self, _id: CatalogItemId, diff --git a/src/storage-types/src/connections/iceberg_credentials.rs b/src/storage-types/src/connections/iceberg_credentials.rs new file mode 100644 index 0000000000000..0c446e8890bcc --- /dev/null +++ b/src/storage-types/src/connections/iceberg_credentials.rs @@ -0,0 +1,503 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Use of this software is governed by the Business Source License +// included in the LICENSE file. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0. + +//! Catalog-vended storage credentials for Iceberg sinks. +//! +//! A REST catalog asked for access delegation answers `loadTable` with temporary, table-scoped +//! storage credentials instead of expecting the client to hold its own. Those credentials expire, +//! and OpenDAL has no notion of that: it takes what the `FileIO` was built with and signs every +//! request with it. [`VendedCredentialLoader`] closes that gap by re-fetching from the catalog's +//! `loadCredentials` endpoint, and the rest of this module locates that endpoint. +//! +//! Locating it takes a round trip of its own. The REST specification has servers announce a +//! request prefix from their `config` endpoint (`catalogs/` for Unity Catalog, absent for +//! others), and every resource path carries it. `iceberg-rust` resolves the same value when it +//! builds the catalog but keeps it private, so this module asks the server directly. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::time::{Duration, Instant, SystemTime}; + +use anyhow::{Context, anyhow}; +use iceberg::TableIdent; +use iceberg::io::{S3_ACCESS_KEY_ID, S3_SECRET_ACCESS_KEY, S3_SESSION_TOKEN}; +use iceberg_catalog_rest::{StorageCredential, TokenProvider}; +use iceberg_storage_opendal::{AwsCredential, ProvideCredential}; +use mz_ore::error::ErrorExt; +use reqsign_core::time::Timestamp; +use reqwest::StatusCode; +use serde::Deserialize; +use tokio::sync::Mutex; +use tracing::{debug, warn}; +use url::Url; + +use crate::connections::IcebergAccessDelegation; + +/// The `X-Iceberg-Access-Delegation` header, spelled for requests Materialize issues itself +/// rather than through the catalog client, which takes it as a `header.*` catalog property. +const ICEBERG_ACCESS_DELEGATION_HEADER: &str = "X-Iceberg-Access-Delegation"; + +/// Property name for the most common way catalogs report when vended S3 credentials expire. +/// Catalogs are not required to send it. +const S3_SESSION_TOKEN_EXPIRES_AT_MS: &str = "s3.session-token-expires-at-ms"; + +/// How far ahead of a reported expiry to re-fetch a vended credential. +const VENDED_CREDENTIAL_REFRESH_BUFFER: Duration = Duration::from_secs(900); + +/// How long to trust a vended credential that reports no expiry. +/// +/// A catalog that omits [`S3_SESSION_TOKEN_EXPIRES_AT_MS`] leaves nothing to schedule against, and +/// a credential held past its real lifetime fails every S3 request until the dataflow restarts. So +/// re-fetch on a short interval instead: each one is a single REST call against a credential the +/// sink is already using. +// TODO SS-449: make this a dyncfg +const VENDED_CREDENTIAL_DEFAULT_TTL: Duration = Duration::from_secs(300); + +#[derive(Debug)] +pub(super) struct VendedCredentialLoader { + client: reqwest::Client, + credential_endpoint: Url, + token: Arc, + cached: Mutex>, +} + +impl VendedCredentialLoader { + pub(super) fn new( + client: reqwest::Client, + credential_endpoint: Url, + token: Arc, + ) -> Self { + Self { + client, + credential_endpoint, + token, + cached: Mutex::new(None), + } + } + + /// Fetches a fresh credential from the catalog, paired with the instant at which it should be + /// re-fetched. + async fn fetch(&self) -> reqsign_core::Result<(AwsCredential, Instant)> { + let token = self.token.token().await.map_err(|e| { + reqsign_core::Error::credential_invalid( + "failed to obtain a catalog token for vended Iceberg storage credentials", + ) + .with_source(e) + })?; + + let response = self + .client + .get(self.credential_endpoint.clone()) + .bearer_auth(token) + .header( + ICEBERG_ACCESS_DELEGATION_HEADER, + IcebergAccessDelegation::VendedCredentials.as_header_value(), + ) + .send() + .await + .map_err(|e| { + reqsign_core::Error::unexpected(format!( + "failed to request vended Iceberg storage credentials from {}", + self.credential_endpoint + )) + .with_source(e) + })?; + + let status = response.status(); + if !status.is_success() { + // A rejected token stays rejected until something re-mints it, and nothing else on + // this path does: the REST client's own 401 handling covers catalog requests, not + // ours. Drop it so the next attempt fetches a new one. + if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { + if let Err(e) = self.token.invalidate().await { + warn!( + error = %e.display_with_causes(), + "failed to invalidate catalog token after {status} from the \ + Iceberg credentials endpoint" + ); + } + } + // Safe to surface the body: only success responses carry credentials. + let body = response.text().await.unwrap_or_default(); + return Err(reqsign_core::Error::unexpected(format!( + "Iceberg catalog returned {status} for vended storage credentials at {}: {body}", + self.credential_endpoint + ))); + } + + let response: LoadCredentialsResponse = response.json().await.map_err(|e| { + reqsign_core::Error::unexpected( + "failed to parse the Iceberg catalog's vended storage credentials", + ) + .with_source(e) + })?; + + // `provide_credential` is handed no path, so the longest-prefix match the Iceberg spec + // describes is not available to us. This endpoint is scoped to a single table and in + // practice returns one credential; if a catalog returns several, the most specific + // prefix is the closest thing to a safe default. + if response.storage_credentials.len() > 1 { + debug!( + endpoint = %self.credential_endpoint, + count = response.storage_credentials.len(), + "Iceberg catalog vended multiple storage credentials; using the longest prefix" + ); + } + let credential = response + .storage_credentials + .into_iter() + .max_by_key(|credential| credential.prefix.len()) + .ok_or_else(|| { + reqsign_core::Error::credential_invalid(format!( + "Iceberg catalog vended no storage credentials at {}", + self.credential_endpoint + )) + })?; + + let missing = |prop: &str| { + reqsign_core::Error::credential_invalid(format!( + "vended Iceberg storage credential for prefix {} is missing {prop}", + credential.prefix + )) + }; + let access_key_id = credential + .config + .get(S3_ACCESS_KEY_ID) + .ok_or_else(|| missing(S3_ACCESS_KEY_ID))? + .clone(); + let secret_access_key = credential + .config + .get(S3_SECRET_ACCESS_KEY) + .ok_or_else(|| missing(S3_SECRET_ACCESS_KEY))? + .clone(); + + let expires_in = credential + .config + .get(S3_SESSION_TOKEN_EXPIRES_AT_MS) + .map(|raw| { + let millis = raw.parse::().map_err(|e| { + reqsign_core::Error::credential_invalid(format!( + "vended Iceberg storage credential for prefix {} has an unparseable \ + {S3_SESSION_TOKEN_EXPIRES_AT_MS}", + credential.prefix + )) + .with_source(e) + })?; + Timestamp::from_millisecond(millis) + }) + .transpose()?; + + Ok(( + AwsCredential { + access_key_id, + secret_access_key, + session_token: credential.config.get(S3_SESSION_TOKEN).cloned(), + expires_in, + }, + refresh_deadline(expires_in), + )) + } +} + +impl ProvideCredential for VendedCredentialLoader { + type Credential = AwsCredential; + + async fn provide_credential( + &self, + _ctx: &reqsign_core::Context, + ) -> reqsign_core::Result> { + // The lock is deliberately held across the fetch. `create_operator` builds a fresh + // OpenDAL `Operator` for every file operation, so reqsign's own credential cache never + // outlives a single call and this cache is all that stands between the sink and one + // catalog round trip per S3 request. Serializing here means a stale entry costs one + // refetch rather than one per in-flight operation. + let mut cached = self.cached.lock().await; + + if let Some((credential, refresh_at)) = &*cached + && Instant::now() < *refresh_at + { + return Ok(Some(credential.clone())); + } + + let (credential, refresh_at) = self.fetch().await?; + *cached = Some((credential.clone(), refresh_at)); + Ok(Some(credential)) + } +} + +/// Returns when a vended credential expiring at `expires_in` should be re-fetched. +/// +/// Refreshes [`VENDED_CREDENTIAL_REFRESH_BUFFER`] early to absorb clock skew and the latency of +/// the fetch itself. A credential that expires within the buffer, or has expired already, is +/// re-fetched on the next call. +fn refresh_deadline(expires_in: Option) -> Instant { + let now = Instant::now(); + let Some(expires_in) = expires_in else { + return now + VENDED_CREDENTIAL_DEFAULT_TTL; + }; + let remaining = expires_in + .as_system_time() + .duration_since(SystemTime::now()) + .unwrap_or_default(); + now + remaining.saturating_sub(VENDED_CREDENTIAL_REFRESH_BUFFER) +} + +/// The `loadCredentials` response envelope. `iceberg-rust` models the credential entries but not +/// this wrapper, because it only ever reads them out of a `loadTable` response. +#[derive(Debug, Deserialize)] +struct LoadCredentialsResponse { + #[serde(default, rename = "storage-credentials")] + storage_credentials: Vec, +} + +/// The part of the Iceberg REST `config` response Materialize reads for itself. +#[derive(Debug, Default, Deserialize)] +struct CatalogConfigResponse { + #[serde(default)] + defaults: BTreeMap, + #[serde(default)] + overrides: BTreeMap, +} + +impl CatalogConfigResponse { + /// The request prefix the server wants inserted between `/v1` and the resource path, + /// `catalogs/` for Unity Catalog and absent for catalogs that do not use one. + /// + /// Overrides win over defaults, matching how the catalog client merges the two. + fn announced_prefix(&self) -> Option<&str> { + self.overrides + .get("prefix") + .or_else(|| self.defaults.get("prefix")) + .map(String::as_str) + } +} + +/// Builds the catalog's `config` endpoint, which is where a server announces its request prefix. +/// +/// Unlike every other REST path, this one is not prefixed: the prefix is what it returns. +fn catalog_config_url(uri: &Url, warehouse: Option<&str>) -> Result { + let mut url = uri.clone(); + url.path_segments_mut() + .map_err(|_| anyhow!("Iceberg catalog URI cannot be a base: {uri}"))? + // A configured URI is as likely to be written with a trailing slash as without, and + // that empty last segment would otherwise become `//v1` in the path. + .pop_if_empty() + .extend(["v1", "config"]); + if let Some(warehouse) = warehouse { + url.query_pairs_mut().append_pair("warehouse", warehouse); + } + Ok(url) +} + +/// Builds the endpoint that vends storage credentials for `table`. +fn table_credentials_url( + uri: &Url, + prefix: Option<&str>, + table: &TableIdent, +) -> Result { + let mut url = uri.clone(); + { + let mut segments = url + .path_segments_mut() + .map_err(|_| anyhow!("Iceberg catalog URI cannot be a base: {uri}"))?; + segments.pop_if_empty().push("v1"); + // A prefix is a path fragment rather than a single segment, so it is split out instead + // of pushed whole, which would percent-encode its separators. + if let Some(prefix) = prefix { + segments.extend(prefix.split('/').filter(|s| !s.is_empty())); + } + // `to_url_string` joins multi-level namespaces with the unit separator the + // specification mandates; pushing it as one segment percent-encodes it to `%1F`. + segments + .push("namespaces") + .push(&table.namespace().to_url_string()) + .push("tables") + .push(&table.name) + .push("credentials"); + } + Ok(url) +} + +/// Resolves the REST endpoint that vends storage credentials for `table`. +/// +/// Takes a round trip to the catalog's `config` endpoint, because the resource path carries a +/// request prefix that only the server knows. +pub(super) async fn table_credentials_endpoint( + uri: &Url, + client: &reqwest::Client, + token: &Arc, + warehouse: Option<&str>, + table: &TableIdent, +) -> Result { + let config_endpoint = catalog_config_url(uri, warehouse)?; + + let bearer = token + .token() + .await + .map_err(|e| anyhow!("failed to obtain an Iceberg catalog token: {e}"))?; + let response = client + .get(config_endpoint.clone()) + .bearer_auth(bearer) + .send() + .await + .with_context(|| { + format!("failed to request Iceberg catalog config at {config_endpoint}") + })?; + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + return Err(anyhow!( + "Iceberg catalog returned {status} for its config at {config_endpoint}: {body}" + )); + } + let config: CatalogConfigResponse = response.json().await.with_context(|| { + format!("failed to parse Iceberg catalog config from {config_endpoint}") + })?; + + table_credentials_url(uri, config.announced_prefix(), table) +} + +#[cfg(test)] +mod tests { + use iceberg::NamespaceIdent; + + use super::*; + + fn table(namespace: &[&str], name: &str) -> TableIdent { + TableIdent::new( + NamespaceIdent::from_strs(namespace).expect("valid namespace"), + name.to_string(), + ) + } + + fn config_with(defaults: &[(&str, &str)], overrides: &[(&str, &str)]) -> CatalogConfigResponse { + let to_map = |pairs: &[(&str, &str)]| { + pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect::>() + }; + CatalogConfigResponse { + defaults: to_map(defaults), + overrides: to_map(overrides), + } + } + + #[mz_ore::test] + fn test_announced_prefix() { + // Overrides win, matching how the catalog client merges the two. + let config = config_with( + &[("prefix", "from-defaults")], + &[("prefix", "from-overrides")], + ); + assert_eq!(config.announced_prefix(), Some("from-overrides")); + + // Defaults are consulted only when overrides is silent, and an unrelated override does + // not suppress them. + let config = config_with(&[("prefix", "from-defaults")], &[("warehouse", "wh")]); + assert_eq!(config.announced_prefix(), Some("from-defaults")); + + // A catalog that uses no prefix announces none. + assert_eq!(config_with(&[], &[]).announced_prefix(), None); + } + + #[mz_ore::test] + fn test_catalog_config_url() { + let url = |uri: &str, warehouse: Option<&str>| { + catalog_config_url(&Url::parse(uri).expect("valid URI"), warehouse) + .expect("URI is a base") + .to_string() + }; + + // The config endpoint is never prefixed: the prefix is what it returns. + assert_eq!( + url("https://catalog.example/api", None), + "https://catalog.example/api/v1/config" + ); + + // A trailing slash must not leave an empty segment behind as `//v1`. + assert_eq!( + url("https://catalog.example/api/", None), + "https://catalog.example/api/v1/config" + ); + assert_eq!( + url("https://catalog.example", None), + "https://catalog.example/v1/config" + ); + + // The warehouse rides along as a query parameter, percent-encoded. + assert_eq!( + url("https://catalog.example", Some("my catalog")), + "https://catalog.example/v1/config?warehouse=my+catalog" + ); + + // A URI that cannot be a base has no path segments to extend. + assert!( + catalog_config_url(&Url::parse("mailto:nobody@example.com").unwrap(), None).is_err() + ); + } + + #[mz_ore::test] + fn test_table_credentials_url() { + let url = |uri: &str, prefix: Option<&str>, t: &TableIdent| { + table_credentials_url(&Url::parse(uri).expect("valid URI"), prefix, t) + .expect("URI is a base") + .to_string() + }; + + // A catalog announcing no prefix puts the resource path directly under `/v1`. + assert_eq!( + url( + "https://catalog.example", + None, + &table(&["sales"], "orders") + ), + "https://catalog.example/v1/namespaces/sales/tables/orders/credentials" + ); + + // Unity Catalog's multi-segment prefix keeps its separator rather than being encoded + // as one segment, and the base path of the configured URI is preserved. + assert_eq!( + url( + "https://dbc.cloud.databricks.com/api/2.1/unity-catalog/iceberg-rest", + Some("catalogs/sink-catalog"), + &table(&["sink-namespace"], "mz_append_test") + ), + "https://dbc.cloud.databricks.com/api/2.1/unity-catalog/iceberg-rest/v1/\ + catalogs/sink-catalog/namespaces/sink-namespace/tables/mz_append_test/credentials" + ); + + // A trailing slash on either the URI or the prefix must not produce an empty segment. + assert_eq!( + url( + "https://catalog.example/", + Some("/catalogs/main/"), + &table(&["sales"], "orders") + ), + "https://catalog.example/v1/catalogs/main/namespaces/sales/tables/orders/credentials" + ); + + // Multi-level namespaces are one segment joined by the unit separator the specification + // mandates, which percent-encodes to `%1F`. + assert_eq!( + url( + "https://catalog.example", + None, + &table(&["sales", "eu"], "orders") + ), + "https://catalog.example/v1/namespaces/sales%1Feu/tables/orders/credentials" + ); + + // Names that would otherwise change the path are escaped, not injected into it. + assert_eq!( + url("https://catalog.example", None, &table(&["a/b"], "c?d#e")), + "https://catalog.example/v1/namespaces/a%2Fb/tables/c%3Fd%23e/credentials" + ); + } +}