apacheGH-41095: [C++][FS][Azure] Add support for CopyFile with hierar…

…chical namespace support We need to add SAS (Shared Access Signatures) token for source URL.
kou · Apr 18, 2024 · 2d0b183 · 2d0b183
1 parent be5f5b8
commit 2d0b183
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 8 deletions.
diff --git a/cpp/src/arrow/filesystem/azurefs.cc b/cpp/src/arrow/filesystem/azurefs.cc
@@ -277,6 +277,13 @@ std::string BuildBaseUrl(const std::string& scheme, const std::string& authority
   url += "/";
   return url;
 }
+
+template <typename... PrefixArgs>
+Status ExceptionToStatus(const Storage::StorageException& exception,
+                         PrefixArgs&&... prefix_args) {
+  return Status::IOError(std::forward<PrefixArgs>(prefix_args)..., " Azure Error: [",
+                         exception.ErrorCode, "] ", exception.what());
+}
 }  // namespace
 
 std::string AzureOptions::AccountBlobUrl(const std::string& account_name) const {
@@ -381,6 +388,24 @@ AzureOptions::MakeDataLakeServiceClient() const {
   return Status::Invalid("AzureOptions doesn't contain a valid auth configuration");
 }
 
+Result<std::string> AzureOptions::GenerateSASToken(
+    Storage::Sas::BlobSasBuilder* builder) const {
+  if (storage_shared_key_credential_) {
+    return builder->GenerateSasToken(*storage_shared_key_credential_);
+  } else {
+    // This part isn't tested. This may not work.
+    ARROW_ASSIGN_OR_RAISE(auto client, MakeBlobServiceClient());
+    try {
+      auto delegation_key_response = client->GetUserDelegationKey(builder->ExpiresOn);
+
+      return builder->GenerateSasToken(delegation_key_response.Value, account_name);
+    } catch (const Storage::StorageException& exception) {
+      return ExceptionToStatus(exception, "GetUserDelegationKey failed for '",
+                               client->GetUrl(), "'.");
+    }
+  }
+}
+
 namespace {
 
 // An AzureFileSystem represents an Azure storage account. An AzureLocation describes a
@@ -450,13 +475,6 @@ struct AzureLocation {
   }
 };
 
-template <typename... PrefixArgs>
-Status ExceptionToStatus(const Storage::StorageException& exception,
-                         PrefixArgs&&... prefix_args) {
-  return Status::IOError(std::forward<PrefixArgs>(prefix_args)..., " Azure Error: [",
-                         exception.ErrorCode, "] ", exception.what());
-}
-
 Status PathNotFound(const AzureLocation& location) {
   return ::arrow::fs::internal::PathNotFound(location.all);
 }
@@ -2868,8 +2886,19 @@ class AzureFileSystem::Impl {
     if (src == dest) {
       return Status::OK();
     }
+    std::string sas_token;
+    {
+      Storage::Sas::BlobSasBuilder builder;
+      std::chrono::seconds available_period(60);
+      builder.ExpiresOn = std::chrono::system_clock::now() + available_period;
+      builder.BlobContainerName = src.container;
+      builder.BlobName = src.path;
+      builder.Resource = Storage::Sas::BlobSasResource::Blob;
+      builder.SetPermissions(Storage::Sas::BlobSasPermissions::Read);
+      ARROW_ASSIGN_OR_RAISE(sas_token, options_.GenerateSASToken(&builder));
+    }
+    auto src_url = GetBlobClient(src.container, src.path).GetUrl() + sas_token;
     auto dest_blob_client = GetBlobClient(dest.container, dest.path);
-    auto src_url = GetBlobClient(src.container, src.path).GetUrl();
     if (!dest.path.empty()) {
       auto dest_parent = dest.parent();
       if (!dest_parent.path.empty()) {

diff --git a/cpp/src/arrow/filesystem/azurefs.h b/cpp/src/arrow/filesystem/azurefs.h
@@ -37,6 +37,10 @@ namespace Azure::Storage::Blobs {
 class BlobServiceClient;
 }
 
+namespace Azure::Storage::Sas {
+class BlobSasBuilder;
+}
+
 namespace Azure::Storage::Files::DataLake {
 class DataLakeFileSystemClient;
 class DataLakeServiceClient;
@@ -196,6 +200,9 @@ struct ARROW_EXPORT AzureOptions {
 
   Result<std::unique_ptr<Azure::Storage::Files::DataLake::DataLakeServiceClient>>
   MakeDataLakeServiceClient() const;
+
+  Result<std::string> GenerateSASToken(
+      Azure::Storage::Sas::BlobSasBuilder* builder) const;
 };
 
 /// \brief FileSystem implementation backed by Azure Blob Storage (ABS) [1] and