From 142a8b64dbaa732f4efb0d5b12beaf8c9e194d7e Mon Sep 17 00:00:00 2001
From: Thomas Newton <thomas.w.newton@gmail.com>
Date: Wed, 20 Dec 2023 13:08:28 +0000
Subject: [PATCH] Add managed identity

---
 cpp/src/arrow/filesystem/azurefs.cc      | 8 ++++++++
 cpp/src/arrow/filesystem/azurefs.h       | 3 +++
 cpp/src/arrow/filesystem/azurefs_test.cc | 9 +++++++++
 3 files changed, 20 insertions(+)

diff --git a/cpp/src/arrow/filesystem/azurefs.cc b/cpp/src/arrow/filesystem/azurefs.cc
index 1aa3e86a6f926..122fdcd8b7a20 100644
--- a/cpp/src/arrow/filesystem/azurefs.cc
+++ b/cpp/src/arrow/filesystem/azurefs.cc
@@ -117,6 +117,14 @@ Status AzureOptions::ConfigureDefaultCredential(const std::string& account_name)
   return Status::OK();
 }
 
+Status AzureOptions::ConfigureManagedIdentityCredential(const std::string& account_name,
+                                                        std::string const& clientId) {
+  credential_kind_ = CredentialKind::kTokenCredential;
+  token_credential_ =
+      std::make_shared<Azure::Identity::ManagedIdentityCredential>(clientId);
+  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..85aa6bd71c284 100644
--- a/cpp/src/arrow/filesystem/azurefs.h
+++ b/cpp/src/arrow/filesystem/azurefs.h
@@ -102,6 +102,9 @@ struct ARROW_EXPORT AzureOptions {
 
   Status ConfigureDefaultCredential(const std::string& account_name);
 
+  Status ConfigureManagedIdentityCredential(const std::string& account_name,
+                                            std::string const& clientId = std::string());
+
   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..f5ebe98340275 100644
--- a/cpp/src/arrow/filesystem/azurefs_test.cc
+++ b/cpp/src/arrow/filesystem/azurefs_test.cc
@@ -276,6 +276,15 @@ TEST(AzureFileSystem, InitializeFilesystemWithDefaultCredential) {
   EXPECT_OK_AND_ASSIGN(auto default_credential_fs, AzureFileSystem::Make(options));
 }
 
+TEST(AzureFileSystem, InitializeFilesystemWithManagedIdentityCredential) {
+  AzureOptions options;
+  ARROW_EXPECT_OK(options.ConfigureManagedIdentityCredential("dummy-account-name"));
+  EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options));
+
+  ARROW_EXPECT_OK(options.ConfigureManagedIdentityCredential("dummy-account-name", "specific-client-id"));
+  EXPECT_OK_AND_ASSIGN(fs, AzureFileSystem::Make(options));
+}
+
 TEST(AzureFileSystem, OptionsCompare) {
   AzureOptions options;
   EXPECT_TRUE(options.Equals(options));