From 91090ae3b477bb1eda0b45a22b1bb139f02a766d Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Wed, 20 Dec 2023 12:48:54 +0000 Subject: [PATCH 1/3] 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(); + return Status::OK(); +} + Result> 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 { 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) { From 0f6aab3ca40bdcbff3eaa006d4eaf1a275709829 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Wed, 20 Dec 2023 12:57:14 +0000 Subject: [PATCH 2/3] Auto-format --- cpp/src/arrow/filesystem/azurefs.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/src/arrow/filesystem/azurefs.cc b/cpp/src/arrow/filesystem/azurefs.cc index 8b22c795226de..fcaf26f31d05e 100644 --- a/cpp/src/arrow/filesystem/azurefs.cc +++ b/cpp/src/arrow/filesystem/azurefs.cc @@ -117,7 +117,8 @@ Status AzureOptions::ConfigureDefaultCredential(const std::string& account_name) return Status::OK(); } -Status AzureOptions::ConfigureWorkloadIdentityCredential(const std::string& account_name) { +Status AzureOptions::ConfigureWorkloadIdentityCredential( + const std::string& account_name) { credential_kind_ = CredentialKind::kTokenCredential; token_credential_ = std::make_shared(); return Status::OK(); From a81fc3468fda17ffaf23777a7b8edc9c742b3b73 Mon Sep 17 00:00:00 2001 From: Thomas Newton Date: Wed, 20 Dec 2023 13:07:04 +0000 Subject: [PATCH 3/3] Test the right auth --- cpp/src/arrow/filesystem/azurefs_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/arrow/filesystem/azurefs_test.cc b/cpp/src/arrow/filesystem/azurefs_test.cc index 0086cdaad6d41..130bb995eed6d 100644 --- a/cpp/src/arrow/filesystem/azurefs_test.cc +++ b/cpp/src/arrow/filesystem/azurefs_test.cc @@ -278,7 +278,7 @@ TEST(AzureFileSystem, InitializeFilesystemWithDefaultCredential) { TEST(AzureFileSystem, InitializeFilesystemWithWorkloadIdentityCredential) { AzureOptions options; - ARROW_EXPECT_OK(options.ConfigureDefaultCredential("dummy-account-name")); + ARROW_EXPECT_OK(options.ConfigureWorkloadIdentityCredential("dummy-account-name")); EXPECT_OK_AND_ASSIGN(auto fs, AzureFileSystem::Make(options)); }