diff --git a/Cargo.lock b/Cargo.lock index 76484079..34cf4af5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1925,6 +1925,7 @@ dependencies = [ "async-trait", "aws-config", "aws-credential-types", + "aws-sdk-s3tables", "bytes", "chrono", "dashmap", diff --git a/Cargo.toml b/Cargo.toml index 5ad311ae..61de8d06 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -34,7 +34,8 @@ queries = { path = "crates/queries" } async-trait = { version = "0.1.84" } aws-config = { version = "1.5.17" } aws-sdk-dynamodb = { version = "1.100.0" } -aws-credential-types = { version = "1.2.1", features = ["hardcoded-credentials"]} +aws-sdk-s3tables = { version = "1.47.0" } +aws-credential-types = { version = "1.2.11", features = ["hardcoded-credentials"]} axum = { version = "0.8.1", features = ["multipart", "macros"] } axum-macros = "0.5" bytes = { version = "1.8.0" } diff --git a/crates/catalog-metastore/Cargo.toml b/crates/catalog-metastore/Cargo.toml index 3b409c0b..17b26eab 100644 --- a/crates/catalog-metastore/Cargo.toml +++ b/crates/catalog-metastore/Cargo.toml @@ -10,7 +10,8 @@ error-stack = { path = "../error-stack" } async-trait = { workspace = true } aws-config = { workspace = true } -aws-credential-types = { version = "1.2.11" } +aws-sdk-s3tables = { workspace = true } +aws-credential-types = { workspace = true } chrono = { workspace = true } dashmap = { workspace = true } iceberg-rust = { workspace = true } diff --git a/crates/catalog-metastore/src/metastore_bootstrap_config.rs b/crates/catalog-metastore/src/metastore_bootstrap_config.rs index c50243e8..ae6f2579 100644 --- a/crates/catalog-metastore/src/metastore_bootstrap_config.rs +++ b/crates/catalog-metastore/src/metastore_bootstrap_config.rs @@ -3,7 +3,10 @@ use crate::{ SchemaIdent, TableFormat, TableIdent, Volume, VolumeIdent, VolumeType, }; use aws_config::meta::credentials::CredentialsProviderChain; -use aws_credential_types::provider::ProvideCredentials; +use aws_config::{BehaviorVersion, Region}; +use aws_credential_types::Credentials; +use aws_credential_types::provider::{ProvideCredentials, SharedCredentialsProvider}; +use aws_sdk_s3tables::Client as S3TablesClient; use iceberg_rust::spec::table_metadata::TableMetadata; use iceberg_rust::spec::util::strip_prefix; use serde::Deserialize; @@ -426,8 +429,14 @@ async fn load_volume_from_env() -> Result, ConfigError> { "s3tables" | "s3_tables" | "s3-tables" => { let arn = env::var("VOLUME_ARN").map_err(|_| missing_var_error("VOLUME_ARN"))?; - let credentials = - credentials_from_env_or_provider("VOLUME_ACCESS_KEY", "VOLUME_SECRET_KEY").await?; + let credentials = credentials_from_env_or_provider( + "VOLUME_ACCESS_KEY", + "VOLUME_SECRET_KEY", + "VOLUME_AWS_SESSION_TOKEN", + ) + .await?; + + validate_s3tables_credentials(&arn, &credentials).await?; VolumeType::S3Tables(S3TablesVolume { endpoint: None, @@ -445,6 +454,7 @@ async fn load_volume_from_env() -> Result, ConfigError> { let credentials = AwsCredentials::AccessKey(AwsAccessKeyCredentials { aws_access_key_id: access_key, aws_secret_access_key: secret_key, + aws_session_token: None, }); VolumeType::S3( @@ -479,11 +489,14 @@ async fn load_volume_from_env() -> Result, ConfigError> { async fn credentials_from_env_or_provider( access_key_env: &str, secret_key_env: &str, + session_token_env: &str, ) -> Result { if let (Ok(access_key), Ok(secret_key)) = (env::var(access_key_env), env::var(secret_key_env)) { + let session_token = env::var(session_token_env).ok(); return Ok(AwsCredentials::AccessKey(AwsAccessKeyCredentials { aws_access_key_id: access_key, aws_secret_access_key: secret_key, + aws_session_token: session_token, })); } @@ -506,5 +519,54 @@ async fn credentials_from_env_or_provider( Ok(AwsCredentials::AccessKey(AwsAccessKeyCredentials { aws_access_key_id: creds.access_key_id().to_string(), aws_secret_access_key: creds.secret_access_key().to_string(), + aws_session_token: creds.session_token().map(std::string::ToString::to_string), })) } + +async fn validate_s3tables_credentials( + arn: &str, + credentials: &AwsCredentials, +) -> Result<(), ConfigError> { + let (access_key, secret_key, token) = match credentials { + AwsCredentials::AccessKey(creds) => ( + creds.aws_access_key_id.clone(), + creds.aws_secret_access_key.clone(), + creds.aws_session_token.clone(), + ), + AwsCredentials::Token(_) => { + return Err(ConfigError::EnvConfig { + reason: "S3 Tables validation requires access key credentials".to_string(), + }); + } + }; + + let region = arn + .split(':') + .nth(3) + .filter(|value| !value.trim().is_empty()) + .unwrap_or("us-east-1") + .to_string(); + + let config = aws_config::defaults(BehaviorVersion::latest()) + .credentials_provider(SharedCredentialsProvider::new(Credentials::from_keys( + access_key, secret_key, token, + ))) + .region(Region::new(region)) + .load() + .await; + let client = S3TablesClient::new(&config); + + client + .get_table_bucket() + .table_bucket_arn(arn) + .send() + .await + .map_err(|error| ConfigError::EnvConfig { + reason: format!( + "Failed to validate S3 Tables credentials for {arn}: {:?}", + error.as_service_error() + ), + })?; + + Ok(()) +} diff --git a/crates/catalog-metastore/src/models/volumes.rs b/crates/catalog-metastore/src/models/volumes.rs index aa277d2e..08e12568 100644 --- a/crates/catalog-metastore/src/models/volumes.rs +++ b/crates/catalog-metastore/src/models/volumes.rs @@ -53,6 +53,7 @@ pub struct AwsAccessKeyCredentials { pub aws_access_key_id: String, #[validate(regex(path = aws_secret_access_key_regex_func(), message = "AWS Secret access key is expected to be 40 chars Base64-like string with uppercase, lowercase, digits, and +/= .\n"))] pub aws_secret_access_key: String, + pub aws_session_token: Option, } impl std::fmt::Display for AwsAccessKeyCredentials { @@ -70,6 +71,7 @@ impl std::fmt::Debug for AwsAccessKeyCredentials { f.debug_struct("AwsAccessKeyCredentials") .field("aws_access_key_id", &self.aws_access_key_id) .field("aws_secret_access_key", &"**********") + .field("aws_session_token", &"**********") .finish() } } diff --git a/crates/catalog/src/catalog_list.rs b/crates/catalog/src/catalog_list.rs index 42b83de8..3a257598 100644 --- a/crates/catalog/src/catalog_list.rs +++ b/crates/catalog/src/catalog_list.rs @@ -251,7 +251,7 @@ impl EmbucketCatalogList { AwsCredentials::AccessKey(ref creds) => ( Some(creds.aws_access_key_id.clone()), Some(creds.aws_secret_access_key.clone()), - None, + creds.aws_session_token.clone(), ), AwsCredentials::Token(ref t) => (None, None, Some(t.clone())), }; diff --git a/crates/executor/src/query.rs b/crates/executor/src/query.rs index cd7e25bd..a5420a41 100644 --- a/crates/executor/src/query.rs +++ b/crates/executor/src/query.rs @@ -1500,6 +1500,7 @@ impl UserQuery { credentials: AwsCredentials::AccessKey(AwsAccessKeyCredentials { aws_access_key_id: key_id, aws_secret_access_key: secret_key, + aws_session_token: None, }), arn: params.aws_access_point_arn.unwrap_or_default(), client_options: None, @@ -1517,6 +1518,7 @@ impl UserQuery { let aws_credentials = AwsCredentials::AccessKey(AwsAccessKeyCredentials { aws_access_key_id: key_id, aws_secret_access_key: secret_key, + aws_session_token: None, }); Volume::new( @@ -2844,6 +2846,7 @@ impl UserQuery { Some(AwsCredentials::AccessKey(AwsAccessKeyCredentials { aws_access_key_id: access_key.to_string(), aws_secret_access_key: secret_key.to_string(), + aws_session_token: None, })) };