[arrow-rs] branch master updated: Build AmazonS3builder from environment variables (#2361) (#2536)

[tu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new 4949a3d02 Build AmazonS3builder from environment variables (#2361) (#2536)
4949a3d02 is described below

commit 4949a3d026bc15dd8320063d9e06fe048d26c6e8
Author: Quentin <am...@gmail.com>
AuthorDate: Sat Aug 20 17:53:50 2022 +0200

    Build AmazonS3builder from environment variables (#2361) (#2536)
    
    * Build AmazonS3Builder from env (#2361)
    
    * clippy: use Self instead of AmazonS3Builder
    
    * rustdoc: remove unnecessary bare_url
    
    * Save the current environment setup
    
    * fmt: missing use in documentation
---
 object_store/src/aws/mod.rs | 78 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 78 insertions(+)

diff --git a/object_store/src/aws/mod.rs b/object_store/src/aws/mod.rs
index 06d20ccc9..2a1308346 100644
--- a/object_store/src/aws/mod.rs
+++ b/object_store/src/aws/mod.rs
@@ -347,6 +347,48 @@ impl AmazonS3Builder {
         Default::default()
     }
 
+    /// Fill the [`AmazonS3Builder`] with regular AWS environment variables
+    ///
+    /// Variables extracted from environment:
+    /// * AWS_ACCESS_KEY_ID -> access_key_id
+    /// * AWS_SECRET_ACCESS_KEY -> secret_access_key
+    /// * AWS_DEFAULT_REGION -> region
+    /// * AWS_ENDPOINT -> endpoint
+    /// * AWS_SESSION_TOKEN -> token
+    /// # Example
+    /// ```
+    /// use object_store::aws::AmazonS3Builder;
+    ///
+    /// let s3 = AmazonS3Builder::from_env()
+    ///     .with_bucket_name("foo")
+    ///     .build();
+    /// ```
+    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);
+        }
+
+        builder
+    }
+
     /// Set the AWS Access Key (required)
     pub fn with_access_key_id(mut self, access_key_id: impl Into<String>) -> Self {
         self.access_key_id = Some(access_key_id.into());
@@ -574,6 +616,42 @@ mod tests {
         }};
     }
 
+    #[test]
+    fn s3_test_config_from_env() {
+        let aws_access_key_id = env::var("AWS_ACCESS_KEY_ID")
+            .unwrap_or("object_store:fake_access_key_id".into());
+        let aws_secret_access_key = env::var("AWS_SECRET_ACCESS_KEY")
+            .unwrap_or("object_store:fake_secret_key".into());
+
+        let aws_default_region = env::var("AWS_DEFAULT_REGION")
+            .unwrap_or("object_store:fake_default_region".into());
+
+        let aws_endpoint =
+            env::var("AWS_ENDPOINT").unwrap_or("object_store:fake_endpoint".into());
+        let aws_session_token = env::var("AWS_SESSION_TOKEN")
+            .unwrap_or("object_store:fake_session_token".into());
+
+        // required
+        env::set_var("AWS_ACCESS_KEY_ID", &aws_access_key_id);
+        env::set_var("AWS_SECRET_ACCESS_KEY", &aws_secret_access_key);
+        env::set_var("AWS_DEFAULT_REGION", &aws_default_region);
+
+        // optional
+        env::set_var("AWS_ENDPOINT", &aws_endpoint);
+        env::set_var("AWS_SESSION_TOKEN", &aws_session_token);
+
+        let builder = AmazonS3Builder::from_env();
+        assert_eq!(builder.access_key_id.unwrap(), aws_access_key_id.as_str());
+        assert_eq!(
+            builder.secret_access_key.unwrap(),
+            aws_secret_access_key.as_str()
+        );
+        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!();