[fix](backup) support aws_token for backup/restore (apache#26242)

xinyiZzz · Nov 3, 2023 · 03ab822 · 03ab822
1 parent 91946dc
commit 03ab822
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/be/src/util/s3_util.cpp b/be/src/util/s3_util.cpp
@@ -152,6 +152,9 @@ std::shared_ptr<Aws::S3::S3Client> S3ClientFactory::create(const S3Conf& s3_conf
 
     Aws::Auth::AWSCredentials aws_cred(s3_conf.ak, s3_conf.sk);
     DCHECK(!aws_cred.IsExpiredOrEmpty());
+    if (!s3_conf.token.empty()) {
+        aws_cred.SetSessionToken(s3_conf.token);
+    }
 
     Aws::Client::ClientConfiguration aws_config = S3ClientFactory::getClientConfiguration();
     aws_config.endpointOverride = s3_conf.endpoint;
@@ -189,6 +192,9 @@ Status S3ClientFactory::convert_properties_to_s3_conf(
     StringCaseMap<std::string> properties(prop.begin(), prop.end());
     s3_conf->ak = properties.find(S3_AK)->second;
     s3_conf->sk = properties.find(S3_SK)->second;
+    if (properties.find(S3_TOKEN) != properties.end()) {
+        s3_conf->token = properties.find(S3_TOKEN)->second;
+    }
     s3_conf->endpoint = properties.find(S3_ENDPOINT)->second;
     s3_conf->region = properties.find(S3_REGION)->second;
 

diff --git a/be/src/util/s3_util.h b/be/src/util/s3_util.h
@@ -69,6 +69,7 @@ const static std::string S3_CONN_TIMEOUT_MS = "AWS_CONNECTION_TIMEOUT_MS";
 struct S3Conf {
     std::string ak;
     std::string sk;
+    std::string token;
     std::string endpoint;
     std::string region;
     std::string bucket;
@@ -80,16 +81,18 @@ struct S3Conf {
 
     std::string to_string() const {
         return fmt::format(
-                "(ak={}, sk=*, endpoint={}, region={}, bucket={}, prefix={}, max_connections={}, "
-                "request_timeout_ms={}, connect_timeout_ms={}, use_virtual_addressing={})",
-                ak, endpoint, region, bucket, prefix, max_connections, request_timeout_ms,
+                "(ak={}, sk=*, token={}, endpoint={}, region={}, bucket={}, prefix={}, "
+                "max_connections={}, request_timeout_ms={}, connect_timeout_ms={}, "
+                "use_virtual_addressing={})",
+                ak, token, endpoint, region, bucket, prefix, max_connections, request_timeout_ms,
                 connect_timeout_ms, use_virtual_addressing);
     }
 
     uint64_t get_hash() const {
         uint64_t hash_code = 0;
         hash_code += Fingerprint(ak);
         hash_code += Fingerprint(sk);
+        hash_code += Fingerprint(token);
         hash_code += Fingerprint(endpoint);
         hash_code += Fingerprint(region);
         hash_code += Fingerprint(bucket);