|
| 1 | +/*! Utils for converting standard Iceberg config formats to equivalent `object_store` options |
| 2 | +*/ |
| 3 | + |
| 4 | +use crate::error::Error; |
| 5 | +use object_store::aws::{AmazonS3Builder, AmazonS3ConfigKey}; |
| 6 | +use object_store::gcp::{GcpCredential, GoogleCloudStorageBuilder, GoogleConfigKey}; |
| 7 | +use object_store::{parse_url_opts, ObjectStore, ObjectStoreScheme, StaticCredentialProvider}; |
| 8 | +use std::collections::HashMap; |
| 9 | +use std::sync::Arc; |
| 10 | +use url::Url; |
| 11 | + |
| 12 | +/// AWS configs |
| 13 | +const CLIENT_REGION: &str = "client.region"; |
| 14 | +const AWS_ACCESS_KEY_ID: &str = "s3.access-key-id"; |
| 15 | +const AWS_SECRET_ACCESS_KEY: &str = "s3.secret-access-key"; |
| 16 | +const AWS_SESSION_TOKEN: &str = "s3.session-token"; |
| 17 | +const AWS_REGION: &str = "s3.region"; |
| 18 | +const AWS_ENDPOINT: &str = "s3.endpoint"; |
| 19 | +const AWS_ALLOW_ANONYMOUS: &str = "s3.allow-anonymous"; |
| 20 | + |
| 21 | +/// GCP configs |
| 22 | +const GCS_BUCKET: &str = "gcs.bucket"; |
| 23 | +const GCS_CREDENTIALS_JSON: &str = "gcs.credentials-json"; |
| 24 | +const GCS_TOKEN: &str = "gcs.oauth2.token"; |
| 25 | + |
| 26 | +/// Parse the url and Iceberg format of variuos storage options into the equivalent `object_store` |
| 27 | +/// options and build the corresponding `ObjectStore`. |
| 28 | +pub fn object_store_from_config( |
| 29 | + url: Url, |
| 30 | + config: HashMap<String, String>, |
| 31 | +) -> Result<Arc<dyn ObjectStore>, Error> { |
| 32 | + let store = match ObjectStoreScheme::parse(&url).map_err(object_store::Error::from)? { |
| 33 | + (ObjectStoreScheme::AmazonS3, _) => { |
| 34 | + let mut builder = AmazonS3Builder::new().with_url(url); |
| 35 | + for (key, option) in config { |
| 36 | + let s3_key = match key.as_str() { |
| 37 | + AWS_ACCESS_KEY_ID => AmazonS3ConfigKey::AccessKeyId, |
| 38 | + AWS_SECRET_ACCESS_KEY => AmazonS3ConfigKey::SecretAccessKey, |
| 39 | + AWS_SESSION_TOKEN => AmazonS3ConfigKey::Token, |
| 40 | + CLIENT_REGION | AWS_REGION => AmazonS3ConfigKey::Region, |
| 41 | + AWS_ENDPOINT => { |
| 42 | + if option.starts_with("http://") { |
| 43 | + // This is mainly used for testing, e.g. against MinIO |
| 44 | + builder = builder.with_allow_http(true); |
| 45 | + } |
| 46 | + AmazonS3ConfigKey::Endpoint |
| 47 | + } |
| 48 | + AWS_ALLOW_ANONYMOUS => AmazonS3ConfigKey::SkipSignature, |
| 49 | + _ => continue, |
| 50 | + }; |
| 51 | + builder = builder.with_config(s3_key, option); |
| 52 | + } |
| 53 | + Arc::new(builder.build()?) as Arc<dyn ObjectStore> |
| 54 | + } |
| 55 | + |
| 56 | + (ObjectStoreScheme::GoogleCloudStorage, _) => { |
| 57 | + let mut builder = GoogleCloudStorageBuilder::new().with_url(url); |
| 58 | + for (key, option) in config { |
| 59 | + let gcs_key = match key.as_str() { |
| 60 | + GCS_CREDENTIALS_JSON => GoogleConfigKey::ServiceAccountKey, |
| 61 | + GCS_BUCKET => GoogleConfigKey::Bucket, |
| 62 | + GCS_TOKEN => { |
| 63 | + let credential = GcpCredential { bearer: option }; |
| 64 | + let credential_provider = |
| 65 | + Arc::new(StaticCredentialProvider::new(credential)) as _; |
| 66 | + builder = builder.with_credentials(credential_provider); |
| 67 | + continue; |
| 68 | + } |
| 69 | + _ => continue, |
| 70 | + }; |
| 71 | + builder = builder.with_config(gcs_key, option); |
| 72 | + } |
| 73 | + Arc::new(builder.build()?) as Arc<dyn ObjectStore> |
| 74 | + } |
| 75 | + |
| 76 | + _ => { |
| 77 | + let (store, _path) = parse_url_opts(&url, config)?; |
| 78 | + store.into() |
| 79 | + } |
| 80 | + }; |
| 81 | + |
| 82 | + Ok(store) |
| 83 | +} |
| 84 | + |
| 85 | +#[cfg(test)] |
| 86 | +mod tests { |
| 87 | + use super::*; |
| 88 | + use serde_json::json; |
| 89 | + use std::collections::HashMap; |
| 90 | + use url::Url; |
| 91 | + |
| 92 | + #[test] |
| 93 | + fn test_s3_config_basic() { |
| 94 | + let url = Url::parse("s3://test-bucket/path").unwrap(); |
| 95 | + let mut config = HashMap::new(); |
| 96 | + config.insert(AWS_ACCESS_KEY_ID.to_string(), "test-key".to_string()); |
| 97 | + config.insert(AWS_SECRET_ACCESS_KEY.to_string(), "test-secret".to_string()); |
| 98 | + config.insert(AWS_SESSION_TOKEN.to_string(), "test-session".to_string()); |
| 99 | + config.insert(AWS_REGION.to_string(), "us-east-1".to_string()); |
| 100 | + |
| 101 | + let store = object_store_from_config(url, config).unwrap(); |
| 102 | + let store_repr = format!("{:?}", store); |
| 103 | + |
| 104 | + assert!(store_repr.contains("region: \"us-east-1\"")); |
| 105 | + assert!(store_repr.contains("bucket: \"test-bucket\"")); |
| 106 | + assert!(store_repr.contains("key_id: \"test-key\"")); |
| 107 | + assert!(store_repr.contains("secret_key: \"test-secret\"")); |
| 108 | + assert!(store_repr.contains("token: Some(\"test-session\")")); |
| 109 | + assert!(store_repr.contains("endpoint: None")); |
| 110 | + assert!(store_repr.contains("allow_http: Parsed(false)")); |
| 111 | + assert!(store_repr.contains("skip_signature: false")); |
| 112 | + } |
| 113 | + |
| 114 | + #[test] |
| 115 | + fn test_s3_config_with_http_endpoint() { |
| 116 | + let url = Url::parse("s3://test-bucket/").unwrap(); |
| 117 | + let mut config = HashMap::new(); |
| 118 | + config.insert( |
| 119 | + AWS_ENDPOINT.to_string(), |
| 120 | + "http://localhost:9000".to_string(), |
| 121 | + ); |
| 122 | + config.insert(AWS_ALLOW_ANONYMOUS.to_string(), "true".to_string()); |
| 123 | + |
| 124 | + let store = object_store_from_config(url, config).unwrap(); |
| 125 | + let store_repr = format!("{:?}", store); |
| 126 | + |
| 127 | + assert!(store_repr.contains("region: \"us-east-1\"")); |
| 128 | + assert!(store_repr.contains("bucket: \"test-bucket\"")); |
| 129 | + assert!(!store_repr.contains("key_id: ")); |
| 130 | + assert!(!store_repr.contains("secret_key: ")); |
| 131 | + assert!(!store_repr.contains("token: ")); |
| 132 | + assert!(store_repr.contains("endpoint: Some(\"http://localhost:9000\")")); |
| 133 | + assert!(store_repr.contains("allow_http: Parsed(true)")); |
| 134 | + assert!(store_repr.contains("skip_signature: true")); |
| 135 | + } |
| 136 | + |
| 137 | + #[test] |
| 138 | + fn test_gcs_config_with_service_account() { |
| 139 | + let url = Url::parse("gs://test-bucket/").unwrap(); |
| 140 | + let mut config = HashMap::new(); |
| 141 | + config.insert( |
| 142 | + GCS_CREDENTIALS_JSON.to_string(), |
| 143 | + json!( |
| 144 | + { |
| 145 | + "disable_oauth": true, "client_email": "", "private_key": "", "private_key_id": "" |
| 146 | + } |
| 147 | + ) |
| 148 | + .to_string(), |
| 149 | + ); |
| 150 | + config.insert(GCS_BUCKET.to_string(), "test-bucket".to_string()); |
| 151 | + |
| 152 | + let store = object_store_from_config(url, config).unwrap(); |
| 153 | + let store_repr = format!("{:?}", store); |
| 154 | + |
| 155 | + assert!(store_repr.contains("bearer: \"\"")); |
| 156 | + assert!(store_repr.contains("bucket_name: \"test-bucket\"")); |
| 157 | + } |
| 158 | + |
| 159 | + #[test] |
| 160 | + fn test_gcs_config_with_oauth_token() { |
| 161 | + let url = Url::parse("gs://test-bucket/").unwrap(); |
| 162 | + let mut config = HashMap::new(); |
| 163 | + config.insert(GCS_TOKEN.to_string(), "oauth-token-123".to_string()); |
| 164 | + config.insert(GCS_BUCKET.to_string(), "test-bucket".to_string()); |
| 165 | + |
| 166 | + let store = object_store_from_config(url, config).unwrap(); |
| 167 | + let store_repr = format!("{:?}", store); |
| 168 | + |
| 169 | + assert!(store_repr.contains("bearer: \"oauth-token-123\"")); |
| 170 | + assert!(store_repr.contains("bucket_name: \"test-bucket\"")); |
| 171 | + } |
| 172 | +} |
0 commit comments