Add managed identity

Tom-Newton · Dec 20, 2023 · 142a8b6 · 142a8b6
1 parent 87865b5
commit 142a8b6
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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));