Skip to content
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.

Sign up for GitHub

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

Jump to bottom

object_store: builder configuration api #3436

Merged
merged 10 commits into from
Jan 4, 2023
Merged

object_store: builder configuration api #3436

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
0597445
feat: draf configuration api for azure
roeap Jan 3, 2023
360c9db
feat: add configuration for aws and gcp
roeap Jan 3, 2023
3ab6ac8
fix: clippy
roeap Jan 3, 2023
ac638a7
feat: allow passing typed config keys
roeap Jan 4, 2023
2f5e5cc
refactor: implement try_from for config keys
roeap Jan 4, 2023
f7ebee1
chore: PR feedback
roeap Jan 4, 2023
a9cf97f
refactor: make options api fallible
roeap Jan 4, 2023
786f469
fix: docs errors
roeap Jan 4, 2023
15668bd
chore: remove helpers
roeap Jan 4, 2023
7dea6c8
test: test sas key splitting and un-nit nits
roeap Jan 4, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion object_store/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ walkdir = "2"

# Cloud storage support
base64 = { version = "0.20", default-features = false, features = ["std"], optional = true }
once_cell = { version = "1.12.0", optional = true }
quick-xml = { version = "0.27.0", features = ["serialize"], optional = true }
serde = { version = "1.0", default-features = false, features = ["derive"], optional = true }
serde_json = { version = "1.0", default-features = false, optional = true }
Expand All @@ -57,7 +58,7 @@ aws-types = { version = "0.52", optional = true }
aws-config = { version = "0.52", optional = true }

