GH-39345: [C++][FS][Azure] Add support for environment credential (#4…

…1715) ### Rationale for this change Maybe be useful to support explicit environment credential (currently environment credential can be used as part of the Azure default credential flow). ### What changes are included in this PR? ### Are these changes tested? There are new unittests but no integration tests that we can actually authenticate successfully. We are relying on the Azure C++ SDK to abstracting that away. ### Are there any user-facing changes? Yes, environment credential is now available. * GitHub Issue: #39345 Authored-by: Thomas Newton <thomas.w.newton@gmail.com> Signed-off-by: Sutou Kouhei <kou@clear-code.com>
apache · Jun 3, 2024 · 02585cd · 02585cd
1 parent 44070eb
commit 02585cd
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 3 deletions.
diff --git a/cpp/src/arrow/filesystem/azurefs.cc b/cpp/src/arrow/filesystem/azurefs.cc
@@ -119,6 +119,8 @@ Status AzureOptions::ExtractFromUriQuery(const Uri& uri) {
         credential_kind = CredentialKind::kAnonymous;
       } else if (kv.second == "workload_identity") {
         credential_kind = CredentialKind::kWorkloadIdentity;
+      } else if (kv.second == "environment") {
+        credential_kind = CredentialKind::kEnvironment;
       } else {
         // Other credential kinds should be inferred from the given
         // parameters automatically.
@@ -171,6 +173,9 @@ Status AzureOptions::ExtractFromUriQuery(const Uri& uri) {
       case CredentialKind::kWorkloadIdentity:
         RETURN_NOT_OK(ConfigureWorkloadIdentityCredential());
         break;
+      case CredentialKind::kEnvironment:
+        RETURN_NOT_OK(ConfigureEnvironmentCredential());
+        break;
       default:
         // Default credential
         break;
@@ -252,6 +257,7 @@ bool AzureOptions::Equals(const AzureOptions& other) const {
     case CredentialKind::kClientSecret:
     case CredentialKind::kManagedIdentity:
     case CredentialKind::kWorkloadIdentity:
+    case CredentialKind::kEnvironment:
       return token_credential_->GetCredentialName() ==
              other.token_credential_->GetCredentialName();
   }
@@ -337,6 +343,12 @@ Status AzureOptions::ConfigureWorkloadIdentityCredential() {
   return Status::OK();
 }
 
+Status AzureOptions::ConfigureEnvironmentCredential() {
+  credential_kind_ = CredentialKind::kEnvironment;
+  token_credential_ = std::make_shared<Azure::Identity::EnvironmentCredential>();
+  return Status::OK();
+}
+
 Result<std::unique_ptr<Blobs::BlobServiceClient>> AzureOptions::MakeBlobServiceClient()
     const {
   if (account_name.empty()) {
@@ -353,6 +365,7 @@ Result<std::unique_ptr<Blobs::BlobServiceClient>> AzureOptions::MakeBlobServiceC
     case CredentialKind::kClientSecret:
     case CredentialKind::kManagedIdentity:
     case CredentialKind::kWorkloadIdentity:
+    case CredentialKind::kEnvironment:
       return std::make_unique<Blobs::BlobServiceClient>(AccountBlobUrl(account_name),
                                                         token_credential_);
     case CredentialKind::kStorageSharedKey:
@@ -379,6 +392,7 @@ AzureOptions::MakeDataLakeServiceClient() const {
     case CredentialKind::kClientSecret:
     case CredentialKind::kManagedIdentity:
     case CredentialKind::kWorkloadIdentity:
+    case CredentialKind::kEnvironment:
       return std::make_unique<DataLake::DataLakeServiceClient>(
           AccountDfsUrl(account_name), token_credential_);
     case CredentialKind::kStorageSharedKey:

diff --git a/cpp/src/arrow/filesystem/azurefs.h b/cpp/src/arrow/filesystem/azurefs.h
@@ -120,6 +120,7 @@ struct ARROW_EXPORT AzureOptions {
     kClientSecret,
     kManagedIdentity,
     kWorkloadIdentity,
+    kEnvironment,
   } credential_kind_ = CredentialKind::kDefault;
 
   std::shared_ptr<Azure::Storage::StorageSharedKeyCredential>
@@ -160,11 +161,13 @@ struct ARROW_EXPORT AzureOptions {
   /// * dfs_storage_authority: Set AzureOptions::dfs_storage_authority
   /// * enable_tls: If it's "false" or "0", HTTP not HTTPS is used.
   /// * credential_kind: One of "default", "anonymous",
-  ///   "workload_identity". If "default" is specified, it's just
-  ///   ignored.  If "anonymous" is specified,
+  ///   "workload_identity" or "environment". If "default" is specified, it's
+  ///   just ignored.  If "anonymous" is specified,
   ///   AzureOptions::ConfigureAnonymousCredential() is called. If
   ///   "workload_identity" is specified,
-  ///   AzureOptions::ConfigureWorkloadIdentityCredential() is called.
+  ///   AzureOptions::ConfigureWorkloadIdentityCredential() is called, If
+  ///   "environment" is specified,
+  ///   AzureOptions::ConfigureEnvironmentCredential() is called.
   /// * tenant_id: You must specify "client_id" and "client_secret"
   ///   too. AzureOptions::ConfigureClientSecretCredential() is called.
   /// * client_id: If you don't specify "tenant_id" and
@@ -188,6 +191,7 @@ struct ARROW_EXPORT AzureOptions {
                                          const std::string& client_secret);
   Status ConfigureManagedIdentityCredential(const std::string& client_id = std::string());
   Status ConfigureWorkloadIdentityCredential();
+  Status ConfigureEnvironmentCredential();
 
   bool Equals(const AzureOptions& other) const;
 

diff --git a/cpp/src/arrow/filesystem/azurefs_test.cc b/cpp/src/arrow/filesystem/azurefs_test.cc
@@ -528,6 +528,13 @@ TEST(AzureFileSystem, InitializeWithWorkloadIdentityCredential) {
   EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
 }
 
+TEST(AzureFileSystem, InitializeWithEnvironmentCredential) {
+  AzureOptions options;
+  options.account_name = "dummy-account-name";
+  ARROW_EXPECT_OK(options.ConfigureEnvironmentCredential());
+  EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
+}
+
 TEST(AzureFileSystem, OptionsCompare) {
   AzureOptions options;
   EXPECT_TRUE(options.Equals(options));
@@ -669,6 +676,15 @@ class TestAzureOptions : public ::testing::Test {
     ASSERT_EQ(options.credential_kind_, AzureOptions::CredentialKind::kWorkloadIdentity);
   }
 
+  void TestFromUriCredentialEnvironment() {
+    ASSERT_OK_AND_ASSIGN(
+        auto options,
+        AzureOptions::FromUri("abfs://account.blob.core.windows.net/container/dir/blob?"
+                              "credential_kind=environment",
+                              nullptr));
+    ASSERT_EQ(options.credential_kind_, AzureOptions::CredentialKind::kEnvironment);
+  }
+
   void TestFromUriCredentialInvalid() {
     ASSERT_RAISES(Invalid, AzureOptions::FromUri(
                                "abfs://file_system@account.dfs.core.windows.net/dir/file?"
@@ -720,6 +736,9 @@ TEST_F(TestAzureOptions, FromUriCredentialManagedIdentity) {
 TEST_F(TestAzureOptions, FromUriCredentialWorkloadIdentity) {
   TestFromUriCredentialWorkloadIdentity();
 }
+TEST_F(TestAzureOptions, FromUriCredentialEnvironment) {
+  TestFromUriCredentialEnvironment();
+}
 TEST_F(TestAzureOptions, FromUriCredentialInvalid) { TestFromUriCredentialInvalid(); }
 TEST_F(TestAzureOptions, FromUriBlobStorageAuthority) {
   TestFromUriBlobStorageAuthority();