From 91090ae3b477bb1eda0b45a22b1bb139f02a766d Mon Sep 17 00:00:00 2001
From: Thomas Newton <thomas.w.newton@gmail.com>
Date: Wed, 20 Dec 2023 12:48:54 +0000
Subject: [PATCH] Add workload identity credential

---
 cpp/src/arrow/filesystem/azurefs.cc      | 6 ++++++
 cpp/src/arrow/filesystem/azurefs.h       | 2 ++
 cpp/src/arrow/filesystem/azurefs_test.cc | 8 +++++++-
 3 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/cpp/src/arrow/filesystem/azurefs.cc b/cpp/src/arrow/filesystem/azurefs.cc
index 1aa3e86a6f926..8b22c795226de 100644
--- a/cpp/src/arrow/filesystem/azurefs.cc
+++ b/cpp/src/arrow/filesystem/azurefs.cc
@@ -117,6 +117,12 @@ Status AzureOptions::ConfigureDefaultCredential(const std::string& account_name)
   return Status::OK();
 }
 
+Status AzureOptions::ConfigureWorkloadIdentityCredential(const std::string& account_name) {
+  credential_kind_ = CredentialKind::kTokenCredential;
+  token_credential_ = std::make_shared<Azure::Identity::WorkloadIdentityCredential>();
+  return Status::OK();
+}
+
 Result<std::unique_ptr<Blobs::BlobServiceClient>> AzureOptions::MakeBlobServiceClient()
     const {
   switch (credential_kind_) {
diff --git a/cpp/src/arrow/filesystem/azurefs.h b/cpp/src/arrow/filesystem/azurefs.h
index 35c140b1097c7..dafc8be44969c 100644
--- a/cpp/src/arrow/filesystem/azurefs.h
+++ b/cpp/src/arrow/filesystem/azurefs.h
@@ -102,6 +102,8 @@ struct ARROW_EXPORT AzureOptions {
 
   Status ConfigureDefaultCredential(const std::string& account_name);
 
+  Status ConfigureWorkloadIdentityCredential(const std::string& account_name);
+
   Status ConfigureAccountKeyCredential(const std::string& account_name,
                                        const std::string& account_key);
 
diff --git a/cpp/src/arrow/filesystem/azurefs_test.cc b/cpp/src/arrow/filesystem/azurefs_test.cc
index 8a39c4c554897..0086cdaad6d41 100644
--- a/cpp/src/arrow/filesystem/azurefs_test.cc
+++ b/cpp/src/arrow/filesystem/azurefs_test.cc
@@ -273,7 +273,13 @@ class AzureHierarchicalNSEnv : public AzureEnvImpl<AzureHierarchicalNSEnv> {
 TEST(AzureFileSystem, InitializeFilesystemWithDefaultCredential) {
   AzureOptions options;
   ARROW_EXPECT_OK(options.ConfigureDefaultCredential("dummy-account-name"));
-  EXPECT_OK_AND_ASSIGN(auto default_credential_fs, AzureFileSystem::Make(options));
+  EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
+}
+
+TEST(AzureFileSystem, InitializeFilesystemWithWorkloadIdentityCredential) {
+  AzureOptions options;
+  ARROW_EXPECT_OK(options.ConfigureDefaultCredential("dummy-account-name"));
+  EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
 }
 
 TEST(AzureFileSystem, OptionsCompare) {