[features]
cloud = ["serde", "serde_json", "quick-xml", "reqwest", "reqwest/json", "reqwest/stream", "chrono/serde", "base64", "rand", "ring"]
cloud = ["serde", "serde_json", "quick-xml", "reqwest", "reqwest/json", "reqwest/stream", "chrono/serde", "base64", "rand", "ring", "once_cell"]
azure = ["cloud"]
gcp = ["cloud", "rustls-pemfile"]
aws = ["cloud"]
Expand Down
256 changes: 233 additions & 23 deletions object_store/src/aws/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,8 +37,9 @@ use chrono::{DateTime, Utc};
use futures::stream::BoxStream;
use futures::TryStreamExt;
use itertools::Itertools;
use once_cell::sync::Lazy;
use snafu::{OptionExt, ResultExt, Snafu};
use std::collections::BTreeSet;
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::ops::Range;
use std::sync::Arc;
use tokio::io::AsyncWrite;
Expand All @@ -51,6 +52,7 @@ use crate::aws::credential::{
StaticCredentialProvider, WebIdentityProvider,
};
use crate::multipart::{CloudMultiPartUpload, CloudMultiPartUploadImpl, UploadPart};
use crate::util::str_is_truthy;
use crate::{
ClientOptions, GetResult, ListResult, MultipartId, ObjectMeta, ObjectStore, Path,
Result, RetryConfig, StreamExt,
Expand Down Expand Up @@ -379,6 +381,162 @@ pub struct AmazonS3Builder {
client_options: ClientOptions,
}

#[derive(PartialEq, Eq)]
enum AmazonS3ConfigKey {
/// AWS Access Key
///
/// See [`AmazonS3Builder::with_access_key_id`] for details.
///
/// Supported keys:
/// - `aws_access_key_id`
/// - `access_key_id`
AccessKeyId,

/// Secret Access Key
///
/// See [`AmazonS3Builder::with_secret_access_key`] for details.
///
/// Supported keys:
/// - `aws_secret_access_key`
/// - `secret_access_key`
SecretAccessKey,

/// Region
///
/// See [`AmazonS3Builder::with_region`] for details.
///
/// Supported keys:
/// - `aws_region`
/// - `region`
Region,

/// Default region
///
/// See [`AmazonS3Builder::with_region`] for details.
///
/// Supported keys:
/// - `aws_default_region`
/// - `default_region`
DefaultRegion,

/// Bucket name
///
/// See [`AmazonS3Builder::with_bucket_name`] for details.
///
/// Supported keys:
/// - `aws_bucket`
/// - `aws_bucket_name`
/// - `bucket`
/// - `bucket_name`
Bucket,

/// Sets custom endpoint for communicating with AWS S3.
///
/// See [`AmazonS3Builder::with_endpoint`] for details.
///
/// Supported keys:
/// - `aws_endpoint`
/// - `aws_endpoint_url`
/// - `endpoint`
/// - `endpoint_url`
Endpoint,

/// Token to use for requests (passed to underlying provider)
///
/// See [`AmazonS3Builder::with_token`] for details.
///
/// Supported keys:
/// - `aws_session_token`
/// - `aws_token`
/// - `session_token`
/// - `token`
Token,

/// Fall back to ImdsV1
///
/// See [`AmazonS3Builder::with_imdsv1_fallback`] for details.
///
/// Supported keys:
/// - `aws_imdsv1_fallback`
/// - `imdsv1_fallback`
ImdsV1Fallback,

/// If virtual hosted style request has to be used
///
/// See [`AmazonS3Builder::with_virtual_hosted_style_request`] for details.
///
/// Supported keys:
/// - `aws_virtual_hosted_style_request`
/// - `virtual_hosted_style_request`
VirtualHostedStyleRequest,

/// Set the instance metadata endpoint
///
/// See [`AmazonS3Builder::with_metadata_endpoint`] for details.
///
/// Supported keys:
/// - `aws_metadata_endpoint`
/// - `metadata_endpoint`
MetadataEndpoint,

/// AWS profile name
///
/// Supported keys:
/// - `aws_profile`
/// - `profile`
Profile,
}

static ALIAS_MAP: Lazy<BTreeMap<&'static str, AmazonS3ConfigKey>> = Lazy::new(|| {
BTreeMap::from([
// access key id
("aws_access_key_id", AmazonS3ConfigKey::AccessKeyId),
("access_key_id", AmazonS3ConfigKey::AccessKeyId),
// secret access key
("aws_secret_access_key", AmazonS3ConfigKey::SecretAccessKey),
("secret_access_key", AmazonS3ConfigKey::SecretAccessKey),
// default region
("aws_default_region", AmazonS3ConfigKey::DefaultRegion),
("default_region", AmazonS3ConfigKey::DefaultRegion),
// region
("aws_region", AmazonS3ConfigKey::Region),
("region", AmazonS3ConfigKey::Region),
// bucket
("aws_bucket", AmazonS3ConfigKey::Bucket),
("aws_bucket_name", AmazonS3ConfigKey::Bucket),
("bucket_name", AmazonS3ConfigKey::Bucket),
("bucket", AmazonS3ConfigKey::Bucket),
// custom S3 endpoint
("aws_endpoint_url", AmazonS3ConfigKey::Endpoint),
("aws_endpoint", AmazonS3ConfigKey::Endpoint),
("endpoint_url", AmazonS3ConfigKey::Endpoint),
("endpoint", AmazonS3ConfigKey::Endpoint),
// session token
("aws_session_token", AmazonS3ConfigKey::Token),
("aws_token", AmazonS3ConfigKey::Token),
("session_token", AmazonS3ConfigKey::Token),
("token", AmazonS3ConfigKey::Token),
// virtual hosted style request
(
"aws_virtual_hosted_style_request",
AmazonS3ConfigKey::VirtualHostedStyleRequest,
),
(
"virtual_hosted_style_request",
AmazonS3ConfigKey::VirtualHostedStyleRequest,
),
// profile
("aws_profile", AmazonS3ConfigKey::Profile),
("profile", AmazonS3ConfigKey::Profile),
// imds v1 fallback
("aws_imdsv1_fallback", AmazonS3ConfigKey::ImdsV1Fallback),
("imdsv1_fallback", AmazonS3ConfigKey::ImdsV1Fallback),
// metadata endpoint
("aws_metadata_endpoint", AmazonS3ConfigKey::MetadataEndpoint),
("metadata_endpoint", AmazonS3ConfigKey::MetadataEndpoint),
])
});

