implement

ByteYue · Jun 25, 2024 · bff8048 · bff8048
1 parent 4b975ae
commit bff8048
Show file tree

Hide file tree

Showing 6 changed files with 32 additions and 7 deletions.
diff --git a/be/src/io/fs/azure_obj_storage_client.cpp b/be/src/io/fs/azure_obj_storage_client.cpp
@@ -26,15 +26,18 @@
 #include <azure/storage/blobs.hpp>
 #include <azure/storage/blobs/blob_client.hpp>
 #include <azure/storage/blobs/blob_container_client.hpp>
+#include <azure/storage/blobs/blob_sas_builder.hpp>
 #include <azure/storage/blobs/rest_client.hpp>
 #include <azure/storage/common/storage_credential.hpp>
 #include <azure/storage/common/storage_exception.hpp>
+#include <chrono>
 #include <exception>
 #include <iterator>
 
 #include "common/logging.h"
 #include "common/status.h"
 #include "io/fs/obj_storage_client.h"
+#include "util/s3_util.h"
 
 namespace {
 std::string wrap_object_storage_path_msg(const doris::io::ObjectStoragePathOptions& opts) {
@@ -46,6 +49,8 @@ auto base64_encode_part_num(int part_num) {
     return Aws::Utils::HashingUtils::Base64Encode(
             {reinterpret_cast<unsigned char*>(&part_num), sizeof(part_num)});
 }
+
+constexpr char SAS_TOKEN_URL_TEMPLATE[] = "https://{}.blob.core.windows.net/{}/{}{}";
 } // namespace
 
 namespace doris::io {
@@ -251,4 +256,21 @@ ObjectStorageResponse AzureObjStorageClient::delete_objects_recursively(
     }
     return {};
 }
+
+std::string AzureObjStorageClient::generate_presigned_url(const ObjectStoragePathOptions& opts,
+                                                          int64_t expiration_secs,
+                                                          const S3ClientConf& conf) {
+    Azure::Storage::Sas::BlobSasBuilder sas_builder;
+    sas_builder.ExpiresOn =
+            std::chrono::system_clock::now() + std::chrono::seconds(expiration_secs);
+    sas_builder.BlobContainerName = opts.bucket;
+    sas_builder.BlobName = opts.key;
+    sas_builder.Resource = Azure::Storage::Sas::BlobSasResource::Blob;
+    sas_builder.SetPermissions(Azure::Storage::Sas::BlobSasPermissions::Write);
+
+    std::string sasToken = sas_builder.GenerateSasToken(
+            Azure::Storage::StorageSharedKeyCredential(conf.ak, conf.sk));
+
+    return fmt::format(SAS_TOKEN_URL_TEMPLATE, conf.ak, conf.bucket, opts.key, sasToken);
+}
 } // namespace doris::io
diff --git a/be/src/io/fs/azure_obj_storage_client.h b/be/src/io/fs/azure_obj_storage_client.h
@@ -51,14 +51,13 @@ class AzureObjStorageClient final : public ObjStorageClient {
                                          std::vector<std::string> objs) override;
     ObjectStorageResponse delete_object(const ObjectStoragePathOptions& opts) override;
     ObjectStorageResponse delete_objects_recursively(const ObjectStoragePathOptions& opts) override;
-    // TODO(ByteYue) : to be implemented
     std::string generate_presigned_url(const ObjectStoragePathOptions& opts,
-                                       int64_t expiration_secs) override {
-        return "http://azure.to.be.implenmented";
-    };
+                                       int64_t expiration_secs, const S3ClientConf& conf) override;
 
 private:
     std::shared_ptr<Azure::Storage::Blobs::BlobContainerClient> _client;
+    std::string account_name;
+    std::string account_key;
 };
 
 } // namespace doris::io
diff --git a/be/src/io/fs/obj_storage_client.h b/be/src/io/fs/obj_storage_client.h
@@ -23,6 +23,7 @@
 #include "io/fs/path.h"
 namespace doris {
 class Status;
+class S3ClientConf;
 namespace io {
 
 // Names are in lexico order.
@@ -118,7 +119,8 @@ class ObjStorageClient {
             const ObjectStoragePathOptions& opts) = 0;
     // Return a presigned URL for users to access the object
     virtual std::string generate_presigned_url(const ObjectStoragePathOptions& opts,
-                                               int64_t expiration_secs) = 0;
+                                               int64_t expiration_secs,
+                                               const S3ClientConf& conf) = 0;
 };
 } // namespace io
 } // namespace doris
diff --git a/be/src/io/fs/s3_obj_storage_client.cpp b/be/src/io/fs/s3_obj_storage_client.cpp
@@ -393,7 +393,8 @@ ObjectStorageResponse S3ObjStorageClient::delete_objects_recursively(
 }
 
 std::string S3ObjStorageClient::generate_presigned_url(const ObjectStoragePathOptions& opts,
-                                                       int64_t expiration_secs) {
+                                                       int64_t expiration_secs,
+                                                       const S3ClientConf&) {
     return _client->GeneratePresignedUrl(opts.bucket, opts.key, Aws::Http::HttpMethod::HTTP_GET,
                                          expiration_secs);
 }

diff --git a/be/src/io/fs/s3_obj_storage_client.h b/be/src/io/fs/s3_obj_storage_client.h
@@ -54,7 +54,7 @@ class S3ObjStorageClient final : public ObjStorageClient {
     ObjectStorageResponse delete_object(const ObjectStoragePathOptions& opts) override;
     ObjectStorageResponse delete_objects_recursively(const ObjectStoragePathOptions& opts) override;
     std::string generate_presigned_url(const ObjectStoragePathOptions& opts,
-                                       int64_t expiration_secs) override;
+                                       int64_t expiration_secs, const S3ClientConf&) override;
 
 private:
     std::shared_ptr<Aws::S3::S3Client> _client;

diff --git a/be/src/util/s3_util.h b/be/src/util/s3_util.h
@@ -107,6 +107,7 @@ struct S3ClientConf {
         hash_code ^= request_timeout_ms;
         hash_code ^= connect_timeout_ms;
         hash_code ^= use_virtual_addressing;
+        hash_code ^= static_cast<int>(provider);
         return hash_code;
     }