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

mergify[bot] · web-flow · commit 7367522126ff · 2022-09-21T03:25:58.000Z
feat(query): add security token support for AWS S3
diff --git a/docs/doc/21-load-data/01-s3.md b/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);
 ```
diff --git a/src/common/storage/src/config.rs b/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()
     }
diff --git a/src/common/storage/src/location.rs b/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,
diff --git a/src/common/storage/tests/location.rs b/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,
diff --git a/src/meta/proto-conv/src/config_from_to_protobuf_impl.rs b/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(),
diff --git a/src/meta/proto-conv/src/util.rs b/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;
diff --git a/src/meta/proto-conv/tests/it/user_proto_conv.rs b/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()
             }),
diff --git a/src/meta/proto-conv/tests/it/user_stage.rs b/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:
diff --git a/src/meta/protos/proto/config.proto b/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 {
diff --git a/src/query/config/src/outer_v0.rs b/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,
diff --git a/src/query/service/tests/it/configs.rs b/src/query/service/tests/it/configs.rs
@@ -117,6 +117,7 @@ region = ""
 endpoint_url = "https://s3.amazonaws.com"
 access_key_id = ""
 secret_access_key = ""
+security_token = ""
 bucket = ""
 root = ""
 master_key = ""
diff --git a/src/query/service/tests/it/interpreters/interpreter_user_stage.rs b/src/query/service/tests/it/interpreters/interpreter_user_stage.rs
@@ -51,11 +51,11 @@ async fn test_user_stage_interpreter() -> Result<()> {
 
         common_datablocks::assert_blocks_eq(
             vec![
-                "+------------+------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------+-----------------+--------------------+---------+",
-                "| name       | stage_type | stage_params                                                                                                                                                                                                                                                                          | copy_options                                                | file_format_options                                                                                                | number_of_files | creator            | comment |",
-                "+------------+------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------+-----------------+--------------------+---------+",
-                "| test_stage | External   | StageParams { storage: S3(StorageS3Config { endpoint_url: \"https://s3.amazonaws.com\", region: \"\", bucket: \"load\", root: \"/files/\", disable_credential_loader: true, enable_virtual_host_style: false, access_key_id: \"******b3c\", secret_access_key: \"******y6z\", master_key: \"\" }) } | CopyOptions { on_error: None, size_limit: 0, purge: false } | FileFormatOptions { format: Csv, skip_header: 0, field_delimiter: \",\", record_delimiter: \"\\n\", compression: None } | NULL            | 'root'@'127.0.0.1' |         |",
-                "+------------+------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------+-----------------+--------------------+---------+",
+                "+------------+------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------+-----------------+--------------------+---------+",
+                "| name       | stage_type | stage_params                                                                                                                                                                                                                                                                                              | copy_options                                                | file_format_options                                                                                                | number_of_files | creator            | comment |",
+                "+------------+------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------+-----------------+--------------------+---------+",
+                "| test_stage | External   | StageParams { storage: S3(StorageS3Config { endpoint_url: \"https://s3.amazonaws.com\", region: \"\", bucket: \"load\", root: \"/files/\", disable_credential_loader: true, enable_virtual_host_style: false, access_key_id: \"******b3c\", secret_access_key: \"******y6z\", security_token: \"\", master_key: \"\" }) } | CopyOptions { on_error: None, size_limit: 0, purge: false } | FileFormatOptions { format: Csv, skip_header: 0, field_delimiter: \",\", record_delimiter: \"\\n\", compression: None } | NULL            | 'root'@'127.0.0.1' |         |",
+                "+------------+------------+-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------------+-----------------+--------------------+---------+",
             ],
             &blocks,
         );
diff --git a/src/query/service/tests/it/storages/system/configs_table.rs b/src/query/service/tests/it/storages/system/configs_table.rs
@@ -133,6 +133,7 @@ async fn test_configs_table() -> Result<()> {
         "| storage | s3.region                            |                                |             |",
         "| storage | s3.root                              |                                |             |",
         "| storage | s3.secret_access_key                 |                                |             |",
+        "| storage | s3.security_token                    |                                |             |",
         "| storage | type                                 | fs                             |             |",
         "+---------+--------------------------------------+--------------------------------+-------------+",
     ];
@@ -266,6 +267,7 @@ async fn test_configs_table_redact() -> Result<()> {
         "| storage | s3.region                            | us-east-2                      |             |",
         "| storage | s3.root                              |                                |             |",
         "| storage | s3.secret_access_key                 | ******key                      |             |",
+        "| storage | s3.security_token                    |                                |             |",
         "| storage | type                                 | s3                             |             |",
         "+---------+--------------------------------------+--------------------------------+-------------+",
     ];

Original file line number	Diff line number	Diff line change
`@@ -29,6 +29,7 @@ message S3StorageConfig {`
`29`	`29`	`string master_key = 7;`
`30`	`30`	`bool disable_credential_loader = 8;`
`31`	`31`	`bool enable_virtual_host_style = 9;`
	`32`	`+ string security_token = 10;`
`32`	`33`	`}`
`33`	`34`
`34`	`35`	`message FsStorageConfig {`