Merge pull request #7758 from ClSlaid/handle-session-token-copy-into

feat(query): add security token support for AWS S3

Author: mergify[bot]

docs/doc/21-load-data/01-s3.md

@@ -82,7 +82,7 @@ To use the COPY data loading, you will need the following information:
 * Your AWS account’s access keys, such as:
   * Access Key ID: *your-access-key-id*
   * Secret Access Key: *your-secret-access-key*
-
+  * Security Token (Optional): *your-aws-temporary-access-token*
 
 ### Step 5: Copy Data into the Target Tables
 
@@ -95,7 +95,7 @@ Using this URI and keys, execute the following statement, replacing the placehol
 ```sql
 COPY INTO books
   FROM 's3://databend-bohu/data/'
-  credentials=(aws_key_id='<your-access-key-id>' aws_secret_key='<your-secret-access-key>')
+  credentials=(aws_key_id='<your-access-key-id>' aws_secret_key='<your-secret-access-key>' [aws-token='<your-aws-temporary-access-token>'])
   pattern ='.*[.]csv'
   file_format = (type = 'CSV' field_delimiter = ','  record_delimiter = '\n' skip_header = 0);
 ```

src/common/storage/src/config.rs

@@ -200,6 +200,13 @@ pub struct StorageS3Config {
     pub bucket: String,
     pub access_key_id: String,
     pub secret_access_key: String,
+    /// Temporary security token used for authentications
+    ///
+    /// This recommended to use since users don't need to store their permanent credentials in their
+    /// scripts or worksheets.
+    ///
+    /// refer to [documentations](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_temp.html) for details.
+    pub security_token: String,
     pub master_key: String,
     pub root: String,
     /// This flag is used internally to control whether databend load
@@ -220,6 +227,7 @@ impl Default for StorageS3Config {
             bucket: "".to_string(),
             access_key_id: "".to_string(),
             secret_access_key: "".to_string(),
+            security_token: "".to_string(),
             master_key: "".to_string(),
             root: "".to_string(),
             disable_credential_loader: false,
@@ -242,6 +250,7 @@ impl Debug for StorageS3Config {
                 "secret_access_key",
                 &mask_string(&self.secret_access_key, 3),
             )
+            .field("security_token", &mask_string(&self.security_token, 3))
             .field("master_key", &mask_string(&self.master_key, 3))
             .finish()
     }

src/common/storage/src/location.rs

@@ -115,6 +115,12 @@ pub fn parse_uri_location(l: &UriLocation) -> Result<(StorageParams, String)> {
                 .or_else(|| l.connection.get("aws_secret_key"))
                 .cloned()
                 .unwrap_or_default(),
+            security_token: l
+                .connection
+                .get("security_token")
+                .or_else(|| l.connection.get("aws_token"))
+                .cloned()
+                .unwrap_or_default(),
             master_key: l.connection.get("master_key").cloned().unwrap_or_default(),
             root: root.to_string(),
             disable_credential_loader: true,

src/common/storage/tests/location.rs

@@ -48,6 +48,7 @@ fn test_parse_uri_location() -> Result<()> {
                     bucket: "test".to_string(),
                     access_key_id: "access_key_id".to_string(),
                     secret_access_key: "secret_access_key".to_string(),
+                    security_token: "".to_string(),
                     master_key: "".to_string(),
                     root: "/tmp/".to_string(),
                     disable_credential_loader: true,
@@ -65,6 +66,7 @@ fn test_parse_uri_location() -> Result<()> {
                 connection: vec![
                     ("aws_key_id", "access_key_id"),
                     ("aws_secret_key", "secret_access_key"),
+                    ("security_token", "security_token"),
                 ]
                 .iter()
                 .map(|(k, v)| (k.to_string(), v.to_string()))
@@ -77,6 +79,38 @@ fn test_parse_uri_location() -> Result<()> {
                     bucket: "test".to_string(),
                     access_key_id: "access_key_id".to_string(),
                     secret_access_key: "secret_access_key".to_string(),
+                    security_token: "security_token".to_string(),
+                    master_key: "".to_string(),
+                    root: "/tmp/".to_string(),
+                    disable_credential_loader: true,
+                    enable_virtual_host_style: false,
+                }),
+                "/".to_string(),
+            ),
+        ),
+        (
+            "s3_with_aws_security_token",
+            UriLocation {
+                protocol: "s3".to_string(),
+                name: "test".to_string(),
+                path: "/tmp/".to_string(),
+                connection: vec![
+                    ("aws_key_id", "access_key_id"),
+                    ("aws_secret_key", "secret_access_key"),
+                    ("aws_token", "security_token"),
+                ]
+                .iter()
+                .map(|(k, v)| (k.to_string(), v.to_string()))
+                .collect(),
+            },
+            (
+                StorageParams::S3(StorageS3Config {
+                    endpoint_url: STORAGE_S3_DEFAULT_ENDPOINT.to_string(),
+                    region: "".to_string(),
+                    bucket: "test".to_string(),
+                    access_key_id: "access_key_id".to_string(),
+                    secret_access_key: "secret_access_key".to_string(),
+                    security_token: "security_token".to_string(),
                     master_key: "".to_string(),
                     root: "/tmp/".to_string(),
                     disable_credential_loader: true,

src/meta/proto-conv/src/config_from_to_protobuf_impl.rs

@@ -35,6 +35,7 @@ impl FromToProto for StorageS3Config {
             endpoint_url: p.endpoint_url,
             access_key_id: p.access_key_id,
             secret_access_key: p.secret_access_key,
+            security_token: p.security_token,
             bucket: p.bucket,
             root: p.root,
             master_key: p.master_key,
@@ -51,6 +52,7 @@ impl FromToProto for StorageS3Config {
             endpoint_url: self.endpoint_url.clone(),
             access_key_id: self.access_key_id.clone(),
             secret_access_key: self.secret_access_key.clone(),
+            security_token: self.security_token.clone(),
             bucket: self.bucket.clone(),
             root: self.root.clone(),
             master_key: self.master_key.clone(),

src/meta/proto-conv/src/util.rs

@@ -43,6 +43,10 @@ const META_CHANGE_LOG: &[(u64, &str)] = &[
         "2022-09-09: Add: table.proto/{TableCopiedFileInfo,TableCopiedFileLock} type",
     ),
     (8, "2022-09-16: Add: users.proto/StageFile::entity_tag"),
+    (
+        9,
+        "2022-09-20: Add: config.proto/S3StorageConfig::security_token",
+    ),
 ];
 
 pub const VER: u64 = META_CHANGE_LOG.last().unwrap().0;

src/meta/proto-conv/tests/it/user_proto_conv.rs

@@ -102,6 +102,7 @@ pub(crate) fn test_s3_stage_info() -> mt::UserStageInfo {
                 root: "/data/files".to_string(),
                 access_key_id: "my_key_id".to_string(),
                 secret_access_key: "my_secret_key".to_string(),
+                security_token: "my_security_token".to_string(),
                 master_key: "my_master_key".to_string(),
                 ..Default::default()
             }),

src/meta/proto-conv/tests/it/user_stage.rs

@@ -43,6 +43,55 @@ fn test_user_stage_gcs_latest() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_user_stage_s3_v8() -> anyhow::Result<()> {
+    // Encoded data of version 8 of user_stage_s3:
+    // It is generated with common::test_pb_from_to.
+    let user_stage_s3_v8 = vec![
+        10, 24, 115, 51, 58, 47, 47, 109, 121, 98, 117, 99, 107, 101, 116, 47, 100, 97, 116, 97,
+        47, 102, 105, 108, 101, 115, 16, 1, 26, 100, 10, 98, 10, 96, 18, 24, 104, 116, 116, 112,
+        115, 58, 47, 47, 115, 51, 46, 97, 109, 97, 122, 111, 110, 97, 119, 115, 46, 99, 111, 109,
+        26, 9, 109, 121, 95, 107, 101, 121, 95, 105, 100, 34, 13, 109, 121, 95, 115, 101, 99, 114,
+        101, 116, 95, 107, 101, 121, 42, 8, 109, 121, 98, 117, 99, 107, 101, 116, 50, 11, 47, 100,
+        97, 116, 97, 47, 102, 105, 108, 101, 115, 58, 13, 109, 121, 95, 109, 97, 115, 116, 101,
+        114, 95, 107, 101, 121, 160, 6, 8, 168, 6, 1, 34, 20, 8, 1, 16, 128, 8, 26, 1, 124, 34, 2,
+        47, 47, 40, 2, 160, 6, 8, 168, 6, 1, 42, 10, 10, 3, 32, 154, 5, 16, 142, 8, 24, 1, 50, 4,
+        116, 101, 115, 116, 160, 6, 8, 168, 6, 1,
+    ];
+
+    let want = mt::UserStageInfo {
+        stage_name: "s3://mybucket/data/files".to_string(),
+        stage_type: mt::StageType::External,
+        stage_params: mt::StageParams {
+            storage: StorageParams::S3(StorageS3Config {
+                bucket: "mybucket".to_string(),
+                root: "/data/files".to_string(),
+                access_key_id: "my_key_id".to_string(),
+                secret_access_key: "my_secret_key".to_string(),
+                master_key: "my_master_key".to_string(),
+                ..Default::default()
+            }),
+        },
+        file_format_options: mt::FileFormatOptions {
+            format: mt::StageFileFormatType::Json,
+            skip_header: 1024,
+            field_delimiter: "|".to_string(),
+            record_delimiter: "//".to_string(),
+            compression: mt::StageFileCompression::Bz2,
+        },
+        copy_options: mt::CopyOptions {
+            on_error: mt::OnErrorMode::SkipFileNum(666),
+            size_limit: 1038,
+            purge: true,
+        },
+        comment: "test".to_string(),
+        ..Default::default()
+    };
+
+    common::test_load_old(func_name!(), user_stage_s3_v8.as_slice(), want)?;
+    Ok(())
+}
+
 #[test]
 fn test_user_stage_fs_v6() -> anyhow::Result<()> {
     // Encoded data of version 6 of user_stage_fs:

src/meta/protos/proto/config.proto

@@ -29,6 +29,7 @@ message S3StorageConfig {
   string master_key = 7;
   bool disable_credential_loader = 8;
   bool enable_virtual_host_style = 9;
+  string security_token = 10;
 }
 
 message FsStorageConfig {

src/query/config/src/outer_v0.rs

@@ -402,6 +402,12 @@ pub struct S3StorageConfig {
     #[clap(long = "storage-s3-secret-access-key", default_value_t)]
     pub secret_access_key: String,
 
+    /// Security token for S3 storage
+    ///
+    /// Check out [documents](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_temp.html) for details
+    #[clap(long = "storage-s3-security-token", default_value_t)]
+    pub security_token: String,
+
     /// S3 Bucket to use for storage
     #[clap(long = "storage-s3-bucket", default_value_t)]
     pub bucket: String,
@@ -449,6 +455,7 @@ impl From<InnerStorageS3Config> for S3StorageConfig {
             endpoint_url: inner.endpoint_url,
             access_key_id: inner.access_key_id,
             secret_access_key: inner.secret_access_key,
+            security_token: inner.security_token,
             bucket: inner.bucket,
             root: inner.root,
             master_key: inner.master_key,
@@ -467,6 +474,7 @@ impl TryInto<InnerStorageS3Config> for S3StorageConfig {
             bucket: self.bucket,
             access_key_id: self.access_key_id,
             secret_access_key: self.secret_access_key,
+            security_token: self.security_token,
             master_key: self.master_key,
             root: self.root,
             disable_credential_loader: false,