Add managed identity

apache · Dec 21, 2023 · bdb0a00 · bdb0a00
1 parent 2308cdf
commit bdb0a00
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 0 deletions.
diff --git a/cpp/src/arrow/filesystem/azurefs.cc b/cpp/src/arrow/filesystem/azurefs.cc
@@ -119,6 +119,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();
+}
+
 Status AzureOptions::ConfigureWorkloadIdentityCredential(
     const std::string& account_name) {
   credential_kind_ = CredentialKind::kTokenCredential;

diff --git a/cpp/src/arrow/filesystem/azurefs.h b/cpp/src/arrow/filesystem/azurefs.h
@@ -103,6 +103,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 ConfigureWorkloadIdentityCredential(const std::string& account_name);
 
   Status ConfigureAccountKeyCredential(const std::string& account_name,

diff --git a/cpp/src/arrow/filesystem/azurefs_test.cc b/cpp/src/arrow/filesystem/azurefs_test.cc
@@ -278,6 +278,16 @@ TEST(AzureFileSystem, InitializeFilesystemWithDefaultCredential) {
   EXPECT_OK_AND_ASSIGN(auto 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, InitializeFilesystemWithWorkloadIdentityCredential) {
   AzureOptions options;
   ARROW_EXPECT_OK(options.ConfigureWorkloadIdentityCredential("dummy-account-name"));