impl AmazonS3Builder {
/// Create a new [`AmazonS3Builder`] with default values.
pub fn new() -> Self {
Expand Down Expand Up @@ -407,28 +565,12 @@ impl AmazonS3Builder {
pub fn from_env() -> Self {
let mut builder: Self = Default::default();

if let Ok(access_key_id) = std::env::var("AWS_ACCESS_KEY_ID") {
builder.access_key_id = Some(access_key_id);
}

if let Ok(secret_access_key) = std::env::var("AWS_SECRET_ACCESS_KEY") {
builder.secret_access_key = Some(secret_access_key);
}

if let Ok(secret) = std::env::var("AWS_DEFAULT_REGION") {
builder.region = Some(secret);
}

if let Ok(endpoint) = std::env::var("AWS_ENDPOINT") {
builder.endpoint = Some(endpoint);
}

if let Ok(token) = std::env::var("AWS_SESSION_TOKEN") {
builder.token = Some(token);
}

if let Ok(profile) = std::env::var("AWS_PROFILE") {
builder.profile = Some(profile);
for (key, _) in ALIAS_MAP.iter() {
if key.starts_with("aws_") {
if let Ok(value) = std::env::var(key.to_ascii_uppercase()) {
builder = builder.with_option(*key, value)
}
}
}

// This env var is set in ECS
Expand Down Expand Up @@ -472,6 +614,49 @@ impl AmazonS3Builder {
self
}

/// Set an option on the builder via a key - value pair.
pub fn with_option(
mut self,
key: impl Into<String>,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@roeap thanks for the super quick turn around time on this one. This is pretty close to what I am thinking, I would like to suggest one more change?

As an user I would like to have my configuration fully checked by the compiler. So I would like to either be able to pass in one of the AmazonS3ConfigKeys as the key or have a helper function that returns one of the string representations. This way in my code I never have to type a string constant for the key.

Let me know what you think....

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I ported my PR to the code so far. I had to toggle back and forth between the 2 PRs to make sure I get these constants correctly.
https://github.com/pola-rs/polars/pull/5972/files#diff-28ad3830b06d49511094951431beea1edcf6a567f646250482e9f00f769ace27R16

This is error prone specially since you have a type safe Enum sitting right there...

Copy link
Contributor Author

@roeap roeap Jan 4, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We now have some conversions implemented for the config key enums, so you should be able to pass in the variants directly, or constuct a map with the variants as keys. Is that what you are looking for here?

roeap marked this conversation as resolved.
Show resolved Hide resolved
value: impl Into<String>,
) -> Self {
let raw = key.into();
if let Some(key) = ALIAS_MAP.get(&*raw.to_ascii_lowercase()) {
match key {
AmazonS3ConfigKey::AccessKeyId => self.access_key_id = Some(value.into()),
AmazonS3ConfigKey::SecretAccessKey => {
self.secret_access_key = Some(value.into())
}
AmazonS3ConfigKey::Region => self.region = Some(value.into()),
AmazonS3ConfigKey::Bucket => self.bucket_name = Some(value.into()),
AmazonS3ConfigKey::Endpoint => self.endpoint = Some(value.into()),
AmazonS3ConfigKey::Token => self.token = Some(value.into()),
AmazonS3ConfigKey::ImdsV1Fallback => {
self.imdsv1_fallback = str_is_truthy(&value.into())
}
AmazonS3ConfigKey::VirtualHostedStyleRequest => {
self.virtual_hosted_style_request = str_is_truthy(&value.into())
}
AmazonS3ConfigKey::DefaultRegion => {
self.region = self.region.or_else(|| Some(value.into()))
}
AmazonS3ConfigKey::MetadataEndpoint => {
self.metadata_endpoint = Some(value.into())
}
AmazonS3ConfigKey::Profile => self.profile = Some(value.into()),
};
}
self
}

/// Hydrate builder from key value pairs
pub fn with_options(mut self, options: &HashMap<String, String>) -> Self {
for (key, value) in options {
self = self.with_option(key, value);
}
self
}

/// Sets properties on this builder based on a URL
///
/// This is a separate member function to allow fallible computation to
Expand Down Expand Up @@ -915,6 +1100,31 @@ mod tests {
assert_eq!(builder.metadata_endpoint.unwrap(), metadata_uri);
}

#[test]
fn s3_test_config_from_map() {
let aws_access_key_id = "object_store:fake_access_key_id".to_string();
let aws_secret_access_key = "object_store:fake_secret_key".to_string();
let aws_default_region = "object_store:fake_default_region".to_string();
let aws_endpoint = "object_store:fake_endpoint".to_string();
let aws_session_token = "object_store:fake_session_token".to_string();
let options = HashMap::from([
("aws_access_key_id".to_string(), aws_access_key_id.clone()),
("aws_secret_access_key".to_string(), aws_secret_access_key),
("aws_default_region".to_string(), aws_default_region.clone()),
("aws_endpoint".to_string(), aws_endpoint.clone()),
("aws_session_token".to_string(), aws_session_token.clone()),
]);

let builder = AmazonS3Builder::new()
.with_options(&options)
.with_option("aws_secret_access_key", "new-secret-key");
assert_eq!(builder.access_key_id.unwrap(), aws_access_key_id.as_str());
assert_eq!(builder.secret_access_key.unwrap(), "new-secret-key");
assert_eq!(builder.region.unwrap(), aws_default_region);
assert_eq!(builder.endpoint.unwrap(), aws_endpoint);
assert_eq!(builder.token.unwrap(), aws_session_token);
}

#[tokio::test]
async fn s3_test() {
let config = maybe_skip_integration!();
Expand Down
Loading